From 0c5eccc8b9ab2c597513810c21153a57d880cdc5 Mon Sep 17 00:00:00 2001
From: zhi <thecoderzhi@gmail.com>
Date: Tue, 11 Oct 2022 09:44:35 -0700
Subject: [PATCH] [protocol] add PreCommit() interface

---
 action/protocol/protocol.go |  5 +++++
 state/factory/util.go       | 14 ++++++++++++++
 state/factory/workingset.go |  4 ++++
 3 files changed, 23 insertions(+)

diff --git a/action/protocol/protocol.go b/action/protocol/protocol.go
index 1a14a2dba1..66e452729b 100644
--- a/action/protocol/protocol.go
+++ b/action/protocol/protocol.go
@@ -52,6 +52,11 @@ type PreStatesCreator interface {
 	CreatePreStates(context.Context, StateManager) error
 }
 
+// PreCommitter performs pre-commit action of the protocol
+type PreCommitter interface {
+	PreCommit(context.Context, StateManager) error
+}
+
 // Committer performs commit action of the protocol
 type Committer interface {
 	Commit(context.Context, StateManager) error
diff --git a/state/factory/util.go b/state/factory/util.go
index ab416ad696..d524c24f23 100644
--- a/state/factory/util.go
+++ b/state/factory/util.go
@@ -81,6 +81,20 @@ func generateWorkingSetCacheKey(blkHeader block.Header, producerAddr string) has
 	return hash.Hash256b(sum)
 }
 
+func protocolPreCommit(ctx context.Context, sr protocol.StateManager) error {
+	if reg, ok := protocol.GetRegistry(ctx); ok {
+		for _, p := range reg.All() {
+			post, ok := p.(protocol.PreCommitter)
+			if ok && sr.ProtocolDirty(p.Name()) {
+				if err := post.PreCommit(ctx, sr); err != nil {
+					return err
+				}
+			}
+		}
+	}
+	return nil
+}
+
 func protocolCommit(ctx context.Context, sr protocol.StateManager) error {
 	if reg, ok := protocol.GetRegistry(ctx); ok {
 		for _, p := range reg.All() {
diff --git a/state/factory/workingset.go b/state/factory/workingset.go
index 048c84d532..cc80966dad 100644
--- a/state/factory/workingset.go
+++ b/state/factory/workingset.go
@@ -216,10 +216,14 @@ func (ws *workingSet) ResetSnapshots() {
 
 // Commit persists all changes in RunActions() into the DB
 func (ws *workingSet) Commit(ctx context.Context) error {
+	if err := protocolPreCommit(ctx, ws); err != nil {
+		return err
+	}
 	if err := ws.store.Commit(); err != nil {
 		return err
 	}
 	if err := protocolCommit(ctx, ws); err != nil {
+		// TODO (zhi): wrap the error and eventually panic it in caller side
 		return err
 	}
 	ws.Reset()