From 170a96ebbef82c0202e13c9c2fe2912a15e13f6a Mon Sep 17 00:00:00 2001
From: Till Faelligen <2353100+S7evinK@users.noreply.github.com>
Date: Tue, 8 Nov 2022 15:15:59 +0100
Subject: [PATCH] Don't drop presence for remote servers immediately

---
 federationapi/consumers/presence.go | 46 +++++++++++++++++++----------
 1 file changed, 30 insertions(+), 16 deletions(-)

diff --git a/federationapi/consumers/presence.go b/federationapi/consumers/presence.go
index 153fc40b5..1c410f1b6 100644
--- a/federationapi/consumers/presence.go
+++ b/federationapi/consumers/presence.go
@@ -18,6 +18,7 @@ import (
 	"context"
 	"encoding/json"
 	"strconv"
+	"time"
 
 	"github.com/matrix-org/dendrite/federationapi/queue"
 	"github.com/matrix-org/dendrite/federationapi/storage"
@@ -32,6 +33,9 @@ import (
 	log "github.com/sirupsen/logrus"
 )
 
+const maxRetries = 3
+const retryDelay = time.Second
+
 // OutputReceiptConsumer consumes events that originate in the clientapi.
 type OutputPresenceConsumer struct {
 	ctx                     context.Context
@@ -93,16 +97,6 @@ func (t *OutputPresenceConsumer) onMessage(ctx context.Context, msgs []*nats.Msg
 		return true
 	}
 
-	var queryRes roomserverAPI.QueryRoomsForUserResponse
-	err = t.rsAPI.QueryRoomsForUser(t.ctx, &roomserverAPI.QueryRoomsForUserRequest{
-		UserID:         userID,
-		WantMembership: "join",
-	}, &queryRes)
-	if err != nil {
-		log.WithError(err).Error("failed to calculate joined rooms for user")
-		return true
-	}
-
 	presence := msg.Header.Get("presence")
 
 	ts, err := strconv.Atoi(msg.Header.Get("last_active_ts"))
@@ -110,13 +104,33 @@ func (t *OutputPresenceConsumer) onMessage(ctx context.Context, msgs []*nats.Msg
 		return true
 	}
 
-	// send this presence to all servers who share rooms with this user.
-	joined, err := t.db.GetJoinedHostsForRooms(t.ctx, queryRes.RoomIDs, true)
-	if err != nil {
-		log.WithError(err).Error("failed to get joined hosts")
-		return true
-	}
+	var joined []gomatrixserverlib.ServerName
+	// We're trying to get joined rooms for this user for 3 seconds (3 retries, 1s delay)
+	// If we fail to get joined hosts, if we fail to, we discard the presence event.
+	for i := 0; i < maxRetries; i++ {
+		var queryRes roomserverAPI.QueryRoomsForUserResponse
+		err = t.rsAPI.QueryRoomsForUser(t.ctx, &roomserverAPI.QueryRoomsForUserRequest{
+			UserID:         userID,
+			WantMembership: "join",
+		}, &queryRes)
+		if err != nil {
+			log.WithError(err).Error("failed to calculate joined rooms for user")
+			return true
+		}
 
+		joined, err = t.db.GetJoinedHostsForRooms(t.ctx, queryRes.RoomIDs, true)
+		if err != nil {
+			log.WithError(err).Error("failed to get joined hosts")
+			return true
+		}
+
+		if len(joined) == 0 {
+			time.Sleep(retryDelay)
+			continue
+		}
+		break
+	}
+	// If we still have no joined hosts, discard the event.
 	if len(joined) == 0 {
 		return true
 	}