From 6e060d9ce9981a21270af0f6f1fbc1a06f744810 Mon Sep 17 00:00:00 2001
From: Jussi Maki <jussi@isovalent.com>
Date: Mon, 26 Jan 2026 16:40:20 +0100
Subject: [PATCH 1/2] part: Test MapTxn reusability with non-singleton

The MapTxn reuse test only tested for the singleton case.
Extend to also test when both previous and new maps are not
singletons.

Signed-off-by: Jussi Maki <jussi@isovalent.com>
---
 part/map_test.go | 27 +++++++++++++++++++++++----
 1 file changed, 23 insertions(+), 4 deletions(-)

diff --git a/part/map_test.go b/part/map_test.go
index 55e23fe..85a0f91 100644
--- a/part/map_test.go
+++ b/part/map_test.go
@@ -368,19 +368,38 @@ func TestMapTxn(t *testing.T) {
 	assert.Equal(t, 88, v)
 	tree = txn.Commit()
 
-	// Transaction can be reused after commit
+	// Transaction can be reused after commit (singleton case)
+	assert.Nil(t, tree.tree.root)
+	assert.False(t, tree.hasTree)
 	txn.Set("baz", 99)
-	mp = maps.Collect(txn.All())
-	assert.Len(t, mp, 2)
-	assert.Equal(t, map[string]int{"bar": 88, "baz": 99}, mp)
 	mp = maps.Collect(tree.All())
 	assert.Len(t, mp, 1)
 	assert.Equal(t, map[string]int{"bar": 88}, mp)
+	mp = maps.Collect(txn.All())
+	assert.Len(t, mp, 2)
+	assert.Equal(t, map[string]int{"bar": 88, "baz": 99}, mp)
 
+	// Set again since All() bumps txnID
+	txn.Set("baz", 99)
 	tree = txn.Commit()
+
+	// Transaction can be reused after commit (non-singleton case)
+	assert.NotNil(t, tree.tree.root)
+	assert.True(t, tree.hasTree)
+	txn.Set("baz", 999)
+
 	mp = maps.Collect(tree.All())
 	assert.Len(t, mp, 2)
 	assert.Equal(t, map[string]int{"bar": 88, "baz": 99}, mp)
+
+	mp = maps.Collect(txn.All())
+	assert.Len(t, mp, 2)
+	assert.Equal(t, map[string]int{"bar": 88, "baz": 999}, mp)
+
+	tree = txn.Commit()
+	mp = maps.Collect(tree.All())
+	assert.Len(t, mp, 2)
+	assert.Equal(t, map[string]int{"bar": 88, "baz": 999}, mp)
 }
 
 func TestUint64Map(t *testing.T) {

From 73f0877946d4240993706732e072455aa966379a Mon Sep 17 00:00:00 2001
From: Jussi Maki <jussi@isovalent.com>
Date: Mon, 26 Jan 2026 16:41:11 +0100
Subject: [PATCH 2/2] part: Fix MapTxn reusability

Now that txnIDs are the mechanism to detect if a node can be mutated or not
we need to bump txnID in when committing to allow MapTxn to keep reusing part.Txn.

Rename 'prevTxnID' to 'nextTxnID' and move incrementing the txnID from 'Tree.Txn()'
into 'Txn.Commit()'.

Signed-off-by: Jussi Maki <jussi@isovalent.com>
---
 part/tree.go | 4 ++--
 part/txn.go  | 5 +++--
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/part/tree.go b/part/tree.go
index 7134549..24ea5e6 100644
--- a/part/tree.go
+++ b/part/tree.go
@@ -19,7 +19,7 @@ type Tree[T any] struct {
 	size      int // the number of objects in the tree
 	opts      options
 	prevTxn   *atomic.Pointer[Txn[T]] // the previous txn for reusing the allocation
-	prevTxnID uint64                  // the transaction ID that produced this tree
+	nextTxnID uint64                  // the next transaction ID to use
 }
 
 // New constructs a new tree.
@@ -74,7 +74,7 @@ func (t *Tree[T]) Txn() *Txn[T] {
 	txn.rootWatch = t.rootWatch
 	txn.size = t.size
 	txn.prevTxn = t.prevTxn
-	txn.txnID = t.prevTxnID + 1
+	txn.txnID = t.nextTxnID
 	return txn
 }
 
diff --git a/part/txn.go b/part/txn.go
index b6f484e..4d4936d 100644
--- a/part/txn.go
+++ b/part/txn.go
@@ -65,7 +65,7 @@ func (txn *Txn[T]) Clone() Tree[T] {
 		rootWatch: txn.rootWatch,
 		size:      txn.size,
 		prevTxn:   txn.prevTxn,
-		prevTxnID: txn.txnID,
+		nextTxnID: txn.txnID,
 	}
 }
 
@@ -185,13 +185,14 @@ func (txn *Txn[T]) Commit() Tree[T] {
 		validateTree(txn.oldRoot, nil, nil, txn.txnID)
 		validateTree(txn.root, nil, txn.watches, txn.txnID)
 	}
+	txn.txnID++
 	t := Tree[T]{
 		opts:      txn.opts,
 		root:      txn.root,
 		rootWatch: newRootWatch,
 		size:      txn.size,
 		prevTxn:   txn.prevTxn,
-		prevTxnID: txn.txnID,
+		nextTxnID: txn.txnID,
 	}
 	// Store this txn in the tree to reuse the allocation next time.
 	t.prevTxn.Store(txn)