Merge pull request #1848 from molecula/fb-828

[FB-828] Fix RBF WAL size check

rbf/db_test.go

@@ -3,6 +3,7 @@ package rbf_test
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"math/rand"
 	"net"
@@ -46,6 +47,35 @@ func TestDB_WAL(t *testing.T) {
 		}
 	})
 
+	t.Run("ErrTxTooLargeWithBitmap", func(t *testing.T) {
+		config := rbfcfg.NewDefaultConfig()
+		config.MaxWALSize = 5 * rbf.PageSize
+
+		db := MustOpenDB(t, config)
+		defer MustCloseDB(t, db)
+
+		tx := MustBegin(t, db, true)
+		defer tx.Rollback()
+
+		if err := tx.CreateBitmap("x"); err != nil {
+			t.Fatal(err)
+		}
+
+		// Fill array until it has the maximum number of elements.
+		for i := uint64(0); i < rbf.ArrayMaxSize; i++ {
+			if _, err := tx.Add("x", i); err != nil {
+				t.Fatal(err)
+			}
+		}
+
+		// Issuing one more item to a full array should convert it to a bitmap
+		// page and cause the write to return "tx too large". Previous to the
+		// FB-828 fix, this would write past the mmap size so it was inaccessible.
+		if _, err := tx.Add("x", rbf.ArrayMaxSize); err == nil || !errors.Is(err, rbf.ErrTxTooLarge) {
+			t.Fatalf("unexpected error: %#v", err)
+		}
+	})
+
 	t.Run("Halt", func(t *testing.T) {
 		if testing.Short() {
 			t.Skip("-short enabled, skipping")

rbf/tx.go

@@ -1163,7 +1163,8 @@ func (tx *Tx) writeBitmapPage(pgno uint32, page []byte) error {
 }
 
 func (tx *Tx) checkTxSize() error {
-	if (tx.walPageN+tx.dirtyN())*PageSize >= len(tx.db.wal) {
+	pageN := tx.walPageN + len(tx.dirtyPages) + (len(tx.dirtyBitmapPages) * 2)
+	if pageN*PageSize >= len(tx.db.wal) {
 		return ErrTxTooLarge
 	}
 	return nil