trace context fence

backport cosmos#11117

store/v2/multi/store.go

@@ -125,9 +125,10 @@ type SchemaBuilder struct {
 
 // Mixin type that to compose trace & listen state into each root store variant type
 type traceListenMixin struct {
-	listeners    map[types.StoreKey][]types.WriteListener
-	TraceWriter  io.Writer
-	TraceContext types.TraceContext
+	listeners         map[types.StoreKey][]types.WriteListener
+	TraceWriter       io.Writer
+	TraceContext      types.TraceContext
+	traceContextMutex sync.RWMutex
 }
 
 func newTraceListenMixin() *traceListenMixin {
@@ -902,12 +903,35 @@ func (tlm *traceListenMixin) SetTracer(w io.Writer) {
 	tlm.TraceWriter = w
 }
 func (tlm *traceListenMixin) SetTracingContext(tc types.TraceContext) {
-	tlm.TraceContext = tc
+	tlm.traceContextMutex.Lock()
+	defer tlm.traceContextMutex.Unlock()
+	if tlm.TraceContext != nil {
+		for k, v := range tc {
+			tlm.TraceContext[k] = v
+		}
+	} else {
+		tlm.TraceContext = tc
+	}
+}
+
+func (tlm *traceListenMixin) getTracingContext() types.TraceContext {
+	tlm.traceContextMutex.Lock()
+	defer tlm.traceContextMutex.Unlock()
+
+	if tlm.TraceContext == nil {
+		return nil
+	}
+
+	ctx := types.TraceContext{}
+	for k, v := range tlm.TraceContext {
+		ctx[k] = v
+	}
+	return ctx
 }
 
 func (tlm *traceListenMixin) wrapTraceListen(store types.KVStore, skey types.StoreKey) types.KVStore {
 	if tlm.TracingEnabled() {
-		store = tracekv.NewStore(store, tlm.TraceWriter, tlm.TraceContext)
+		store = tracekv.NewStore(store, tlm.TraceWriter, tlm.getTracingContext())
 	}
 	if tlm.ListeningEnabled(skey) {
 		store = listenkv.NewStore(store, skey, tlm.listeners[skey])

store/v2/multi/store_test.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"math"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 
@@ -942,6 +943,53 @@ func TestTrace(t *testing.T) {
 	require.NoError(t, store.Close())
 }
 
+func TestTraceConcurrency(t *testing.T) {
+	db := memdb.NewDB()
+	opts := storeConfig123(t)
+	store, err := NewStore(db, opts)
+	require.NoError(t, err)
+
+	b := &bytes.Buffer{}
+	tc := types.TraceContext(map[string]interface{}{"blockHeight": 64})
+
+	store.SetTracer(b)
+	store.SetTracingContext(tc)
+
+	cms := store.CacheWrap()
+	s1 := cms.GetKVStore(skey_1)
+	require.NotNil(t, s1)
+
+	stop := make(chan struct{})
+	stopW := make(chan struct{})
+
+	go func(stop chan struct{}) {
+		for {
+			select {
+			case <-stop:
+				return
+			default:
+				s1.Set([]byte{1}, []byte{1})
+				cms.Write()
+			}
+		}
+	}(stop)
+
+	go func(stop chan struct{}) {
+		for {
+			select {
+			case <-stop:
+				return
+			default:
+				store.SetTracingContext(tc)
+			}
+		}
+	}(stopW)
+
+	time.Sleep(1 * time.Second)
+	stop <- struct{}{}
+	stopW <- struct{}{}
+}
+
 func TestListeners(t *testing.T) {
 	kvPairs := []types.KVPair{
 		{Key: []byte{1}, Value: []byte("v1")},