Merge pull request #1860 from molecula/rbfPages

RBF page/cursor management improvements

etcd/embed.go

@@ -210,7 +210,16 @@ func (e *Etcd) retryClient(fn func(cli *clientv3.Client) error) (err error) {
 			e.cli = cli
 			e.cliMu.Unlock()
 			break
-		case etcdserver.ErrTimeout:
+		case nil:
+			return nil
+		default:
+			msg := err.Error()
+			if !strings.HasPrefix(msg, "etcdserver: request timed out") {
+				// not a known error, also not a wrapped timeout
+				return errors.Wrap(err, "non-retryable error")
+			}
+			fallthrough // treat this as being like a timeout error
+		case etcdserver.ErrTimeout, etcdserver.ErrTimeoutDueToLeaderFail, etcdserver.ErrTimeoutDueToConnectionLost, etcdserver.ErrTimeoutLeaderTransfer:
 			// sporadic timeouts are concerning but not necessarily fatal
 			// and can usually be retried.
 			elapsed := time.Since(start)
@@ -225,9 +234,6 @@ func (e *Etcd) retryClient(fn func(cli *clientv3.Client) error) (err error) {
 			// from spamming these.
 			time.Sleep(100 * time.Millisecond)
 			break
-		default:
-			// nil, or an error we don't know about
-			return errors.Wrap(err, "non-retryable error")
 		}
 	}
 	// if we got here, we got a total of three of some combination of

rbf/cursor.go

@@ -526,7 +526,7 @@ func (c *Cursor) putLeafCellFast(in leafCell, isInsert bool) (err error) {
 	}
 
 	// Write page header.
-	dst := allocPage() // make([]byte, PageSize)
+	dst := allocPage()
 	writePageNo(dst, readPageNo(src))
 	writeFlags(dst, PageTypeLeaf)
 	writeCellN(dst, dstCellN)
@@ -616,7 +616,7 @@ func (c *Cursor) deleteLeafCell(key uint64) (err error) {
 	cells = cells[:len(cells)-1]
 
 	// Write cells to page.
-	buf := make([]byte, PageSize)
+	buf := allocPage()
 	writePageNo(buf[:], elem.pgno)
 	writeFlags(buf[:], PageTypeLeaf)
 	writeCellN(buf[:], len(cells))
@@ -800,7 +800,7 @@ func (c *Cursor) deleteBranchCell(stackIndex int, key uint64) (err error) {
 			return err
 		}
 
-		buf := make([]byte, PageSize)
+		buf := allocPage()
 		copy(buf, target)
 		writePageNo(buf[:], elem.pgno)
 

rbf/db.go

@@ -11,6 +11,7 @@ import (
 	"sort"
 	"sync"
 	"syscall"
+	"unsafe"
 
 	"github.com/benbjohnson/immutable"
 	"github.com/molecula/featurebase/v2/logger"
@@ -68,6 +69,8 @@ type DB struct {
 
 	// Path represents the path to the database file.
 	Path string
+
+	freelistCursor Cursor // cursor to reuse for freelist operations
 }
 
 // NewDB returns a new instance of DB.
@@ -560,7 +563,7 @@ func (db *DB) init() error {
 // initMetaPage initializes the meta page.
 func (db *DB) initMetaPage() error {
 
-	page := make([]byte, PageSize)
+	page := allocPage()
 	writeMetaMagic(page)
 	writeMetaPageN(page, 3)
 	writeMetaRootRecordPageNo(page, 1)
@@ -572,7 +575,7 @@ func (db *DB) initMetaPage() error {
 // initRootRecordPage initializes the initial root record page.
 func (db *DB) initRootRecordPage() error {
 
-	page := make([]byte, PageSize)
+	page := allocPage()
 	writePageNo(page, 1)
 	writeFlags(page, PageTypeRootRecord)
 	_, err := db.file.WriteAt(page, 1*PageSize)
@@ -582,7 +585,7 @@ func (db *DB) initRootRecordPage() error {
 // initFreelistPage initializes the initial freelist btree page.
 func (db *DB) initFreelistPage() error {
 
-	page := make([]byte, PageSize)
+	page := allocPage()
 	writePageNo(page, 2)
 	writeFlags(page, PageTypeLeaf)
 	_, err := db.file.WriteAt(page, 2*PageSize)
@@ -807,6 +810,11 @@ func (db *DB) readMetaPage() ([]byte, error) {
 	return db.readDBPage(0)
 }
 
+// getCursor returns a cursor which has not been zeroed. The only thing
+// a caller should need to do is set c.stack's top correctly (it should be
+// 0, and the [0] elem should be the root page to start on).
+//
+// TODO: Should this do anything about c.buffered?
 func (db *DB) getCursor(tx *Tx) *Cursor {
 	c := cursorSyncPool.Get().(*Cursor)
 	c.tx = tx
@@ -827,20 +835,40 @@ type DebugInfo struct {
 	Txs  []*TxDebugInfo `json:"txs"`
 }
 
-// Shared pool for in-memory database pages.
-// These are used before being flushed to disk.
-var pagePool = &sync.Pool{
-	New: func() interface{} {
-		page := make([]byte, PageSize)
-		return &page
-	},
+// when we want a cursor to access a free list, we are always doing this in
+// a context specific to a write transaction, of which any DB can only have
+// one at a time, and the operations modifying the free list don't recurse,
+// because that would corrupt the list (see tx.freelistCleanup for the hairy
+// details), which means that there is only ever one cursor being used for the
+// free list, but also we use that cursor very often, and if we have to allocate
+// it or zero it we end up with a lot of excess allocations and zeroing.
+func (db *DB) getFreelistCursor(tx *Tx) *Cursor {
+	c := &db.freelistCursor
+	c.tx = tx
+	c.stack.elems[0] = stackElem{pgno: readMetaFreelistPageNo(tx.meta[:])}
+	c.stack.top = 0
+	c.buffered = false
+	return c
 }
 
+// Shared pool for in-memory database pages.
+// These are used before being flushed to disk.
+var pagePool = &sync.Pool{}
+
 func allocPage() []byte {
-	page := pagePool.Get().(*[]byte)
-	return *page
+	existing := pagePool.Get()
+	if existing == nil {
+		return make([]byte, PageSize)
+	}
+	// zero the existing page before returning it
+	page := existing.(*[PageSize]byte)[:]
+	for i := range page {
+		page[i] = 0
+	}
+	return page
 }
 
 func freePage(page []byte) {
-	pagePool.Put(&page)
+	data := (*[PageSize]byte)(unsafe.Pointer(&page[0]))
+	pagePool.Put(data)
 }

rbf/db_test.go

@@ -339,7 +339,7 @@ func TestDB_MultiTx(t *testing.T) {
 	}
 
 	// Continuously set/clear bits while readers are executing.
-	for i := 0; i < 1000; i++ {
+	for i := 0; i < 100; i++ {
 		func() {
 			tx, err := db.Begin(true)
 			if err != nil {

rbf/tx.go

@@ -236,7 +236,7 @@ func (tx *Tx) createBitmap(name string) error {
 	}
 
 	// Write root page.
-	page := make([]byte, PageSize)
+	page := allocPage()
 	writePageNo(page, pgno)
 	writeFlags(page, PageTypeLeaf)
 	writeCellN(page, 0)
@@ -449,7 +449,7 @@ func (tx *Tx) writeRootRecordPages(records *immutable.SortedMap) (err error) {
 	// Write new root record pages.
 	for itr := records.Iterator(); !itr.Done(); {
 		// Initialize page & write as many records as will fit.
-		page := make([]byte, PageSize)
+		page := allocPage()
 		writePageNo(page, pgno)
 		writeFlags(page, PageTypeRootRecord)
 
@@ -984,6 +984,10 @@ func (tx *Tx) walkTree(pgno, parent uint32, fn func(pgno, parent, typ uint32, er
 // about that removing things from the free list, because the add logic
 // already just uses new pages rather than trying to use the free list
 // when it knows the free list is involved.
+//
+// Because this is expected to be used in a defer, instead of returning an
+// error, it will set the error it got the address of to a new error if it
+// encounters one and there wasn't one already.
 func (tx *Tx) freelistCleanup(outErr *error) {
 	defer func() {
 		// no matter what, we're done with this after this, but we still
@@ -994,8 +998,7 @@ func (tx *Tx) freelistCleanup(outErr *error) {
 	if len(tx.pendingFreelistAdds) == 0 {
 		return
 	}
-	c := Cursor{tx: tx}
-	c.stack.elems[0] = stackElem{pgno: readMetaFreelistPageNo(tx.meta[:])}
+	c := tx.db.getFreelistCursor(tx)
 	for len(tx.pendingFreelistAdds) > 0 {
 		var pass []uint32
 		pass, tx.pendingFreelistAdds = tx.pendingFreelistAdds, nil
@@ -1006,7 +1009,7 @@ func (tx *Tx) freelistCleanup(outErr *error) {
 				}
 				return
 			} else if !changed {
-				vprint.PanicOn(fmt.Sprintf("rbf.Tx.freePgno(): double free: %d", tx.pendingFreelistAdds))
+				vprint.PanicOn(fmt.Sprintf("rbf.Tx.freelistCleanup(): double free: %d", pass))
 			}
 		}
 	}
@@ -1015,44 +1018,25 @@ func (tx *Tx) freelistCleanup(outErr *error) {
 // allocatePgno returns a page number for a new available page. This page may be
 // pulled from the free list or, if no free pages are available, it will be
 // created by extending the file size.
+//
+// allocatePgno uses the freelist cursor (a shared db-wide thing), and sets
+// the "modifyingFreelist" flag while it's running. If for some reason a
+// modification to the freelist would require a new allocation or free,
+// allocations always just create a new page, and frees are processed later
+// by a separate call through a deferred tx.freelistCleanup().
 func (tx *Tx) allocatePgno() (_ uint32, outErr error) {
 	if tx.modifyingFreelist {
 		return tx.allocateNewPgno(), nil
 	}
-	// Attempt to find page in freelist.
-	pgno, err := tx.nextFreelistPageNo()
-
-	if err != nil {
-		return 0, err
-	} else if pgno != 0 {
-		tx.modifyingFreelist = true
-		defer tx.freelistCleanup(&outErr)
-		c := Cursor{tx: tx}
-		c.stack.elems[0] = stackElem{pgno: readMetaFreelistPageNo(tx.meta[:])}
-		if changed, err := c.Remove(uint64(pgno)); err != nil {
-			return 0, err
-		} else if !changed {
-			vprint.PanicOn(fmt.Sprintf("tx.Tx.allocatePgno(): double alloc: %d", pgno))
-		}
-		return pgno, nil
-	}
-	// no freelist pages, fall back
-	return tx.allocateNewPgno(), nil
-}
-
-// allocateNewPgno requests a new page unconditionally, ignoring the free list.
-func (tx *Tx) allocateNewPgno() uint32 {
-	// Increment the total page count by one and return the last page.
-	pgno := readMetaPageN(tx.meta[:])
-	writeMetaPageN(tx.meta[:], pgno+1)
-	return pgno
-}
-
-func (tx *Tx) nextFreelistPageNo() (uint32, error) {
-	c := Cursor{tx: tx}
-	c.stack.elems[0] = stackElem{pgno: readMetaFreelistPageNo(tx.meta[:])}
+	// this serves as a precaution against double-use of the freelist cursor
+	// used database-wide. we don't have actual synchronization here because
+	// only one write Tx should exist at once and it's not safe to use its
+	// write-capable ops concurrently anyway.
+	tx.modifyingFreelist = true
+	defer tx.freelistCleanup(&outErr)
+	c := tx.db.getFreelistCursor(tx)
 	if err := c.First(); err == io.EOF {
-		return 0, nil
+		return tx.allocateNewPgno(), nil
 	} else if err != nil {
 		return 0, err
 	}
@@ -1067,17 +1051,30 @@ func (tx *Tx) nextFreelistPageNo() (uint32, error) {
 	v := cell.firstValue(tx)
 
 	pgno := uint32((cell.Key << 16) | uint64(v))
+
+	if changed, err := c.Remove(uint64(pgno)); err != nil {
+		return 0, err
+	} else if !changed {
+		vprint.PanicOn(fmt.Sprintf("tx.Tx.allocatePgno(): double alloc: %d", pgno))
+	}
 	return pgno, nil
 }
 
+// allocateNewPgno requests a new page unconditionally, ignoring the free list.
+func (tx *Tx) allocateNewPgno() uint32 {
+	// Increment the total page count by one and return the last page.
+	pgno := readMetaPageN(tx.meta[:])
+	writeMetaPageN(tx.meta[:], pgno+1)
+	return pgno
+}
+
 // deallocate releases a page number to the freelist.
 func (tx *Tx) freePgno(pgno uint32) (outErr error) {
 	if tx.modifyingFreelist {
 		tx.pendingFreelistAdds = append(tx.pendingFreelistAdds, pgno)
 		return nil
 	}
-	c := Cursor{tx: tx}
-	c.stack.elems[0] = stackElem{pgno: readMetaFreelistPageNo(tx.meta[:])}
+	c := tx.db.getFreelistCursor(tx)
 
 	tx.modifyingFreelist = true
 	defer tx.freelistCleanup(&outErr)
@@ -1800,9 +1797,16 @@ func (tx *Tx) flush() error {
 	}
 
 	// Write bitmap headers & pages to WAL.
+	//
+	// We need to write a bitmap header before each such page. We only allocate
+	// one header, and we reuse it, because each write is flushing it out to
+	// disk, and it doesn't get stored in-memory.
+	var hdr []byte
+	if len(tx.dirtyBitmapPages) > 0 {
+		hdr = allocPage()
+	}
 	for _, pgno := range dirtyPageMapKeys(tx.dirtyBitmapPages) {
 		// Write header page.
-		hdr := make([]byte, PageSize)
 		writePageNo(hdr[:], pgno)
 		writeFlags(hdr[:], PageTypeBitmapHeader)
 		if _, err := tx.writeToWAL(w, hdr); err != nil {

rbf/tx_test.go

@@ -248,6 +248,27 @@ func TestTx_DeallocateTree(t *testing.T) {
 	}
 }
 
+func arraySizedChunk() []uint16 {
+	v := make([]uint16, rbf.ArrayMaxSize)
+	for i := range v {
+		v[i] = uint16(i)
+	}
+	return v
+}
+
+var convenientPrepopulatedArray = arraySizedChunk()
+
+// populateBitmapWithArrays
+func populateBitmapWithArrays(tb testing.TB, tx *rbf.Tx, n int, name string) {
+	c := roaring.NewContainerArray(convenientPrepopulatedArray)
+	for i := 0; i < n; i++ {
+		err := tx.PutContainer(name, uint64(i), c)
+		if err != nil {
+			tb.Fatal(err)
+		}
+	}
+}
+
 func TestTx_RecreateBitmap(t *testing.T) {
 	db := MustOpenDB(t)
 	defer MustCloseDB(t, db)
@@ -258,14 +279,8 @@ func TestTx_RecreateBitmap(t *testing.T) {
 	if err := tx.CreateBitmap("x"); err != nil {
 		t.Fatal(err)
 	}
-	const N = 825000
-	slots := make([]uint64, N)
-	for i := range slots {
-		slots[i] = uint64(i) << 20
-	}
-	if _, err := tx.Add("x", slots...); err != nil {
-		t.Fatal(err)
-	}
+	const N = 825
+	populateBitmapWithArrays(t, tx, N, "x")
 	err := tx.Commit()
 	if err != nil {
 		t.Fatal(err)
@@ -291,9 +306,7 @@ func TestTx_RecreateBitmap(t *testing.T) {
 	if err := tx.CreateBitmap("x"); err != nil {
 		t.Fatal(err)
 	}
-	if _, err := tx.Add("x", slots...); err != nil {
-		t.Fatal(err)
-	}
+	populateBitmapWithArrays(t, tx, N, "x")
 	err = tx.Commit()
 	if err != nil {
 		t.Fatal(err)
@@ -374,20 +387,14 @@ func TestTx_DeallocateToFreeList(t *testing.T) {
 	if err = tx.CreateBitmap("y"); err != nil {
 		t.Fatal(err)
 	}
-	const N = 12274831
-	slots := make([]uint64, N)
-	for i := range slots {
-		slots[i] = uint64(i) << 10
-	}
-	bm := roaring.NewBitmap(slots...)
-	if _, err = tx.AddRoaring("x", bm); err != nil {
-		t.Fatal(err)
-	}
+	// Insert large array values.
+	populateBitmapWithArrays(t, tx, 4080, "x")
+
 	if err = tx.Check(); err != nil {
 		t.Fatal(err)
 	}
 	for i := 0; i < 500; i++ {
-		if _, err := tx.Add("y", uint64(i)<<16); err != nil {
+		if _, err := tx.Add("y", uint64(i)<<16+32768); err != nil {
 			t.Fatal(err)
 		}
 	}
@@ -426,9 +433,8 @@ func TestTx_DeallocateToFreeList(t *testing.T) {
 	if err := tx.CreateBitmap("x"); err != nil {
 		t.Fatal(err)
 	}
-	if _, err := tx.AddRoaring("x", bm); err != nil {
-		t.Fatal(err)
-	}
+	populateBitmapWithArrays(t, tx, 4080, "x")
+
 	if err = tx.Check(); err != nil {
 		t.Fatal(err)
 	}
@@ -451,17 +457,7 @@ func TestTx_Remove(t *testing.T) {
 	}
 
 	// Insert large array values.
-	var values []uint64
-	for i := 0; i < 1000; i++ {
-		for j := 0; j < rbf.ArrayMaxSize; j++ {
-			v := uint64((i << 16) + j)
-			values = append(values, v)
-
-			if _, err := tx.Add("x", v); err != nil {
-				t.Fatalf("Add(%d) err=%q", v, err)
-			}
-		}
-	}
+	populateBitmapWithArrays(t, tx, 500, "x")
 
 	if err := tx.Commit(); err != nil {
 		t.Fatal(err)
@@ -471,12 +467,17 @@ func TestTx_Remove(t *testing.T) {
 	defer tx.Rollback()
 
 	// Remove all array values.
-	for _, i := range rand.Perm(len(values)) {
-		v := values[i]
-		if _, err := tx.Remove("x", v); err != nil {
-			t.Fatalf("Remove(%d) err=%q", v, err)
+	for i := 0; i < 500; i++ {
+		err := tx.RemoveContainer("x", uint64(i))
+		if err != nil {
+			t.Fatal(err)
 		}
 	}
+	// This triggered a different panic without the relevant patch.
+	err := tx.RemoveContainer("x", 500)
+	if err != nil {
+		t.Fatal(err)
+	}
 
 	if err := tx.Commit(); err != nil {
 		t.Fatal(err)