Merge pull request #202 from matrix-org/kegan/fast-lookup

Use the rooms table initially when querying latest nids

state/event_table.go

@@ -317,10 +317,10 @@ func (t *EventTable) LatestEventInRooms(txn *sqlx.Tx, roomIDs []string, highestN
 	return
 }
 
-func (t *EventTable) LatestEventNIDInRooms(roomIDs []string, highestNID int64) (roomToNID map[string]int64, err error) {
+func (t *EventTable) LatestEventNIDInRooms(txn *sqlx.Tx, roomIDs []string, highestNID int64) (roomToNID map[string]int64, err error) {
 	// the position (event nid) may be for a random different room, so we need to find the highest nid <= this position for this room
 	var events []Event
-	err = t.db.Select(
+	err = txn.Select(
 		&events,
 		`SELECT event_nid, room_id FROM syncv3_events
 		WHERE event_nid IN (SELECT max(event_nid) FROM syncv3_events WHERE event_nid <= $1 AND room_id = ANY($2) GROUP BY room_id)`,

state/event_table_test.go

@@ -951,7 +951,11 @@ func TestLatestEventNIDInRooms(t *testing.T) {
 		},
 	}
 	for _, tc := range testCases {
-		gotRoomToNID, err := table.LatestEventNIDInRooms(tc.roomIDs, int64(tc.highestNID))
+		var gotRoomToNID map[string]int64
+		err = sqlutil.WithTransaction(table.db, func(txn *sqlx.Tx) error {
+			gotRoomToNID, err = table.LatestEventNIDInRooms(txn, tc.roomIDs, int64(tc.highestNID))
+			return err
+		})
 		assertNoError(t, err)
 		want := make(map[string]int64) // map event IDs to nids
 		for roomID, eventID := range tc.wantMap {

state/storage.go

@@ -838,6 +838,43 @@ func (s *Storage) AllJoinedMembers(txn *sqlx.Tx, tempTableName string) (result m
 	return result, metadata, nil
 }
 
+func (s *Storage) LatestEventNIDInRooms(roomIDs []string, highestNID int64) (roomToNID map[string]int64, err error) {
+	roomToNID = make(map[string]int64)
+	err = sqlutil.WithTransaction(s.Accumulator.db, func(txn *sqlx.Tx) error {
+		// Pull out the latest nids for all the rooms. If they are < highestNID then use them, else we need to query the
+		// events table (slow) for the latest nid in this room which is < highestNID.
+		fastRoomToLatestNIDs, err := s.Accumulator.roomsTable.LatestNIDs(txn, roomIDs)
+		if err != nil {
+			return err
+		}
+		var slowRooms []string
+		for _, roomID := range roomIDs {
+			nid := fastRoomToLatestNIDs[roomID]
+			if nid > 0 && nid <= highestNID {
+				roomToNID[roomID] = nid
+			} else {
+				// we need to do a slow query for this
+				slowRooms = append(slowRooms, roomID)
+			}
+		}
+
+		if len(slowRooms) == 0 {
+			return nil // no work to do
+		}
+		logger.Warn().Int("slow_rooms", len(slowRooms)).Msg("LatestEventNIDInRooms: pos value provided is far behind the database copy, performance degraded")
+
+		slowRoomToLatestNIDs, err := s.EventsTable.LatestEventNIDInRooms(txn, slowRooms, highestNID)
+		if err != nil {
+			return err
+		}
+		for roomID, nid := range slowRoomToLatestNIDs {
+			roomToNID[roomID] = nid
+		}
+		return nil
+	})
+	return roomToNID, err
+}
+
 // Returns a map from joined room IDs to EventMetadata, which is nil iff a non-nil error
 // is returned.
 func (s *Storage) JoinedRoomsAfterPosition(userID string, pos int64) (

sync3/caches/global.go

@@ -158,7 +158,7 @@ func (c *GlobalCache) LoadJoinedRooms(ctx context.Context, userID string) (
 		i++
 	}
 
-	latestNIDs, err = c.store.EventsTable.LatestEventNIDInRooms(roomIDs, initialLoadPosition)
+	latestNIDs, err = c.store.LatestEventNIDInRooms(roomIDs, initialLoadPosition)
 	if err != nil {
 		return 0, nil, nil, nil, err
 	}