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/cursor_test.go b/rbf/cursor_test.go index 9c5603a3a..7de5f7d77 100644 --- a/rbf/cursor_test.go +++ b/rbf/cursor_test.go @@ -2,6 +2,7 @@ package rbf_test import ( + "bytes" "io" "math/bits" "math/rand" @@ -852,8 +853,10 @@ func TestDumpDot(t *testing.T) { if err != nil { t.Fatal(err) } - rbf.Dumpdot(tx, 0, " ", os.Stdout) + var b bytes.Buffer + rbf.Dumpdot(tx, 0, " ", &b) } + func TestCursor_UpdateBranchCells(t *testing.T) { db := MustOpenDB(t) defer MustCloseDB(t, db) 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) { diff --git a/rbf/rbf_test.go b/rbf/rbf_test.go index 3ff80350b..0a2432128 100644 --- a/rbf/rbf_test.go +++ b/rbf/rbf_test.go @@ -164,79 +164,6 @@ func GenerateValues(rand *rand.Rand, n int) []uint64 { return a } -var _ = ToRows - -// ToRows returns a sorted list of rows from a set of values. -func ToRows(values []uint64) []*Row { - m := make(map[uint64][]uint64) - for _, v := range values { - id := v / rbf.ShardWidth - m[id] = append(m[id], v&rbf.RowValueMask) - } - - a := make([]*Row, 0, len(m)) - for id, values := range m { - a = append(a, &Row{ID: id, Values: values}) - } - sort.Slice(a, func(i, j int) bool { return a[i].ID < a[j].ID }) - return a -} - -var _ = Row{} - -type Row struct { - ID uint64 - Values []uint64 -} - -func (r *Row) Bitmap() []uint64 { - a := make([]uint64, rbf.ShardWidth/64) - for _, v := range r.Values { - a[v/64] |= 1 << (v % 64) - } - return a -} - -// Union returns the union of r and other's values. -func (r *Row) Union(other *Row) []uint64 { - m := make(map[uint64]struct{}) - for _, v := range r.Values { - m[v] = struct{}{} - } - for _, v := range other.Values { - m[v] = struct{}{} - } - - a := make([]uint64, 0, len(m)) - for v := range m { - a = append(a, v) - } - sort.Slice(a, func(i, j int) bool { return a[i] < a[j] }) - return a -} - -// Intersect returns the intersection of r & other's values. -func (r *Row) Intersect(other *Row) []uint64 { - m := make(map[uint64]struct{}) - for _, v := range r.Values { - m[v] = struct{}{} - } - - a := make([]uint64, 0) - used := make(map[uint64]struct{}) - for _, v := range other.Values { - if _, ok := used[v]; ok { - continue - } - if _, ok := m[v]; ok { - used[v] = struct{}{} - a = append(a, v) - } - } - sort.Slice(a, func(i, j int) bool { return a[i] < a[j] }) - return a -} - // QuickCheck executes fn multiple times with a different PRNG. func QuickCheck(t *testing.T, fn func(t *testing.T, rand *rand.Rand)) { for i := 0; i < *quickCheckN; i++ { diff --git a/rbf/tx_test.go b/rbf/tx_test.go index e9e493afc..a02361315 100644 --- a/rbf/tx_test.go +++ b/rbf/tx_test.go @@ -2,6 +2,7 @@ package rbf_test import ( + "bytes" "encoding/binary" "fmt" "math/rand" @@ -742,7 +743,11 @@ func TestTx_DeleteBitmapsWithPrefix(t *testing.T) { t.Fatal(err) } } - checkInfos := func() { + var b bytes.Buffer + pBuf := func(msg string, args ...interface{}) (int, error) { + return fmt.Fprintf(&b, msg, args...) + } + checkInfos := func(pf func(string, ...interface{}) (int, error)) { tx := MustBegin(t, db, false) defer tx.Rollback() infos, err := tx.PageInfos() @@ -750,34 +755,34 @@ func TestTx_DeleteBitmapsWithPrefix(t *testing.T) { for pgno, info := range infos { switch info := info.(type) { case *rbf.MetaPageInfo: - fmt.Printf("%-8d ", pgno) - fmt.Printf("%-10s ", "meta") - fmt.Printf("pageN=%d,walid=%d,rootrec=%d,freelist=%d\n", info.PageN, info.WALID, info.RootRecordPageNo, info.FreelistPageNo) + pf("%-8d ", pgno) + pf("%-10s ", "meta") + pf("pageN=%d,walid=%d,rootrec=%d,freelist=%d\n", info.PageN, info.WALID, info.RootRecordPageNo, info.FreelistPageNo) case *rbf.RootRecordPageInfo: - fmt.Printf("%-8d ", pgno) - fmt.Printf("%-10s ", "rootrec") - fmt.Printf("next=%d\n", info.Next) + pf("%-8d ", pgno) + pf("%-10s ", "rootrec") + pf("next=%d\n", info.Next) case *rbf.LeafPageInfo: - fmt.Printf("%-8d ", pgno) - fmt.Printf("%-10s ", "leaf") - fmt.Printf("flags=x%x,celln=%d\n", info.Flags, info.CellN) + pf("%-8d ", pgno) + pf("%-10s ", "leaf") + pf("flags=x%x,celln=%d\n", info.Flags, info.CellN) case *rbf.BranchPageInfo: - fmt.Printf("%-8d ", pgno) - fmt.Printf("%-10s ", "branch") - fmt.Printf("flags=x%x,celln=%d\n", info.Flags, info.CellN) + pf("%-8d ", pgno) + pf("%-10s ", "branch") + pf("flags=x%x,celln=%d\n", info.Flags, info.CellN) case *rbf.BitmapPageInfo: - fmt.Printf("%-8d ", pgno) - fmt.Printf("%-10s ", "bitmap") - fmt.Printf("-\n") + pf("%-8d ", pgno) + pf("%-10s ", "bitmap") + pf("-\n") case *rbf.FreePageInfo: - fmt.Printf("%-8d ", pgno) - fmt.Printf("%-10s ", "free") - fmt.Printf("-\n") + pf("%-8d ", pgno) + pf("%-10s ", "free") + pf("-\n") default: t.Fatal(fmt.Sprintf("unexpected page info type %T", info)) @@ -806,19 +811,19 @@ func TestTx_DeleteBitmapsWithPrefix(t *testing.T) { ifError(tx.Commit()) } - checkInfos() + checkInfos(pBuf) populate() - checkInfos() + checkInfos(pBuf) ifError(db.Check()) tx := MustBegin(t, db, true) tx.DeleteBitmapsWithPrefix(prefix) ifError(tx.Commit()) ifError(db.Check()) - checkInfos() + checkInfos(pBuf) populate() ifError(db.Check()) - checkInfos() + checkInfos(pBuf) }