One NATS instance per `BaseDendrite` (#2438)

* One NATS instance per `BaseDendrite`

* Fix roomserver

2 files changed, 17 insertions, 22 deletions

diff --git a/setup/base/base.go b/setup/base/base.go
index ef449cc3..0e7528a0 100644
--- a/setup/base/base.go
+++ b/setup/base/base.go
@@ -41,6 +41,7 @@ import (
 	"golang.org/x/net/http2/h2c"
 
 	"github.com/matrix-org/dendrite/internal"
+	"github.com/matrix-org/dendrite/setup/jetstream"
 	"github.com/matrix-org/dendrite/setup/process"
 
 	"github.com/gorilla/mux"
@@ -77,6 +78,7 @@ type BaseDendrite struct {
 	InternalAPIMux         *mux.Router
 	DendriteAdminMux       *mux.Router
 	SynapseAdminMux        *mux.Router
+	NATS                   *jetstream.NATSInstance
 	UseHTTPAPIs            bool
 	apiHttpClient          *http.Client
 	Cfg                    *config.Dendrite
@@ -240,6 +242,7 @@ func NewBaseDendrite(cfg *config.Dendrite, componentName string, options ...Base
 		InternalAPIMux:         mux.NewRouter().SkipClean(true).PathPrefix(httputil.InternalPathPrefix).Subrouter().UseEncodedPath(),
 		DendriteAdminMux:       mux.NewRouter().SkipClean(true).PathPrefix(httputil.DendriteAdminPathPrefix).Subrouter().UseEncodedPath(),
 		SynapseAdminMux:        mux.NewRouter().SkipClean(true).PathPrefix(httputil.SynapseAdminPathPrefix).Subrouter().UseEncodedPath(),
+		NATS:                   &jetstream.NATSInstance{},
 		apiHttpClient:          &apiClient,
 		Database:               db,     // set if monolith with global connection pool only
 		DatabaseWriter:         writer, // set if monolith with global connection pool only
diff --git a/setup/jetstream/nats.go b/setup/jetstream/nats.go
index 8d528969..426f02bb 100644
--- a/setup/jetstream/nats.go
+++ b/setup/jetstream/nats.go
@@ -13,31 +13,23 @@ import (
 	"github.com/sirupsen/logrus"
 
 	natsserver "github.com/nats-io/nats-server/v2/server"
-	"github.com/nats-io/nats.go"
 	natsclient "github.com/nats-io/nats.go"
 )
 
-var natsServer *natsserver.Server
-var natsServerMutex sync.Mutex
-
-func PrepareForTests() (*process.ProcessContext, nats.JetStreamContext, *nats.Conn) {
-	cfg := &config.Dendrite{}
-	cfg.Defaults(true)
-	cfg.Global.JetStream.InMemory = true
-	pc := process.NewProcessContext()
-	js, jc := Prepare(pc, &cfg.Global.JetStream)
-	return pc, js, jc
+type NATSInstance struct {
+	*natsserver.Server
+	sync.Mutex
 }
 
-func Prepare(process *process.ProcessContext, cfg *config.JetStream) (natsclient.JetStreamContext, *natsclient.Conn) {
+func (s *NATSInstance) Prepare(process *process.ProcessContext, cfg *config.JetStream) (natsclient.JetStreamContext, *natsclient.Conn) {
 	// check if we need an in-process NATS Server
 	if len(cfg.Addresses) != 0 {
 		return setupNATS(process, cfg, nil)
 	}
-	natsServerMutex.Lock()
-	if natsServer == nil {
+	s.Lock()
+	if s.Server == nil {
 		var err error
-		natsServer, err = natsserver.NewServer(&natsserver.Options{
+		s.Server, err = natsserver.NewServer(&natsserver.Options{
 			ServerName:      "monolith",
 			DontListen:      true,
 			JetStream:       true,
@@ -49,23 +41,23 @@ func Prepare(process *process.ProcessContext, cfg *config.JetStream) (natsclient
 		if err != nil {
 			panic(err)
 		}
-		natsServer.ConfigureLogger()
+		s.ConfigureLogger()
 		go func() {
 			process.ComponentStarted()
-			natsServer.Start()
+			s.Start()
 		}()
 		go func() {
 			<-process.WaitForShutdown()
-			natsServer.Shutdown()
-			natsServer.WaitForShutdown()
+			s.Shutdown()
+			s.WaitForShutdown()
 			process.ComponentFinished()
 		}()
 	}
-	natsServerMutex.Unlock()
-	if !natsServer.ReadyForConnections(time.Second * 10) {
+	s.Unlock()
+	if !s.ReadyForConnections(time.Second * 10) {
 		logrus.Fatalln("NATS did not start in time")
 	}
-	nc, err := natsclient.Connect("", natsclient.InProcessServer(natsServer))
+	nc, err := natsclient.Connect("", natsclient.InProcessServer(s))
 	if err != nil {
 		logrus.Fatalln("Failed to create NATS client")
 	}