Review comments

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
+	stateKeyNIDMap map[string]int64
-	state       stateEntryMap
+	state          stateEntryMap
-	events      eventMap
+	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 result.stateKeyNIDMap, err = db.EventStateKeyNIDs(eventStateKeys); err != nil {
-	if 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 result.events, err = db.Events(eventNIDs); err != nil {
-	if 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
+// Map from event type, state key tuple to numeric event ID.
-
+// Implemented using binary search on a sorted array.
-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
-}
-
 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.

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)" +
+	" WHERE event_state_key = ANY($1)"
-	" ORDER BY event_state_key ASC"
 
 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))
+	result := make(map[string]int64, len(eventStateKeys))
-	i := 0
 	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() {

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)
 }