fix: Close data race in InMemoryIndex Add/Evict with RWMutex

Signed-off-by: Guangya Liu <gyliu513@gmail.com>

Author: gyliu513

pkg/kvcache/kvblock/in_memory.go

@@ -76,6 +76,9 @@ func NewInMemoryIndex(cfg *InMemoryIndexConfig) (*InMemoryIndex, error) {
 
 // InMemoryIndex is an in-memory implementation of the Index interface.
 type InMemoryIndex struct {
+	// mu protects multi-step read-modify-write sequences on data/engineToRequestKeys
+	// from racing with concurrent Add/Evict operations.
+	mu sync.RWMutex
 	// data holds the mapping of requestKeys to sets of pod identifiers.
 	data *lru.Cache[BlockHash, *PodCache]
 	// engineToRequestKeys holds the mapping of engineKeys to requestKeys.
@@ -89,10 +92,7 @@ var _ Index = &InMemoryIndex{}
 // PodCache represents a cache for pod entries.
 type PodCache struct {
 	// cache is an LRU cache that maps PodEntry to their last access time.
-	// thread-safe.
 	cache *lru.Cache[PodEntry, struct{}]
-	// mu protects the cache from concurrent access during check-and-set operations.
-	mu sync.Mutex
 }
 
 // Lookup receives a list of requestKeys and a set of pod identifiers,
@@ -112,6 +112,9 @@ func (m *InMemoryIndex) Lookup(ctx context.Context, requestKeys []BlockHash,
 
 	traceLogger := log.FromContext(ctx).V(logging.TRACE).WithName("kvblock.InMemoryIndex.Lookup")
 
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+
 	podsPerKey := make(map[BlockHash][]PodEntry)
 	highestHitIdx := 0
 
@@ -157,6 +160,9 @@ func (m *InMemoryIndex) Add(ctx context.Context, engineKeys, requestKeys []Block
 
 	traceLogger := log.FromContext(ctx).V(logging.TRACE).WithName("kvblock.InMemoryIndex.Add")
 
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
 	for i, requestKey := range requestKeys {
 		engineKey := engineKeys[i]
 
@@ -197,11 +203,9 @@ func (m *InMemoryIndex) Add(ctx context.Context, engineKeys, requestKeys []Block
 			}
 		}
 
-		podCache.mu.Lock()
 		for _, entry := range entries {
 			podCache.cache.Add(entry, struct{}{})
 		}
-		podCache.mu.Unlock()
 
 		traceLogger.Info("added pods to key", "requestKey", requestKey, "engineKey", engineKey, "pods", entries)
 	}
@@ -217,6 +221,9 @@ func (m *InMemoryIndex) Evict(ctx context.Context, engineKey BlockHash, entries
 
 	traceLogger := log.FromContext(ctx).V(logging.TRACE).WithName("kvblock.InMemoryIndex.Evict")
 
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
 	requestKey, found := m.engineToRequestKeys.Get(engineKey)
 	if !found {
 		traceLogger.Info("engineKey not found in index, nothing to evict", "engineKey", engineKey)
@@ -230,28 +237,17 @@ func (m *InMemoryIndex) Evict(ctx context.Context, engineKey BlockHash, entries
 		return nil
 	}
 
-	podCache.mu.Lock()
 	for _, entry := range entries {
 		podCache.cache.Remove(entry)
 	}
 
-	isEmpty := podCache.cache.Len() == 0
-	podCache.mu.Unlock()
-
 	traceLogger.Info("evicted pods from key", "requestKey", requestKey, "engineKey", engineKey, "pods", entries)
 
 	// Remove key from main cache if empty
-	if isEmpty {
-		// Re-fetch and hold the lock through removal to prevent racing with Add
-		if currentCache, stillExists := m.data.Get(requestKey); stillExists && currentCache != nil {
-			currentCache.mu.Lock()
-			if currentCache.cache.Len() == 0 {
-				m.data.Remove(requestKey)
-				m.engineToRequestKeys.Remove(engineKey)
-				traceLogger.Info("removed requestKey from index as no pods remain", "requestKey", requestKey, "engineKey", engineKey)
-			}
-			currentCache.mu.Unlock()
-		}
+	if podCache.cache.Len() == 0 {
+		m.data.Remove(requestKey)
+		m.engineToRequestKeys.Remove(engineKey)
+		traceLogger.Info("removed requestKey from index as no pods remain", "requestKey", requestKey, "engineKey", engineKey)
 	}
 
 	return nil
@@ -260,6 +256,9 @@ func (m *InMemoryIndex) Evict(ctx context.Context, engineKey BlockHash, entries
 // GetRequestKey returns the requestKey associated with the given engineKey.
 // Returns an error if the engineKey mapping is missing (e.g., already evicted).
 func (m *InMemoryIndex) GetRequestKey(ctx context.Context, engineKey BlockHash) (BlockHash, error) {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+
 	requestKey, found := m.engineToRequestKeys.Get(engineKey)
 	if !found {
 		return EmptyBlockHash, fmt.Errorf("engine key not found: %s", engineKey.String())