fix(storage): share aloneInBucket flag per bucket via atomic.Bool #1255
Conversation
…Factory The previous subquery-based approach in newScopedSelect executed a COUNT(*) on _system.ledgers for every scoped query. Replace it with a cached *atomic.Bool per bucket, shared across all stores through the DefaultFactory. When a new ledger is created, SetAloneInBucket(false) immediately propagates to every store in that bucket, ensuring the WHERE ledger=? predicate is never incorrectly skipped. Also adds CountLedgersInBucket to the system store interface so CreateLedger and OpenLedger can seed the flag correctly.
WalkthroughThe changes implement a per-bucket optimization flag for ledgers. When creating or opening a ledger, the code counts ledgers in the target bucket and sets an atomic flag if the ledger is the sole ledger in its bucket. This flag is shared across all stores in the same bucket via the factory, enabling efficient query optimization without repeated database counts. Changes
Sequence Diagram(s)sequenceDiagram
participant Driver
participant SystemStore
participant Factory
participant Store
Driver->>SystemStore: CountLedgersInBucket(bucket)
SystemStore-->>Driver: count
Driver->>Factory: Create(ctx, ledger)
Factory->>Factory: Get/Create atomic.Bool for bucket
Factory->>Store: newStore(...)
Store-->>Factory: Store instance
Factory->>Store: SetAloneInBucket(count == 1)
Factory-->>Driver: Store with flag set
Note over Store: Uses aloneInBucket flag<br/>in newScopedSelect for<br/>query optimization
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Comment |
![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
🤖 Fix all issues with AI agents
In `@internal/storage/driver/driver.go`:
- Around line 78-84: The CreateLedger path currently uses
systemStore.CountLedgersInBucket and may incorrectly set
ret.SetAloneInBucket(count == 1) under READ COMMITTED; change CreateLedger to
always call ret.SetAloneInBucket(false) (since creating a ledger cannot make it
alone) and remove/ignore the COUNT-based logic here; rely on OpenLedger to
compute and promote aloneInBucket=true when it can safely observe committed
state, or alternatively move the COUNT call to after the transaction commit if
you prefer the other approach.
🧹 Nitpick comments (1)
internal/storage/system/store.go (1)
73-82: Consider filtering out soft-deleted ledgers.
GetLedger(Line 198) andDeleteBucketapply adeleted_at IS NULLfilter, butCountLedgersInBucketdoes not. This means soft-deleted ledgers inflate the count, causing thealoneInBucketoptimization to be missed when a bucket contains one active ledger and one or more soft-deleted ledgers. The behavior is conservative (safe), but suboptimal.Suggested fix
func (d *DefaultStore) CountLedgersInBucket(ctx context.Context, bucket string) (int, error) { count, err := d.db.NewSelect(). Model(&ledger.Ledger{}). Where("bucket = ?", bucket). + Where("deleted_at IS NULL"). Count(ctx) if err != nil { return 0, postgres.ResolveError(err) } return count, nil }
| count, err := systemStore.CountLedgersInBucket(ctx, l.Bucket) | ||
| if err != nil { | ||
| return fmt.Errorf("counting ledgers in bucket: %w", err) | ||
| } | ||
|
|
||
| ret = d.ledgerStoreFactory.Create(b, *l) | ||
| ret.SetAloneInBucket(count == 1) |
There was a problem hiding this comment.
TOCTOU race: concurrent CreateLedger calls for the same bucket can leave aloneInBucket = true with 2 ledgers.
Under PostgreSQL READ COMMITTED (the default with nil tx options on Line 44), each transaction's COUNT only sees rows committed before that statement began plus its own inserts. When two CreateLedger calls for the same bucket overlap, both can observe count == 1 (each seeing only its own uncommitted row), leading both to call SetAloneInBucket(true). After both commit, the bucket has 2 ledgers but the shared flag is true, causing newScopedSelect to omit the WHERE ledger = ? predicate and potentially leaking cross-ledger data.
This is the exact class of bug described in the PR root-cause. The flag self-corrects on the next OpenLedger, but the window could be significant if stores are cached.
Simplest safe fix: always set aloneInBucket = false inside CreateLedger (since creating a ledger can only increase the count), and rely on OpenLedger to promote the flag to true when appropriate.
Suggested fix
- count, err := systemStore.CountLedgersInBucket(ctx, l.Bucket)
- if err != nil {
- return fmt.Errorf("counting ledgers in bucket: %w", err)
- }
-
ret = d.ledgerStoreFactory.Create(b, *l)
- ret.SetAloneInBucket(count == 1)
+ // Always conservative: a concurrent create may not yet be visible.
+ // OpenLedger will promote to true when the count is actually 1.
+ ret.SetAloneInBucket(false)
return nilAlternatively, move the count + set after the transaction commits so the read sees all committed rows:
Alternative fix — count after commit
if err != nil {
return nil, postgres.ResolveError(err)
}
+ count, err := d.systemStoreFactory.Create(d.db).CountLedgersInBucket(ctx, l.Bucket)
+ if err != nil {
+ return nil, fmt.Errorf("counting ledgers in bucket: %w", err)
+ }
+ ret.SetAloneInBucket(count == 1)
+
return ret, nil📝 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.
| count, err := systemStore.CountLedgersInBucket(ctx, l.Bucket) | |
| if err != nil { | |
| return fmt.Errorf("counting ledgers in bucket: %w", err) | |
| } | |
| ret = d.ledgerStoreFactory.Create(b, *l) | |
| ret.SetAloneInBucket(count == 1) | |
| ret = d.ledgerStoreFactory.Create(b, *l) | |
| // Always conservative: a concurrent create may not yet be visible. | |
| // OpenLedger will promote to true when the count is actually 1. | |
| ret.SetAloneInBucket(false) | |
| return nil |
🤖 Prompt for AI Agents
In `@internal/storage/driver/driver.go` around lines 78 - 84, The CreateLedger
path currently uses systemStore.CountLedgersInBucket and may incorrectly set
ret.SetAloneInBucket(count == 1) under READ COMMITTED; change CreateLedger to
always call ret.SetAloneInBucket(false) (since creating a ledger cannot make it
alone) and remove/ignore the COUNT-based logic here; rely on OpenLedger to
compute and promote aloneInBucket=true when it can safely observe committed
state, or alternatively move the COUNT call to after the transaction commit if
you prefer the other approach.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1255 +/- ##
==========================================
- Coverage 82.20% 82.16% -0.05%
==========================================
Files 198 198
Lines 10224 10243 +19
==========================================
+ Hits 8405 8416 +11
- Misses 1326 1330 +4
- Partials 493 497 +4 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
2 participants
Summary
newScopedSelectwith a shared*atomic.Boolper bucket, managed by theDefaultFactorySetAloneInBucket(false)now propagates immediately to all existing stores in that bucketCountLedgersInBucketto the system store interface soCreateLedgerandOpenLedgercan seed the flag correctlyRoot cause
The previous subquery-based approach (
SELECT count = 1 FROM _system.ledgers ...) innewScopedSelectadded overhead to every query. More critically, when stores were cached, a per-store boolean could miss bucket membership changes: existing stores kept believing they were alone and skipped theWHERE ledger = ?predicate, causing cross-ledger data leaks (e.g.TestAccountsListfailures).Changes
internal/storage/ledger/store.go:aloneInBucket *atomic.Boolfield +SetAloneInBucket/newScopedSelectusing atomic loadinternal/storage/ledger/factory.go:bucketFlags map[string]*atomic.Boolshared across all stores of the same bucketinternal/storage/driver/driver.go:CreateLedgerandOpenLedgercallCountLedgersInBucketandSetAloneInBucketinternal/storage/system/store.go: newCountLedgersInBucketmethod on interface + implementationinternal/storage/driver/store.go:CountLedgersInBucketadded to localSystemStoreinterfaceinternal/storage/driver/store_generated_test.go: mock updated for new methodTest plan
earthly -P +testspasses (all unit, integration, e2e, stress tests)TestAccountsListno longer flaky with multiple ledgers in the same bucket