Roomserver API changes (#1118)

* s/QueryBackfill/PerformBackfill/g

* OutputEvent now includes AddStateEvents which contain the full event of extra state events

* Only include adds not the current event

* Get adding state right

roomserver/internal/input_latest_events.go

@@ -19,6 +19,7 @@ package internal
 import (
 	"bytes"
 	"context"
+	"fmt"
 
 	"github.com/matrix-org/dendrite/internal"
 	"github.com/matrix-org/dendrite/roomserver/api"
@@ -310,24 +311,11 @@ func (u *latestEventsUpdater) makeOutputNewRoomEvent() (*api.OutputEvent, error)
 		TransactionID:   u.transactionID,
 	}
 
-	var stateEventNIDs []types.EventNID
-	for _, entry := range u.added {
-		stateEventNIDs = append(stateEventNIDs, entry.EventNID)
-	}
-	for _, entry := range u.removed {
-		stateEventNIDs = append(stateEventNIDs, entry.EventNID)
-	}
-	for _, entry := range u.stateBeforeEventRemoves {
-		stateEventNIDs = append(stateEventNIDs, entry.EventNID)
-	}
-	for _, entry := range u.stateBeforeEventAdds {
-		stateEventNIDs = append(stateEventNIDs, entry.EventNID)
-	}
-	stateEventNIDs = stateEventNIDs[:util.SortAndUnique(eventNIDSorter(stateEventNIDs))]
-	eventIDMap, err := u.api.DB.EventIDs(u.ctx, stateEventNIDs)
+	eventIDMap, err := u.stateEventMap()
 	if err != nil {
 		return nil, err
 	}
+
 	for _, entry := range u.added {
 		ore.AddsStateEventIDs = append(ore.AddsStateEventIDs, eventIDMap[entry.EventNID])
 	}
@@ -342,12 +330,60 @@ func (u *latestEventsUpdater) makeOutputNewRoomEvent() (*api.OutputEvent, error)
 	}
 	ore.SendAsServer = u.sendAsServer
 
+	// include extra state events if they were added as nearly every downstream component will care about it
+	// and we'd rather not have them all hit QueryEventsByID at the same time!
+	if len(ore.AddsStateEventIDs) > 0 {
+		ore.AddStateEvents, err = u.extraEventsForIDs(roomVersion, ore.AddsStateEventIDs)
+		if err != nil {
+			return nil, fmt.Errorf("failed to load add_state_events from db: %w", err)
+		}
+	}
+
 	return &api.OutputEvent{
 		Type:         api.OutputTypeNewRoomEvent,
 		NewRoomEvent: &ore,
 	}, nil
 }
 
+// extraEventsForIDs returns the full events for the event IDs given, but does not include the current event being
+// updated.
+func (u *latestEventsUpdater) extraEventsForIDs(roomVersion gomatrixserverlib.RoomVersion, eventIDs []string) ([]gomatrixserverlib.HeaderedEvent, error) {
+	var extraEventIDs []string
+	for _, e := range eventIDs {
+		if e == u.event.EventID() {
+			continue
+		}
+		extraEventIDs = append(extraEventIDs, e)
+	}
+	if len(extraEventIDs) == 0 {
+		return nil, nil
+	}
+	extraEvents, err := u.api.DB.EventsFromIDs(u.ctx, extraEventIDs)
+	if err != nil {
+		return nil, err
+	}
+	var h []gomatrixserverlib.HeaderedEvent
+	for _, e := range extraEvents {
+		h = append(h, e.Headered(roomVersion))
+	}
+	return h, nil
+}
+
+// retrieve an event nid -> event ID map for all events that need updating
+func (u *latestEventsUpdater) stateEventMap() (map[types.EventNID]string, error) {
+	var stateEventNIDs []types.EventNID
+	var allStateEntries []types.StateEntry
+	allStateEntries = append(allStateEntries, u.added...)
+	allStateEntries = append(allStateEntries, u.removed...)
+	allStateEntries = append(allStateEntries, u.stateBeforeEventRemoves...)
+	allStateEntries = append(allStateEntries, u.stateBeforeEventAdds...)
+	for _, entry := range allStateEntries {
+		stateEventNIDs = append(stateEventNIDs, entry.EventNID)
+	}
+	stateEventNIDs = stateEventNIDs[:util.SortAndUnique(eventNIDSorter(stateEventNIDs))]
+	return u.api.DB.EventIDs(u.ctx, stateEventNIDs)
+}
+
 type eventNIDSorter []types.EventNID
 
 func (s eventNIDSorter) Len() int           { return len(s) }

roomserver/internal/query.go

@@ -441,11 +441,11 @@ func (r *RoomserverInternalAPI) QueryMissingEvents(
 	return err
 }
 
-// QueryBackfill implements api.RoomServerQueryAPI
-func (r *RoomserverInternalAPI) QueryBackfill(
+// PerformBackfill implements api.RoomServerQueryAPI
+func (r *RoomserverInternalAPI) PerformBackfill(
 	ctx context.Context,
-	request *api.QueryBackfillRequest,
-	response *api.QueryBackfillResponse,
+	request *api.PerformBackfillRequest,
+	response *api.PerformBackfillResponse,
 ) error {
 	// if we are requesting the backfill then we need to do a federation hit
 	// TODO: we could be more sensible and fetch as many events we already have then request the rest
@@ -489,7 +489,7 @@ func (r *RoomserverInternalAPI) QueryBackfill(
 	return err
 }
 
-func (r *RoomserverInternalAPI) backfillViaFederation(ctx context.Context, req *api.QueryBackfillRequest, res *api.QueryBackfillResponse) error {
+func (r *RoomserverInternalAPI) backfillViaFederation(ctx context.Context, req *api.PerformBackfillRequest, res *api.PerformBackfillResponse) error {
 	roomVer, err := r.DB.GetRoomVersionForRoom(ctx, req.RoomID)
 	if err != nil {
 		return fmt.Errorf("backfillViaFederation: unknown room version for room %s : %w", req.RoomID, err)
@@ -647,7 +647,7 @@ func (r *RoomserverInternalAPI) scanEventTree(
 	var pre string
 
 	// TODO: add tests for this function to ensure it meets the contract that callers expect (and doc what that is supposed to be)
-	// Currently, callers like QueryBackfill will call scanEventTree with a pre-populated `visited` map, assuming that by doing
+	// Currently, callers like PerformBackfill will call scanEventTree with a pre-populated `visited` map, assuming that by doing
 	// so means that the events in that map will NOT be returned from this function. That is not currently true, resulting in
 	// duplicate events being sent in response to /backfill requests.
 	initialIgnoreList := make(map[string]bool, len(visited))