From 921fa574a69d7cb7dee0c034ec6e57740931395d Mon Sep 17 00:00:00 2001
From: ucwong <ucwong@126.com>
Date: Sat, 7 Sep 2024 04:18:22 +0800
Subject: [PATCH] core/state/snapshot : use mark stale in different layer

---
 core/state/snapshot/difflayer.go | 8 ++++++++
 core/state/snapshot/snapshot.go  | 8 ++------
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/core/state/snapshot/difflayer.go b/core/state/snapshot/difflayer.go
index 779c1ea98c..d51cd3f0c3 100644
--- a/core/state/snapshot/difflayer.go
+++ b/core/state/snapshot/difflayer.go
@@ -541,3 +541,11 @@ func (dl *diffLayer) StorageList(accountHash common.Hash) ([]common.Hash, bool)
 	dl.memory += uint64(len(dl.storageList)*common.HashLength + common.HashLength)
 	return storageList, destructed
 }
+
+// markStale sets the stale flag as true.
+func (dl *diffLayer) markStale() {
+	dl.lock.Lock()
+	defer dl.lock.Unlock()
+
+	dl.stale.Store(true)
+}
diff --git a/core/state/snapshot/snapshot.go b/core/state/snapshot/snapshot.go
index d1ffb5cc2d..b2274ec443 100644
--- a/core/state/snapshot/snapshot.go
+++ b/core/state/snapshot/snapshot.go
@@ -266,9 +266,7 @@ func (t *Tree) Disable() {
 
 		case *diffLayer:
 			// If the layer is a simple diff, simply mark as stale
-			layer.lock.Lock()
-			layer.stale.Store(true)
-			layer.lock.Unlock()
+			layer.markStale()
 
 		default:
 			panic(fmt.Sprintf("unknown layer type: %T", layer))
@@ -725,9 +723,7 @@ func (t *Tree) Rebuild(root common.Hash) {
 
 		case *diffLayer:
 			// If the layer is a simple diff, simply mark as stale
-			layer.lock.Lock()
-			layer.stale.Store(true)
-			layer.lock.Unlock()
+			layer.markStale()
 
 		default:
 			panic(fmt.Sprintf("unknown layer type: %T", layer))