From 1907d9b830e28fef0d2a0691028889f4646fc88b Mon Sep 17 00:00:00 2001
From: Marco Iorio <marco.iorio@isovalent.com>
Date: Thu, 10 Sep 2026 17:48:46 +0200
Subject: [PATCH] db: introduce support for commit hooks

Introduce a new [statedb.New] option that allows to register commit
hooks that get invoked every time that a transaction is committed,
before releasing the associated locks. The hooks are strictly read
only, and cannot abort the transaction. Additionally provide the
possibility of registering the hooks through [statedb.Cell].

One possible use-case for the hooks is to validate properties of
the entries of specific tables, for testing purposes. However,
they may also unlock other use-cases, including the possibility
of synchronously dumping the state to disk, and so on.

Signed-off-by: Marco Iorio <marco.iorio@isovalent.com>
---
 cell.go      | 14 +++++++++++---
 db.go        | 23 ++++++++++++++++++++++-
 db_test.go   | 35 +++++++++++++++++++++++++++++++++++
 write_txn.go |  5 +++++
 4 files changed, 73 insertions(+), 4 deletions(-)

diff --git a/cell.go b/cell.go
index 9a233950..72d61daa 100644
--- a/cell.go
+++ b/cell.go
@@ -20,15 +20,23 @@ var Cell = cell.Module(
 	),
 )
 
+// CommitHookOut registers a [CommitHook].
+type CommitHookOut struct {
+	cell.Out
+
+	CommitHook CommitHook `group:"statedb-commit-hooks"`
+}
+
 type params struct {
 	cell.In
 
-	Lifecycle cell.Lifecycle
-	Metrics   Metrics `optional:"true"`
+	Lifecycle   cell.Lifecycle
+	Metrics     Metrics      `optional:"true"`
+	CommitHooks []CommitHook `group:"statedb-commit-hooks"`
 }
 
 func newHiveDB(p params) *DB {
-	db := New(WithMetrics(p.Metrics))
+	db := New(WithMetrics(p.Metrics), WithCommitHooks(p.CommitHooks...))
 	p.Lifecycle.Append(
 		cell.Hook{
 			OnStart: func(cell.HookContext) error {
diff --git a/db.go b/db.go
index bc70559a..8de2b040 100644
--- a/db.go
+++ b/db.go
@@ -97,6 +97,7 @@ type dbState struct {
 	gcRateLimitInterval time.Duration
 	metrics             Metrics
 	writeTxnPool        sync.Pool
+	commitHooks         []CommitHook
 }
 
 type dbRoot = []*tableEntry
@@ -104,7 +105,8 @@ type dbRoot = []*tableEntry
 type Option func(*opts)
 
 type opts struct {
-	metrics Metrics
+	metrics     Metrics
+	commitHooks []CommitHook
 }
 
 func WithMetrics(m Metrics) Option {
@@ -113,6 +115,24 @@ func WithMetrics(m Metrics) Option {
 	}
 }
 
+// CommitHook is a commit hook that can be registered through [WithCommitHooks].
+// Commit hooks are run synchronously, must be strictly read-only, and fast.
+// The given list of tables is only valid until the hook terminates, and must
+// explicitly copied if longer retention is needed.
+type CommitHook func(txn ReadTxn, tables []string)
+
+// WithCommitHooks registers hooks that get invoked every time that a transaction
+// is committed, before releasing the associated locks. They get passed the read
+// transaction for that snapshot, and the list of tables locked by the transaction.
+// The hooks must be strictly read-only, and cannot abort the transaction.
+func WithCommitHooks(hooks ...CommitHook) Option {
+	return func(o *opts) {
+		o.commitHooks = append(o.commitHooks,
+			slices.DeleteFunc(hooks, func(hook CommitHook) bool { return hook == nil })...,
+		)
+	}
+}
+
 // New creates a new database.
 //
 // The created database must be started and stopped!
@@ -130,6 +150,7 @@ func New(options ...Option) *DB {
 		dbState: &dbState{
 			metrics:             opts.metrics,
 			gcRateLimitInterval: defaultGCRateLimitInterval,
+			commitHooks:         opts.commitHooks,
 		},
 	}
 	db.updateWriteTxnPoolLocked(0)
diff --git a/db_test.go b/db_test.go
index bfdf068f..207ddf5b 100644
--- a/db_test.go
+++ b/db_test.go
@@ -1520,6 +1520,41 @@ func TestDB_DeleteEmptySecondaryKey(t *testing.T) {
 	require.Empty(t, Collect(table.List(db.ReadTxn(), tagsIndex.Query(""))))
 }
 
+func TestDB_CommitHooks(t *testing.T) {
+	t.Parallel()
+
+	var (
+		cnt  int
+		last struct {
+			txn    ReadTxn
+			tables []string
+		}
+
+		hook = func(txn ReadTxn, tables []string) {
+			cnt, last.txn, last.tables = cnt+1, txn, slices.Clone(tables)
+		}
+
+		db   = New(WithCommitHooks(hook, nil, hook)) // nil hooks are ignored
+		tbl1 = newTestObjectTable(t, db, "foo")
+		tbl2 = newTestObjectTable(t, db, "bar")
+		tbl3 = newTestObjectTable(t, db, "baz")
+	)
+
+	db.WriteTxn(tbl1, tbl3).Commit()
+
+	require.Equal(t, 2, cnt, "The commit hook should have been invoked exactly twice (registered twice)")
+	require.Equal(t, db.ReadTxn(), last.txn, "The commit hook should be given the correct read transaction")
+	require.ElementsMatch(t, last.tables, []string{"foo", "baz"},
+		"The commit hook should be given the correct list of tables")
+
+	db.WriteTxn(tbl2).Commit()
+
+	require.Equal(t, 4, cnt, "The commit hook should have been invoked two more times")
+	require.Equal(t, db.ReadTxn(), last.txn, "The commit hook should be given the correct read transaction")
+	require.ElementsMatch(t, last.tables, []string{"bar"},
+		"The commit hook should be given the correct list of tables")
+}
+
 func TestWriteJSON(t *testing.T) {
 	t.Parallel()
 
diff --git a/write_txn.go b/write_txn.go
index 684d7972..2004382a 100644
--- a/write_txn.go
+++ b/write_txn.go
@@ -454,6 +454,11 @@ func (handle *writeTxnHandle) Commit() ReadTxn {
 		txn.notify()
 	}
 
+	// Invoke commit hooks, if any.
+	for _, hook := range db.commitHooks {
+		hook((*readTxn)(&root), txn.tableNames)
+	}
+
 	// With the root pointer updated, we can now release the tables for the next write transaction.
 	txn.smus.Unlock()