db: introduce support for commit hooks

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 {

db.go

@@ -97,14 +97,16 @@ type dbState struct {
 	gcRateLimitInterval time.Duration
 	metrics             Metrics
 	writeTxnPool        sync.Pool
+	commitHooks         []CommitHook
 }
 
 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)

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()
 

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()