From c71680ded880653c8032a24d7becafea95972344 Mon Sep 17 00:00:00 2001
From: Fridrik Asmundsson <fridrik01@gmail.com>
Date: Thu, 23 Mar 2023 13:14:49 +0000
Subject: [PATCH] Use TipSetKey as key in cache and return copies

---
 chain/stmgr/execute.go | 27 ++++++++++++++++++++++++---
 chain/stmgr/stmgr.go   |  9 +++++----
 2 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/chain/stmgr/execute.go b/chain/stmgr/execute.go
index 3c2644550..761398928 100644
--- a/chain/stmgr/execute.go
+++ b/chain/stmgr/execute.go
@@ -128,8 +128,16 @@ func (sm *StateManager) ExecutionTraceWithMonitor(ctx context.Context, ts *types
 }
 
 func (sm *StateManager) ExecutionTrace(ctx context.Context, ts *types.TipSet) (cid.Cid, []*api.InvocResult, error) {
-	if entry, ok := sm.execTraceCache.Get(ts); ok {
-		return entry.cid, entry.invocTrace, nil
+	tsKey := ts.Key()
+
+	{
+		// check if we have the trace for this tipset in the cache
+		sm.execTraceCacheLock.Lock()
+		defer sm.execTraceCacheLock.Unlock()
+		if entry, ok := sm.execTraceCache.Get(tsKey); ok {
+			// we have to make a deep copy since caller can modify the invocTrace
+			return entry.postStateRoot, makeDeepCopy(entry.invocTrace), nil
+		}
 	}
 
 	var invocTrace []*api.InvocResult
@@ -138,7 +146,20 @@ func (sm *StateManager) ExecutionTrace(ctx context.Context, ts *types.TipSet) (c
 		return cid.Undef, nil, err
 	}
 
-	sm.execTraceCache.Add(ts, tipSetCacheEntry{st, invocTrace})
+	sm.execTraceCache.Add(tsKey, tipSetCacheEntry{st, makeDeepCopy(invocTrace)})
 
 	return st, invocTrace, nil
 }
+
+func makeDeepCopy(invocTrace []*api.InvocResult) []*api.InvocResult {
+	c := make([]*api.InvocResult, len(invocTrace))
+	for i, ir := range invocTrace {
+		if ir == nil {
+			continue
+		}
+		tmp := *ir
+		c[i] = &tmp
+	}
+
+	return c
+}
diff --git a/chain/stmgr/stmgr.go b/chain/stmgr/stmgr.go
index 86de96376..d7eabaaec 100644
--- a/chain/stmgr/stmgr.go
+++ b/chain/stmgr/stmgr.go
@@ -142,7 +142,8 @@ type StateManager struct {
 
 	// We keep a small cache for calls to ExecutionTrace which helps improve
 	// performance for node operators like exchanges and block explorers
-	execTraceCache *lru.ARCCache[*types.TipSet, tipSetCacheEntry]
+	execTraceCache     *lru.ARCCache[types.TipSetKey, tipSetCacheEntry]
+	execTraceCacheLock sync.Mutex
 }
 
 // Caches a single state tree
@@ -152,8 +153,8 @@ type treeCache struct {
 }
 
 type tipSetCacheEntry struct {
-	cid        cid.Cid
-	invocTrace []*api.InvocResult
+	postStateRoot cid.Cid
+	invocTrace    []*api.InvocResult
 }
 
 func NewStateManager(cs *store.ChainStore, exec Executor, sys vm.SyscallBuilder, us UpgradeSchedule, beacon beacon.Schedule, metadataDs dstore.Batching) (*StateManager, error) {
@@ -195,7 +196,7 @@ func NewStateManager(cs *store.ChainStore, exec Executor, sys vm.SyscallBuilder,
 		}
 	}
 
-	execTraceCache, err := lru.NewARC[*types.TipSet, tipSetCacheEntry](execTraceCacheSize)
+	execTraceCache, err := lru.NewARC[types.TipSetKey, tipSetCacheEntry](execTraceCacheSize)
 	if err != nil {
 		return nil, err
 	}