Add pinecone demo toggle for dis/enabling relaying for other nodes

3 files changed, 23 insertions, 3 deletions

diff --git a/relayapi/internal/api.go b/relayapi/internal/api.go
index 3ff8c2ad..3a5762fb 100644
--- a/relayapi/internal/api.go
+++ b/relayapi/internal/api.go
@@ -15,6 +15,8 @@
 package internal
 
 import (
+	"sync"
+
 	fedAPI "github.com/matrix-org/dendrite/federationapi/api"
 	"github.com/matrix-org/dendrite/federationapi/producers"
 	"github.com/matrix-org/dendrite/relayapi/storage"
@@ -30,6 +32,8 @@ type RelayInternalAPI struct {
 	producer               *producers.SyncAPIProducer
 	presenceEnabledInbound bool
 	serverName             gomatrixserverlib.ServerName
+	relayingEnabledMutex   sync.Mutex
+	relayingEnabled        bool
 }
 
 func NewRelayInternalAPI(
@@ -40,6 +44,7 @@ func NewRelayInternalAPI(
 	producer *producers.SyncAPIProducer,
 	presenceEnabledInbound bool,
 	serverName gomatrixserverlib.ServerName,
+	relayingEnabled bool,
 ) *RelayInternalAPI {
 	return &RelayInternalAPI{
 		db:                     db,
@@ -49,5 +54,6 @@ func NewRelayInternalAPI(
 		producer:               producer,
 		presenceEnabledInbound: presenceEnabledInbound,
 		serverName:             serverName,
+		relayingEnabled:        relayingEnabled,
 	}
 }
diff --git a/relayapi/internal/perform.go b/relayapi/internal/perform.go
index 59429933..c00d0d0e 100644
--- a/relayapi/internal/perform.go
+++ b/relayapi/internal/perform.go
@@ -24,6 +24,20 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
+// SetRelayingEnabled implements api.RelayInternalAPI
+func (r *RelayInternalAPI) SetRelayingEnabled(enabled bool) {
+	r.relayingEnabledMutex.Lock()
+	defer r.relayingEnabledMutex.Unlock()
+	r.relayingEnabled = enabled
+}
+
+// RelayingEnabled implements api.RelayInternalAPI
+func (r *RelayInternalAPI) RelayingEnabled() bool {
+	r.relayingEnabledMutex.Lock()
+	defer r.relayingEnabledMutex.Unlock()
+	return r.relayingEnabled
+}
+
 // PerformRelayServerSync implements api.RelayInternalAPI
 func (r *RelayInternalAPI) PerformRelayServerSync(
 	ctx context.Context,
diff --git a/relayapi/internal/perform_test.go b/relayapi/internal/perform_test.go
index fb71b7d0..5673b199 100644
--- a/relayapi/internal/perform_test.go
+++ b/relayapi/internal/perform_test.go
@@ -72,7 +72,7 @@ func TestPerformRelayServerSync(t *testing.T) {
 
 	fedClient := &testFedClient{}
 	relayAPI := NewRelayInternalAPI(
-		&db, fedClient, nil, nil, nil, false, "",
+		&db, fedClient, nil, nil, nil, false, "", true,
 	)
 
 	err = relayAPI.PerformRelayServerSync(context.Background(), *userID, gomatrixserverlib.ServerName("relay"))
@@ -92,7 +92,7 @@ func TestPerformRelayServerSyncFedError(t *testing.T) {
 
 	fedClient := &testFedClient{shouldFail: true}
 	relayAPI := NewRelayInternalAPI(
-		&db, fedClient, nil, nil, nil, false, "",
+		&db, fedClient, nil, nil, nil, false, "", true,
 	)
 
 	err = relayAPI.PerformRelayServerSync(context.Background(), *userID, gomatrixserverlib.ServerName("relay"))
@@ -112,7 +112,7 @@ func TestPerformRelayServerSyncRunsUntilQueueEmpty(t *testing.T) {
 
 	fedClient := &testFedClient{queueDepth: 2}
 	relayAPI := NewRelayInternalAPI(
-		&db, fedClient, nil, nil, nil, false, "",
+		&db, fedClient, nil, nil, nil, false, "", true,
 	)
 
 	err = relayAPI.PerformRelayServerSync(context.Background(), *userID, gomatrixserverlib.ServerName("relay"))