Merge pull request #1928 from molecula/fb1211

[FB-1211] don't crash on close during reads

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 {

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)

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 {

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) {

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++ {

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)
 
 }