Skip to content

Conversation

@flemzord
Copy link
Member

@flemzord flemzord commented Feb 6, 2026

Summary

  • Replace the per-query subquery approach in newScopedSelect with a shared *atomic.Bool per bucket, managed by the DefaultFactory
  • When a second ledger is created in the same bucket, SetAloneInBucket(false) now propagates immediately to all existing stores in that bucket
  • Add CountLedgersInBucket to the system store interface so CreateLedger and OpenLedger can seed the flag correctly

Root cause

The previous subquery-based approach (SELECT count = 1 FROM _system.ledgers ...) in newScopedSelect added 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 the WHERE ledger = ? predicate, causing cross-ledger data leaks (e.g. TestAccountsList failures).

Changes

  • internal/storage/ledger/store.go: aloneInBucket *atomic.Bool field + SetAloneInBucket / newScopedSelect using atomic load
  • internal/storage/ledger/factory.go: bucketFlags map[string]*atomic.Bool shared across all stores of the same bucket
  • internal/storage/driver/driver.go: CreateLedger and OpenLedger call CountLedgersInBucket and SetAloneInBucket
  • internal/storage/system/store.go: new CountLedgersInBucket method on interface + implementation
  • internal/storage/driver/store.go: CountLedgersInBucket added to local SystemStore interface
  • internal/storage/driver/store_generated_test.go: mock updated for new method

Test plan

  • earthly -P +tests passes (all unit, integration, e2e, stress tests)
  • TestAccountsList no longer flaky with multiple ledgers in the same bucket

…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.
@flemzord flemzord requested a review from a team as a code owner February 6, 2026 12:58
@coderabbitai
Copy link
Contributor

coderabbitai bot commented Feb 6, 2026

Walkthrough

This change introduces a per-bucket optimization flag to track whether a ledger is alone in its bucket. The implementation counts ledgers in each bucket during creation and opening, propagates the flag through driver and factory layers, and leverages it to simplify query optimization logic.

Changes

Cohort / File(s) Summary
Driver Behavioral Logic
internal/storage/driver/driver.go
Modified CreateLedger and OpenLedger to count ledgers in the ledger's bucket using a new SystemStore method, set the AloneInBucket flag based on count equality to 1, and propagate counting errors.
Driver Storage Interface
internal/storage/driver/store.go
Added CountLedgersInBucket(ctx, bucket) method to SystemStore interface for querying ledger counts per bucket.
Driver Storage Test Mocks
internal/storage/driver/store_generated_test.go
Added mock implementations of CountLedgersInBucket for both MockSystemStore and MockSystemStoreMockRecorder to enable testing.
Ledger Factory Synchronization
internal/storage/ledger/factory.go
Introduced mutex and per-bucket atomic flag map (bucketFlags) in DefaultFactory to track and manage thread-safe per-bucket state; Create method now lazily initializes and propagates per-bucket atomic flags to Store instances.
Ledger Store Query Optimization
internal/storage/ledger/store.go
Added aloneInBucket field (*atomic.Bool) to Store for per-bucket optimization flag; simplified newScopedSelect logic to conditionally omit ledger filter when alone-in-bucket is true; added SetAloneInBucket setter method.
System Store Implementation
internal/storage/system/store.go
Implemented CountLedgersInBucket method on DefaultStore to query Ledger records by bucket and return the count with error propagation via postgres.ResolveError.

Sequence Diagram

sequenceDiagram
    participant Client
    participant Driver as Driver Layer
    participant Factory as Ledger Factory
    participant SystemStore as System Store
    participant Database as Database
    
    rect rgba(100, 150, 200, 0.5)
    note over Client,Database: Ledger Creation/Opening Flow
    Client->>Driver: CreateLedger() or OpenLedger()
    Driver->>SystemStore: CountLedgersInBucket(bucket)
    SystemStore->>Database: Query Ledger count by bucket
    Database-->>SystemStore: Count result
    SystemStore-->>Driver: Return count
    
    Driver->>Factory: Create store & set flag
    Factory->>Factory: Acquire mutex
    Factory->>Factory: Get/create per-bucket atomic flag
    Factory->>Factory: Release mutex
    Factory->>Factory: Assign atomic flag to store.aloneInBucket
    Factory-->>Driver: Store with flag set
    
    Driver-->>Client: Return ledger store
    
    rect rgba(150, 200, 150, 0.5)
    note over Factory: Query Optimization Later
    Client->>Factory: Query operations on store
    Factory->>Factory: Check store.aloneInBucket
    alt Alone in bucket (true)
        Factory->>Database: Query without ledger filter
    else Not alone (false/nil)
        Factory->>Database: Query with ledger = ? filter
    end
    Database-->>Factory: Results
    Factory-->>Client: Return data
    end
    end
Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Poem

🐰 A flag per bucket, now we see—
Which ledgers hop so solo and free!
Optimization through atomic care,
Threads and queries dancing fair!

🚥 Pre-merge checks | ✅ 2 | ❌ 1
❌ Failed checks (1 warning)
Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (2 passed)
Check name Status Explanation
Title check ✅ Passed The title accurately summarizes the main change: replacing a per-query subquery approach with a shared atomic.Bool per bucket to fix data leakage issues.
Description check ✅ Passed The description is directly related to the changeset, providing clear context on the root cause, changes, and test plan for the atomic.Bool bucket flag optimization.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing touches
  • 📝 Generate docstrings
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Post copyable unit tests in a comment
  • Commit unit tests in branch fix/v2.3/shared-alone-in-bucket-flag

Comment @coderabbitai help to get the list of available commands and usage tips.

Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🤖 Fix all issues with AI agents
In `@internal/storage/driver/driver.go`:
- Around line 93-112: OpenLedger currently reads the ledger count then calls
store.SetAloneInBucket(count == 1), which can race with concurrent CreateLedger
and mistakenly set the shared alone flag back to true; change OpenLedger to
never set the flag to true — only clear it when you observe count > 1 (i.e.,
call SetAloneInBucket(false) when count > 1) or rely on CreateLedger/factory to
initialize true; alternatively make SetAloneInBucket one-directional (only allow
true→false, never false→true) so OpenLedger cannot revert a concurrent false to
true; update references in OpenLedger and ensure CreateLedger remains the only
path that can set true.

Comment on lines 93 to 112
func (d *Driver) OpenLedger(ctx context.Context, name string) (*ledgerstore.Store, *ledger.Ledger, error) {
// todo: keep the ledger in cache somewhere to avoid read the ledger at each request, maybe in the factory
ret, err := d.systemStoreFactory.Create(d.db).GetLedger(ctx, name)
// NOTE: the aloneInBucket flag is now shared per bucket via the Factory,
// so all stores in the same bucket see updates immediately.
systemStore := d.systemStoreFactory.Create(d.db)

ret, err := systemStore.GetLedger(ctx, name)
if err != nil {
return nil, nil, err
}

count, err := systemStore.CountLedgersInBucket(ctx, ret.Bucket)
if err != nil {
return nil, nil, fmt.Errorf("counting ledgers in bucket: %w", err)
}

store := d.ledgerStoreFactory.Create(d.bucketFactory.Create(ret.Bucket), *ret)
store.SetAloneInBucket(count == 1)

return store, ret, err
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

TOCTOU race: OpenLedger can overwrite a concurrent CreateLedger's flag update.

OpenLedger reads the count outside a transaction and then calls SetAloneInBucket(count == 1). If a concurrent CreateLedger commits between the count read and the Store call, the sequence can be:

  1. OpenLedger reads count=1 (2nd ledger not yet committed)
  2. CreateLedger commits → sets shared flag to false
  3. OpenLedger calls SetAloneInBucket(true) → overwrites back to true

All stores in that bucket now incorrectly skip WHERE ledger = ?, potentially returning data from the wrong ledger.

Consider making the flag transition one-directional: SetAloneInBucket should only allow true → false, never false → true, or use CompareAndSwap. The false → true transition should only happen in a controlled path (e.g., ledger deletion, if applicable), not in OpenLedger.

Proposed fix in store.go
 func (store *Store) SetAloneInBucket(alone bool) {
 	if store.aloneInBucket != nil {
-		store.aloneInBucket.Store(alone)
+		if alone {
+			// Only upgrade from false→true is unsafe from OpenLedger;
+			// use CAS to avoid overwriting a concurrent false.
+			store.aloneInBucket.CompareAndSwap(false, true)
+		} else {
+			store.aloneInBucket.Store(false)
+		}
 	}
 }

Actually, even the CAS above doesn't fully resolve this—it still races. A safer approach: OpenLedger should never set the flag to true; it should only set it to false when count > 1. Leave the true initialization to CreateLedger (which runs transactionally) or to the factory's lazy init.

Alternative: only propagate false from OpenLedger
 	store := d.ledgerStoreFactory.Create(d.bucketFactory.Create(ret.Bucket), *ret)
-	store.SetAloneInBucket(count == 1)
+	if count > 1 {
+		store.SetAloneInBucket(false)
+	}
 
 	return store, ret, err
📝 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.

Suggested change
func (d *Driver) OpenLedger(ctx context.Context, name string) (*ledgerstore.Store, *ledger.Ledger, error) {
// todo: keep the ledger in cache somewhere to avoid read the ledger at each request, maybe in the factory
ret, err := d.systemStoreFactory.Create(d.db).GetLedger(ctx, name)
// NOTE: the aloneInBucket flag is now shared per bucket via the Factory,
// so all stores in the same bucket see updates immediately.
systemStore := d.systemStoreFactory.Create(d.db)
ret, err := systemStore.GetLedger(ctx, name)
if err != nil {
return nil, nil, err
}
count, err := systemStore.CountLedgersInBucket(ctx, ret.Bucket)
if err != nil {
return nil, nil, fmt.Errorf("counting ledgers in bucket: %w", err)
}
store := d.ledgerStoreFactory.Create(d.bucketFactory.Create(ret.Bucket), *ret)
store.SetAloneInBucket(count == 1)
return store, ret, err
func (d *Driver) OpenLedger(ctx context.Context, name string) (*ledgerstore.Store, *ledger.Ledger, error) {
// todo: keep the ledger in cache somewhere to avoid read the ledger at each request, maybe in the factory
// NOTE: the aloneInBucket flag is now shared per bucket via the Factory,
// so all stores in the same bucket see updates immediately.
systemStore := d.systemStoreFactory.Create(d.db)
ret, err := systemStore.GetLedger(ctx, name)
if err != nil {
return nil, nil, err
}
count, err := systemStore.CountLedgersInBucket(ctx, ret.Bucket)
if err != nil {
return nil, nil, fmt.Errorf("counting ledgers in bucket: %w", err)
}
store := d.ledgerStoreFactory.Create(d.bucketFactory.Create(ret.Bucket), *ret)
if count > 1 {
store.SetAloneInBucket(false)
}
return store, ret, err
}
🤖 Prompt for AI Agents
In `@internal/storage/driver/driver.go` around lines 93 - 112, OpenLedger
currently reads the ledger count then calls store.SetAloneInBucket(count == 1),
which can race with concurrent CreateLedger and mistakenly set the shared alone
flag back to true; change OpenLedger to never set the flag to true — only clear
it when you observe count > 1 (i.e., call SetAloneInBucket(false) when count >
1) or rely on CreateLedger/factory to initialize true; alternatively make
SetAloneInBucket one-directional (only allow true→false, never false→true) so
OpenLedger cannot revert a concurrent false to true; update references in
OpenLedger and ensure CreateLedger remains the only path that can set true.

@codecov
Copy link

codecov bot commented Feb 6, 2026

Codecov Report

❌ Patch coverage is 83.33333% with 6 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.69%. Comparing base (42bbfb5) to head (0b8547f).
⚠️ Report is 1 commits behind head on release/v2.3.

Files with missing lines Patch % Lines
internal/storage/driver/driver.go 60.00% 2 Missing and 2 partials ⚠️
internal/storage/system/store.go 75.00% 1 Missing and 1 partial ⚠️
Additional details and impacted files
@@               Coverage Diff                @@
##           release/v2.3    #1254      +/-   ##
================================================
- Coverage         81.78%   81.69%   -0.10%     
================================================
  Files               187      187              
  Lines              9059     9078      +19     
================================================
+ Hits               7409     7416       +7     
- Misses             1215     1222       +7     
- Partials            435      440       +5     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@flemzord flemzord merged commit 526163a into release/v2.3 Feb 6, 2026
12 of 14 checks passed
@flemzord flemzord deleted the fix/v2.3/shared-alone-in-bucket-flag branch February 6, 2026 13:42
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants