From d7ff2b245c82c576128602fa741f2710787b9068 Mon Sep 17 00:00:00 2001
From: Alex Peters <alpe@users.noreply.github.com>
Date: Fri, 2 Dec 2022 11:04:18 +0100
Subject: [PATCH] Add feat: add store name in tracekv-emitted store traces
 #11646

---
 store/cachemulti/store.go | 10 +++++++++-
 store/rootmulti/store.go  |  4 +---
 store/types/store.go      | 23 +++++++++++++++++++++++
 3 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/store/cachemulti/store.go b/store/cachemulti/store.go
index d64dc03aca12..deb1d46272dd 100644
--- a/store/cachemulti/store.go
+++ b/store/cachemulti/store.go
@@ -13,6 +13,10 @@ import (
 	"github.com/cosmos/cosmos-sdk/store/types"
 )
 
+// storeNameCtxKey is the TraceContext metadata key that identifies
+// the store which emitted a given trace.
+const storeNameCtxKey = "store_name"
+
 //----------------------------------------
 // Store
 
@@ -55,7 +59,11 @@ func NewFromKVStore(
 
 	for key, store := range stores {
 		if cms.TracingEnabled() {
-			store = tracekv.NewStore(store.(types.KVStore), cms.traceWriter, cms.traceContext)
+			tctx := cms.traceContext.Clone().Merge(types.TraceContext{
+				storeNameCtxKey: key.Name(),
+			})
+
+			store = tracekv.NewStore(store.(types.KVStore), cms.traceWriter, tctx)
 		}
 		if cms.ListeningEnabled(key) {
 			store = listenkv.NewStore(store.(types.KVStore), key, listeners[key])
diff --git a/store/rootmulti/store.go b/store/rootmulti/store.go
index 7e5d032850e3..b6f9fceb2a0f 100644
--- a/store/rootmulti/store.go
+++ b/store/rootmulti/store.go
@@ -332,9 +332,7 @@ func (rs *Store) SetTracingContext(tc types.TraceContext) types.MultiStore {
 	rs.traceContextMutex.Lock()
 	defer rs.traceContextMutex.Unlock()
 	if rs.traceContext != nil {
-		for k, v := range tc {
-			rs.traceContext[k] = v
-		}
+		rs.traceContext = rs.traceContext.Merge(tc)
 	} else {
 		rs.traceContext = tc
 	}
diff --git a/store/types/store.go b/store/types/store.go
index 8bb7037ea9f0..7e7db85722ed 100644
--- a/store/types/store.go
+++ b/store/types/store.go
@@ -418,6 +418,29 @@ type KVPair kv.Pair
 // every trace operation.
 type TraceContext map[string]interface{}
 
+// Clone clones tc into another instance of TraceContext.
+func (tc TraceContext) Clone() TraceContext {
+	ret := TraceContext{}
+	for k, v := range tc {
+		ret[k] = v
+	}
+
+	return ret
+}
+
+// Merge merges value of newTc into tc.
+func (tc TraceContext) Merge(newTc TraceContext) TraceContext {
+	if tc == nil {
+		tc = TraceContext{}
+	}
+
+	for k, v := range newTc {
+		tc[k] = v
+	}
+
+	return tc
+}
+
 // MultiStorePersistentCache defines an interface which provides inter-block
 // (persistent) caching capabilities for multiple CommitKVStores based on StoreKeys.
 type MultiStorePersistentCache interface {