Merge pull request #1819 from molecula/rbf-check-empty-branch

[FB-1105] Add rbf check for empty branch pages

rbf/db.go

@@ -265,6 +265,14 @@ func (db *DB) methodicalWALPageN(pageN int) (lastMeta int, err error) {
 	return lastMeta, nil
 }
 
+// Checkpoint performs a manual checkpoint. This is not necessary except for tests.
+func (db *DB) Checkpoint() error {
+	db.mu.Lock()
+	defer db.mu.Unlock()
+	db.rwmu.Lock()
+	return db.checkpoint()
+}
+
 // checkpoint moves all WAL pages to the main DB file. Must be called
 // while holding both db.mu and db.rwmu. Should release db.rwmu, but not
 // db.mu.

rbf/rbf_test.go

@@ -86,7 +86,14 @@ func MustCloseDB(tb testing.TB, db *rbf.DB) {
 	tb.Helper()
 	if err := db.Check(); err != nil && err != rbf.ErrClosed {
 		tb.Fatal(err)
-	} else if n := db.TxN(); n != 0 {
+	}
+	MustCloseDBNoCheck(tb, db)
+}
+
+// MustCloseDBNoCheck closes db. On error, fail test.
+func MustCloseDBNoCheck(tb testing.TB, db *rbf.DB) {
+	tb.Helper()
+	if n := db.TxN(); n != 0 {
 		tb.Fatalf("db still has %d active transactions; must closed before closing db", n)
 	} else if err := db.Close(); err != nil && err != rbf.ErrClosed {
 		tb.Fatal(err)

rbf/tx.go

@@ -741,6 +741,27 @@ func (tx *Tx) Check() error {
 	return nil
 }
 
+func (tx *Tx) checkPage(pgno, parent, typ uint32) error {
+	switch typ {
+	case PageTypeBranch:
+		return tx.checkBranchPage(pgno, parent, typ)
+	default:
+		return nil
+	}
+}
+
+func (tx *Tx) checkBranchPage(pgno, parent, typ uint32) error {
+	page, _, err := tx.readPage(pgno)
+	if err != nil {
+		return err
+	}
+
+	if readCellN(page) == 0 {
+		return fmt.Errorf("branch page %d is empty", pgno)
+	}
+	return nil
+}
+
 // checkPageAllocations ensures that all pages are either in-use or on the freelist.
 func (tx *Tx) checkPageAllocations() error {
 	freePageSet, err := tx.freePageSet()
@@ -830,7 +851,7 @@ func (tx *Tx) inusePageSet() (map[uint32]struct{}, error) {
 	// Traverse freelist and mark pages as in-use.
 	if err := tx.walkTree(readMetaFreelistPageNo(tx.meta[:]), 0, func(pgno, parent, typ uint32) error {
 		m[pgno] = struct{}{}
-		return nil
+		return tx.checkPage(pgno, parent, typ)
 	}); err != nil {
 		return m, err
 	}
@@ -846,7 +867,8 @@ func (tx *Tx) inusePageSet() (map[uint32]struct{}, error) {
 
 		if err := tx.walkTree(pgno.(uint32), 0, func(pgno, parent, typ uint32) error {
 			m[pgno] = struct{}{}
-			return nil
+
+			return tx.checkPage(pgno, parent, typ)
 		}); err != nil {
 			return m, err
 		}

rbf/tx_test.go

@@ -2,8 +2,11 @@
 package rbf_test
 
 import (
+	"encoding/binary"
 	"fmt"
 	"math/rand"
+	"os"
+	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -816,3 +819,83 @@ func TestTx_DeleteBitmapsWithPrefix(t *testing.T) {
 	checkInfos()
 
 }
+
+func TestTx_Check(t *testing.T) {
+	t.Run("EmptyBranchPage", func(t *testing.T) {
+		t.Parallel()
+
+		db := MustOpenDB(t)
+		defer MustCloseDBNoCheck(t, db)
+		tx := MustBegin(t, db, true)
+		defer tx.Rollback()
+
+		if err := tx.CreateBitmap("x"); err != nil {
+			t.Fatal(err)
+		}
+
+		// Insert enough array containers to split page.
+		for i := 0; i < 1000; i++ {
+			if _, err := tx.Add("x", uint64(i<<16)); err != nil {
+				t.Fatalf("Add(%d) err=%q", i<<16, err)
+			}
+		}
+
+		// Read page types for all pages.
+		infos, err := tx.PageInfos()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		// Commit & checkpoint to flush to the data file.
+		if err := tx.Commit(); err != nil {
+			t.Fatal(err)
+		} else if err := db.Checkpoint(); err != nil {
+			t.Fatal(err)
+		}
+
+		// Corrupt first branch page found by zeroing out the cell count.
+		var pgno uint32
+		for _, info := range infos {
+			if info, ok := info.(*rbf.BranchPageInfo); ok {
+				pgno = info.Pgno
+				page := mustReadPage(t, db.DataPath(), pgno)
+				binary.BigEndian.PutUint16(page[8:10], 0) // zero cell count
+				mustWritePage(t, db.DataPath(), pgno, page)
+				break
+			}
+		}
+
+		// Verify that check now returns an error.
+		if err := db.Check(); err == nil || !strings.Contains(err.Error(), fmt.Sprintf("branch page %d is empty", pgno)) {
+			t.Fatalf("unexpected error: %#v", err)
+		}
+	})
+}
+
+func mustReadPage(tb testing.TB, path string, pgno uint32) []byte {
+	tb.Helper()
+	f, err := os.Open(path)
+	if err != nil {
+		tb.Fatal(err)
+	}
+	defer f.Close()
+
+	buf := make([]byte, rbf.PageSize)
+	if _, err := f.ReadAt(buf, int64(pgno)*rbf.PageSize); err != nil {
+		tb.Fatal(err)
+	}
+	return buf
+}
+
+func mustWritePage(tb testing.TB, path string, pgno uint32, buf []byte) {
+	tb.Helper()
+	f, err := os.OpenFile(path, os.O_WRONLY, 0666)
+	if err != nil {
+		tb.Fatal(err)
+	}
+	defer f.Close()
+
+	if _, err := f.WriteAt(buf, int64(pgno)*rbf.PageSize); err != nil {
+		tb.Fatal(err)
+	}
+}