Make tests more reliable (#2948)

When using `testrig.CreateBase` and then using that base for other
`NewInternalAPI` calls, we never actually shutdown the components.
`testrig.CreateBase` returns a `close` function, which only removes the
database, so still running components have issues connecting to the
database, since we ripped it out underneath it - which can result in
"Disk I/O" or "pq deadlock detected" issues.

5 files changed, 28 insertions, 14 deletions

diff --git a/federationapi/federationapi.go b/federationapi/federationapi.go
index 87eb751f..ce0ce98e 100644
--- a/federationapi/federationapi.go
+++ b/federationapi/federationapi.go
@@ -85,10 +85,7 @@ func AddPublicRoutes(
 	}
 
 	routing.Setup(
-		base.PublicFederationAPIMux,
-		base.PublicKeyAPIMux,
-		base.PublicWellKnownAPIMux,
-		cfg,
+		base,
 		rsAPI, f, keyRing,
 		federation, userAPI, keyAPI, mscCfg,
 		servers, producer,
diff --git a/federationapi/federationapi_test.go b/federationapi/federationapi_test.go
index 68a06a03..7009230c 100644
--- a/federationapi/federationapi_test.go
+++ b/federationapi/federationapi_test.go
@@ -273,12 +273,12 @@ func TestRoomsV3URLEscapeDoNot404(t *testing.T) {
 	cfg.Global.ServerName = gomatrixserverlib.ServerName("localhost")
 	cfg.Global.PrivateKey = privKey
 	cfg.Global.JetStream.InMemory = true
-	base := base.NewBaseDendrite(cfg, "Monolith")
+	b := base.NewBaseDendrite(cfg, "Monolith", base.DisableMetrics)
 	keyRing := &test.NopJSONVerifier{}
 	// TODO: This is pretty fragile, as if anything calls anything on these nils this test will break.
 	// Unfortunately, it makes little sense to instantiate these dependencies when we just want to test routing.
-	federationapi.AddPublicRoutes(base, nil, nil, keyRing, nil, &internal.FederationInternalAPI{}, nil, nil)
-	baseURL, cancel := test.ListenAndServe(t, base.PublicFederationAPIMux, true)
+	federationapi.AddPublicRoutes(b, nil, nil, keyRing, nil, &internal.FederationInternalAPI{}, nil, nil)
+	baseURL, cancel := test.ListenAndServe(t, b.PublicFederationAPIMux, true)
 	defer cancel()
 	serverName := gomatrixserverlib.ServerName(strings.TrimPrefix(baseURL, "https://"))
 
diff --git a/federationapi/routing/routing.go b/federationapi/routing/routing.go
index 0a3ab7a8..04eb3d06 100644
--- a/federationapi/routing/routing.go
+++ b/federationapi/routing/routing.go
@@ -32,6 +32,7 @@ import (
 	keyserverAPI "github.com/matrix-org/dendrite/keyserver/api"
 	"github.com/matrix-org/dendrite/roomserver/api"
 	roomserverAPI "github.com/matrix-org/dendrite/roomserver/api"
+	"github.com/matrix-org/dendrite/setup/base"
 	"github.com/matrix-org/dendrite/setup/config"
 	userapi "github.com/matrix-org/dendrite/userapi/api"
 	"github.com/matrix-org/gomatrixserverlib"
@@ -49,8 +50,7 @@ import (
 // applied:
 // nolint: gocyclo
 func Setup(
-	fedMux, keyMux, wkMux *mux.Router,
-	cfg *config.FederationAPI,
+	base *base.BaseDendrite,
 	rsAPI roomserverAPI.FederationRoomserverAPI,
 	fsAPI *fedInternal.FederationInternalAPI,
 	keys gomatrixserverlib.JSONVerifier,
@@ -61,9 +61,16 @@ func Setup(
 	servers federationAPI.ServersInRoomProvider,
 	producer *producers.SyncAPIProducer,
 ) {
-	prometheus.MustRegister(
-		pduCountTotal, eduCountTotal,
-	)
+	fedMux := base.PublicFederationAPIMux
+	keyMux := base.PublicKeyAPIMux
+	wkMux := base.PublicWellKnownAPIMux
+	cfg := &base.Cfg.FederationAPI
+
+	if base.EnableMetrics {
+		prometheus.MustRegister(
+			pduCountTotal, eduCountTotal,
+		)
+	}
 
 	v2keysmux := keyMux.PathPrefix("/v2").Subrouter()
 	v1fedmux := fedMux.PathPrefix("/v1").Subrouter()
diff --git a/setup/base/base.go b/setup/base/base.go
index d3adbf53..ff38209f 100644
--- a/setup/base/base.go
+++ b/setup/base/base.go
@@ -264,6 +264,8 @@ func NewBaseDendrite(cfg *config.Dendrite, componentName string, options ...Base
 
 // Close implements io.Closer
 func (b *BaseDendrite) Close() error {
+	b.ProcessContext.ShutdownDendrite()
+	b.ProcessContext.WaitForShutdown()
 	return b.tracerCloser.Close()
 }
 
diff --git a/test/testrig/base.go b/test/testrig/base.go
index 7bc26a5c..52e6ef5f 100644
--- a/test/testrig/base.go
+++ b/test/testrig/base.go
@@ -62,7 +62,12 @@ func CreateBaseDendrite(t *testing.T, dbType test.DBType) (*base.BaseDendrite, f
 			MaxIdleConnections:     2,
 			ConnMaxLifetimeSeconds: 60,
 		}
-		return base.NewBaseDendrite(&cfg, "Test", base.DisableMetrics), close
+		base := base.NewBaseDendrite(&cfg, "Test", base.DisableMetrics)
+		return base, func() {
+			base.ShutdownDendrite()
+			base.WaitForShutdown()
+			close()
+		}
 	case test.DBTypeSQLite:
 		cfg.Defaults(config.DefaultOpts{
 			Generate:   true,
@@ -72,7 +77,10 @@ func CreateBaseDendrite(t *testing.T, dbType test.DBType) (*base.BaseDendrite, f
 		// use a distinct prefix else concurrent postgres/sqlite runs will clash since NATS will use
 		// the file system event with InMemory=true :(
 		cfg.Global.JetStream.TopicPrefix = fmt.Sprintf("Test_%d_", dbType)
-		return base.NewBaseDendrite(&cfg, "Test", base.DisableMetrics), func() {
+		base := base.NewBaseDendrite(&cfg, "Test", base.DisableMetrics)
+		return base, func() {
+			base.ShutdownDendrite()
+			base.WaitForShutdown()
 			// cleanup db files. This risks getting out of sync as we add more database strings :(
 			dbFiles := []config.DataSource{
 				cfg.FederationAPI.Database.ConnectionString,