From 01cf2ad6ff5bdfb41207cb1233a731810f761f99 Mon Sep 17 00:00:00 2001 From: Glenn Rempe Date: Mon, 9 Nov 2026 13:27:51 -0500 Subject: [PATCH] storage: close what Tessera's drivers open Tessera's storage drivers open database connections and never close them: the AWS driver's MySQL pool, the Spanner client the GCP driver creates when none is passed in, and the pools, clients and databases of the AWS, GCP and Badger antispams. A process that opens a log more than once keeps every one of them, and Badger keeps its directory lock, so its database can't be opened again in that process. Construction that fails part way leaks what it opened too. Each appender now closes what it opened once its context is done and its background jobs have returned, following the documented shutdown: call the shutdown function, then cancel the context. The AWS driver closes its MySQL pool. The GCP driver closes a Spanner client it created, but not one passed in Config.SpannerClient, and no longer keeps the client it creates in its Config, so that other lifecycles from the same driver don't share a client it closes. The GCS client is left open, as its Close makes any later use panic. Construction that fails closes what it opened. The antispams are created by the caller, so each gets a Close method. The GCP antispam closes only a client it created. The Badger antispam stops its compaction, then closes its database, releasing the lock. There is no way yet to wait for the cleanup: it happens shortly after the context ends. --- storage/aws/antispam/aws.go | 11 +++ storage/aws/antispam/aws_test.go | 51 +++++++++++ storage/aws/aws.go | 41 ++++++++- storage/aws/aws_test.go | 117 ++++++++++++++++++++++++++ storage/gcp/antispam/gcp.go | 26 ++++-- storage/gcp/antispam/gcp_test.go | 48 +++++++++++ storage/gcp/gcp.go | 54 +++++++++--- storage/gcp/gcp_test.go | 55 ++++++++++++ storage/posix/antispam/badger.go | 25 +++++- storage/posix/antispam/badger_test.go | 47 +++++++++++ 10 files changed, 451 insertions(+), 24 deletions(-) diff --git a/storage/aws/antispam/aws.go b/storage/aws/antispam/aws.go index 57b6d4818..83f0c36a5 100644 --- a/storage/aws/antispam/aws.go +++ b/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 ( diff --git a/storage/aws/antispam/aws_test.go b/storage/aws/antispam/aws_test.go index c0cd67014..35e2abba3 100644 --- a/storage/aws/antispam/aws_test.go +++ b/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. diff --git a/storage/aws/aws.go b/storage/aws/aws.go index 960510940..d6989c3a1 100644 --- a/storage/aws/aws.go +++ b/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. // diff --git a/storage/aws/aws_test.go b/storage/aws/aws_test.go index 4069a8b6c..a45ae64a7 100644 --- a/storage/aws/aws_test.go +++ b/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) { diff --git a/storage/gcp/antispam/gcp.go b/storage/gcp/antispam/gcp.go index 81cfe5d65..52adde1a4 100644 --- a/storage/gcp/antispam/gcp.go +++ b/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. diff --git a/storage/gcp/antispam/gcp_test.go b/storage/gcp/antispam/gcp_test.go index b2f5a6a7e..df7b43b9e 100644 --- a/storage/gcp/antispam/gcp_test.go +++ b/storage/gcp/antispam/gcp_test.go @@ -200,6 +200,54 @@ func TestAntispamSharedClientWrongDatabase(t *testing.T) { } } +// TestAntispamClose checks that Close closes a Spanner client that NewAntispam created, and leaves +// open one that was passed in. +func TestAntispamClose(t *testing.T) { + const db = "projects/p/instances/i/databases/d" + read := func(c *spanner.Client) error { + ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) + defer cancel() + _, err := c.Single().ReadRow(ctx, "FollowCoord", spanner.Key{0}, []string{"nextIdx"}) + return err + } + t.Run("its own client", func(t *testing.T) { + closeDB := newSpannerDB(t) + defer closeDB() + as, err := NewAntispam(t.Context(), db, AntispamOpts{}) + if err != nil { + t.Fatalf("NewAntispam: %v", err) + } + if err := read(as.dbPool); err != nil { + t.Fatalf("read before Close: %v", err) + } + if err := as.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + if err := read(as.dbPool); err == nil { + t.Error("read after Close succeeded, want the client closed") + } + }) + t.Run("a client passed in", func(t *testing.T) { + closeDB := newSpannerDB(t) + defer closeDB() + c, err := spanner.NewClient(t.Context(), db) + if err != nil { + t.Fatalf("spanner.NewClient: %v", err) + } + defer c.Close() + as, err := NewAntispam(t.Context(), db, AntispamOpts{SpannerClient: c}) + if err != nil { + t.Fatalf("NewAntispam: %v", err) + } + if err := as.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + if err := read(c); err != nil { + t.Errorf("read with the caller's client after Close: %v, want it still open", err) + } + }) +} + func TestAntispamPushbackRecovers(t *testing.T) { for _, test := range []struct { name string diff --git a/storage/gcp/gcp.go b/storage/gcp/gcp.go index 7ee211b09..a9db2b899 100644 --- a/storage/gcp/gcp.go +++ b/storage/gcp/gcp.go @@ -250,24 +250,30 @@ func (s *Storage) Appender(ctx context.Context, opts *tessera.AppendOptions) (*t bucketPrefix: s.cfg.BucketPrefix, } - var err error - if s.cfg.SpannerClient == nil { - s.cfg.SpannerClient, err = spanner.NewClient(ctx, s.cfg.Spanner) + spannerClient, closeClient := s.cfg.SpannerClient, func() {} + if spannerClient == nil { + var err error + spannerClient, err = spanner.NewClient(ctx, s.cfg.Spanner) if err != nil { return nil, nil, fmt.Errorf("failed to connect to Spanner: %v", err) } + // This appender created the client, so it closes it. + closeClient = spannerClient.Close } table := func(t string) string { return s.cfg.SpannerTablePrefix + t } - if err := initDB(ctx, s.cfg.Spanner, s.cfg.SpannerClient, table); err != nil { + if err := initDB(ctx, s.cfg.Spanner, spannerClient, table); err != nil { + closeClient() return nil, nil, fmt.Errorf("failed to verify/init Spanner schema: %v", err) } - seq, err := newSpannerCoordinator(ctx, s.cfg.SpannerClient, table, uint64(opts.PushbackMaxOutstanding())) + seq, err := newSpannerCoordinator(ctx, spannerClient, table, uint64(opts.PushbackMaxOutstanding())) if err != nil { + closeClient() return nil, nil, fmt.Errorf("failed to create Spanner coordinator: %v", err) } + seq.closeClient = closeClient a, lr, err := s.newAppender(ctx, gs, seq, opts) if err != nil { @@ -279,7 +285,15 @@ func (s *Storage) Appender(ctx context.Context, opts *tessera.AppendOptions) (*t } // newAppender creates and initialises a tessera.Appender struct with the provided underlying storage implementations. -func (s *Storage) newAppender(ctx context.Context, o objStore, seq *spannerCoordinator, 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 *spannerCoordinator, 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) } @@ -314,7 +328,6 @@ func (s *Storage) newAppender(ctx context.Context, o objStore, seq *spannerCoord }, nextIndex: a.sequencer.nextIndex, } - var err error a.newCP, err = opts.CheckpointPublisherContext(ctx, reader, s.cfg.HTTPClient) if err != nil { return nil, nil, fmt.Errorf("failed to create checkpoint publisher: %v", err) @@ -324,12 +337,21 @@ func (s *Storage) newAppender(ctx context.Context, o objStore, seq *spannerCoord return nil, nil, fmt.Errorf("failed to initialise log storage: %v", err) } - go a.integrateEntriesJob(ctx) - go a.publishCheckpointJob(ctx, opts.CheckpointInterval(), opts.CheckpointRepublishInterval(), opts.CheckpointPublicationTimeout()) + a.jobs.Go(func() { a.integrateEntriesJob(ctx) }) + a.jobs.Go(func() { + a.publishCheckpointJob(ctx, opts.CheckpointInterval(), opts.CheckpointRepublishInterval(), opts.CheckpointPublicationTimeout()) + }) if i := opts.GarbageCollectionInterval(); i > 0 { - go a.garbageCollectorJob(ctx, i) + a.jobs.Go(func() { a.garbageCollectorJob(ctx, i) }) } + // Once ctx is done the appender is finished, and its jobs return: then close the coordinator. + go func() { + <-ctx.Done() + a.jobs.Wait() + seq.close() + }() + return a, reader, nil } @@ -344,6 +366,9 @@ type Appender struct { cpUpdated chan struct{} entriesAssigned chan struct{} + + // jobs tracks the background jobs, which return once the appender's context is done. + jobs sync.WaitGroup } // Add is the entrypoint for adding entries to a sequencing log. @@ -812,6 +837,8 @@ func (a *Appender) updateEntryBundles(ctx context.Context, fromSeq uint64, entri // a durable and thread/multi-process safe sequencer. type spannerCoordinator struct { dbPool *spanner.Client + // closeClient, if set, closes dbPool. It's set only when Tessera created dbPool. + closeClient func() // table returns the provided table name with the configured table prefix, if any, prepended. table func(string) string maxOutstanding uint64 @@ -835,6 +862,13 @@ func newSpannerCoordinator(ctx context.Context, dbPool *spanner.Client, table fu return r, nil } +// close closes the coordinator's Spanner client, if Tessera created it. +func (s *spannerCoordinator) close() { + if s.closeClient != nil { + s.closeClient() + } +} + // initDB ensures that the coordination DB is initialised correctly. // // The database schema consists of 5 tables: diff --git a/storage/gcp/gcp_test.go b/storage/gcp/gcp_test.go index 66aa4442a..002639cec 100644 --- a/storage/gcp/gcp_test.go +++ b/storage/gcp/gcp_test.go @@ -504,6 +504,61 @@ func makeTile(t *testing.T, size uint64) *api.HashTile { return r } +// TestAppenderClosesSpannerClient checks that an appender closes a Spanner client that Tessera +// created once its context is done, and when newAppender fails, and leaves a caller's client open. +func TestAppenderClosesSpannerClient(t *testing.T) { + ctx := t.Context() + sk, _ := mustGenerateKeys(t) + + for _, test := range []struct { + desc string + interval time.Duration + ownClient bool + wantErr bool + }{ + {desc: "its own client, when its context is done", interval: minCheckpointInterval, ownClient: true}, + {desc: "its own client, when newAppender fails", interval: minCheckpointInterval / 2, ownClient: true, wantErr: true}, + {desc: "a caller's client", interval: minCheckpointInterval}, + } { + t.Run(test.desc, func(t *testing.T) { + db, closeDB := newSpannerDB(t) + defer closeDB() + defer db.Close() + seq, err := newSpannerCoordinator(ctx, db, prefixTable(""), 1000) + if err != nil { + t.Fatalf("newSpannerCoordinator: %v", err) + } + closed := make(chan struct{}) + if test.ownClient { + seq.closeClient = func() { close(closed) } + } + actx, cancel := context.WithCancel(ctx) + defer cancel() + opts := tessera.NewAppendOptions().WithCheckpointInterval(test.interval).WithCheckpointSigner(sk) + a, _, 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 { + cancel() + } + if !test.ownClient { + a.jobs.Wait() + time.Sleep(100 * time.Millisecond) + if _, err := db.Single().ReadRow(ctx, "Tessera", spanner.Key{0}, []string{"compatibilityVersion"}); err != nil { + t.Errorf("read with the caller's client after the appender finished: %v, want it still open", err) + } + return + } + select { + case <-closed: + case <-time.After(5 * time.Second): + t.Fatal("the appender did not close its Spanner client") + } + }) + } +} + func TestTileRoundtrip(t *testing.T) { ctx := t.Context() m := newMemObjStore() diff --git a/storage/posix/antispam/badger.go b/storage/posix/antispam/badger.go index 8eca9932e..d1bb6896f 100644 --- a/storage/posix/antispam/badger.go +++ b/storage/posix/antispam/badger.go @@ -111,19 +111,23 @@ func NewAntispam(ctx context.Context, badgerPath string, opts AntispamOpts) (*An return nil, fmt.Errorf("failed to open badger: %v", err) } + gcCtx, stopGC := context.WithCancel(ctx) r := &AntispamStorage{ - opts: opts, - db: db, + opts: opts, + db: db, + stopGC: stopGC, + gcDone: make(chan struct{}), } go func() { + defer close(r.gcDone) // runsInTick tracks the number of GC runs per tick. var runsInTick int64 ticker := time.NewTicker(opts.CompactionInterval) defer ticker.Stop() for { select { - case <-ctx.Done(): + case <-gcCtx.Done(): return case <-ticker.C: } @@ -158,11 +162,26 @@ func NewAntispam(ctx context.Context, badgerPath string, opts AntispamOpts) (*An return r, nil } +// Close stops the antispam's compaction, then closes its Badger database, which releases the +// database's directory lock so that it can be opened again. +// +// 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.stopGC() + <-s.gcDone + return s.db.Close() +} + type AntispamStorage struct { opts AntispamOpts db *badger.DB + // stopGC stops the compaction goroutine, which closes gcDone when it returns. + stopGC context.CancelFunc + gcDone chan struct{} + // pushBack is used to prevent the follower from getting too far underwater. // Populate dynamically will set this to true/false based on how far behind the follower is from the // currently integrated tree size. diff --git a/storage/posix/antispam/badger_test.go b/storage/posix/antispam/badger_test.go index 6043299ca..544ed02b7 100644 --- a/storage/posix/antispam/badger_test.go +++ b/storage/posix/antispam/badger_test.go @@ -18,6 +18,8 @@ import ( "context" "crypto/sha256" "fmt" + "runtime" + "strings" "testing" "time" @@ -242,6 +244,51 @@ func TestAntispamFollowerErrorBeforeStream(t *testing.T) { cancel() } +// TestAntispamClose checks that Close releases the database, so it can be opened again in the +// same process, and stops the antispam's goroutines and Badger's. +func TestAntispamClose(t *testing.T) { + dir := t.TempDir() + before := antispamGoroutines() + for i := range 3 { + as, err := NewAntispam(t.Context(), dir, AntispamOpts{CompactionInterval: time.Millisecond}) + if err != nil { + t.Fatalf("NewAntispam, opening %d: %v", i+1, err) + } + time.Sleep(10 * time.Millisecond) // let compaction run + if err := as.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + } + for deadline := time.Now().Add(5 * time.Second); ; time.Sleep(10 * time.Millisecond) { + got := antispamGoroutines() + if got == before { + break + } + if time.Now().After(deadline) { + t.Fatalf("%d antispam and Badger goroutines running after Close, want %d", got, before) + } + } +} + +// antispamGoroutines returns the number of goroutines started by NewAntispam, or running in +// Badger and its cache. +func antispamGoroutines() int { + buf := make([]byte, 1<<20) + for { + n := runtime.Stack(buf, true) + if n < len(buf) { + count := 0 + for _, g := range strings.Split(string(buf[:n]), "\n\n") { + if strings.Contains(g, "github.com/dgraph-io/") || strings.Contains(g, "storage/posix/antispam.NewAntispam") { + count++ + } + } + return count + } + buf = make([]byte, 2*len(buf)) + } +} + type failingLogReader struct { tessera.LogReader }