statedb: Fix regression in Modify()

The change to not call merge() if the object did not exist was broken
as merge() wasn't called for the object inserted into secondary indexes.

Fix the issue by returning the merged new object from tableIndexTxn.modify
and inserting that into the secondary indexes.

Signed-off-by: Jussi Maki <jussi.maki@isovalent.com>

Author: joamaki

lpm_index.go

@@ -312,7 +312,7 @@ func (l *lpmIndexTxn) insert(key index.Key, obj object) (old object, hadOld bool
 }
 
 // modify implements tableIndexTxn.
-func (l *lpmIndexTxn) modify(key index.Key, obj object, mod func(old, new object) object) (old object, hadOld bool, watch <-chan struct{}) {
+func (l *lpmIndexTxn) modify(key index.Key, obj object, mod func(old, new object) object) (old object, newObj object, hadOld bool, watch <-chan struct{}) {
 	panic("LPM index cannot be the primary index")
 }
 

part/part_test.go

@@ -672,7 +672,7 @@ func Test_modify(t *testing.T) {
 
 	txn := tree.Txn()
 	for i := range 1000 {
-		old, hadOld := txn.Modify(key, 123, func(x, _ int) int { return x + 1 })
+		old, _, hadOld := txn.Modify(key, 123, func(x, _ int) int { return x + 1 })
 		require.True(t, hadOld)
 		require.Equal(t, i+1, old)
 	}

part/tree.go

@@ -127,7 +127,7 @@ func (t *Tree[T]) Insert(key []byte, value T) (old T, hadOld bool, tree Tree[T])
 // Returns the old value if it exists.
 func (t *Tree[T]) Modify(key []byte, value T, mod func(T, T) T) (old T, hadOld bool, tree Tree[T]) {
 	txn := t.Txn()
-	old, hadOld = txn.Modify(key, value, mod)
+	old, _, hadOld = txn.Modify(key, value, mod)
 	tree = txn.CommitAndNotify()
 	return
 }

part/txn.go

@@ -80,7 +80,7 @@ func (txn *Txn[T]) Insert(key []byte, value T) (old T, hadOld bool) {
 // Returns the old value if it exists and a watch channel that closes when the
 // key changes again.
 func (txn *Txn[T]) InsertWatch(key []byte, value T) (old T, hadOld bool, watch <-chan struct{}) {
-	old, hadOld, watch, txn.root = txn.insert(txn.root, key, value)
+	old, _, hadOld, watch, txn.root = txn.insert(txn.root, key, value)
 	validateTree(txn.root, nil, txn.watches, txn.txnID)
 	if !hadOld {
 		txn.size++
@@ -93,19 +93,19 @@ func (txn *Txn[T]) InsertWatch(key []byte, value T) (old T, hadOld bool, watch <
 
 // Modify a value in the tree. It is up to the
 // caller to not mutate the value in-place and to return a clone.
-// Returns the old value if it exists.
-func (txn *Txn[T]) Modify(key []byte, value T, mod func(T, T) T) (old T, hadOld bool) {
-	old, hadOld, _ = txn.ModifyWatch(key, value, mod)
+// Returns the old value (if it exists) and the new possibly merged value.
+func (txn *Txn[T]) Modify(key []byte, value T, mod func(T, T) T) (old T, newValue T, hadOld bool) {
+	old, newValue, hadOld, _ = txn.ModifyWatch(key, value, mod)
 	return
 }
 
 // Modify a value in the tree. If the key does not exist the modify
 // function is called with the zero value for T. It is up to the
 // caller to not mutate the value in-place and to return a clone.
-// Returns the old value if it exists and a watch channel that closes
-// when the key changes again.
-func (txn *Txn[T]) ModifyWatch(key []byte, value T, mod func(T, T) T) (old T, hadOld bool, watch <-chan struct{}) {
-	old, hadOld, watch, txn.root = txn.modify(txn.root, key, value, mod)
+// Returns the old value (if it exists) and the new possibly merged value,
+// and a watch channel that closes when the key changes again.
+func (txn *Txn[T]) ModifyWatch(key []byte, value T, mod func(T, T) T) (old T, newValue T, hadOld bool, watch <-chan struct{}) {
+	old, newValue, hadOld, watch, txn.root = txn.modify(txn.root, key, value, mod)
 	validateTree(txn.root, nil, txn.watches, txn.txnID)
 	if !hadOld {
 		txn.size++
@@ -243,17 +243,18 @@ func (txn *Txn[T]) cloneNode(n *header[T]) *header[T] {
 	return n
 }
 
-func (txn *Txn[T]) insert(root *header[T], key []byte, value T) (oldValue T, hadOld bool, watch <-chan struct{}, newRoot *header[T]) {
+func (txn *Txn[T]) insert(root *header[T], key []byte, value T) (oldValue T, newValue T, hadOld bool, watch <-chan struct{}, newRoot *header[T]) {
 	return txn.modify(root, key, value, nil)
 }
 
-func (txn *Txn[T]) modify(root *header[T], key []byte, newValue T, mod func(T, T) T) (oldValue T, hadOld bool, watch <-chan struct{}, newRoot *header[T]) {
+func (txn *Txn[T]) modify(root *header[T], key []byte, newValue T, mod func(T, T) T) (oldValue T, newValueOut T, hadOld bool, watch <-chan struct{}, newRoot *header[T]) {
 	txn.dirty = true
 	fullKey := key
+	newValueOut = newValue
 
 	if root == nil {
 		leaf := newLeaf(txn.opts, key, fullKey, newValue)
-		return oldValue, false, leaf.watch, leaf.self()
+		return oldValue, newValueOut, false, leaf.watch, leaf.self()
 	}
 
 	// Start recursing from the root to find the insertion point.
@@ -323,7 +324,8 @@ func (txn *Txn[T]) modify(root *header[T], key []byte, newValue T, mod func(T, T
 			}
 			watch = leaf.watch
 			if mod != nil {
-				leaf.value = mod(oldValue, newValue)
+				newValueOut = mod(oldValue, newValue)
+				leaf.value = newValueOut
 			} else {
 				leaf.value = newValue
 			}

part_index.go

@@ -321,7 +321,7 @@ func (r *partIndexTxn) len() int {
 }
 
 // modify implements tableIndexTxn.
-func (r *partIndexTxn) modify(key index.Key, obj object, mod func(old, new object) object) (old object, hadOld bool, watch <-chan struct{}) {
+func (r *partIndexTxn) modify(key index.Key, obj object, mod func(old, new object) object) (old object, newObj object, hadOld bool, watch <-chan struct{}) {
 	return r.tx.ModifyWatch(key, obj, mod)
 }
 

table.go

@@ -486,7 +486,7 @@ func (t *genTable[Obj]) InsertWatch(txn WriteTxn, obj Obj) (oldObj Obj, hadOld b
 
 func (t *genTable[Obj]) Modify(txn WriteTxn, obj Obj, merge func(old, new Obj) Obj) (oldObj Obj, hadOld bool, err error) {
 	mergeObjects := func(old object, new object) object {
-		new.data = merge(old.data.(Obj), obj)
+		new.data = merge(old.data.(Obj), new.data.(Obj))
 		return new
 	}
 	var old object

types.go

@@ -414,7 +414,7 @@ type tableIndexTxn interface {
 	tableIndex
 
 	insert(key index.Key, obj object) (old object, hadOld bool, watch <-chan struct{})
-	modify(key index.Key, obj object, mod func(old, new object) object) (old object, hadOld bool, watch <-chan struct{})
+	modify(key index.Key, obj object, mod func(old, new object) object) (old object, new object, hadOld bool, watch <-chan struct{})
 	delete(key index.Key) (old object, hadOld bool)
 	reindex(primaryKey index.Key, old object, new object)
 }

write_txn.go

@@ -139,7 +139,9 @@ func (txn *writeTxnState) modify(meta TableMeta, guardRevision Revision, newData
 	if merge == nil {
 		oldObj, oldExists, watch = idIndexTxn.insert(idKey, obj)
 	} else {
-		oldObj, oldExists, watch = idIndexTxn.modify(idKey, obj, merge)
+		// Insert the object into the primary index. This returns the merged new
+		// object which we'll then insert into the secondary indexes.
+		oldObj, obj, oldExists, watch = idIndexTxn.modify(idKey, obj, merge)
 	}
 
 	// Sanity check: is the same object being inserted back and thus the