Merge pull request #857 from molecula/rbf-fix-direct-write-corruption

Fix RBF write corruption during direct write.

rbf/cursor.go

@@ -322,14 +322,25 @@ func toPgno(val []byte) uint32 {
 	return binary.LittleEndian.Uint32(val)
 }
 func (c *Cursor) putLeafCell(in leafCell) (err error) {
-	cells := readLeafCells(c.leafPage, c.leafCells[:])
+	// Copy target page if we are using direct writes because a split will
+	// cause the source data to be overwritten after the first page is written.
+	leafPage := c.leafPage
+	if c.tx.exclusive {
+		leafPage = make([]byte, PageSize)
+		copy(leafPage, c.leafPage)
+	}
+
+	cells := readLeafCells(leafPage, c.leafCells[:])
 	elem := &c.stack.elems[c.stack.index]
 	cell := in
 	if elem.index >= len(cells) || c.Key() != cell.Key {
 		//new cell
 		if in.Type == ContainerTypeBitmap {
 			//allocated bitmap()
-			bitmapPgno, _ := c.tx.allocate()
+			bitmapPgno, err := c.tx.allocate()
+			if err != nil {
+				return err
+			}
 			cell.Data = fromPgno(bitmapPgno)
 			cell.Type = ContainerTypeBitmapPtr
 		}
@@ -494,6 +505,15 @@ func (c *Cursor) putBranchCells(stackIndex int, newCells []branchCell) (err erro
 	if err != nil {
 		return err
 	}
+
+	// Copy target page if we are using direct writes because a split will
+	// cause the source data to be overwritten after the first page is written.
+	if c.tx.exclusive {
+		tmp := make([]byte, PageSize)
+		copy(tmp, page)
+		page = tmp
+	}
+
 	cells := readBranchCells(page)
 
 	// Update current cell & insert additional cells after it.
@@ -522,7 +542,7 @@ func (c *Cursor) putBranchCells(stackIndex int, newCells []branchCell) (err erro
 			parent.Pgno = origPgno
 		} else {
 			if parent.Pgno, err = c.tx.allocate(); err != nil {
-				return fmt.Errorf("cannot allocate leaf: %w", err)
+				return fmt.Errorf("cannot allocate branch: %w", err)
 			}
 		}
 		parents = append(parents, parent)
@@ -544,6 +564,8 @@ func (c *Cursor) putBranchCells(stackIndex int, newCells []branchCell) (err erro
 		}
 	}
 
+	// TODO(BBJ): Check if key on page changes and update parent if so.
+
 	// TODO(BBJ): Update page in buffer & cursor stack.
 
 	// If this is not a split, then exit now.

rbf/db.go

@@ -441,7 +441,6 @@ func (db *DB) readWALPage(walID int64) ([]byte, error) {
 }
 
 func (db *DB) writeWALPage(page []byte, isMeta bool) (walID int64, err error) {
-
 	if err := db.ensureWritableWALSegment(); err != nil {
 		return 0, err
 	}
@@ -829,7 +828,6 @@ func (db *DB) Check() error {
 
 // writePage writes a page to the data file.
 func (db *DB) writePage(pgno uint32, page []byte) error {
-
 	_, err := db.file.WriteAt(page, int64(pgno)*PageSize)
 	return err
 }

rbf/tx.go

@@ -720,10 +720,12 @@ func (tx *Tx) checkPageAllocations() error {
 		if isInuse && isFree {
 			return fmt.Errorf("page in-use & free: pgno=%d", pgno)
 		} else if !isInuse && !isFree {
-			page, _ := tx.readPage(pgno)
+			page, err := tx.readPage(pgno)
+			if err != nil {
+				return err
+			}
 			flags := readFlags(page)
 			if flags == PageTypeBranch || flags == PageTypeLeaf {
-
 				return fmt.Errorf("page not in-use & not free: pgno=%d", pgno)
 			}
 			//assuming its a bitmap so its ok TODO ben?