testdb: Fix stuck PgBouncer connections by killing via admin command

The PgBouncer-based test job is really flaky, where the symptom of the
flakiness is that we fail to close all connections to the database on
test cleanup. Seemingly, PgBouncer will in some situations keep around
the connection without us being able to close them in the 20 second
timeframe we allow the conection termination to take.

While we could try to increase the timeout even further, this is not a
good way to guarantee that terminations do indeed connect. Actually, we
probably cannot require PgBouncer to really terminate connections only
because we kill them in the Postgres backend itself: PgBouncer may
reconnect automatically or anything else, so we should treat it as a
black box which just does some automagic things we cannot control by
issuing commands against Postgres.

Luckily, PgBouncer has an admin console which allows us to take control
over what it does directly. Instead of trying to terminate connections
via Postgres, we can thus ask PgBouncer to terminate connections by
issuing a `KILL` command, which will immediately terminate all client
and server connections.

Convert the test cleanup logic to do so. Given that we have test setups
where the same database configuration (and thus database) is reused, we
also need to issue a `RESUME` command such that subsequent test cases
will be able to reconnect to the database.

1 files changed, 22 insertions, 1 deletions

diff --git a/internal/testhelper/testdb/db.go b/internal/testhelper/testdb/db.go
index c3d855425..529ebb4a0 100644
--- a/internal/testhelper/testdb/db.go
+++ b/internal/testhelper/testdb/db.go
@@ -19,6 +19,7 @@ import (
 	"github.com/stretchr/testify/require"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/config"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/datastore/glsql"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/testhelper"
 )
 
 const (
@@ -268,9 +269,29 @@ func initPraefectDB(t testing.TB, database string) *sql.DB {
 	require.NoErrorf(t, err, "failed to create %q database", praefectTemplateDatabase)
 
 	t.Cleanup(func() {
+		if _, ok := getDatabaseEnvironment(t)["PGHOST_PGBOUNCER"]; ok {
+			pgbouncerCfg := dbCfg
+			// This database name will connect us to the special admin console.
+			pgbouncerCfg.DBName = "pgbouncer"
+
+			// We cannot use `requireSQLOpen()` because it would ping the database,
+			// which is not supported by the PgBouncer admin console.
+			pgbouncerDB, err := sql.Open("postgres", glsql.DSN(pgbouncerCfg, false))
+			require.NoError(t, err)
+			defer testhelper.MustClose(t, pgbouncerDB)
+
+			// Trying to release connections like we do with the "normal" Postgres
+			// database regularly results in flaky tests with PgBouncer given that the
+			// connections are seemingly never released. Instead, we kill PgBouncer
+			// connections by connecting to its admin console and using the KILL
+			// command, which instructs it to kill all client and server connections.
+			_, err = pgbouncerDB.Exec("KILL " + database)
+			require.NoError(t, err, "killing PgBouncer connections")
+		}
+
 		dbCfg.DBName = "postgres"
 		postgresDB := requireSQLOpen(t, dbCfg, true)
-		defer func() { require.NoErrorf(t, postgresDB.Close(), "release connection to the %q database", dbCfg.DBName) }()
+		defer testhelper.MustClose(t, postgresDB)
 
 		// We need to force-terminate open connections as for the tasks that use PgBouncer
 		// the actual client connected to the database is a PgBouncer and not a test that is