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

9 files changed, 130 insertions, 67 deletions

diff --git a/roomserver/api/api.go b/roomserver/api/api.go
index 3a2ad059..967f58ba 100644
--- a/roomserver/api/api.go
+++ b/roomserver/api/api.go
@@ -89,10 +89,10 @@ type RoomserverInternalAPI interface {
 	) error
 
 	// Query a given amount (or less) of events prior to a given set of events.
-	QueryBackfill(
+	PerformBackfill(
 		ctx context.Context,
-		request *QueryBackfillRequest,
-		response *QueryBackfillResponse,
+		request *PerformBackfillRequest,
+		response *PerformBackfillResponse,
 	) error
 
 	// Asks for the default room version as preferred by the server.
diff --git a/roomserver/api/output.go b/roomserver/api/output.go
index 92a468a9..2bbd97af 100644
--- a/roomserver/api/output.go
+++ b/roomserver/api/output.go
@@ -63,6 +63,13 @@ type OutputNewRoomEvent struct {
 	// Together with RemovesStateEventIDs this allows the receiver to keep an up to date
 	// view of the current state of the room.
 	AddsStateEventIDs []string `json:"adds_state_event_ids"`
+	// All extra newly added state events. This is only set if there are *extra* events
+	// other than `Event`. This can happen when forks get merged because state resolution
+	// may decide a bunch of state events on one branch are now valid, so they will be
+	// present in this list. This is useful when trying to maintain the current state of a room
+	// as to do so you need to include both these events and `Event`.
+	AddStateEvents []gomatrixserverlib.HeaderedEvent `json:"adds_state_events"`
+
 	// The state event IDs that were removed from the state of the room by this event.
 	RemovesStateEventIDs []string `json:"removes_state_event_ids"`
 	// The ID of the event that was output before this event.
@@ -112,6 +119,26 @@ type OutputNewRoomEvent struct {
 	TransactionID *TransactionID `json:"transaction_id"`
 }
 
+// AddsState returns all added state events from this event.
+//
+// This function is needed because `AddStateEvents` will not include a copy of
+// the original event to save space, so you cannot use that slice alone.
+// Instead, use this function which will add the original event if it is present
+// in `AddsStateEventIDs`.
+func (ore *OutputNewRoomEvent) AddsState() []gomatrixserverlib.HeaderedEvent {
+	includeOutputEvent := false
+	for _, id := range ore.AddsStateEventIDs {
+		if id == ore.Event.EventID() {
+			includeOutputEvent = true
+			break
+		}
+	}
+	if !includeOutputEvent {
+		return ore.AddStateEvents
+	}
+	return append(ore.AddStateEvents, ore.Event)
+}
+
 // An OutputNewInviteEvent is written whenever an invite becomes active.
 // Invite events can be received outside of an existing room so have to be
 // tracked separately from the room events themselves.
diff --git a/roomserver/api/perform.go b/roomserver/api/perform.go
index 1cf54144..3e5cae1b 100644
--- a/roomserver/api/perform.go
+++ b/roomserver/api/perform.go
@@ -2,6 +2,7 @@ package api
 
 import (
 	"github.com/matrix-org/gomatrixserverlib"
+	"github.com/matrix-org/util"
 )
 
 type PerformJoinRequest struct {
@@ -22,3 +23,31 @@ type PerformLeaveRequest struct {
 
 type PerformLeaveResponse struct {
 }
+
+// PerformBackfillRequest is a request to PerformBackfill.
+type PerformBackfillRequest struct {
+	// The room to backfill
+	RoomID string `json:"room_id"`
+	// A map of backwards extremity event ID to a list of its prev_event IDs.
+	BackwardsExtremities map[string][]string `json:"backwards_extremities"`
+	// The maximum number of events to retrieve.
+	Limit int `json:"limit"`
+	// The server interested in the events.
+	ServerName gomatrixserverlib.ServerName `json:"server_name"`
+}
+
+// PrevEventIDs returns the prev_event IDs of all backwards extremities, de-duplicated in a lexicographically sorted order.
+func (r *PerformBackfillRequest) PrevEventIDs() []string {
+	var prevEventIDs []string
+	for _, pes := range r.BackwardsExtremities {
+		prevEventIDs = append(prevEventIDs, pes...)
+	}
+	prevEventIDs = util.UniqueStrings(prevEventIDs)
+	return prevEventIDs
+}
+
+// PerformBackfillResponse is a response to PerformBackfill.
+type PerformBackfillResponse struct {
+	// Missing events, arbritrary order.
+	Events []gomatrixserverlib.HeaderedEvent `json:"events"`
+}
diff --git a/roomserver/api/query.go b/roomserver/api/query.go
index c9a46ae9..b1525342 100644
--- a/roomserver/api/query.go
+++ b/roomserver/api/query.go
@@ -18,7 +18,6 @@ package api
 
 import (
 	"github.com/matrix-org/gomatrixserverlib"
-	"github.com/matrix-org/util"
 )
 
 // QueryLatestEventsAndStateRequest is a request to QueryLatestEventsAndState
@@ -204,34 +203,6 @@ type QueryStateAndAuthChainResponse struct {
 	AuthChainEvents []gomatrixserverlib.HeaderedEvent `json:"auth_chain_events"`
 }
 
-// QueryBackfillRequest is a request to QueryBackfill.
-type QueryBackfillRequest struct {
-	// The room to backfill
-	RoomID string `json:"room_id"`
-	// A map of backwards extremity event ID to a list of its prev_event IDs.
-	BackwardsExtremities map[string][]string `json:"backwards_extremities"`
-	// The maximum number of events to retrieve.
-	Limit int `json:"limit"`
-	// The server interested in the events.
-	ServerName gomatrixserverlib.ServerName `json:"server_name"`
-}
-
-// PrevEventIDs returns the prev_event IDs of all backwards extremities, de-duplicated in a lexicographically sorted order.
-func (r *QueryBackfillRequest) PrevEventIDs() []string {
-	var prevEventIDs []string
-	for _, pes := range r.BackwardsExtremities {
-		prevEventIDs = append(prevEventIDs, pes...)
-	}
-	prevEventIDs = util.UniqueStrings(prevEventIDs)
-	return prevEventIDs
-}
-
-// QueryBackfillResponse is a response to QueryBackfill.
-type QueryBackfillResponse struct {
-	// Missing events, arbritrary order.
-	Events []gomatrixserverlib.HeaderedEvent `json:"events"`
-}
-
 // QueryRoomVersionCapabilitiesRequest asks for the default room version
 type QueryRoomVersionCapabilitiesRequest struct{}
 
diff --git a/roomserver/internal/input_latest_events.go b/roomserver/internal/input_latest_events.go
index aea85ca9..e69307ad 100644
--- a/roomserver/internal/input_latest_events.go
+++ b/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) }
diff --git a/roomserver/internal/query_backfill.go b/roomserver/internal/perform_backfill.go
index 23ae9455..23ae9455 100644
--- a/roomserver/internal/query_backfill.go
+++ b/roomserver/internal/perform_backfill.go
diff --git a/roomserver/internal/query.go b/roomserver/internal/query.go
index aea93388..375ddc23 100644
--- a/roomserver/internal/query.go
+++ b/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))
diff --git a/roomserver/inthttp/client.go b/roomserver/inthttp/client.go
index 6f5e882e..1244300d 100644
--- a/roomserver/inthttp/client.go
+++ b/roomserver/inthttp/client.go
@@ -24,8 +24,9 @@ const (
 	RoomserverInputRoomEventsPath = "/roomserver/inputRoomEvents"
 
 	// Perform operations
-	RoomserverPerformJoinPath  = "/roomserver/performJoin"
-	RoomserverPerformLeavePath = "/roomserver/performLeave"
+	RoomserverPerformJoinPath     = "/roomserver/performJoin"
+	RoomserverPerformLeavePath    = "/roomserver/performLeave"
+	RoomserverPerformBackfillPath = "/roomserver/performBackfill"
 
 	// Query operations
 	RoomserverQueryLatestEventsAndStatePath    = "/roomserver/queryLatestEventsAndState"
@@ -36,7 +37,6 @@ const (
 	RoomserverQueryServerAllowedToSeeEventPath = "/roomserver/queryServerAllowedToSeeEvent"
 	RoomserverQueryMissingEventsPath           = "/roomserver/queryMissingEvents"
 	RoomserverQueryStateAndAuthChainPath       = "/roomserver/queryStateAndAuthChain"
-	RoomserverQueryBackfillPath                = "/roomserver/queryBackfill"
 	RoomserverQueryRoomVersionCapabilitiesPath = "/roomserver/queryRoomVersionCapabilities"
 	RoomserverQueryRoomVersionForRoomPath      = "/roomserver/queryRoomVersionForRoom"
 )
@@ -274,16 +274,16 @@ func (h *httpRoomserverInternalAPI) QueryStateAndAuthChain(
 	return internalHTTP.PostJSON(ctx, span, h.httpClient, apiURL, request, response)
 }
 
-// QueryBackfill implements RoomServerQueryAPI
-func (h *httpRoomserverInternalAPI) QueryBackfill(
+// PerformBackfill implements RoomServerQueryAPI
+func (h *httpRoomserverInternalAPI) PerformBackfill(
 	ctx context.Context,
-	request *api.QueryBackfillRequest,
-	response *api.QueryBackfillResponse,
+	request *api.PerformBackfillRequest,
+	response *api.PerformBackfillResponse,
 ) error {
-	span, ctx := opentracing.StartSpanFromContext(ctx, "QueryBackfill")
+	span, ctx := opentracing.StartSpanFromContext(ctx, "PerformBackfill")
 	defer span.Finish()
 
-	apiURL := h.roomserverURL + RoomserverQueryBackfillPath
+	apiURL := h.roomserverURL + RoomserverPerformBackfillPath
 	return internalHTTP.PostJSON(ctx, span, h.httpClient, apiURL, request, response)
 }
 
diff --git a/roomserver/inthttp/server.go b/roomserver/inthttp/server.go
index 3a13ce37..8ac815f3 100644
--- a/roomserver/inthttp/server.go
+++ b/roomserver/inthttp/server.go
@@ -165,14 +165,14 @@ func AddRoutes(r api.RoomserverInternalAPI, internalAPIMux *mux.Router) {
 		}),
 	)
 	internalAPIMux.Handle(
-		RoomserverQueryBackfillPath,
-		internal.MakeInternalAPI("QueryBackfill", func(req *http.Request) util.JSONResponse {
-			var request api.QueryBackfillRequest
-			var response api.QueryBackfillResponse
+		RoomserverPerformBackfillPath,
+		internal.MakeInternalAPI("PerformBackfill", func(req *http.Request) util.JSONResponse {
+			var request api.PerformBackfillRequest
+			var response api.PerformBackfillResponse
 			if err := json.NewDecoder(req.Body).Decode(&request); err != nil {
 				return util.ErrorResponse(err)
 			}
-			if err := r.QueryBackfill(req.Context(), &request, &response); err != nil {
+			if err := r.PerformBackfill(req.Context(), &request, &response); err != nil {
 				return util.ErrorResponse(err)
 			}
 			return util.JSONResponse{Code: http.StatusOK, JSON: &response}