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 }