Merge pull request #2010 from molecula/fb-1229

[FB-1229] Shrink RBF freelist & truncate data file on checkpoint

rbf/db.go

@@ -235,12 +235,13 @@ func (db *DB) openWAL() (err error) {
 			break
 		}
 	}
+
 	if fileSize != int64(pageN*PageSize) {
-		if err := db.walFile.Truncate(int64(pageN * PageSize)); err != nil {
+		if err := db.walFile.Truncate(int64(pageN) * PageSize); err != nil {
 			return fmt.Errorf("wal truncate: %w", err)
 		}
 	}
-	if _, err := db.walFile.Seek(int64(pageN*PageSize), io.SeekStart); err != nil {
+	if _, err := db.walFile.Seek(int64(pageN)*PageSize, io.SeekStart); err != nil {
 		return fmt.Errorf("wal seek: %w", err)
 	}
 	db.walPageN = pageN
@@ -310,6 +311,7 @@ func (db *DB) checkpoint() (err error) {
 	}()
 
 	// Copy the pages from the WAL back to the database outside of the lock.
+	var pageN uint32
 	if err := func() error {
 		db.mu.Unlock() // This is intentionally reversed so run w/o lock
 		defer db.mu.Lock()
@@ -367,6 +369,11 @@ func (db *DB) checkpoint() (err error) {
 				return fmt.Errorf("reading page %d [page number %d]: %v", walID, pgno, err)
 			}
 
+			// Determine new database size from the page size in meta page.
+			if pgno == 0 {
+				pageN = readMetaPageN(page)
+			}
+
 			// Write data to the data file.
 			if err = db.writeDBPage(pgno, page); err != nil {
 				return fmt.Errorf("writing page %d: %v", pgno, err)
@@ -404,6 +411,14 @@ func (db *DB) checkpoint() (err error) {
 			db.logger.Errorf("seek wal file: %w", err)
 		}
 
+		// Truncate data file if it has shrunk.
+		if fi, err := db.file.Stat(); err != nil {
+			db.logger.Errorf("stat db file: %w", err)
+		} else if sz := int64(pageN) * PageSize; sz > 0 && fi.Size() > sz {
+			if err := db.file.Truncate(sz); err != nil {
+				db.logger.Errorf("truncate db file: %w", err)
+			}
+		}
 	})
 
 	return nil
@@ -556,7 +571,7 @@ func (db *DB) WALSize() int64 {
 }
 
 func (db *DB) walSize() int64 {
-	return int64(db.walPageN * PageSize)
+	return int64(db.walPageN) * PageSize
 }
 
 // init initializes a new database file.
@@ -634,7 +649,7 @@ func (db *DB) Begin(writable bool) (_ *Tx, err error) {
 	// Wait for WAL size to be below threshold, if we're going to write.
 	// Reads don't care.
 	if writable {
-		for int64(db.walPageN*PageSize) > db.cfg.MaxWALCheckpointSize {
+		for int64(db.walPageN)*PageSize > db.cfg.MaxWALCheckpointSize {
 			if db.isDead != nil {
 				err := db.isDead
 				cleanup()

rbf/tx.go

@@ -103,6 +103,11 @@ func (tx *Tx) Commit() error {
 		return ErrTxClosed
 	}
 
+	// Remove any free pages off the end of the file and update the size.
+	if err := tx.truncateFreelist(); err != nil {
+		return err
+	}
+
 	// If any pages have been written, ensure we write a new meta page with
 	// the commit flag to mark the end of the transaction.
 	if tx.dirty() {
@@ -134,6 +139,59 @@ func (tx *Tx) Commit() error {
 	return tx.db.removeTx(tx)
 }
 
+// truncateFreelist removes any free pages off the end of the file and updates
+// the size of the database. This allows the data file to be resized on checkpoint.
+func (tx *Tx) truncateFreelist() error {
+	for {
+		if truncated, err := tx.truncateLastFreePage(); err != nil {
+			return err
+		} else if !truncated {
+			return nil // no more free pages at end of file, exit
+		}
+	}
+}
+
+// truncateLastFreePage removes the last page from the file if it is a free page.
+// The page count is then decremented to move the high water mark to remove the page.
+// Returns true if a page was removed, otherwise returns false.
+func (tx *Tx) truncateLastFreePage() (truncated bool, outErr error) {
+	tx.modifyingFreelist = true
+	defer tx.freelistCleanup(&outErr)
+
+	c := tx.db.getFreelistCursor(tx)
+	if err := c.Last(); err == io.EOF {
+		return false, nil
+	} else if err != nil {
+		return false, err
+	}
+
+	elem := &c.stack.elems[c.stack.top]
+	leafPage, _, err := c.tx.readPage(elem.pgno)
+	if err != nil {
+		return false, err
+	}
+	cell := readLeafCell(leafPage, elem.index)
+
+	// If page number is not the last page then exit.
+	pgno := uint32((cell.Key << 16) | uint64(cell.lastValue(tx)))
+	pageN := readMetaPageN(tx.meta[:])
+	if pgno < pageN-1 {
+		return false, nil
+	}
+
+	// Otherwise remove it from the freelist.
+	if changed, err := c.Remove(uint64(pgno)); err != nil {
+		return false, err
+	} else if !changed {
+		vprint.PanicOn(fmt.Sprintf("tx.Tx.truncateLastFreePage(): double alloc: %d", pgno))
+	}
+
+	// Decrement the page count in the database.
+	writeMetaPageN(tx.meta[:], pageN-1)
+
+	return true, nil
+}
+
 func (tx *Tx) Rollback() { tx.rollback(false) }
 
 func (tx *Tx) rollback(hasDBLock bool) {

rbf/tx_test.go

@@ -936,6 +936,88 @@ func setArray(tb testing.TB, key, num int, c *rbf.Cursor) {
 	}
 }
 
+// FB-1229: This test verifies that the database will be truncated as pages at
+// the end of the file are pushed to the freelist.
+func TestTx_ReclaimAfterDelete(t *testing.T) {
+	const containerN = 5000
+	const batchSize = 500
+
+	db := MustOpenDB(t)
+	defer MustCloseDB(t, db)
+
+	keys := rand.New(rand.NewSource(0)).Perm(containerN)
+	inserted := make(map[uint64]struct{})
+
+	for i := 0; i < len(keys); i += batchSize {
+		func() {
+			tx := MustBegin(t, db, true)
+			defer tx.Rollback()
+
+			if err := tx.CreateBitmapIfNotExists("x"); err != nil {
+				t.Fatal(err)
+			}
+
+			// Insert a bunch of containers.
+			c := roaring.NewContainerArray(convenientPrepopulatedArray)
+			for j := 0; j < batchSize; j++ {
+				key := uint64(keys[i+j])
+				if err := tx.PutContainer("x", key, c); err != nil {
+					t.Fatal(err)
+				}
+				inserted[key] = struct{}{}
+			}
+
+			// Insert some already inserted containers.
+			var deleted int
+			for k := range inserted {
+				if err := tx.RemoveContainer("x", uint64(k)); err != nil {
+					t.Fatal(err)
+				}
+				delete(inserted, k)
+
+				if deleted++; deleted > 200 {
+					break
+				}
+			}
+
+			if err := tx.Commit(); err != nil {
+				t.Fatal(err)
+			}
+		}()
+	}
+
+	fi, err := os.Stat(db.DataPath())
+	if err != nil {
+		t.Fatal(err)
+	}
+	origSize := fi.Size()
+
+	// Delete all containers.
+	func() {
+		tx := MustBegin(t, db, true)
+		defer tx.Rollback()
+
+		for _, key := range keys {
+			if err := tx.RemoveContainer("x", uint64(key)); err != nil {
+				t.Fatal(err)
+			}
+		}
+
+		if err := tx.Commit(); err != nil {
+			t.Fatal(err)
+		}
+	}()
+
+	// Verify database has shrunk after checkpoint.
+	if err := db.Checkpoint(); err != nil {
+		t.Fatal(err)
+	} else if fi, err := os.Stat(db.DataPath()); err != nil {
+		t.Fatal(err)
+	} else if fi.Size() >= origSize {
+		t.Fatalf("size did not shrink: originally %d bytes, ended with %d bytes", fi.Size(), origSize)
+	}
+}
+
 func BenchmarkTx_Add(b *testing.B) {
 	for _, n := range []int{1, 10, 1000} {
 		b.Run(fmt.Sprint(n), func(b *testing.B) {