Rollback alloc map: remove all page ids which are allocated by the txid

Signed-off-by: Benjamin Wang <benjamin.ahrtr@gmail.com>

freelist.go

@@ -252,6 +252,14 @@ func (f *freelist) rollback(txid txid) {
 	}
 	// Remove pages from pending list and mark as free if allocated by txid.
 	delete(f.pending, txid)
+
+	// Remove pgids which are allocated by this txid
+	for pgid, tid := range f.allocs {
+		if tid == txid {
+			delete(f.allocs, pgid)
+		}
+	}
+
 	f.mergeSpans(m)
 }
 

freelist_test.go

@@ -7,6 +7,8 @@ import (
 	"sort"
 	"testing"
 	"unsafe"
+
+	"github.com/stretchr/testify/require"
 )
 
 // TestFreelistType is used as a env variable for test to indicate the backend type
@@ -253,6 +255,38 @@ func TestFreelistArray_allocate(t *testing.T) {
 	}
 }
 
+func Test_Freelist_Rollback(t *testing.T) {
+	f := newTestFreelist()
+
+	f.readIDs([]pgid{3, 5, 6, 7, 12, 13})
+
+	f.free(100, &page{
+		id:       20,
+		flags:    0,
+		count:    0,
+		overflow: 1,
+	})
+	f.allocate(100, 3)
+	f.free(100, &page{
+		id:       25,
+		flags:    0,
+		count:    0,
+		overflow: 0,
+	})
+	f.allocate(100, 2)
+
+	require.Equal(t, map[pgid]txid{5: 100, 12: 100}, f.allocs)
+	require.Equal(t, map[txid]*txPending{100: {
+		ids:     []pgid{20, 21, 25},
+		alloctx: []txid{0, 0, 0},
+	}}, f.pending)
+
+	f.rollback(100)
+
+	require.Equal(t, map[pgid]txid{}, f.allocs)
+	require.Equal(t, map[txid]*txPending{}, f.pending)
+}
+
 // Ensure that a freelist can deserialize from a freelist page.
 func TestFreelist_read(t *testing.T) {
 	// Create a page.