From d960248ea8eef3c07a17477c106d01e1089ed206 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Fri, 19 Mar 2021 14:31:37 +0000
Subject: [PATCH] Retry room directory lookups if they fail

---
 cmd/dendrite-demo-pinecone/rooms/rooms.go | 22 ++++++++++++++++------
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/cmd/dendrite-demo-pinecone/rooms/rooms.go b/cmd/dendrite-demo-pinecone/rooms/rooms.go
index 5cf5e5280..6b1410d13 100644
--- a/cmd/dendrite-demo-pinecone/rooms/rooms.go
+++ b/cmd/dendrite-demo-pinecone/rooms/rooms.go
@@ -29,6 +29,8 @@ import (
 	pineconeSessions "github.com/matrix-org/pinecone/sessions"
 )
 
+const pineconeRoomAttempts = 3
+
 type PineconeRoomProvider struct {
 	r         *pineconeRouter.Router
 	s         *pineconeSessions.Sessions
@@ -84,12 +86,20 @@ func bulkFetchPublicRoomsFromServers(
 		go func(homeserverDomain gomatrixserverlib.ServerName) {
 			defer wg.Done()
 			util.GetLogger(ctx).WithField("hs", homeserverDomain).Info("Querying HS for public rooms")
-			fres, err := fedClient.GetPublicRooms(ctx, homeserverDomain, int(limit), "", false, "")
-			if err != nil {
-				util.GetLogger(ctx).WithError(err).WithField("hs", homeserverDomain).Warn(
-					"bulkFetchPublicRoomsFromServers: failed to query hs",
-				)
-				return
+			var fres gomatrixserverlib.RespPublicRooms
+			var err error
+			for i := 0; i < pineconeRoomAttempts; i++ {
+				fres, err = fedClient.GetPublicRooms(ctx, homeserverDomain, int(limit), "", false, "")
+				if err != nil {
+					util.GetLogger(ctx).WithError(err).WithField("hs", homeserverDomain).Warn(
+						"bulkFetchPublicRoomsFromServers: failed to query hs",
+					)
+					if i == pineconeRoomAttempts-1 {
+						return
+					}
+				} else {
+					break
+				}
 			}
 			for _, room := range fres.Chunk {
 				// atomically send a room or stop