Feat/shard memory storage - #4401
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
💤 Files with no reviewable changes (2)
WalkthroughReplaces single global mutex maps with fixed-size sharded maps for internal memory storage and the idempotency locker; operations are routed by an FNV-1a-style hash to per-shard maps protected by shard-local locks, and GC/connection snapshots operate per-shard. ChangesSharded concurrency pattern for storage and idempotency
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces sharding to both the in-memory storage and the idempotency locker to reduce lock contention and improve concurrent performance. While the sharding implementation is a solid improvement, the reviewer identifies several opportunities to further optimize performance and reduce heap allocations. Specifically, the reviewer suggests implementing an inline, allocation-free FNV-1a hash function to avoid allocations from fnv.New32a() and string-to-byte conversions. Additionally, the reviewer recommends refactoring the Reset and Keys methods to process shards sequentially rather than spawning concurrent goroutines, which eliminates unnecessary scheduling overhead and slice allocations.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| func getHash(key string) uint32 { | ||
| h := fnv.New32a() | ||
| h.Write([]byte(key)) | ||
| return h.Sum32() | ||
| } |
There was a problem hiding this comment.
The current getHash implementation uses fnv.New32a() which allocates a new hash.Hash32 interface/struct on the heap on every call, and converting the string to a byte slice []byte(key) can also cause allocations. Since this is a high-throughput memory storage, we can implement an inline, allocation-free FNV-1a hash function directly on the string to completely eliminate heap allocations.
| func getHash(key string) uint32 { | |
| h := fnv.New32a() | |
| h.Write([]byte(key)) | |
| return h.Sum32() | |
| } | |
| func getHash(key string) uint32 { | |
| const ( | |
| offset32 = 2166136261 | |
| prime32 = 16777619 | |
| ) | |
| hash := uint32(offset32) | |
| for i := 0; i < len(key); i++ { | |
| hash ^= uint32(key[i]) | |
| hash *= prime32 | |
| } | |
| return hash | |
| } |
| func (l *MemoryLock) getShard(key string) *lockerShard { | ||
| h := fnv.New32a() | ||
| h.Write([]byte(key)) | ||
| return l.shards[h.Sum32()%numShards] | ||
| } |
There was a problem hiding this comment.
The current getShard implementation uses fnv.New32a() which allocates a new hash.Hash32 interface/struct on the heap on every call, and converting the string to a byte slice []byte(key) can also cause allocations. Since this is a high-throughput memory storage, we can implement an inline, allocation-free FNV-1a hash function directly on the string to completely eliminate heap allocations.
| func (l *MemoryLock) getShard(key string) *lockerShard { | |
| h := fnv.New32a() | |
| h.Write([]byte(key)) | |
| return l.shards[h.Sum32()%numShards] | |
| } | |
| func (l *MemoryLock) getShard(key string) *lockerShard { | |
| const ( | |
| offset32 = 2166136261 | |
| prime32 = 16777619 | |
| ) | |
| hash := uint32(offset32) | |
| for i := 0; i < len(key); i++ { | |
| hash ^= uint32(key[i]) | |
| hash *= prime32 | |
| } | |
| return l.shards[hash%numShards] | |
| } |
| func (s *Storage) Reset() error { | ||
| ndb := make(map[string]Entry) | ||
| s.mux.Lock() | ||
| s.db = ndb | ||
| s.mux.Unlock() | ||
| wg := &sync.WaitGroup{} | ||
|
|
||
| for _, shard := range s.shards { | ||
| wg.Add(1) | ||
| go func(shrd *Shard) { | ||
| defer wg.Done() | ||
|
|
||
| shrd.mux.Lock() | ||
| shrd.db = make(map[string]Entry) | ||
| shrd.mux.Unlock() | ||
| }(shard) | ||
| } | ||
|
|
||
| wg.Wait() | ||
| return nil | ||
| } |
There was a problem hiding this comment.
Spawning 32 goroutines and using a sync.WaitGroup to reset 32 shards introduces unnecessary scheduling overhead and allocations. Since numShards is small (32), resetting them sequentially in a simple loop is much faster, cleaner, and avoids spawning goroutines entirely.
func (s *Storage) Reset() error {
for _, shard := range s.shards {
shard.mux.Lock()
shard.db = make(map[string]Entry)
shard.mux.Unlock()
}
return nil
}| func (s *Storage) Keys() ([][]byte, error) { | ||
| s.mux.RLock() | ||
| defer s.mux.RUnlock() | ||
| wg := &sync.WaitGroup{} | ||
| var keysLen = 0 | ||
| for _, shard := range s.shards { | ||
| shard.mux.RLock() | ||
| keysLen += len(shard.db) | ||
| shard.mux.RUnlock() | ||
| } | ||
|
|
||
| if len(s.db) == 0 { | ||
| // check if no valid keys were found | ||
| if keysLen == 0 { | ||
| return nil, nil | ||
| } | ||
|
|
||
| localKeys := make([][][]byte, numShards) | ||
| ts := utils.Timestamp() | ||
| keys := make([][]byte, 0, len(s.db)) | ||
| for key, v := range s.db { | ||
| // Filter out the expired keys | ||
| if v.expiry == 0 || v.expiry > ts { | ||
| keys = append(keys, []byte(key)) | ||
| } | ||
| for i, shard := range s.shards { | ||
| wg.Add(1) | ||
| go func(idx int, shrd *Shard) { | ||
| defer wg.Done() | ||
| shrd.mux.RLock() | ||
| defer shrd.mux.RUnlock() | ||
| for key, v := range shrd.db { | ||
| // Filter out the expired keys | ||
| if v.expiry == 0 || v.expiry > ts { | ||
| localKeys[idx] = append(localKeys[idx], []byte(key)) | ||
| } | ||
| } | ||
| }(i, shard) | ||
|
|
||
| } | ||
|
|
||
| wg.Wait() | ||
|
|
||
| keys := make([][]byte, 0, keysLen) | ||
| for _, shardKeys := range localKeys { | ||
| keys = append(keys, shardKeys...) | ||
| } | ||
| // Double check if no valid keys were found | ||
| if len(keys) == 0 { | ||
| return nil, nil | ||
| } | ||
|
|
||
| return keys, nil | ||
| } |
There was a problem hiding this comment.
Spawning 32 goroutines in Keys() to collect keys concurrently introduces significant overhead and multiple slice allocations (a slice of slices localKeys, plus individual slices per shard).
By refactoring this to a sequential loop, we can pre-allocate a single keys slice with the exact capacity needed (keysLen) and append directly to it. This completely eliminates goroutine scheduling overhead and reduces heap allocations to a single slice allocation.
func (s *Storage) Keys() ([][]byte, error) {
var keysLen = 0
for _, shard := range s.shards {
shard.mux.RLock()
keysLen += len(shard.db)
shard.mux.RUnlock()
}
if keysLen == 0 {
return nil, nil
}
keys := make([][]byte, 0, keysLen)
ts := utils.Timestamp()
for _, shard := range s.shards {
shard.mux.RLock()
for key, v := range shard.db {
if v.expiry == 0 || v.expiry > ts {
keys = append(keys, []byte(key))
}
}
shard.mux.RUnlock()
}
if len(keys) == 0 {
return nil, nil
}
return keys, nil
}![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@internal/storage/memory/memory.go`:
- Around line 235-254: The Conn() comment is stale: after the sharding refactor
Conn() builds and returns a fresh mergedMaps snapshot instead of a live shared
map; update the docstring for Storage.Conn to state it returns a point-in-time
copy (a new map built from s.shards/shard.db) and warn callers that the returned
map is not kept synchronized with storage and must not be used for concurrent
mutations or to observe live updates.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 1efbd556-7e84-4ad6-8bce-da2f8d8cf894
📒 Files selected for processing (2)
internal/storage/memory/memory.gomiddleware/idempotency/locker.go
| // Conn returns the underlying storage map. The returned map remains shared with | ||
| // the storage, so callers must not modify it and must synchronize any access | ||
| // that overlaps with other storage operations. | ||
| func (s *Storage) Conn() map[string]Entry { | ||
| s.mux.RLock() | ||
| defer s.mux.RUnlock() | ||
| return s.db | ||
| var allocatedMapLen = 0 | ||
| for _, shard := range s.shards { | ||
| shard.mux.RLock() | ||
| allocatedMapLen += len(shard.db) | ||
| shard.mux.RUnlock() | ||
| } | ||
| mergedMaps := make(map[string]Entry, allocatedMapLen) | ||
| for _, shard := range s.shards { | ||
| shard.mux.RLock() | ||
| for k, v := range shard.db { | ||
| mergedMaps[k] = v | ||
| } | ||
| shard.mux.RUnlock() | ||
| } | ||
| return mergedMaps | ||
| } |
There was a problem hiding this comment.
Documentation is outdated after sharding refactor.
The comment states the returned map "remains shared with the storage," but the implementation now creates a fresh merged snapshot each call. Callers expecting live updates via the returned map would be surprised.
If the snapshot behavior is intentional, update the comment to reflect that this returns a point-in-time copy.
📝 Suggested comment update
-// Conn returns the underlying storage map. The returned map remains shared with
-// the storage, so callers must not modify it and must synchronize any access
-// that overlaps with other storage operations.
+// Conn returns a merged snapshot of all shard maps. The returned map is a copy;
+// modifications to it do not affect the storage. Note that the snapshot is not
+// atomic across shards—concurrent writes may result in a partially stale view.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@internal/storage/memory/memory.go` around lines 235 - 254, The Conn() comment
is stale: after the sharding refactor Conn() builds and returns a fresh
mergedMaps snapshot instead of a live shared map; update the docstring for
Storage.Conn to state it returns a point-in-time copy (a new map built from
s.shards/shard.db) and warn callers that the returned map is not kept
synchronized with storage and must not be used for concurrent mutations or to
observe live updates.
|
Thanks @dima11223432 for picking this up, the implementation work and the careful per-shard refactoring across both files is real engineering effort and I appreciate the time you put into it. I want to be upfront about why we are not going to merge this in the current form: after looking at the trade-off and at the underlying issue (#4361), the sharding pattern is not the right fit for these two components in this codebase. The short version:
Separate from the underlying question, several things on the change itself would need work before it could land:
If at some point we see profile data from a real workload showing the storage or locker mutex as the bottleneck, we will reopen this discussion with that evidence in hand. Until then, the maintenance and correctness cost of carrying this design exceeds the benefit for the workloads the in-process implementation targets. Closing #4361 alongside this PR with the same reasoning. Thanks again for the effort, and sorry for the negative outcome on a substantial chunk of work. |
|
Thanks for your feedback, Rene! It is incredibly important and educational for me. Just to clarify and lock down the learning experience: so the underlying issue basically described a theoretical problem rather than a real-world bottleneck, meaning the lock contention it anticipated doesn't actually impact performance in realistic workloads? I completely understand the reasoning behind the trade-offs regarding simplicity and hot-key patterns now. Thank you for taking the time to write such a detailed design review! |
|
That's right; unfortunately, I didn't realize that until I looked at the resulting code and did some research. |
2 participants
Description
This PR eliminates global mutex contention under high concurrent loads (>10K QPS) by implementing the Sharded Map pattern across two internal components:
middleware/idempotency/locker.go): Replaced the single globalsync.Mutexwith a 32-shard structure. Implemented a robust double-checked locking mechanism to safely handle thread creation, deletion, and reference counting without concurrent map writes or deadlocks.internal/storage/memory/memory.go): Converted the singlesync.RWMutexprotecting the entire database map into a 32-shard segmented array ([]*Shard). Each shard manages its own lock, isolated map, and garbage collection sequence.Benefits
Set/Delete) in one shard no longer block read operations (Get) in other shards.Fixes #4361
Changes introduced
Type of change
Checklist