storage: close what Tessera's drivers open

storage/aws/antispam/aws.go

@@ -126,6 +126,7 @@ func NewAntispam(ctx context.Context, dsn string, opts AntispamOpts) (*AntispamS
 	}
 
 	if err := dbPool.Ping(); err != nil {
+		_ = dbPool.Close()
 		return nil, fmt.Errorf("failed to ping MySQL db: %v", err)
 	}
 
@@ -135,14 +136,24 @@ func NewAntispam(ctx context.Context, dsn string, opts AntispamOpts) (*AntispamS
 	}
 
 	if err := r.initDB(ctx); err != nil {
+		_ = dbPool.Close()
 		return nil, fmt.Errorf("failed to initDB: %v", err)
 	}
 	if err := r.checkDataCompatibility(ctx); err != nil {
+		_ = dbPool.Close()
 		return nil, fmt.Errorf("schema is not compatible with this version of the Tessera library: %v", err)
 	}
 	return r, nil
 }
 
+// Close closes the antispam's connection pool.
+//
+// Call it once nothing uses the antispam: after the appender it was given to has been shut down,
+// and that appender's context is done.
+func (s *AntispamStorage) Close() error {
+	return s.dbPool.Close()
+}
+
 func (s *AntispamStorage) initDB(ctx context.Context) error {
 	if _, err := s.dbPool.ExecContext(ctx,
 		`CREATE TABLE IF NOT EXISTS AntispamMeta (

storage/aws/antispam/aws_test.go

@@ -20,6 +20,8 @@ import (
 	"database/sql"
 	"flag"
 	"os"
+	"runtime"
+	"strings"
 	"testing"
 	"time"
 
@@ -230,6 +232,55 @@ func TestAntispamIdleConnections(t *testing.T) {
 	}
 }
 
+func TestAntispamClose(t *testing.T) {
+	ctx := t.Context()
+	if canSkipMySQLTest(t, ctx) {
+		slog.WarnContext(ctx, "MySQL not available, skipping", slog.String("name", t.Name()))
+		t.Skip("MySQL not available, skipping test")
+	}
+	as, err := NewAntispam(ctx, *mySQLURI, AntispamOpts{})
+	if err != nil {
+		t.Fatalf("NewAntispam: %v", err)
+	}
+	if err := as.Close(); err != nil {
+		t.Fatalf("Close: %v", err)
+	}
+	if err := as.dbPool.PingContext(ctx); err == nil || !strings.Contains(err.Error(), "database is closed") {
+		t.Errorf("Ping after Close: %v, want the pool closed", err)
+	}
+}
+
+// TestNewAntispamClosesPoolOnFailure checks that NewAntispam closes the pool it opened when it
+// fails, as there is then nothing the caller could close.
+func TestNewAntispamClosesPoolOnFailure(t *testing.T) {
+	before := openPools()
+	for range 3 {
+		if _, err := NewAntispam(t.Context(), "root:root@tcp(127.0.0.1:1)/tessera", AntispamOpts{}); err == nil {
+			t.Fatal("NewAntispam succeeded with nothing listening")
+		}
+	}
+	for deadline := time.Now().Add(5 * time.Second); ; time.Sleep(10 * time.Millisecond) {
+		got := openPools()
+		if got == before {
+			break
+		}
+		if time.Now().After(deadline) {
+			t.Fatalf("%d pools open after 3 failures, want %d", got, before)
+		}
+	}
+}
+
+// openPools returns the number of open database/sql pools: each runs a connectionOpener goroutine.
+func openPools() int {
+	buf := make([]byte, 1<<20)
+	for {
+		if n := runtime.Stack(buf, true); n < len(buf) {
+			return strings.Count(string(buf[:n]), "database/sql.(*DB).connectionOpener(")
+		}
+		buf = make([]byte, 2*len(buf))
+	}
+}
+
 // canSkipMySQLTest checks if the test MySQL db is available and, if not, if the test can be skipped.
 //
 // Use this method before every MySQL test, and if it returns true, skip the test.

storage/aws/aws.go

@@ -132,6 +132,9 @@ type sequencer interface {
 
 	// garbageCollect coordinates the removal of unneeded partial tiles/entry bundles for the provided tree size, up to a maximum number of deletes per invocation.
 	garbageCollect(ctx context.Context, treeSize uint64, maxDeletes uint, removePrefix func(ctx context.Context, prefix string) error, entriesPath func(uint64, uint8) string) error
+
+	// close releases the sequencer's resources. Nothing may use the sequencer after it.
+	close() error
 }
 
 // consumeFunc is the signature of a function which can consume entries from the sequencer.
@@ -220,7 +223,15 @@ func (s *Storage) Appender(ctx context.Context, opts *tessera.AppendOptions) (*t
 }
 
 // newAppender creates and initialises an Appender struct with the provided underlying storage implementations.
-func (s *Storage) newAppender(ctx context.Context, o objStore, seq sequencer, opts *tessera.AppendOptions) (*Appender, tessera.LogReader, error) {
+//
+// The Appender owns seq: newAppender closes it if it fails, and otherwise the Appender closes it
+// once ctx is done and its background jobs have returned.
+func (s *Storage) newAppender(ctx context.Context, o objStore, seq sequencer, opts *tessera.AppendOptions) (_ *Appender, _ tessera.LogReader, err error) {
+	defer func() {
+		if err != nil {
+			_ = seq.close()
+		}
+	}()
 	if opts.CheckpointInterval() < minCheckpointInterval {
 		return nil, nil, fmt.Errorf("requested CheckpointInterval (%v) is less than minimum permitted %v", opts.CheckpointInterval(), minCheckpointInterval)
 	}
@@ -267,15 +278,26 @@ func (s *Storage) newAppender(ctx context.Context, o objStore, seq sequencer, op
 	}
 
 	// Kick off go-routine which handles the integration of entries.
-	go r.integrateEntriesJob(ctx)
+	r.jobs.Go(func() { r.integrateEntriesJob(ctx) })
 
 	// Kick off go-routine which handles the publication of checkpoints.
-	go r.publishCheckpointJob(ctx, opts.CheckpointInterval(), opts.CheckpointRepublishInterval(), opts.CheckpointPublicationTimeout())
+	r.jobs.Go(func() {
+		r.publishCheckpointJob(ctx, opts.CheckpointInterval(), opts.CheckpointRepublishInterval(), opts.CheckpointPublicationTimeout())
+	})
 
 	if i := opts.GarbageCollectionInterval(); i > 0 {
-		go r.garbageCollectorJob(ctx, i)
+		r.jobs.Go(func() { r.garbageCollectorJob(ctx, i) })
 	}
 
+	// Once ctx is done the appender is finished, and its jobs return: then close the sequencer.
+	go func() {
+		<-ctx.Done()
+		r.jobs.Wait()
+		if err := seq.close(); err != nil {
+			slog.WarnContext(ctx, "Failed to close the MySQL sequencer", slog.Any("error", err))
+		}
+	}()
+
 	return r, r.logStore, nil
 }
 
@@ -290,6 +312,9 @@ type Appender struct {
 
 	treeUpdated     chan struct{}
 	entriesAssigned chan struct{}
+
+	// jobs tracks the background jobs, which return once the appender's context is done.
+	jobs sync.WaitGroup
 }
 
 // integrateEntriesJob periodically appends newly sequenced entries to the log.
@@ -916,6 +941,7 @@ func newMySQLSequencer(ctx context.Context, dsn string, maxOutstanding uint64, m
 	}
 
 	if err := dbPool.Ping(); err != nil {
+		_ = dbPool.Close()
 		return nil, fmt.Errorf("failed to ping MySQL db: %v", err)
 	}
 
@@ -925,9 +951,11 @@ func newMySQLSequencer(ctx context.Context, dsn string, maxOutstanding uint64, m
 	}
 
 	if err := r.initDB(ctx); err != nil {
+		_ = dbPool.Close()
 		return nil, fmt.Errorf("failed to initDB: %v", err)
 	}
 	if err := r.checkDataCompatibility(ctx); err != nil {
+		_ = dbPool.Close()
 		return nil, fmt.Errorf("schema is not compatible with this version of the Tessera library: %v", err)
 	}
 	return r, nil
@@ -1284,6 +1312,11 @@ func (s *mySQLSequencer) nextIndex(ctx context.Context) (uint64, error) {
 	return nextSeq, nil
 }
 
+// close closes the sequencer's connection pool.
+func (s *mySQLSequencer) close() error {
+	return s.dbPool.Close()
+}
+
 // publishCheckpoint checks when the last checkpoint was published, and if it was more than minAge ago, calls the provided
 // function to publish a new one.
 //

storage/aws/aws_test.go

@@ -34,6 +34,7 @@ import (
 	"net/http/httptest"
 	"os"
 	"reflect"
+	"runtime"
 	"strings"
 	"sync"
 	"testing"
@@ -189,6 +190,122 @@ func TestMySQLSequencerIdleConnections(t *testing.T) {
 	}
 }
 
+// TestMySQLSequencerClosesPoolOnFailure checks that newMySQLSequencer closes the pool it opened
+// when it fails, as there is then nothing the caller could close.
+func TestMySQLSequencerClosesPoolOnFailure(t *testing.T) {
+	ctx := t.Context()
+	t.Run("unreachable", func(t *testing.T) {
+		before := openPools()
+		for range 3 {
+			if _, err := newMySQLSequencer(ctx, "root:root@tcp(127.0.0.1:1)/tessera", 1000, 0, 0); err == nil {
+				t.Fatal("newMySQLSequencer succeeded with nothing listening")
+			}
+		}
+		if got := waitForOpenPools(before); got != before {
+			t.Errorf("%d pools open after 3 failures, want %d", got, before)
+		}
+	})
+	t.Run("incompatible schema", func(t *testing.T) {
+		if canSkipMySQLTest(t, ctx) {
+			t.Skip("MySQL not available, skipping test")
+		}
+		mustDropTables(t, ctx)
+		seq, err := newMySQLSequencer(ctx, *mySQLURI, 1000, 0, 0)
+		if err != nil {
+			t.Fatalf("newMySQLSequencer: %v", err)
+		}
+		setVersion := func(v uint64) {
+			if _, err := seq.dbPool.ExecContext(context.Background(), "UPDATE Tessera SET compatibilityVersion = ? WHERE id = 0", v); err != nil {
+				t.Fatalf("set compatibilityVersion: %v", err)
+			}
+		}
+		setVersion(SchemaCompatibilityVersion + 1)
+		t.Cleanup(func() {
+			setVersion(SchemaCompatibilityVersion)
+			_ = seq.close()
+		})
+		before := openPools()
+		for range 3 {
+			if _, err := newMySQLSequencer(ctx, *mySQLURI, 1000, 0, 0); err == nil {
+				t.Fatal("newMySQLSequencer succeeded with an incompatible schema")
+			}
+		}
+		if got := waitForOpenPools(before); got != before {
+			t.Errorf("%d pools open after 3 failures, want %d", got, before)
+		}
+	})
+}
+
+// TestAppenderClosesSequencer checks that an appender closes its sequencer's pool once its
+// context is done, and when newAppender fails.
+func TestAppenderClosesSequencer(t *testing.T) {
+	ctx := t.Context()
+	if canSkipMySQLTest(t, ctx) {
+		slog.WarnContext(ctx, "MySQL not available, skipping", slog.String("name", t.Name()))
+		t.Skip("MySQL not available, skipping test")
+	}
+	mustDropTables(t, ctx)
+	sk, _ := mustGenerateKeys(t)
+
+	for _, test := range []struct {
+		desc     string
+		interval time.Duration
+		wantErr  bool
+	}{
+		{desc: "when its context is done", interval: time.Second},
+		{desc: "when newAppender fails", interval: minCheckpointInterval / 2, wantErr: true},
+	} {
+		t.Run(test.desc, func(t *testing.T) {
+			seq, err := newMySQLSequencer(ctx, *mySQLURI, 1000, 0, 0)
+			if err != nil {
+				t.Fatalf("newMySQLSequencer: %v", err)
+			}
+			actx, cancel := context.WithCancel(ctx)
+			defer cancel()
+			opts := tessera.NewAppendOptions().WithCheckpointInterval(test.interval).WithCheckpointSigner(sk)
+			_, _, err = (&Storage{}).newAppender(actx, newMemObjStore(), seq, opts)
+			if gotErr := err != nil; gotErr != test.wantErr {
+				t.Fatalf("newAppender: %v, want error %v", err, test.wantErr)
+			}
+			if !test.wantErr {
+				if err := seq.dbPool.PingContext(ctx); err != nil {
+					t.Fatalf("Ping before the context is done: %v", err)
+				}
+				cancel()
+			}
+			for deadline := time.Now().Add(5 * time.Second); ; time.Sleep(10 * time.Millisecond) {
+				err := seq.dbPool.PingContext(ctx)
+				if err != nil && strings.Contains(err.Error(), "database is closed") {
+					break
+				}
+				if time.Now().After(deadline) {
+					t.Fatalf("the sequencer's pool is still open: Ping returns %v", err)
+				}
+			}
+		})
+	}
+}
+
+// openPools returns the number of open database/sql pools: each runs a connectionOpener goroutine.
+func openPools() int {
+	buf := make([]byte, 1<<20)
+	for {
+		if n := runtime.Stack(buf, true); n < len(buf) {
+			return strings.Count(string(buf[:n]), "database/sql.(*DB).connectionOpener(")
+		}
+		buf = make([]byte, 2*len(buf))
+	}
+}
+
+// waitForOpenPools waits up to 5 seconds for the number of open pools to fall to want, and returns it.
+func waitForOpenPools(want int) int {
+	for deadline := time.Now().Add(5 * time.Second); ; time.Sleep(10 * time.Millisecond) {
+		if got := openPools(); got <= want || time.Now().After(deadline) {
+			return got
+		}
+	}
+}
+
 func TestMySQLSequencerPushback(t *testing.T) {
 	ctx := context.Background()
 	if canSkipMySQLTest(t, ctx) {

storage/gcp/antispam/gcp.go

@@ -121,13 +121,14 @@ func NewAntispam(ctx context.Context, spannerDB string, opts AntispamOpts) (*Ant
 		return opts.SpannerTablePrefix + t
 	}
 
-	db := opts.SpannerClient
+	db, closeClient := opts.SpannerClient, func() {}
 	if db == nil {
 		var err error
 		db, err = spanner.NewClient(ctx, spannerDB)
 		if err != nil {
 			return nil, fmt.Errorf("failed to connect to Spanner: %v", err)
 		}
+		closeClient = db.Close
 	}
 
 	// Skip the (slow, even when no-op) DDL if the schema is already present. Keep schemaInitialised in sync with this.
@@ -142,26 +143,37 @@ func NewAntispam(ctx context.Context, spannerDB string, opts AntispamOpts) (*Ant
 				{spanner.Insert(table("FollowCoord"), []string{"id", "nextIdx"}, []any{0, 0})},
 			},
 		); err != nil {
-			if opts.SpannerClient == nil {
-				db.Close()
-			}
+			closeClient()
 			return nil, fmt.Errorf("failed to create tables: %v", err)
 		}
 	}
 
 	r := &AntispamStorage{
-		opts:   opts,
-		dbPool: db,
-		table:  table,
+		opts:        opts,
+		dbPool:      db,
+		closeClient: closeClient,
+		table:       table,
 	}
 
 	return r, nil
 }
 
+// Close closes the antispam's Spanner client if NewAntispam created it. A client passed in
+// AntispamOpts.SpannerClient is left open.
+//
+// Call it once nothing uses the antispam: after the appender it was given to has been shut down,
+// and that appender's context is done.
+func (s *AntispamStorage) Close() error {
+	s.closeClient()
+	return nil
+}
+
 type AntispamStorage struct {
 	opts AntispamOpts
 
 	dbPool *spanner.Client
+	// closeClient closes dbPool if NewAntispam created it, and otherwise does nothing.
+	closeClient func()
 
 	// table returns the provided table name with the configured table prefix, if any,
 	// prepended - see AntispamOpts.SpannerTablePrefix.