Use post commit hook for post-create plugin hooks (#2920)

pkg/txn/hooks.go

@@ -0,0 +1,54 @@
+package txn
+
+import (
+	"context"
+)
+
+type key int
+
+const (
+	hookManagerKey key = iota + 1
+)
+
+type hookManager struct {
+	postCommitHooks   []TxnFunc
+	postRollbackHooks []TxnFunc
+}
+
+func (m *hookManager) register(ctx context.Context) context.Context {
+	return context.WithValue(ctx, hookManagerKey, m)
+}
+
+func hookManagerCtx(ctx context.Context) *hookManager {
+	m, ok := ctx.Value(hookManagerKey).(*hookManager)
+	if !ok {
+		return nil
+	}
+	return m
+}
+
+func executePostCommitHooks(ctx context.Context) {
+	m := hookManagerCtx(ctx)
+	for _, h := range m.postCommitHooks {
+		// ignore errors
+		_ = h(ctx)
+	}
+}
+
+func executePostRollbackHooks(ctx context.Context) {
+	m := hookManagerCtx(ctx)
+	for _, h := range m.postRollbackHooks {
+		// ignore errors
+		_ = h(ctx)
+	}
+}
+
+func AddPostCommitHook(ctx context.Context, hook TxnFunc) {
+	m := hookManagerCtx(ctx)
+	m.postCommitHooks = append(m.postCommitHooks, hook)
+}
+
+func AddPostRollbackHook(ctx context.Context, hook TxnFunc) {
+	m := hookManagerCtx(ctx)
+	m.postRollbackHooks = append(m.postRollbackHooks, hook)
+}

pkg/txn/transaction.go

@@ -11,9 +11,6 @@ type Manager interface {
 	Rollback(ctx context.Context) error
 
 	IsLocked(err error) bool
-
-	AddPostCommitHook(ctx context.Context, hook TxnFunc)
-	AddPostRollbackHook(ctx context.Context, hook TxnFunc)
 }
 
 type DatabaseProvider interface {
@@ -26,7 +23,7 @@ type TxnFunc func(ctx context.Context) error
 // the transaction is rolled back. Otherwise it is committed.
 func WithTxn(ctx context.Context, m Manager, fn TxnFunc) error {
 	var err error
-	ctx, err = m.Begin(ctx)
+	ctx, err = begin(ctx, m)
 	if err != nil {
 		return err
 	}
@@ -34,16 +31,16 @@ func WithTxn(ctx context.Context, m Manager, fn TxnFunc) error {
 	defer func() {
 		if p := recover(); p != nil {
 			// a panic occurred, rollback and repanic
-			_ = m.Rollback(ctx)
+			rollback(ctx, m)
 			panic(p)
 		}
 
 		if err != nil {
 			// something went wrong, rollback
-			_ = m.Rollback(ctx)
+			rollback(ctx, m)
 		} else {
 			// all good, commit
-			err = m.Commit(ctx)
+			err = commit(ctx, m)
 		}
 	}()
 
@@ -51,6 +48,36 @@ func WithTxn(ctx context.Context, m Manager, fn TxnFunc) error {
 	return err
 }
 
+func begin(ctx context.Context, m Manager) (context.Context, error) {
+	var err error
+	ctx, err = m.Begin(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	hm := hookManager{}
+	ctx = hm.register(ctx)
+
+	return ctx, nil
+}
+
+func commit(ctx context.Context, m Manager) error {
+	if err := m.Commit(ctx); err != nil {
+		return err
+	}
+
+	executePostCommitHooks(ctx)
+	return nil
+}
+
+func rollback(ctx context.Context, m Manager) {
+	if err := m.Rollback(ctx); err != nil {
+		return
+	}
+
+	executePostRollbackHooks(ctx)
+}
+
 // WithDatabase executes fn with the context provided by p.WithDatabase.
 // It does not run inside a transaction, so all database operations will be
 // executed in their own transaction.