add locking to cached Put

Author: DracoLi

x/blockdb/cache_db.go

@@ -5,18 +5,25 @@ package blockdb
 
 import (
 	"slices"
+	"sync"
 
 	"go.uber.org/zap"
 
 	"github.com/ava-labs/avalanchego/cache/lru"
 	"github.com/ava-labs/avalanchego/database"
 )
 
+// numShards is the number of mutex shards used to reduce lock contention for
+// concurrent Put operations. Using 256 shards provides a good balance between
+// memory usage (~2KB) and concurrency.
+const numShards = 256
+
 var _ database.HeightIndex = (*cacheDB)(nil)
 
 type cacheDB struct {
-	db    *Database
-	cache *lru.Cache[BlockHeight, BlockData]
+	db     *Database
+	cache  *lru.Cache[BlockHeight, BlockData]
+	shards [numShards]sync.Mutex
 }
 
 func newCacheDB(db *Database, size uint16) *cacheDB {
@@ -46,7 +53,18 @@ func (c *cacheDB) Get(height BlockHeight) (BlockData, error) {
 	return data, nil
 }
 
+// Put writes block data at the specified height to both the underlying database
+// and the cache.
+//
+// Concurrent calls to Put with the same height are serialized using sharded
+// locking to ensure cache consistency with the underlying database.
+// This allows concurrent writes to different heights while preventing race
+// conditions for writes to the same height.
 func (c *cacheDB) Put(height BlockHeight, data BlockData) error {
+	shard := &c.shards[height%numShards]
+	shard.Lock()
+	defer shard.Unlock()
+
 	if err := c.db.Put(height, data); err != nil {
 		return err
 	}

x/blockdb/cache_db_test.go

@@ -110,7 +110,6 @@ func TestCachePutOverridesSameHeight(t *testing.T) {
 	cached, ok = db.cache.Get(height)
 	require.True(t, ok)
 	require.Equal(t, b2, cached)
-	require.NotEqual(t, b1, cached)
 
 	// Get should also return the new block
 	data, err := db.Get(height)