perf: Pool allocations of cachekv stores #24608
Conversation
Ironbird - launch a networkTo use Ironbird, you can use the following commands:
Custom Load Test ConfigurationYou can provide a custom load test configuration using the `--load-test-config=` flag:Use |
go.mod
Outdated
| @@ -204,3 +204,5 @@ retract ( | |||
| // do not use | |||
| v0.43.0 | |||
| ) | |||
|
|
|||
| replace cosmossdk.io/store => github.com/hyphacoop/cosmos-sdk/store v0.0.0-20250428150920-74852268c7e7 | |||
There was a problem hiding this comment.
I can clean this up once I know what the merge plan is, I just don't know what best practices are for this kind of change that affects multiple components
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
baseapp/baseapp.go
Outdated
| return ctx | ||
| } | ||
|
|
||
| type poolingStore interface { | ||
| storetypes.MultiStore | ||
| CacheMultiStorePooled() storetypes.PooledCacheMultiStore | ||
| } | ||
|
|
||
| // cacheTxContext returns a new context based off of the provided context with | ||
| // a branched multi-store. | ||
| func (app *BaseApp) cacheTxContext(ctx sdk.Context, txBytes []byte) (sdk.Context, storetypes.CacheMultiStore) { | ||
| ms := ctx.MultiStore() |
There was a problem hiding this comment.
Change potentially affects state.
Call sequence:
(*github.com/cosmos/cosmos-sdk/baseapp.BaseApp).getContextForTx (baseapp/baseapp.go:677)
(*github.com/cosmos/cosmos-sdk/baseapp.BaseApp).runTx (baseapp/baseapp.go:844)
(*github.com/cosmos/cosmos-sdk/baseapp.BaseApp).deliverTx (baseapp/baseapp.go:772)
(*github.com/cosmos/cosmos-sdk/baseapp.BaseApp).internalFinalizeBlock (baseapp/baseapp.go:705)
(*github.com/cosmos/cosmos-sdk/baseapp.BaseApp).FinalizeBlock (baseapp/baseapp.go:869)
📝 WalkthroughWalkthroughA set of changes introduces object pooling for cache multi-stores and cache key-value stores within the application's storage layer. New interfaces and types are defined to support pooled cache stores, including lifecycle management via Changes
Sequence Diagram(s)sequenceDiagram
participant BaseApp
participant MultiStore (poolingStore)
participant PooledCacheMultiStore
participant RegularCacheMultiStore
BaseApp->>MultiStore: cacheTxContext()
alt MultiStore implements poolingStore
MultiStore->>PooledCacheMultiStore: CacheMultiStorePooled()
BaseApp->>PooledCacheMultiStore: use for tx processing
BaseApp-->>PooledCacheMultiStore: defer Release()
else
MultiStore->>RegularCacheMultiStore: CacheMultiStore()
BaseApp->>RegularCacheMultiStore: use for tx processing
end
sequenceDiagram
participant PooledStorePool
participant PooledStore
participant ParentStore
BaseApp->>PooledStorePool: NewPooledStore(parent)
PooledStorePool->>PooledStore: retrieve or create
PooledStore->>ParentStore: set parent reference
BaseApp->>PooledStore: use for caching
BaseApp->>PooledStore: Release()
PooledStore->>PooledStore: reset caches, clear parent
PooledStore->>PooledStorePool: return to pool
✨ Finishing Touches
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (1)
baseapp/baseapp.go (1)
989-991: Same “defer Release” duplicationSee previous comment – extracting a helper would remove this repetition.
🧹 Nitpick comments (7)
store/cachekv/internal/btree.go (1)
39-41: Receiver should probably be a pointer to avoid silent struct-copy
Clear()takes the receiver by value, which means theBTreestruct (currently one word) is copied before invoking the underlyingbt.tree.Clear().
While harmless today, this has two downsides:
- If more fields are ever added to
BTree, the copy becomes more expensive.- A pointer receiver is the idiomatic choice for mutating methods and makes the intent explicit.
-func (bt BTree) Clear() { +func (bt *BTree) Clear() { bt.tree.Clear() }Call-sites would remain unchanged because the compiler automatically takes the address where needed (
store.sortedCache.Clear()).
Feel free to keep the value receiver if you want to guarantee immutability of the wrapper itself, but please leave a short comment documenting that decision.store/types/store.go (1)
155-158: Document the lifecycle contract ofRelease()The new
PooledCacheMultiStoreinterface introducesRelease(), but there is no doc-comment that explains:
- When callers must invoke it (after
Write()? Always?).- Whether it is idempotent.
- Whether the object is still usable afterwards (it currently is not).
Adding a short comment will avoid misuse and future bugs, e.g.:
// Release returns the instance to its pool. After calling Release the receiver // MUST NOT be used again. Calling Release multiple times is safe/no-op.store/cachekv/store.go (2)
35-37: Consider embedding*Storeinstead ofStoreEmbedding
Storeby value means everyPooledStorecontains its own copy of
theStorestruct. All pointer-receiver methods onStorestill work, but the
extra copy slightly increases the size of each pooled object (two cache maps,
B-tree, mutex, etc.) and forces Go to copy that memory when the pool hands out
an instance.A lighter alternative is:
-type PooledStore struct { Store } +type PooledStore struct { *Store }and constructing it with
&PooledStore{Store: &Store{ ... }}.
Not critical, but worth considering if the pool becomes large.
143-143: Clearing the B-tree may retain large backing arrays
btree.Clear()resets the length to zero but does not free node memory. For
workloads with occasional huge bursts (e.g. Genesis), this can keep tens of MB
alive inside the pool. Consider the same “size threshold” strategy used for
mapcaches: ifstore.sortedCache.Len() > bigN { store.sortedCache = internal.NewBTree() }.baseapp/baseapp.go (2)
706-709: Avoid duplicating public-facing interfaces
poolingStorere-declares behaviour that arguably belongs instore/types.
Long-term, keep the contract in a single place (e.g.storetypes.MultiStoreplus a build‐tagged extension) to prevent diverging expectations and import-cycle work-arounds.[nit]
935-939: DRY – factor the “defer Release” boilerplateThis pattern re-appears several times in
runTx. Consider a tiny helper:func deferIfPooled(cms storetypes.CacheMultiStore) { if p, ok := cms.(storetypes.PooledCacheMultiStore); ok { defer p.Release() } }Then call
deferIfPooled(msCache)to avoid accidental omissions when new call-sites are added.store/cachemulti/store.go (1)
200-202: Minor: method name leaks implementation detail
CacheMultiStorePooledis fine, but for symmetry withCacheMultiStoreconsider
CacheMultiStorePool()orCacheMultiStoreFromPool()to read as an action
instead of an adjective. Optional.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (2)
go.modis excluded by!**/*.modgo.sumis excluded by!**/*.sum,!**/*.sum
📒 Files selected for processing (5)
baseapp/baseapp.go(3 hunks)store/cachekv/internal/btree.go(1 hunks)store/cachekv/store.go(4 hunks)store/cachemulti/store.go(7 hunks)store/types/store.go(1 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (2)
baseapp/baseapp.go (1)
store/types/store.go (3)
MultiStore(115-147)PooledCacheMultiStore(155-158)CacheMultiStore(150-153)
store/cachemulti/store.go (3)
store/types/store.go (9)
Store(17-20)CacheMultiStore(150-153)PooledCacheMultiStore(155-158)KVStore(251-269)StoreKey(371-374)CacheWrapper(308-314)TraceContext(464-464)CacheWrap(297-306)MultiStore(115-147)store/cachekv/store.go (4)
Store(27-33)PooledStore(35-37)NewStore(42-49)NewPooledStore(69-73)store/rootmulti/store.go (2)
Store(59-82)NewStore(93-107)
⏰ Context from checks skipped due to timeout of 90000ms (14)
- GitHub Check: tests (03)
- GitHub Check: tests (01)
- GitHub Check: tests (00)
- GitHub Check: test-sim-nondeterminism
- GitHub Check: test-integration
- GitHub Check: test-e2e
- GitHub Check: Check up to date mocks
- GitHub Check: build (arm64)
- GitHub Check: Analyze
- GitHub Check: build (amd64)
- GitHub Check: Gosec
- GitHub Check: Summary
- GitHub Check: golangci-lint
- GitHub Check: dependency-review
🔇 Additional comments (3)
store/cachekv/store.go (2)
51-61:sync.Poolfactory leaks the innerStoreon grow-only pathsWhen
resetCachesdecides the cache grew past 100 000 entries it allocates
fresh maps, but the capacity of those maps remains captured in the pooled
object forever. Re-using that object therefore retains the high-water memory
usage, partly defeating the goal of lowering allocations.Two quick mitigations:
store.resetCaches() +// Reset heavy internal state that Clear() doesn’t shrink. +if cap(store.cache) > 200_000 { + store.cache = make(map[string]*cValue) +}or periodically create a brand-new
PooledStoreafter N uses.Would you like a micro-benchmark to quantify the retained capacity?
69-73: Ensure pooled instances start with empty caches
NewPooledStoreassumesresetCachesran duringRelease(). If a buggy
caller forgets toRelease, the nextGet()will return an object that still
contains stale state. You can detect this early:if len(store.cache) != 0 { panic("using a pooled Store that was not Released") }or, more leniently, call
store.resetCaches()defensively here.baseapp/baseapp.go (1)
715-720: Gracefully fall back when pooled implementation is missingThe branching logic works, but if a downstream
MultiStoreadvertisesCacheMultiStorePooledyet returnsnil(e.g. pool exhausted or disabled by feature flag) the code panics on the next dereference. A defensive check makes the change safer during rollout:- if msPooled, ok := ms.(poolingStore); ok { - msCache = msPooled.CacheMultiStorePooled() + if msPooled, ok := ms.(poolingStore); ok { + if pooled := msPooled.CacheMultiStorePooled(); pooled != nil { + msCache = pooled + } else { + msCache = ms.CacheMultiStore() + } } else { msCache = ms.CacheMultiStore() }
store/cachekv/store.go
Outdated
| func (store *PooledStore) Release() { | ||
| store.resetCaches() | ||
| store.parent = nil | ||
| storePool.Put(store) | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Reset the mutex to guarantee a clean state before pooling
If any caller accidentally leaves store.mtx in a locked state, Release
will silently pool an object whose mutex can never be locked again, leading to
a hard-to-reproduce deadlock the next time it is checked out.
Guard against this by zeroing the mutex:
func (store *PooledStore) Release() {
store.resetCaches()
+ store.mtx = sync.Mutex{}
store.parent = nil
storePool.Put(store)
}Cheap insurance, even if current call-sites appear correct.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| func (store *PooledStore) Release() { | |
| store.resetCaches() | |
| store.parent = nil | |
| storePool.Put(store) | |
| } | |
| func (store *PooledStore) Release() { | |
| store.resetCaches() | |
| store.mtx = sync.Mutex{} | |
| store.parent = nil | |
| storePool.Put(store) | |
| } |
🤖 Prompt for AI Agents (early access)
In store/cachekv/store.go around lines 63 to 67, the Release method should reset the store's mutex to a zero value before putting the store back into the pool. This prevents pooling a store with a locked mutex, which would cause deadlocks later. Add a line to assign a new zero-value mutex to store.mtx before calling storePool.Put(store).
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
fastfadingviolets
force-pushed
the
perf/pool-cachekv
branch
2 times, most recently
from
4db8b81 to
c4feb72
Compare
fastfadingviolets
force-pushed
the
perf/pool-cachekv
branch
3 times, most recently
from
98af2ef to
e6365dc
Compare
|
@fastfadingviolets Would it be possible to share metrics specifically related to memory and GC as well? I’m curious to see the quantitative improvements, if any. |
absolutely! Let me share some graphs here. I ran two half-hour load tests; you'll see two big humps in my graph. The first one is without the optimization, the second one's with it. First, rates of objects being allocated go down from ~1.4M to ~1.1M:
Allocation rates (i.e. same thing but in MB) goes from ~100MB/s to ~80MB/s:
Annoyingly, the GC histogram seems to show GCs getting longer with the fix, but I think that that's because we're doing fewer garbage collections, and so there's more to collect. This graph doesn't have the 99th percentile:
The 99th percentile goes up pretty significantly (it's likely this is also affected by startup operations, though):
However! If I look at the CPU profile, I can see we're spending less time in garbage collection overall, by about 8%, which would be consistent with the theory that we're just doing fewer GCs:
Lastly, just for the heck of it, we went from cacheTxContext being responsible for 12% of objects allocated to <0.1%:
This translates to about 10% fewer allocated objects in runTx:
Let me know if you'd like me to collect any more metrics! |
This means that hot code that provisions caches, such as runTx, can re-use already allocated memory-space for their cache, without having to suffer the hit of a new allocation.
This prevents a per-tx allocation of the cache.
fastfadingviolets
force-pushed
the
perf/pool-cachekv
branch
from
6fee430 to
a33a22f
Compare
| // a branched multi-store. | ||
| func (app *BaseApp) cacheTxContext(ctx sdk.Context, txBytes []byte) (sdk.Context, storetypes.CacheMultiStore) { | ||
| ms := ctx.MultiStore() | ||
| msCache := ms.CacheMultiStore() | ||
| var msCache storetypes.CacheMultiStore | ||
| if msPooled, ok := ms.(storetypes.PoolingMultiStore); ok { | ||
| msCache = msPooled.CacheMultiStorePooled() | ||
| } else { | ||
| msCache = ms.CacheMultiStore() |
There was a problem hiding this comment.
Change potentially affects state.
Call sequence:
(*github.com/cosmos/cosmos-sdk/baseapp.BaseApp).cacheTxContext (baseapp/baseapp.go:625)
(*github.com/cosmos/cosmos-sdk/baseapp.BaseApp).runTx (baseapp/baseapp.go:760)
(*github.com/cosmos/cosmos-sdk/baseapp.BaseApp).deliverTx (baseapp/baseapp.go:688)
(*github.com/cosmos/cosmos-sdk/baseapp.BaseApp).internalFinalizeBlock (baseapp/baseapp.go:718)
(*github.com/cosmos/cosmos-sdk/baseapp.BaseApp).FinalizeBlock (baseapp/baseapp.go:884)
|
|
||
| // NewFromKVStore creates a new Store object from a mapping of store keys to | ||
| // CacheWrapper objects and a KVStore as the database. Each CacheWrapper store | ||
| // is a branched store. | ||
| func NewFromKVStore( | ||
| store types.KVStore, stores map[types.StoreKey]types.CacheWrapper, | ||
| keys map[string]types.StoreKey, traceWriter io.Writer, traceContext types.TraceContext, | ||
| ) Store { | ||
| cms := Store{ | ||
| ) *Store { |
There was a problem hiding this comment.
why did we need to change this function sig?
There was a problem hiding this comment.
The issue is ultimately that sync.Pool works with pointers--if it didn't you'd end up copying the struct that it's trying to allocate for you--so I needed to be able to return a pointer to the PooledStore in newFromKVStorePooled. Once I'd made that change, interface compliance broke everywhere because the methods had value receivers. It ended up being simplest to just change everything to pointers for consistency.
|
@fastfadingviolets we will need a changelog in both the root CHANGELOG.md and the one in |
|
@aljo242 added changelog stuff & i think i've incorporated all your comments! I'm assuming this is gonna get squashed so I'm just pushing new commits, but lmk if you want me to squash them myself |
fastfadingviolets
force-pushed
the
perf/pool-cachekv
branch
from
abf893a to
092e009
Compare
| for k, v := range cms.stores { | ||
| if pStore, ok := v.(types.PooledCacheKVStore); ok { | ||
| pStore.Release() | ||
| } | ||
| delete(cms.stores, k) | ||
| } |
Check warning
Code scanning / CodeQL
Iteration over map Warning
| for k := range cms.keys { | ||
| delete(cms.keys, k) | ||
| } |
Check warning
Code scanning / CodeQL
Iteration over map Warning
|
adding @aaronc to take a look |
|
@fastfadingviolets we have some conflicts |
Description
Last week, on April 25, we ran a load test involving close to 40 validators and collected performance metrics from them. One thing that came up is that the cachekv that gets allocated to run each transaction is the top allocator:
My suggested perf improvement is to pool the cache objects themselves, as implemented in this PR. I tried a load test on a small set of nodes to see what this would do, and the result is a reduction in failures per second (i.e. "tx broadcasts that fail") from 96.73 in the baseline:
To 80.72 after this fix:
Note that I opened this against
main, but it includes both changes to the sdk itself and to store. Let me know if you'd like this refactored at all, or split into two.I'll throw the changelog together once I know that we're merging this, and how different bits are getting into different components.
Author Checklist
All items are required. Please add a note to the item if the item is not applicable and
please add links to any relevant follow up issues. Your PR will not be merged unless you satisfy
all of these items.
I have...
!in the type prefix if API or client breaking changeCHANGELOG.mdSummary by CodeRabbit
New Features
Performance Improvements