diff --git a/fragment_internal_test.go b/fragment_internal_test.go index 6fd803954..3beea99ec 100644 --- a/fragment_internal_test.go +++ b/fragment_internal_test.go @@ -1570,6 +1570,7 @@ func TestFragment_WriteTo_ReadFrom(t *testing.T) { } // make a read-only Tx after ReadFrom has committed. tx = idx.holder.txf.NewTx(Txo{Write: !writable, Index: idx, Fragment: f1, Shard: f1.shard}) + defer tx.Rollback() // Verify cache is in other fragment. if n := f1.cache.Len(); n != 1 { diff --git a/rbf/db.go b/rbf/db.go index ec75a6064..bd92f2271 100644 --- a/rbf/db.go +++ b/rbf/db.go @@ -411,16 +411,28 @@ func (db *DB) checkpoint() (err error) { // Close closes the database. func (db *DB) Close() (err error) { - // TODO(bbj): Add wait group to hang until last Tx is complete. + // mark db as closed, spawn a thing to wait for existing tx to drain, then + // release the lock so they CAN drain. We do this before getting the + // write lock, so if something else is waiting on rwmu.Lock, and will be + // competing with us, we can ensure that it'll exit out quickly. + db.mu.Lock() + db.opened = false + // wait for transactions to complete + ch := make(chan struct{}) + db.afterCurrentTx(func() { + close(ch) + }) + db.mu.Unlock() + <-ch + // Wait for writer lock. db.rwmu.Lock() defer db.rwmu.Unlock() + // and main DB lock. db.mu.Lock() defer db.mu.Unlock() - db.opened = false - // Close mmap handle. if db.data != nil { if e := syswrap.Munmap(db.data); e != nil && err == nil { diff --git a/rbf/db_test.go b/rbf/db_test.go index 3e6bdd1b1..2b886677c 100644 --- a/rbf/db_test.go +++ b/rbf/db_test.go @@ -139,6 +139,90 @@ func TestDB_WAL(t *testing.T) { t.Fatal(err) } }) + + // initially this is just a cut and paste of the Halt test, except that + // we close the DB while the reads are still running. + t.Run("Close", func(t *testing.T) { + if testing.Short() { + t.Skip("-short enabled, skipping") + } + + config := rbfcfg.NewDefaultConfig() + config.MaxWALSize = 16 * rbf.PageSize + config.MaxWALCheckpointSize = 8 * rbf.PageSize + config.MinWALCheckpointSize = 4 * rbf.PageSize + + db := MustOpenDB(t, config) + + // Continuously run read overlapping transactions. + ctx, cancel := context.WithCancel(context.Background()) + g, ctx := errgroup.WithContext(ctx) + for i := 0; i < 10; i++ { + i := i + g.Go(func() error { + time.Sleep(time.Duration(i) * 10 * time.Millisecond) // stagger + for { + if err := ctx.Err(); err != nil { + return nil + } + + if err := func() error { + tx, err := db.Begin(false) + if err != nil { + return err + } + // give the db time to close between when we opened and + // when we run the Container call + time.Sleep(10 * time.Millisecond) + _, err = tx.Container("x", 0) + if err != nil { + t.Fatalf("requesting container: %v", err) + } + defer tx.Rollback() + return nil + }(); err != nil { + // it's okay to ErrClosed, because we plan to close + // the database out from under us. + if err != rbf.ErrClosed { + return err + } else { + return nil + } + } + } + }) + } + + // Generate updates to the DB/WAL. + for i := 0; i < 100; i++ { + func() { + tx := MustBegin(t, db, true) + defer tx.Rollback() + + if err := tx.CreateBitmapIfNotExists("x"); err != nil { + t.Fatal(err) + } else if _, err := tx.Add("x", uint64(i)); err != nil { + t.Fatal(err) + } else if err := tx.Commit(); err != nil { + t.Fatal(err) + } + time.Sleep(1 * time.Millisecond) + }() + } + // close the db now. + err := db.Close() + if err != nil { + t.Fatalf("closing db: %v", err) + } + // delay a bit to let some readers try to read + time.Sleep(20 * time.Millisecond) + + // Stop read transactions & wait. + cancel() + if err := g.Wait(); err != nil { + t.Fatal(err) + } + }) } func TestDB_Recovery(t *testing.T) {