Inject state blocks from incremental polls into the room timeline

state/accumulator.go

@@ -131,20 +131,37 @@ func (a *Accumulator) roomInfoDelta(roomID string, events []Event) RoomInfo {
 	}
 }
 
+type InitialiseResult struct {
+	// AddedEvents is true iff this call to Initialise added new state events to the DB.
+	AddedEvents bool
+	// SnapshotID is the ID of the snapshot which incorporates all added events.
+	// It has no meaning if AddedEvents is False.
+	SnapshotID int64
+	// PrependTimelineEvents is empty if the room was not initialised prior to this call.
+	// Otherwise, it is an order-preserving subset of the `state` argument to Initialise
+	// containing all events that were not persisted prior to the Initialise call. These
+	// should be prepended to the room timeline by the caller.
+	PrependTimelineEvents []json.RawMessage
+}
+
 // Initialise starts a new sync accumulator for the given room using the given state as a baseline.
-// This will only take effect if this is the first time the v3 server has seen this room, and it wasn't
-// possible to get all events up to the create event (e.g Matrix HQ). Returns true if this call actually
-// added new events, along with the snapshot NID.
 //
+// This will only take effect if this is the first time the v3 server has seen this room, and it wasn't
+// possible to get all events up to the create event (e.g Matrix HQ).
 // This function:
 // - Stores these events
 // - Sets up the current snapshot based on the state list given.
-func (a *Accumulator) Initialise(roomID string, state []json.RawMessage) (bool, int64, error) {
+//
+// If the v3 server has seen this room before, this function
+//   - queries the DB to determine which state events are known to th server,
+//   - returns (via InitialiseResult.PrependTimelineEvents) a slice of unknown state events,
+//
+// and otherwise does nothing.
+func (a *Accumulator) Initialise(roomID string, state []json.RawMessage) (InitialiseResult, error) {
+	var res InitialiseResult
 	if len(state) == 0 {
-		return false, 0, nil
+		return res, nil
 	}
-	addedEvents := false
-	var snapID int64
 	err := sqlutil.WithTransaction(a.db, func(txn *sqlx.Tx) error {
 		// Attempt to short-circuit. This has to be done inside a transaction to make sure
 		// we don't race with multiple calls to Initialise with the same room ID.
@@ -153,8 +170,26 @@ func (a *Accumulator) Initialise(roomID string, state []json.RawMessage) (bool,
 			return fmt.Errorf("error fetching snapshot id for room %s: %s", roomID, err)
 		}
 		if snapshotID > 0 {
-			// we only initialise rooms once
-			logger.Info().Str("room_id", roomID).Int64("snapshot_id", snapshotID).Msg("Accumulator.Initialise called but current snapshot already exists, bailing early")
+			// Poller A has received a gappy sync v2 response with a state block, and
+			// we have seen this room before. If we knew for certain that there is some
+			// other active poller B in this room then we could safely skip this logic.
+
+			// Log at debug for now. If we find an unknown event, we'll return it so
+			// that the poller can log a warning.
+			logger.Debug().Str("room_id", roomID).Int64("snapshot_id", snapshotID).Msg("Accumulator.Initialise called with incremental state but current snapshot already exists.")
+			eventIDs := make([]string, len(state))
+			for i := range state {
+				eventIDs[i] = gjson.ParseBytes(state[i]).Get("event_id").Str
+			}
+			unknownEventIDs, err := a.eventsTable.SelectUnknownEventIDs(txn, eventIDs)
+			if err != nil {
+				return fmt.Errorf("error determing which event IDs are unknown: %s", err)
+			}
+			for i := range state {
+				if _, unknown := unknownEventIDs[eventIDs[i]]; unknown {
+					res.PrependTimelineEvents = append(res.PrependTimelineEvents, state[i])
+				}
+			}
 			return nil
 		}
 
@@ -184,28 +219,19 @@ func (a *Accumulator) Initialise(roomID string, state []json.RawMessage) (bool,
 		}
 
 		// pull out the event NIDs we just inserted
-		eventIDs := make([]string, len(events))
 		membershipEventIDs := make(map[string]struct{}, len(events))
-		for i := range eventIDs {
-			eventIDs[i] = events[i].ID
-			if events[i].Type == "m.room.member" {
-				membershipEventIDs[events[i].ID] = struct{}{}
+		for _, event := range events {
+			if event.Type == "m.room.member" {
+				membershipEventIDs[event.ID] = struct{}{}
 			}
 		}
-		idToNIDs, err := a.eventsTable.SelectNIDsByIDs(txn, eventIDs)
-		if err != nil {
-			return fmt.Errorf("failed to select NIDs for inserted events: %w", err)
-		}
-		if len(idToNIDs) != len(eventIDs) {
-			return fmt.Errorf("missing events just inserted, asked for %v got %v", eventIDs, idToNIDs)
-		}
-		memberNIDs := make([]int64, 0, len(idToNIDs))
-		otherNIDs := make([]int64, 0, len(idToNIDs))
-		for evID, nid := range idToNIDs {
+		memberNIDs := make([]int64, 0, len(eventIDToNID))
+		otherNIDs := make([]int64, 0, len(eventIDToNID))
+		for evID, nid := range eventIDToNID {
 			if _, exists := membershipEventIDs[evID]; exists {
-				memberNIDs = append(memberNIDs, nid)
+				memberNIDs = append(memberNIDs, int64(nid))
 			} else {
-				otherNIDs = append(otherNIDs, nid)
+				otherNIDs = append(otherNIDs, int64(nid))
 			}
 		}
 
@@ -219,7 +245,7 @@ func (a *Accumulator) Initialise(roomID string, state []json.RawMessage) (bool,
 		if err != nil {
 			return fmt.Errorf("failed to insert snapshot: %w", err)
 		}
-		addedEvents = true
+		res.AddedEvents = true
 		latestNID := int64(0)
 		for _, nid := range otherNIDs {
 			if nid > latestNID {
@@ -244,10 +270,10 @@ func (a *Accumulator) Initialise(roomID string, state []json.RawMessage) (bool,
 		// will have an associated state snapshot ID on the event.
 
 		// Set the snapshot ID as the current state
-		snapID = snapshot.SnapshotID
+		res.SnapshotID = snapshot.SnapshotID
 		return a.roomsTable.Upsert(txn, info, snapshot.SnapshotID, latestNID)
 	})
-	return addedEvents, snapID, err
+	return res, err
 }
 
 // Accumulate internal state from a user's sync response. The timeline order MUST be in the order

state/accumulator_test.go

@@ -23,11 +23,11 @@ func TestAccumulatorInitialise(t *testing.T) {
 	db, close := connectToDB(t)
 	defer close()
 	accumulator := NewAccumulator(db)
-	added, initSnapID, err := accumulator.Initialise(roomID, roomEvents)
+	res, err := accumulator.Initialise(roomID, roomEvents)
 	if err != nil {
 		t.Fatalf("falied to Initialise accumulator: %s", err)
 	}
-	if !added {
+	if !res.AddedEvents {
 		t.Fatalf("didn't add events, wanted it to")
 	}
 
@@ -45,8 +45,8 @@ func TestAccumulatorInitialise(t *testing.T) {
 	if snapID == 0 {
 		t.Fatalf("Initialise did not store a current snapshot")
 	}
-	if snapID != initSnapID {
-		t.Fatalf("Initialise returned wrong snapshot ID, got %v want %v", initSnapID, snapID)
+	if snapID != res.SnapshotID {
+		t.Fatalf("Initialise returned wrong snapshot ID, got %v want %v", res.SnapshotID, snapID)
 	}
 
 	// this snapshot should have 1 member event and 2 other events in it
@@ -80,11 +80,11 @@ func TestAccumulatorInitialise(t *testing.T) {
 	}
 
 	// Subsequent calls do nothing and are not an error
-	added, _, err = accumulator.Initialise(roomID, roomEvents)
+	res, err = accumulator.Initialise(roomID, roomEvents)
 	if err != nil {
 		t.Fatalf("falied to Initialise accumulator: %s", err)
 	}
-	if added {
+	if res.AddedEvents {
 		t.Fatalf("added events when it shouldn't have")
 	}
 }
@@ -99,7 +99,7 @@ func TestAccumulatorAccumulate(t *testing.T) {
 	db, close := connectToDB(t)
 	defer close()
 	accumulator := NewAccumulator(db)
-	_, _, err := accumulator.Initialise(roomID, roomEvents)
+	_, err := accumulator.Initialise(roomID, roomEvents)
 	if err != nil {
 		t.Fatalf("failed to Initialise accumulator: %s", err)
 	}
@@ -195,7 +195,7 @@ func TestAccumulatorDelta(t *testing.T) {
 	db, close := connectToDB(t)
 	defer close()
 	accumulator := NewAccumulator(db)
-	_, _, err := accumulator.Initialise(roomID, nil)
+	_, err := accumulator.Initialise(roomID, nil)
 	if err != nil {
 		t.Fatalf("failed to Initialise accumulator: %s", err)
 	}
@@ -244,7 +244,7 @@ func TestAccumulatorMembershipLogs(t *testing.T) {
 	db, close := connectToDB(t)
 	defer close()
 	accumulator := NewAccumulator(db)
-	_, _, err := accumulator.Initialise(roomID, nil)
+	_, err := accumulator.Initialise(roomID, nil)
 	if err != nil {
 		t.Fatalf("failed to Initialise accumulator: %s", err)
 	}
@@ -384,7 +384,7 @@ func TestAccumulatorDupeEvents(t *testing.T) {
 	defer close()
 	accumulator := NewAccumulator(db)
 	roomID := "!buggy:localhost"
-	_, _, err := accumulator.Initialise(roomID, joinRoom.State.Events)
+	_, err := accumulator.Initialise(roomID, joinRoom.State.Events)
 	if err != nil {
 		t.Fatalf("failed to Initialise accumulator: %s", err)
 	}
@@ -426,7 +426,7 @@ func TestAccumulatorMisorderedGraceful(t *testing.T) {
 	accumulator := NewAccumulator(db)
 	roomID := "!TestAccumulatorStateReset:localhost"
 	// Create a room with initial state A,C
-	_, _, err := accumulator.Initialise(roomID, []json.RawMessage{
+	_, err := accumulator.Initialise(roomID, []json.RawMessage{
 		eventA, eventC,
 	})
 	if err != nil {

state/event_table.go

@@ -259,6 +259,34 @@ func (t *EventTable) SelectStrippedEventsByIDs(txn *sqlx.Tx, verifyAll bool, ids
 
 }
 
+// SelectUnknownEventIDs accepts a list of event IDs and returns the subset of those which are not known to the DB.
+func (t *EventTable) SelectUnknownEventIDs(txn *sqlx.Tx, maybeUnknownEventIDs []string) (map[string]struct{}, error) {
+	// Note: in practice, the order of rows returned matches the order of rows of
+	// array entries. But I don't think that's guaranteed. Return an (unordered) set
+	// out of paranoia.
+	queryStr := `
+	WITH maybe_unknown_events(event_id) AS (SELECT unnest($1::text[]))
+	SELECT event_id
+	FROM maybe_unknown_events LEFT JOIN syncv3_events USING(event_id)
+	WHERE event_nid IS NULL;`
+
+	var unknownEventIDs []string
+	var err error
+	if txn != nil {
+		err = txn.Select(&unknownEventIDs, queryStr, pq.StringArray(maybeUnknownEventIDs))
+	} else {
+		err = t.db.Select(&unknownEventIDs, queryStr, pq.StringArray(maybeUnknownEventIDs))
+	}
+	if err != nil {
+		return nil, err
+	}
+	unknownMap := make(map[string]struct{}, len(unknownEventIDs))
+	for _, eventID := range unknownEventIDs {
+		unknownMap[eventID] = struct{}{}
+	}
+	return unknownMap, nil
+}
+
 // UpdateBeforeSnapshotID sets the before_state_snapshot_id field to `snapID` for the given NIDs.
 func (t *EventTable) UpdateBeforeSnapshotID(txn *sqlx.Tx, eventNID, snapID, replacesNID int64) error {
 	_, err := txn.Exec(

state/storage.go

@@ -309,7 +309,7 @@ func (s *Storage) Accumulate(roomID, prevBatch string, timeline []json.RawMessag
 	return s.accumulator.Accumulate(roomID, prevBatch, timeline)
 }
 
-func (s *Storage) Initialise(roomID string, state []json.RawMessage) (bool, int64, error) {
+func (s *Storage) Initialise(roomID string, state []json.RawMessage) (InitialiseResult, error) {
 	return s.accumulator.Initialise(roomID, state)
 }
 

state/storage_test.go

@@ -282,7 +282,7 @@ func TestVisibleEventNIDsBetween(t *testing.T) {
 		},
 	}
 	for roomID, eventMap := range roomIDToEventMap {
-		_, _, err := store.Initialise(roomID, eventMap)
+		_, err := store.Initialise(roomID, eventMap)
 		if err != nil {
 			t.Fatalf("Initialise on %s failed: %s", roomID, err)
 		}
@@ -512,7 +512,7 @@ func TestStorageLatestEventsInRoomsPrevBatch(t *testing.T) {
 		},
 	}
 
-	_, _, err := store.Initialise(roomID, stateEvents)
+	_, err := store.Initialise(roomID, stateEvents)
 	if err != nil {
 		t.Fatalf("failed to initialise: %s", err)
 	}
@@ -616,7 +616,7 @@ func TestGlobalSnapshot(t *testing.T) {
 	store := NewStorage(postgresConnectionString)
 	defer store.Teardown()
 	for roomID, stateEvents := range roomIDToEventMap {
-		_, _, err := store.Initialise(roomID, stateEvents)
+		_, err := store.Initialise(roomID, stateEvents)
 		assertNoError(t, err)
 	}
 	snapshot, err := store.GlobalSnapshot()

sync2/handler2/handler.go

@@ -240,21 +240,20 @@ func (h *Handler) Accumulate(deviceID, roomID, prevBatch string, timeline []json
 	})
 }
 
-func (h *Handler) Initialise(roomID string, state []json.RawMessage) {
-	added, snapID, err := h.Store.Initialise(roomID, state)
+func (h *Handler) Initialise(roomID string, state []json.RawMessage) []json.RawMessage {
+	res, err := h.Store.Initialise(roomID, state)
 	if err != nil {
 		logger.Err(err).Int("state", len(state)).Str("room", roomID).Msg("V2: failed to initialise room")
 		sentry.CaptureException(err)
-		return
+		return nil
 	}
-	if !added {
-		// no new events
-		return
+	if res.AddedEvents {
+		h.v2Pub.Notify(pubsub.ChanV2, &pubsub.V2Initialise{
+			RoomID:      roomID,
+			SnapshotNID: res.SnapshotID,
+		})
 	}
-	h.v2Pub.Notify(pubsub.ChanV2, &pubsub.V2Initialise{
-		RoomID:      roomID,
-		SnapshotNID: snapID,
-	})
+	return res.PrependTimelineEvents
 }
 
 func (h *Handler) SetTyping(roomID string, ephEvent json.RawMessage) {

sync2/poller.go

@@ -3,6 +3,7 @@ package sync2
 import (
 	"context"
 	"encoding/json"
+	"github.com/getsentry/sentry-go"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -23,7 +24,8 @@ type V2DataReceiver interface {
 	// Accumulate data for this room. This means the timeline section of the v2 response.
 	Accumulate(deviceID, roomID, prevBatch string, timeline []json.RawMessage) // latest pos with event nids of timeline entries
 	// Initialise the room, if it hasn't been already. This means the state section of the v2 response.
-	Initialise(roomID string, state []json.RawMessage) // snapshot ID?
+	// If given a state delta from an incremental sync, returns the slice of all state events unknown to the DB.
+	Initialise(roomID string, state []json.RawMessage) []json.RawMessage // snapshot ID?
 	// SetTyping indicates which users are typing.
 	SetTyping(roomID string, ephEvent json.RawMessage)
 	// Sent when there is a new receipt
@@ -196,14 +198,15 @@ func (h *PollerMap) Accumulate(deviceID, roomID, prevBatch string, timeline []js
 	}
 	wg.Wait()
 }
-func (h *PollerMap) Initialise(roomID string, state []json.RawMessage) {
+func (h *PollerMap) Initialise(roomID string, state []json.RawMessage) (result []json.RawMessage) {
 	var wg sync.WaitGroup
 	wg.Add(1)
 	h.executor <- func() {
-		h.callbacks.Initialise(roomID, state)
+		result = h.callbacks.Initialise(roomID, state)
 		wg.Done()
 	}
 	wg.Wait()
+	return
 }
 func (h *PollerMap) SetTyping(roomID string, ephEvent json.RawMessage) {
 	var wg sync.WaitGroup
@@ -496,7 +499,24 @@ func (p *poller) parseRoomsResponse(res *SyncResponse) {
 	for roomID, roomData := range res.Rooms.Join {
 		if len(roomData.State.Events) > 0 {
 			stateCalls++
-			p.receiver.Initialise(roomID, roomData.State.Events)
+			prependStateEvents := p.receiver.Initialise(roomID, roomData.State.Events)
+			if len(prependStateEvents) > 0 {
+				// The poller has just learned of these state events due to an
+				// incremental poller sync; we must have missed the opportunity to see
+				// these down /sync in a timeline. As a workaround, inject these into
+				// the timeline now so that future events are received under the
+				// correct room state.
+				const warnMsg = "parseRoomsResponse: prepending state events to timeline after gappy poll"
+				log.Warn().Str("room_id", roomID).Int("prependStateEvents", len(prependStateEvents)).Msg(warnMsg)
+				sentry.WithScope(func(scope *sentry.Scope) {
+					scope.SetContext("sliding-sync", map[string]interface{}{
+						"room_id":                  roomID,
+						"num_prepend_state_events": len(prependStateEvents),
+					})
+					sentry.CaptureMessage(warnMsg)
+				})
+				roomData.Timeline.Events = append(prependStateEvents, roomData.Timeline.Events...)
+			}
 		}
 		// process typing/receipts before events so we seed the caches correctly for when we return the room
 		for _, ephEvent := range roomData.Ephemeral.Events {

sync2/poller_test.go

@@ -460,14 +460,17 @@ type mockDataReceiver struct {
 func (a *mockDataReceiver) Accumulate(userID, roomID, prevBatch string, timeline []json.RawMessage) {
 	a.timelines[roomID] = append(a.timelines[roomID], timeline...)
 }
-func (a *mockDataReceiver) Initialise(roomID string, state []json.RawMessage) {
+func (a *mockDataReceiver) Initialise(roomID string, state []json.RawMessage) []json.RawMessage {
 	a.states[roomID] = state
 	if a.incomingProcess != nil {
 		a.incomingProcess <- struct{}{}
 	}
 	if a.unblockProcess != nil {
 		<-a.unblockProcess
 	}
+	// The return value is a list of unknown state events to be prepended to the room
+	// timeline. Untested here---return nil for now.
+	return nil
 }
 func (a *mockDataReceiver) SetTyping(roomID string, ephEvent json.RawMessage) {
 }