From 1078a64c0e386f90d6af817bcda05da10271df0b Mon Sep 17 00:00:00 2001
From: dustinxie <dahuaxie@gmail.com>
Date: Wed, 20 Nov 2024 15:34:29 -0800
Subject: [PATCH] address comment

---
 state/factory/factory.go         | 35 +++++++++++++++++++++++---------
 state/factory/workingset.go      |  2 +-
 state/factory/workingset_test.go |  4 +++-
 state/factory/workingsetstore.go | 19 +++++++++++++++--
 4 files changed, 46 insertions(+), 14 deletions(-)

diff --git a/state/factory/factory.go b/state/factory/factory.go
index 9ba8cddd7a..2da378d74a 100644
--- a/state/factory/factory.go
+++ b/state/factory/factory.go
@@ -266,19 +266,35 @@ func (sf *factory) newWorkingSet(ctx context.Context, height uint64) (*workingSe
 	if err != nil {
 		return nil, err
 	}
-	var (
-		rootKey    = ArchiveTrieRootKey
-		createTrie = true
+	store, err := newFactoryWorkingSetStore(sf.protocolView, flusher)
+	if err != nil {
+		return nil, err
+	}
+	return sf.createSfWorkingSet(ctx, height, store)
+}
+
+func (sf *factory) newWorkingSetAtHeight(ctx context.Context, height uint64) (*workingSet, error) {
+	span := tracer.SpanFromContext(ctx)
+	span.AddEvent("factory.newWorkingSet")
+	defer span.End()
+
+	g := genesis.MustExtractGenesisContext(ctx)
+	flusher, err := db.NewKVStoreFlusher(
+		sf.dao,
+		batch.NewCachedBatch(),
+		sf.flusherOptions(!g.IsEaster(height))...,
 	)
-	if height < sf.currentChainHeight {
-		// archive mode
-		rootKey = fmt.Sprintf("%s-%d", ArchiveTrieRootKey, height)
-		createTrie = false
+	if err != nil {
+		return nil, err
 	}
-	store, err := newFactoryWorkingSetStore(sf.protocolView, flusher, rootKey, createTrie)
+	store, err := newFactoryWorkingSetStoreAtHeight(sf.protocolView, flusher, height)
 	if err != nil {
 		return nil, err
 	}
+	return sf.createSfWorkingSet(ctx, height, store)
+}
+
+func (sf *factory) createSfWorkingSet(ctx context.Context, height uint64, store workingSetStore) (*workingSet, error) {
 	if err := store.Start(ctx); err != nil {
 		return nil, err
 	}
@@ -293,7 +309,6 @@ func (sf *factory) newWorkingSet(ctx context.Context, height uint64) (*workingSe
 			}
 		}
 	}
-
 	return newWorkingSet(height, store), nil
 }
 
@@ -403,7 +418,7 @@ func (sf *factory) WorkingSetAtHeight(ctx context.Context, height uint64) (proto
 	if height > sf.currentChainHeight {
 		return nil, errors.Errorf("query height %d is higher than tip height %d", height, sf.currentChainHeight)
 	}
-	return sf.newWorkingSet(ctx, height)
+	return sf.newWorkingSetAtHeight(ctx, height)
 }
 
 // PutBlock persists all changes in RunActions() into the DB
diff --git a/state/factory/workingset.go b/state/factory/workingset.go
index 0a34669078..0f25f293df 100644
--- a/state/factory/workingset.go
+++ b/state/factory/workingset.go
@@ -333,7 +333,7 @@ func (ws *workingSet) Commit(ctx context.Context) error {
 		return err
 	}
 	ws.Reset()
-	return ws.store.Stop(ctx)
+	return nil
 }
 
 // State pulls a state from DB
diff --git a/state/factory/workingset_test.go b/state/factory/workingset_test.go
index fd1e5c1841..78f752ed19 100644
--- a/state/factory/workingset_test.go
+++ b/state/factory/workingset_test.go
@@ -57,7 +57,9 @@ func newFactoryWorkingSet(t testing.TB) *workingSet {
 		genesis.Default,
 	)
 	r.NoError(sf.Start(ctx))
-	// defer r.NoError(sf.Stop(ctx))
+	defer func() {
+		r.NoError(sf.Stop(ctx))
+	}()
 
 	ws, err := sf.(workingSetCreator).newWorkingSet(ctx, 1)
 	r.NoError(err)
diff --git a/state/factory/workingsetstore.go b/state/factory/workingsetstore.go
index dea5d5ffa9..cec2d453e5 100644
--- a/state/factory/workingsetstore.go
+++ b/state/factory/workingsetstore.go
@@ -56,8 +56,23 @@ func newStateDBWorkingSetStore(view protocol.View, flusher db.KVStoreFlusher, re
 	}
 }
 
-func newFactoryWorkingSetStore(view protocol.View, flusher db.KVStoreFlusher, trieRootKey string, createTrie bool) (workingSetStore, error) {
-	tlt, err := newTwoLayerTrie(ArchiveTrieNamespace, flusher.KVStoreWithBuffer(), trieRootKey, createTrie)
+func newFactoryWorkingSetStore(view protocol.View, flusher db.KVStoreFlusher) (workingSetStore, error) {
+	tlt, err := newTwoLayerTrie(ArchiveTrieNamespace, flusher.KVStoreWithBuffer(), ArchiveTrieRootKey, true)
+	if err != nil {
+		return nil, err
+	}
+
+	return &factoryWorkingSetStore{
+		flusher:   flusher,
+		view:      view,
+		tlt:       tlt,
+		trieRoots: make(map[int][]byte),
+	}, nil
+}
+
+func newFactoryWorkingSetStoreAtHeight(view protocol.View, flusher db.KVStoreFlusher, height uint64) (workingSetStore, error) {
+	rootKey := fmt.Sprintf("%s-%d", ArchiveTrieRootKey, height)
+	tlt, err := newTwoLayerTrie(ArchiveTrieNamespace, flusher.KVStoreWithBuffer(), rootKey, false)
 	if err != nil {
 		return nil, err
 	}