From 94db72d652120eb0e2361f55c436333644c8c658 Mon Sep 17 00:00:00 2001
From: Benjamin Wang <benjamin.ahrtr@gmail.com>
Date: Mon, 19 Aug 2024 15:43:43 +0100
Subject: [PATCH] Rollback alloc map: remove all page ids which are allocated
 by the txid

Signed-off-by: Benjamin Wang <benjamin.ahrtr@gmail.com>
---
 freelist.go      |  8 ++++++++
 freelist_test.go | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/freelist.go b/freelist.go
index 61d43f81b..dffc7bc74 100644
--- a/freelist.go
+++ b/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)
 }
 
diff --git a/freelist_test.go b/freelist_test.go
index 6b7861add..7a81db091 100644
--- a/freelist_test.go
+++ b/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.