Relay integration to pinecone demos (#2955)

This extends the dendrite monolith for pinecone to integrate the s&f
features into the mobile apps.
Also makes a few tweaks to federation queueing/statistics to make some
edge cases more robust.

federationapi/api/api.go

@@ -72,12 +72,26 @@ type RoomserverFederationAPI interface {
 }
 
 type P2PFederationAPI interface {
-	// Relay Server sync api used in the pinecone demos.
+	// Get the relay servers associated for the given server.
 	P2PQueryRelayServers(
 		ctx context.Context,
 		request *P2PQueryRelayServersRequest,
 		response *P2PQueryRelayServersResponse,
 	) error
+
+	// Add relay server associations to the given server.
+	P2PAddRelayServers(
+		ctx context.Context,
+		request *P2PAddRelayServersRequest,
+		response *P2PAddRelayServersResponse,
+	) error
+
+	// Remove relay server associations from the given server.
+	P2PRemoveRelayServers(
+		ctx context.Context,
+		request *P2PRemoveRelayServersRequest,
+		response *P2PRemoveRelayServersResponse,
+	) error
 }
 
 // KeyserverFederationAPI is a subset of gomatrixserverlib.FederationClient functions which the keyserver
@@ -256,3 +270,19 @@ type P2PQueryRelayServersRequest struct {
 type P2PQueryRelayServersResponse struct {
 	RelayServers []gomatrixserverlib.ServerName
 }
+
+type P2PAddRelayServersRequest struct {
+	Server       gomatrixserverlib.ServerName
+	RelayServers []gomatrixserverlib.ServerName
+}
+
+type P2PAddRelayServersResponse struct {
+}
+
+type P2PRemoveRelayServersRequest struct {
+	Server       gomatrixserverlib.ServerName
+	RelayServers []gomatrixserverlib.ServerName
+}
+
+type P2PRemoveRelayServersResponse struct {
+}

federationapi/internal/perform.go

@@ -840,6 +840,36 @@ func (r *FederationInternalAPI) P2PQueryRelayServers(
 	return nil
 }
 
+// P2PAddRelayServers implements api.FederationInternalAPI
+func (r *FederationInternalAPI) P2PAddRelayServers(
+	ctx context.Context,
+	request *api.P2PAddRelayServersRequest,
+	response *api.P2PAddRelayServersResponse,
+) error {
+	logrus.Infof("Adding relay servers for: %s", request.Server)
+	err := r.db.P2PAddRelayServersForServer(ctx, request.Server, request.RelayServers)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// P2PRemoveRelayServers implements api.FederationInternalAPI
+func (r *FederationInternalAPI) P2PRemoveRelayServers(
+	ctx context.Context,
+	request *api.P2PRemoveRelayServersRequest,
+	response *api.P2PRemoveRelayServersResponse,
+) error {
+	logrus.Infof("Adding relay servers for: %s", request.Server)
+	err := r.db.P2PRemoveRelayServersForServer(ctx, request.Server, request.RelayServers)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
 func (r *FederationInternalAPI) shouldAttemptDirectFederation(
 	destination gomatrixserverlib.ServerName,
 ) bool {

federationapi/internal/perform_test.go

@@ -123,6 +123,47 @@ func TestQueryRelayServers(t *testing.T) {
 	assert.Equal(t, len(relayServers), len(res.RelayServers))
 }
 
+func TestRemoveRelayServers(t *testing.T) {
+	testDB := test.NewInMemoryFederationDatabase()
+
+	server := gomatrixserverlib.ServerName("wakeup")
+	relayServers := []gomatrixserverlib.ServerName{"relay1", "relay2"}
+	err := testDB.P2PAddRelayServersForServer(context.Background(), server, relayServers)
+	assert.NoError(t, err)
+
+	cfg := config.FederationAPI{
+		Matrix: &config.Global{
+			SigningIdentity: gomatrixserverlib.SigningIdentity{
+				ServerName: "relay",
+			},
+		},
+	}
+	fedClient := &testFedClient{}
+	stats := statistics.NewStatistics(testDB, FailuresUntilBlacklist, FailuresUntilAssumedOffline)
+	queues := queue.NewOutgoingQueues(
+		testDB, process.NewProcessContext(),
+		false,
+		cfg.Matrix.ServerName, fedClient, nil, &stats,
+		nil,
+	)
+	fedAPI := NewFederationInternalAPI(
+		testDB, &cfg, nil, fedClient, &stats, nil, queues, nil,
+	)
+
+	req := api.P2PRemoveRelayServersRequest{
+		Server:       server,
+		RelayServers: []gomatrixserverlib.ServerName{"relay1"},
+	}
+	res := api.P2PRemoveRelayServersResponse{}
+	err = fedAPI.P2PRemoveRelayServers(context.Background(), &req, &res)
+	assert.NoError(t, err)
+
+	finalRelays, err := testDB.P2PGetRelayServersForServer(context.Background(), server)
+	assert.NoError(t, err)
+	assert.Equal(t, 1, len(finalRelays))
+	assert.Equal(t, gomatrixserverlib.ServerName("relay2"), finalRelays[0])
+}
+
 func TestPerformDirectoryLookup(t *testing.T) {
 	testDB := test.NewInMemoryFederationDatabase()
 

federationapi/inthttp/client.go

@@ -25,6 +25,8 @@ const (
 	FederationAPIPerformBroadcastEDUPath           = "/federationapi/performBroadcastEDU"
 	FederationAPIPerformWakeupServers              = "/federationapi/performWakeupServers"
 	FederationAPIQueryRelayServers                 = "/federationapi/queryRelayServers"
+	FederationAPIAddRelayServers                   = "/federationapi/addRelayServers"
+	FederationAPIRemoveRelayServers                = "/federationapi/removeRelayServers"
 
 	FederationAPIGetUserDevicesPath      = "/federationapi/client/getUserDevices"
 	FederationAPIClaimKeysPath           = "/federationapi/client/claimKeys"
@@ -522,3 +524,25 @@ func (h *httpFederationInternalAPI) P2PQueryRelayServers(
 		h.httpClient, ctx, request, response,
 	)
 }
+
+func (h *httpFederationInternalAPI) P2PAddRelayServers(
+	ctx context.Context,
+	request *api.P2PAddRelayServersRequest,
+	response *api.P2PAddRelayServersResponse,
+) error {
+	return httputil.CallInternalRPCAPI(
+		"AddRelayServers", h.federationAPIURL+FederationAPIAddRelayServers,
+		h.httpClient, ctx, request, response,
+	)
+}
+
+func (h *httpFederationInternalAPI) P2PRemoveRelayServers(
+	ctx context.Context,
+	request *api.P2PRemoveRelayServersRequest,
+	response *api.P2PRemoveRelayServersResponse,
+) error {
+	return httputil.CallInternalRPCAPI(
+		"RemoveRelayServers", h.federationAPIURL+FederationAPIRemoveRelayServers,
+		h.httpClient, ctx, request, response,
+	)
+}

federationapi/queue/destinationqueue.go

@@ -410,34 +410,49 @@ func (oq *destinationQueue) nextTransaction(
 	defer cancel()
 
 	relayServers := oq.statistics.KnownRelayServers()
-	if oq.statistics.AssumedOffline() && len(relayServers) > 0 {
-		sendMethod = statistics.SendViaRelay
-		relaySuccess := false
-		logrus.Infof("Sending to relay servers: %v", relayServers)
-		// TODO : how to pass through actual userID here?!?!?!?!
-		userID, userErr := gomatrixserverlib.NewUserID("@user:"+string(oq.destination), false)
-		if userErr != nil {
-			return userErr, sendMethod
-		}
-
-		// Attempt sending to each known relay server.
-		for _, relayServer := range relayServers {
-			_, relayErr := oq.client.P2PSendTransactionToRelay(ctx, *userID, t, relayServer)
-			if relayErr != nil {
-				err = relayErr
-			} else {
-				// If sending to one of the relay servers succeeds, consider the send successful.
-				relaySuccess = true
-			}
-		}
-
-		// Clear the error if sending to any of the relay servers succeeded.
-		if relaySuccess {
-			err = nil
-		}
-	} else {
+	hasRelayServers := len(relayServers) > 0
+	shouldSendToRelays := oq.statistics.AssumedOffline() && hasRelayServers
+	if !shouldSendToRelays {
 		sendMethod = statistics.SendDirect
 		_, err = oq.client.SendTransaction(ctx, t)
+	} else {
+		// Try sending directly to the destination first in case they came back online.
+		sendMethod = statistics.SendDirect
+		_, err = oq.client.SendTransaction(ctx, t)
+		if err != nil {
+			// The destination is still offline, try sending to relays.
+			sendMethod = statistics.SendViaRelay
+			relaySuccess := false
+			logrus.Infof("Sending %q to relay servers: %v", t.TransactionID, relayServers)
+			// TODO : how to pass through actual userID here?!?!?!?!
+			userID, userErr := gomatrixserverlib.NewUserID("@user:"+string(oq.destination), false)
+			if userErr != nil {
+				return userErr, sendMethod
+			}
+
+			// Attempt sending to each known relay server.
+			for _, relayServer := range relayServers {
+				_, relayErr := oq.client.P2PSendTransactionToRelay(ctx, *userID, t, relayServer)
+				if relayErr != nil {
+					err = relayErr
+				} else {
+					// If sending to one of the relay servers succeeds, consider the send successful.
+					relaySuccess = true
+
+					// TODO : what about if the dest comes back online but can't see their relay?
+					// How do I sync with the dest in that case?
+					// Should change the database to have a "relay success" flag on events and if
+					// I see the node back online, maybe directly send through the backlog of events
+					// with "relay success"... could lead to duplicate events, but only those that
+					// I sent. And will lead to a much more consistent experience.
+				}
+			}
+
+			// Clear the error if sending to any of the relay servers succeeded.
+			if relaySuccess {
+				err = nil
+			}
+		}
 	}
 	switch errResponse := err.(type) {
 	case nil:

federationapi/queue/queue_test.go

@@ -923,7 +923,7 @@ func TestSendPDUOnRelaySuccessRemovedFromDB(t *testing.T) {
 	assert.NoError(t, err)
 
 	check := func(log poll.LogT) poll.Result {
-		if fc.txCount.Load() == 1 {
+		if fc.txCount.Load() >= 1 {
 			if fc.txRelayCount.Load() == 1 {
 				data, dbErr := db.GetPendingPDUs(pc.Context(), destination, 100)
 				assert.NoError(t, dbErr)
@@ -962,7 +962,7 @@ func TestSendEDUOnRelaySuccessRemovedFromDB(t *testing.T) {
 	assert.NoError(t, err)
 
 	check := func(log poll.LogT) poll.Result {
-		if fc.txCount.Load() == 1 {
+		if fc.txCount.Load() >= 1 {
 			if fc.txRelayCount.Load() == 1 {
 				data, dbErr := db.GetPendingEDUs(pc.Context(), destination, 100)
 				assert.NoError(t, dbErr)

federationapi/statistics/statistics.go

@@ -164,6 +164,8 @@ func (s *ServerStatistics) Success(method SendMethod) {
 				logrus.WithError(err).Errorf("Failed to remove %q from blacklist", s.serverName)
 			}
 		}
+
+		s.removeAssumedOffline()
 	}
 }