From 6481b4eabe12e3c50ca7c8e735483349484223cb Mon Sep 17 00:00:00 2001 From: Ben Johnson Date: Mon, 20 Dec 2021 13:24:34 -0700 Subject: [PATCH] Add rbf check for empty branch pages --- rbf/db.go | 8 +++++ rbf/rbf_test.go | 9 +++++- rbf/tx.go | 26 ++++++++++++++-- rbf/tx_test.go | 83 +++++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 123 insertions(+), 3 deletions(-) diff --git a/rbf/db.go b/rbf/db.go index f6b9ce5bc..65dd4fbea 100644 --- a/rbf/db.go +++ b/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. diff --git a/rbf/rbf_test.go b/rbf/rbf_test.go index be3461e37..4071a1a95 100644 --- a/rbf/rbf_test.go +++ b/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) diff --git a/rbf/tx.go b/rbf/tx.go index bceabd8f1..5c6c7f7c6 100644 --- a/rbf/tx.go +++ b/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 } diff --git a/rbf/tx_test.go b/rbf/tx_test.go index 1004437a3..345031417 100644 --- a/rbf/tx_test.go +++ b/rbf/tx_test.go @@ -2,8 +2,11 @@ package rbf_test import ( + "encoding/binary" "fmt" "math/rand" + "os" + "strings" "sync" "testing" "time" @@ -770,3 +773,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) + } +}