diff --git a/src/github.com/matrix-org/dendrite/roomserver/input/events.go b/src/github.com/matrix-org/dendrite/roomserver/input/events.go index cad8ede67..7dec91515 100644 --- a/src/github.com/matrix-org/dendrite/roomserver/input/events.go +++ b/src/github.com/matrix-org/dendrite/roomserver/input/events.go @@ -17,7 +17,7 @@ type RoomEventDatabase interface { StateEntriesForEventIDs(eventIDs []string) ([]types.StateEntry, error) // Lookup the numeric IDs for a list of string event state keys. // Returns a sorted list of state entries. - EventStateKeyNIDs(eventStateKeys []string) ([]types.IDPair, error) + EventStateKeyNIDs(eventStateKeys []string) (map[string]int64, error) // Lookup the Events for a list of numeric event IDs. // Returns a sorted list of state entries. Events(eventNIDs []int64) ([]types.Event, error) @@ -99,9 +99,9 @@ func checkAuthEvents(db RoomEventDatabase, event gomatrixserverlib.Event, authEv } type authEvents struct { - stateNIDMap idMap - state stateEntryMap - events eventMap + stateKeyNIDMap map[string]int64 + state stateEntryMap + events eventMap } func (ae *authEvents) Create() (*gomatrixserverlib.Event, error) { @@ -137,7 +137,7 @@ func (ae *authEvents) lookupEventWithEmptyStateKey(typeNID int64) *gomatrixserve } func (ae *authEvents) lookupEvent(typeNID int64, stateKey string) *gomatrixserverlib.Event { - stateKeyNID, ok := ae.stateNIDMap.lookup(stateKey) + stateKeyNID, ok := ae.stateKeyNIDMap[stateKey] if !ok { return nil } @@ -162,32 +162,28 @@ func loadAuthEvents( var eventStateKeys []string eventStateKeys = append(eventStateKeys, needed.Member...) eventStateKeys = append(eventStateKeys, needed.ThirdPartyInvite...) - stateKeyNIDs, err := db.EventStateKeyNIDs(eventStateKeys) - if err != nil { + if result.stateKeyNIDMap, err = db.EventStateKeyNIDs(eventStateKeys); err != nil { return } - result.stateNIDMap = newIDMap(stateKeyNIDs) // Load the events we need. - keysNeeded := stateKeysNeeded(result.stateNIDMap, needed) + result.state = state var eventNIDs []int64 - result.state = newStateEntryMap(state) + keysNeeded := stateKeysNeeded(result.stateKeyNIDMap, needed) for _, keyNeeded := range keysNeeded { eventNID, ok := result.state.lookup(keyNeeded) if ok { eventNIDs = append(eventNIDs, eventNID) } } - events, err := db.Events(eventNIDs) - if err != nil { + if result.events, err = db.Events(eventNIDs); err != nil { return } - result.events = newEventMap(events) return } // stateKeysNeeded works out which numeric state keys we need to authenticate some events. -func stateKeysNeeded(stateNIDMap idMap, stateNeeded gomatrixserverlib.StateNeeded) []types.StateKeyTuple { +func stateKeysNeeded(stateKeyNIDMap map[string]int64, stateNeeded gomatrixserverlib.StateNeeded) []types.StateKeyTuple { var keys []types.StateKeyTuple if stateNeeded.Create { keys = append(keys, types.StateKeyTuple{types.MRoomCreateNID, types.EmptyStateKeyNID}) @@ -199,13 +195,13 @@ func stateKeysNeeded(stateNIDMap idMap, stateNeeded gomatrixserverlib.StateNeede keys = append(keys, types.StateKeyTuple{types.MRoomJoinRulesNID, types.EmptyStateKeyNID}) } for _, member := range stateNeeded.Member { - stateKeyNID, ok := stateNIDMap.lookup(member) + stateKeyNID, ok := stateKeyNIDMap[member] if ok { keys = append(keys, types.StateKeyTuple{types.MRoomMemberNID, stateKeyNID}) } } for _, token := range stateNeeded.ThirdPartyInvite { - stateKeyNID, ok := stateNIDMap.lookup(token) + stateKeyNID, ok := stateKeyNIDMap[token] if ok { keys = append(keys, types.StateKeyTuple{types.MRoomThirdPartyInviteNID, stateKeyNID}) } @@ -213,35 +209,10 @@ func stateKeysNeeded(stateNIDMap idMap, stateNeeded gomatrixserverlib.StateNeede return keys } -type idMap map[string]int64 - -func newIDMap(ids []types.IDPair) idMap { - result := make(map[string]int64) - for _, pair := range ids { - result[pair.ID] = pair.NID - } - return idMap(result) -} - -// lookup an entry in the id map. -func (m idMap) lookup(id string) (nid int64, ok bool) { - // Use a hash map here. - // We could use binary search here like we do for the maps below as it - // would be faster for small lists. - // However the benefits of binary search aren't as strong here and it's - // possible that we could encounter sets of pathological strings since - // the state keys are ultimately controlled by user input. - nid, ok = map[string]int64(m)[id] - return -} - +// Map from event type, state key tuple to numeric event ID. +// Implemented using binary search on a sorted array. type stateEntryMap []types.StateEntry -// newStateEntryMap creates a map from a sorted list of state entries. -func newStateEntryMap(stateEntries []types.StateEntry) stateEntryMap { - return stateEntryMap(stateEntries) -} - // lookup an entry in the event map. func (m stateEntryMap) lookup(stateKey types.StateKeyTuple) (eventNID int64, ok bool) { // Since the list is sorted we can implement this using binary search. @@ -259,13 +230,10 @@ func (m stateEntryMap) lookup(stateKey types.StateKeyTuple) (eventNID int64, ok return } +// Map from numeric event ID to event. +// Implemented using binary search on a sorted array. type eventMap []types.Event -// newEventMap creates a map from a sorted list of events. -func newEventMap(events []types.Event) eventMap { - return eventMap(events) -} - // lookup an entry in the event map. func (m eventMap) lookup(eventNID int64) (event *types.Event, ok bool) { // Since the list is sorted we can implement this using binary search. diff --git a/src/github.com/matrix-org/dendrite/roomserver/storage/sql.go b/src/github.com/matrix-org/dendrite/roomserver/storage/sql.go index 3eaec1859..a93ded15d 100644 --- a/src/github.com/matrix-org/dendrite/roomserver/storage/sql.go +++ b/src/github.com/matrix-org/dendrite/roomserver/storage/sql.go @@ -237,10 +237,11 @@ const insertEventStateKeyNIDSQL = "" + const selectEventStateKeyNIDSQL = "" + "SELECT event_state_key_nid FROM event_state_keys WHERE event_state_key = $1" +// Bulk lookup from string state key to numeric ID for that state key. +// Takes an array of strings as the query parameter. const selectEventStateKeyNIDsSQL = "" + "SELECT event_state_key, event_state_key_nid FROM event_state_keys" + - " WHERE event_state_key = ANY($1)" + - " ORDER BY event_state_key ASC" + " WHERE event_state_key = ANY($1)" func (s *statements) insertEventStateKeyNID(eventStateKey string) (eventStateKeyNID int64, err error) { err = s.insertEventStateKeyNIDStmt.QueryRow(eventStateKey).Scan(&eventStateKeyNID) @@ -252,22 +253,23 @@ func (s *statements) selectEventStateKeyNID(eventStateKey string) (eventStateKey return } -func (s *statements) selectEventStateKeyNIDs(eventStateKeys []string) ([]types.IDPair, error) { +func (s *statements) selectEventStateKeyNIDs(eventStateKeys []string) (map[string]int64, error) { rows, err := s.selectEventStateKeyNIDsStmt.Query(pq.StringArray(eventStateKeys)) if err != nil { return nil, err } defer rows.Close() - results := make([]types.IDPair, len(eventStateKeys)) - i := 0 + result := make(map[string]int64, len(eventStateKeys)) for rows.Next() { - if err := rows.Scan(&results[i].ID, &results[i].NID); err != nil { + var stateKey string + var stateKeyNID int64 + if err := rows.Scan(&stateKey, &stateKeyNID); err != nil { return nil, err } - i++ + result[stateKey] = stateKeyNID } - return results[:i], nil + return result, nil } func (s *statements) prepareRooms(db *sql.DB) (err error) { @@ -339,17 +341,20 @@ CREATE TABLE IF NOT EXISTS events ( -- Needed for setting reference hashes when sending new events. reference_sha256 BYTEA NOT NULL, -- A list of numeric IDs for events that can authenticate this event. - auth_events BIGINT[] NOT NULL, + auth_event_nids BIGINT[] NOT NULL, ); ` const insertEventSQL = "" + - "INSERT INTO events (room_nid, event_type_nid, event_state_key_nid, event_id, reference_sha256, auth_events)" + + "INSERT INTO events (room_nid, event_type_nid, event_state_key_nid, event_id, reference_sha256, auth_event_nids)" + " VALUES ($1, $2, $3, $4, $5, $6)" + " ON CONFLICT ON CONSTRAINT event_id_unique" + " DO UPDATE SET event_id = $1" + " RETURNING event_nid" +// Bulk lookup of events by string ID. +// Sort by the numeric IDs for event type and state key. +// This means we can use binary search to lookup entries by type and state key. const selectStateEventsByIDSQL = "" + "SELECT event_type_nid, event_state_key_nid, event_nid FROM events" + " WHERE event_id = ANY($1)" + @@ -383,12 +388,16 @@ func (s *statements) insertEvent( } func (s *statements) selectStateEventsByID(eventIDs []string) ([]types.StateEntry, error) { - results := make([]types.StateEntry, len(eventIDs)) rows, err := s.selectStateEventsByIDStmt.Query(pq.StringArray(eventIDs)) if err != nil { return nil, err } defer rows.Close() + // We know that we will only get as many results as event IDs + // because of the unique constraint on event IDs. + // So we can allocate an array of the correct size now. + // We might get fewer results than IDs so we adjust the length of the slice before returning it. + results := make([]types.StateEntry, len(eventIDs)) i := 0 for ; rows.Next(); i++ { result := &results[i] @@ -438,6 +447,9 @@ const insertEventJSONSQL = "" + "INSERT INTO event_json (event_nid, event_json) VALUES ($1, $2)" + " ON CONFLICT DO NOTHING" +// Bulk event JSON lookup by numeric event ID. +// Sort by the numeric event ID. +// This means that we can use binary search to lookup by numeric event ID. const selectEventJSONsSQL = "" + "SELECT event_nid, event_json FROM event_json" + " WHERE event_nid = ANY($1)" + @@ -460,6 +472,10 @@ func (s *statements) selectEventJSONs(eventNIDs []int64) ([]eventJSONPair, error } defer rows.Close() + // We know that we will only get as many results as event NIDs + // because of the unique constraint on event NIDs. + // So we can allocate an array of the correct size now. + // We might get fewer results than NIDs so we adjust the length of the slice before returning it. results := make([]eventJSONPair, len(eventNIDs)) i := 0 for rows.Next() { diff --git a/src/github.com/matrix-org/dendrite/roomserver/storage/storage.go b/src/github.com/matrix-org/dendrite/roomserver/storage/storage.go index 65ea09d54..ed220c70e 100644 --- a/src/github.com/matrix-org/dendrite/roomserver/storage/storage.go +++ b/src/github.com/matrix-org/dendrite/roomserver/storage/storage.go @@ -123,7 +123,7 @@ func (d *Database) StateEntriesForEventIDs(eventIDs []string) ([]types.StateEntr } // EventStateKeyNIDs implements input.EventDatabase -func (d *Database) EventStateKeyNIDs(eventStateKeys []string) ([]types.IDPair, error) { +func (d *Database) EventStateKeyNIDs(eventStateKeys []string) (map[string]int64, error) { return d.statements.selectEventStateKeyNIDs(eventStateKeys) }