Change how servers are selected for missing auth/prev events

cmd/dendrite-polylith-multi/personalities/federationapi.go

@@ -33,7 +33,7 @@ func FederationAPI(base *setup.BaseDendrite, cfg *config.Dendrite) {
 		base.PublicFederationAPIMux, base.PublicKeyAPIMux,
 		&base.Cfg.FederationAPI, userAPI, federation, keyRing,
 		rsAPI, fsAPI, base.EDUServerClient(), keyAPI,
-		&base.Cfg.MSCs,
+		&base.Cfg.MSCs, nil,
 	)
 
 	base.SetupAndServeHTTP(

federationapi/federationapi.go

@@ -39,10 +39,12 @@ func AddPublicRoutes(
 	eduAPI eduserverAPI.EDUServerInputAPI,
 	keyAPI keyserverAPI.KeyInternalAPI,
 	mscCfg *config.MSCs,
+	servers routing.ServersInRoomProvider,
 ) {
 	routing.Setup(
 		fedRouter, keyRouter, cfg, rsAPI,
 		eduAPI, federationSenderAPI, keyRing,
 		federation, userAPI, keyAPI, mscCfg,
+		servers,
 	)
 }

federationapi/federationapi_test.go

@@ -31,7 +31,7 @@ func TestRoomsV3URLEscapeDoNot404(t *testing.T) {
 	fsAPI := base.FederationSenderHTTPClient()
 	// TODO: This is pretty fragile, as if anything calls anything on these nils this test will break.
 	// Unfortunately, it makes little sense to instantiate these dependencies when we just want to test routing.
-	federationapi.AddPublicRoutes(base.PublicFederationAPIMux, base.PublicKeyAPIMux, &cfg.FederationAPI, nil, nil, keyRing, nil, fsAPI, nil, nil, &cfg.MSCs)
+	federationapi.AddPublicRoutes(base.PublicFederationAPIMux, base.PublicKeyAPIMux, &cfg.FederationAPI, nil, nil, keyRing, nil, fsAPI, nil, nil, &cfg.MSCs, nil)
 	baseURL, cancel := test.ListenAndServe(t, base.PublicFederationAPIMux, true)
 	defer cancel()
 	serverName := gomatrixserverlib.ServerName(strings.TrimPrefix(baseURL, "https://"))

federationapi/routing/routing.go

@@ -50,6 +50,7 @@ func Setup(
 	userAPI userapi.UserInternalAPI,
 	keyAPI keyserverAPI.KeyInternalAPI,
 	mscCfg *config.MSCs,
+	servers ServersInRoomProvider,
 ) {
 	v2keysmux := keyMux.PathPrefix("/v2").Subrouter()
 	v1fedmux := fedMux.PathPrefix("/v1").Subrouter()
@@ -99,7 +100,7 @@ func Setup(
 		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest, vars map[string]string) util.JSONResponse {
 			return Send(
 				httpReq, request, gomatrixserverlib.TransactionID(vars["txnID"]),
-				cfg, rsAPI, eduAPI, keyAPI, keys, federation, mu,
+				cfg, rsAPI, eduAPI, keyAPI, keys, federation, mu, servers,
 			)
 		},
 	)).Methods(http.MethodPut, http.MethodOptions)

federationapi/routing/send.go

@@ -101,6 +101,7 @@ func Send(
 	keys gomatrixserverlib.JSONVerifier,
 	federation *gomatrixserverlib.FederationClient,
 	mu *internal.MutexByRoom,
+	servers ServersInRoomProvider,
 ) util.JSONResponse {
 	t := txnReq{
 		rsAPI:      rsAPI,
@@ -109,6 +110,7 @@ func Send(
 		federation: federation,
 		hadEvents:  make(map[string]bool),
 		haveEvents: make(map[string]*gomatrixserverlib.HeaderedEvent),
+		servers:    servers,
 		keyAPI:     keyAPI,
 		roomsMu:    mu,
 	}
@@ -158,16 +160,20 @@ func Send(
 	}
 }
 
+type ServersInRoomProvider interface {
+	GetServersForRoom(ctx context.Context, roomID string, event *gomatrixserverlib.Event) []gomatrixserverlib.ServerName
+}
+
 type txnReq struct {
 	gomatrixserverlib.Transaction
-	rsAPI        api.RoomserverInternalAPI
-	eduAPI       eduserverAPI.EDUServerInputAPI
-	keyAPI       keyapi.KeyInternalAPI
-	keys         gomatrixserverlib.JSONVerifier
-	federation   txnFederationClient
-	servers      []gomatrixserverlib.ServerName
-	serversMutex sync.RWMutex
-	roomsMu      *internal.MutexByRoom
+	rsAPI      api.RoomserverInternalAPI
+	eduAPI     eduserverAPI.EDUServerInputAPI
+	keyAPI     keyapi.KeyInternalAPI
+	keys       gomatrixserverlib.JSONVerifier
+	federation txnFederationClient
+	roomsMu    *internal.MutexByRoom
+	// something that can tell us about which servers are in a room right now
+	servers ServersInRoomProvider
 	// a list of events from the auth and prev events which we already had
 	hadEvents map[string]bool
 	// local cache of events for auth checks, etc - this may include events
@@ -466,22 +472,24 @@ func (t *txnReq) processDeviceListUpdate(ctx context.Context, e gomatrixserverli
 	}
 }
 
-func (t *txnReq) getServers(ctx context.Context, roomID string) []gomatrixserverlib.ServerName {
-	t.serversMutex.Lock()
-	defer t.serversMutex.Unlock()
+func (t *txnReq) getServers(ctx context.Context, roomID string, event *gomatrixserverlib.Event) []gomatrixserverlib.ServerName {
+	// The server that sent us the event should be sufficient to tell us about missing
+	// prev and auth events.
+	servers := []gomatrixserverlib.ServerName{t.Origin}
+	// If a specific room-to-server provider exists then use that. This will primarily
+	// be used for the P2P demos.
 	if t.servers != nil {
-		return t.servers
+		return append(servers, t.servers.GetServersForRoom(ctx, roomID, event)...)
 	}
-	t.servers = []gomatrixserverlib.ServerName{t.Origin}
-	serverReq := &api.QueryServerJoinedToRoomRequest{
-		RoomID: roomID,
+	// If the event origin is different to the transaction origin then we can use
+	// this as a last resort. The origin server that created the event would have
+	// had to know the auth and prev events.
+	if event != nil {
+		if origin := event.Origin(); origin != t.Origin {
+			servers = append(servers, origin)
+		}
 	}
-	serverRes := &api.QueryServerJoinedToRoomResponse{}
-	if err := t.rsAPI.QueryServerJoinedToRoom(ctx, serverReq, serverRes); err == nil {
-		t.servers = append(t.servers, serverRes.ServerNames...)
-		util.GetLogger(ctx).Infof("Found %d server(s) to query for missing events in %q", len(t.servers), roomID)
-	}
-	return t.servers
+	return servers
 }
 
 func (t *txnReq) processEvent(ctx context.Context, e *gomatrixserverlib.Event) error {
@@ -566,7 +574,7 @@ func (t *txnReq) retrieveMissingAuthEvents(
 withNextEvent:
 	for missingAuthEventID := range missingAuthEvents {
 	withNextServer:
-		for _, server := range t.getServers(ctx, e.RoomID()) {
+		for _, server := range t.getServers(ctx, e.RoomID(), e) {
 			logger.Infof("Retrieving missing auth event %q from %q", missingAuthEventID, server)
 			tx, err := t.federation.GetEvent(ctx, server, missingAuthEventID)
 			if err != nil {
@@ -948,7 +956,7 @@ func (t *txnReq) getMissingEvents(ctx context.Context, e *gomatrixserverlib.Even
 	}
 
 	var missingResp *gomatrixserverlib.RespMissingEvents
-	servers := t.getServers(ctx, e.RoomID())
+	servers := t.getServers(ctx, e.RoomID(), e)
 	for _, server := range servers {
 		var m gomatrixserverlib.RespMissingEvents
 		if m, err = t.federation.LookupMissingEvents(ctx, server, e.RoomID(), gomatrixserverlib.MissingEvents{
@@ -1220,7 +1228,7 @@ func (t *txnReq) lookupEvent(ctx context.Context, roomVersion gomatrixserverlib.
 	}
 	var event *gomatrixserverlib.Event
 	found := false
-	servers := t.getServers(ctx, roomID)
+	servers := t.getServers(ctx, roomID, nil)
 	for _, serverName := range servers {
 		txn, err := t.federation.GetEvent(ctx, serverName, missingEventID)
 		if err != nil || len(txn.PDUs) == 0 {

setup/monolith.go

@@ -68,7 +68,7 @@ func (m *Monolith) AddAllPublicRoutes(process *process.ProcessContext, csMux, ss
 	federationapi.AddPublicRoutes(
 		ssMux, keyMux, &m.Config.FederationAPI, m.UserAPI, m.FedClient,
 		m.KeyRing, m.RoomserverAPI, m.FederationSenderAPI,
-		m.EDUInternalAPI, m.KeyAPI, &m.Config.MSCs,
+		m.EDUInternalAPI, m.KeyAPI, &m.Config.MSCs, nil,
 	)
 	mediaapi.AddPublicRoutes(mediaMux, &m.Config.MediaAPI, m.UserAPI, m.Client)
 	syncapi.AddPublicRoutes(