fix(bigtable/accelerator): collapse cold-start openHandle bursts via singleflight - #20380
Open
sushanb wants to merge 2 commits into
Open
fix(bigtable/accelerator): collapse cold-start openHandle bursts via singleflight#20380sushanb wants to merge 2 commits into
sushanb wants to merge 2 commits into
Conversation
…singleflight + sync.Map fast path Under high QPS, N concurrent first-touch RPCs for the same resource all called sessionTables.GetOrOpen concurrently, N-1 lost the insertion race, and each loser's Close() ran sessionTable's release path. Against a shared-key session pool that release tore down the pool the winner already memoized and served ErrPoolClosed on every subsequent RPC through that handle. Top-level client.OpenTable never exposed this path because it warms the cache exactly once at construction, so the concurrent-miss race the tests in bigtable/internal/session/pr20366_repro_test.go pin was unreachable from that caller shape — but reachable from the accelerator's per-RPC openHandle. Two additions to *Channel: - handleFastPath sync.Map: pre-lookup for the hot path. Skips the closure allocation, res.Kind switch, resource-name parse, and sessionTables.mu acquire on every cache-hit RPC. Staleness is safe because TableHandle.ReadRow / MutateRow self-heal via dispatch() when the underlying cache entry is TTL-evicted, so a pointer left behind here still routes to the live successor on next use. - openGroup singleflight.Group: collapses the cold-start burst. All N concurrent missers for a resource funnel through one openFn call; the waiters block on singleflight's internal chan and receive the same handle. The loser branch inside sessionTables.GetOrOpen fires zero times from this call site, regardless of the session-layer state. Result: sc.OpenTable is now called exactly once per resource per Channel lifetime, matching the client.OpenTable(t)-then-share-the-shim pattern that never triggered the bug in production. Includes accelerator_singleflight_test.go: fans out 32 concurrent openHandle callers behind a barrier and asserts sc.OpenTable invocation count == 1. Without the fix the same test observes 32 invocations.
product-auto-label
Bot
added
the
api: bigtable
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces a fast-path cache (sync.Map) and a singleflight group to the Channel struct to optimize resource handle retrieval and collapse concurrent cold-start bursts, preventing session-layer race conditions. The review feedback suggests introducing an atomic closed boolean flag to the Channel struct to safely track its lifecycle, ensuring that concurrent or subsequent handle lookups do not store stale handles after Close() has been called.
Comment on lines
+30
to
+33
| "context" | ||
| "io" | ||
| "strings" | ||
| "sync" |
Contributor
There was a problem hiding this comment.
To support tracking the closed state of the channel safely across concurrent goroutines, we need to import the sync/atomic package.
Suggested change
| "context" | |
| "io" | |
| "strings" | |
| "sync" | |
| "context" | |
| "io" | |
| "strings" | |
| "sync" | |
| "sync/atomic" |
Comment on lines
+166
to
167
| openGroup singleflight.Group | ||
| } |
Contributor
There was a problem hiding this comment.
Add a closed atomic boolean field to the Channel struct to safely track whether the channel has been closed. This helps prevent race conditions where stale handles are stored in handleFastPath after Close() has run.
Suggested change
| openGroup singleflight.Group | |
| } | |
| openGroup singleflight.Group | |
| closed atomic.Bool | |
| } |
References
- In Go projects using Go 1.19 or later, prefer using typed atomics (such as
atomic.Int64,atomic.Uint64,atomic.Bool) instead of raw types with package-levelsync/atomicfunctions.
Comment on lines
213
to
275
| func (c *Channel) openHandle(res adapters.Resource) (session.TableAPI, error) { | ||
| var open func() session.TableAPI | ||
| switch res.Kind { | ||
| case adapters.ResourceTable: | ||
| tableID, err := c.parseTableName(res.Name) | ||
| if err != nil { | ||
| return nil, err | ||
| // Fast path — a bare atomic Load. No mutex, no closure allocation, no | ||
| // resource-name parse. Under steady state every RPC after the first for a | ||
| // given resource hits here. | ||
| if v, ok := c.handleFastPath.Load(res.Name); ok { | ||
| return v.(session.TableAPI), nil | ||
| } | ||
|
|
||
| // Slow path — collapsed via singleflight so N concurrent first-touch RPCs | ||
| // for the same resource funnel through ONE sessionTables.GetOrOpen call. | ||
| // See the openGroup field comment for the failure mode this guards against. | ||
| v, err, _ := c.openGroup.Do(res.Name, func() (any, error) { | ||
| // Second cache check inside the singleflight — a prior burst's winner | ||
| // may have published between our fast-path miss and our arrival here. | ||
| if v, ok := c.handleFastPath.Load(res.Name); ok { | ||
| return v, nil | ||
| } | ||
| open = func() session.TableAPI { return c.sc.OpenTable(tableID) } | ||
| case adapters.ResourceAuthorizedView: | ||
| tableID, viewID, err := c.parseAuthorizedViewName(res.Name) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| open = func() session.TableAPI { return c.sc.OpenAuthorizedView(tableID, viewID) } | ||
| case adapters.ResourceMaterializedView: | ||
| viewID, err := c.parseMaterializedViewName(res.Name) | ||
| if err != nil { | ||
| return nil, err | ||
|
|
||
| var open func() session.TableAPI | ||
| switch res.Kind { | ||
| case adapters.ResourceTable: | ||
| tableID, err := c.parseTableName(res.Name) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| open = func() session.TableAPI { return c.sc.OpenTable(tableID) } | ||
| case adapters.ResourceAuthorizedView: | ||
| tableID, viewID, err := c.parseAuthorizedViewName(res.Name) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| open = func() session.TableAPI { return c.sc.OpenAuthorizedView(tableID, viewID) } | ||
| case adapters.ResourceMaterializedView: | ||
| viewID, err := c.parseMaterializedViewName(res.Name) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| open = func() session.TableAPI { return c.sc.OpenMaterializedView(viewID) } | ||
| default: | ||
| // Kind is set by the adapter from the populated V2 name field, so an | ||
| // unrecognized kind is an internal invariant violation, not bad input. | ||
| return nil, status.Errorf(codes.Internal, "accelerator: unknown resource kind %v", res.Kind) | ||
| } | ||
| open = func() session.TableAPI { return c.sc.OpenMaterializedView(viewID) } | ||
| default: | ||
| // Kind is set by the adapter from the populated V2 name field, so an | ||
| // unrecognized kind is an internal invariant violation, not bad input. | ||
| return nil, status.Errorf(codes.Internal, "accelerator: unknown resource kind %v", res.Kind) | ||
| } | ||
|
|
||
| // Key on the full V2 name (the identity Cloud Bigtable uses on the wire), | ||
| // so table / authorized-view / materialized-view keys never collide. A nil | ||
| // handle means the cache has been Closed (Channel.Close), so the channel is | ||
| // no longer usable. | ||
| tbl := c.sessionTables.GetOrOpen(res.Name, open) | ||
| if tbl == nil { | ||
| return nil, status.Error(codes.Unavailable, "accelerator: channel is closed") | ||
| // Key on the full V2 name (the identity Cloud Bigtable uses on the wire), | ||
| // so table / authorized-view / materialized-view keys never collide. A nil | ||
| // handle means the cache has been Closed (Channel.Close), so the channel is | ||
| // no longer usable. | ||
| tbl := c.sessionTables.GetOrOpen(res.Name, open) | ||
| if tbl == nil { | ||
| return nil, status.Error(codes.Unavailable, "accelerator: channel is closed") | ||
| } | ||
| // Publish for future fast-path lookups. Only the singleflight winner | ||
| // reaches this point per resource per cold-start burst, so LoadOrStore | ||
| // here is defensive — a plain Store would also be correct. | ||
| actual, _ := c.handleFastPath.LoadOrStore(res.Name, tbl) | ||
| return actual, nil | ||
| }) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| return tbl, nil | ||
| return v.(session.TableAPI), nil | ||
| } |
Contributor
There was a problem hiding this comment.
Check the closed flag at the beginning of openHandle and right before storing the table handle in handleFastPath. This prevents a race condition where a concurrent Close() clears the map, but openHandle subsequently overwrites it with a stale handle.
func (c *Channel) openHandle(res adapters.Resource) (session.TableAPI, error) {
if c.closed.Load() {
return nil, status.Error(codes.Unavailable, "accelerator: channel is closed")
}
// Fast path — a bare atomic Load. No mutex, no closure allocation, no
// resource-name parse. Under steady state every RPC after the first for a
// given resource hits here.
if v, ok := c.handleFastPath.Load(res.Name); ok {
return v.(session.TableAPI), nil
}
// Slow path — collapsed via singleflight so N concurrent first-touch RPCs
// for the same resource funnel through ONE sessionTables.GetOrOpen call.
// See the openGroup field comment for the failure mode this guards against.
v, err, _ := c.openGroup.Do(res.Name, func() (any, error) {
// Second cache check inside the singleflight — a prior burst's winner
// may have published between our fast-path miss and our arrival here.
if v, ok := c.handleFastPath.Load(res.Name); ok {
return v, nil
}
var open func() session.TableAPI
switch res.Kind {
case adapters.ResourceTable:
tableID, err := c.parseTableName(res.Name)
if err != nil {
return nil, err
}
open = func() session.TableAPI { return c.sc.OpenTable(tableID) }
case adapters.ResourceAuthorizedView:
tableID, viewID, err := c.parseAuthorizedViewName(res.Name)
if err != nil {
return nil, err
}
open = func() session.TableAPI { return c.sc.OpenAuthorizedView(tableID, viewID) }
case adapters.ResourceMaterializedView:
viewID, err := c.parseMaterializedViewName(res.Name)
if err != nil {
return nil, err
}
open = func() session.TableAPI { return c.sc.OpenMaterializedView(viewID) }
default:
// Kind is set by the adapter from the populated V2 name field, so an
// unrecognized kind is an internal invariant violation, not bad input.
return nil, status.Errorf(codes.Internal, "accelerator: unknown resource kind %v", res.Kind)
}
// Key on the full V2 name (the identity Cloud Bigtable uses on the wire),
// so table / authorized-view / materialized-view keys never collide. A nil
// handle means the cache has been Closed (Channel.Close), so the channel is
// no longer usable.
tbl := c.sessionTables.GetOrOpen(res.Name, open)
if tbl == nil {
return nil, status.Error(codes.Unavailable, "accelerator: channel is closed")
}
if c.closed.Load() {
return nil, status.Error(codes.Unavailable, "accelerator: channel is closed")
}
// Publish for future fast-path lookups. Only the singleflight winner
// reaches this point per resource per cold-start burst, so LoadOrStore
// here is defensive — a plain Store would also be correct.
actual, _ := c.handleFastPath.LoadOrStore(res.Name, tbl)
return actual, nil
})
if err != nil {
return nil, err
}
return v.(session.TableAPI), nil
}References
- When implementing lazy initialization of resources alongside a Close or teardown method, use an atomic boolean flag and a guard wrapper around the lazy openers. The guard should check the flag before opening, and re-check it after opening to immediately release the resource if Close was called concurrently, preventing resource leaks.
Comment on lines
356
to
371
| if c.sc == nil { | ||
| return nil | ||
| } | ||
| // Drop fast-path references before sessionTables.Close so a concurrent | ||
| // openHandle racing Close either sees the empty map and takes the slow | ||
| // path (which observes the closed cache and returns Unavailable) or has | ||
| // already snapshotted a stale handle whose next RPC surfaces the | ||
| // closed-pool error. Either outcome is honest; a lingering map entry | ||
| // after Close would just keep the doomed handle reachable from Channel. | ||
| c.handleFastPath.Range(func(k, _ any) bool { | ||
| c.handleFastPath.Delete(k) | ||
| return true | ||
| }) | ||
| c.sessionTables.Close() | ||
| return c.sc.Close() | ||
| } |
Contributor
There was a problem hiding this comment.
Set the closed flag to true at the beginning of Close(). This ensures that any concurrent or subsequent openHandle calls will immediately observe the closed state and avoid storing stale handles in handleFastPath.
Suggested change
| if c.sc == nil { | |
| return nil | |
| } | |
| // Drop fast-path references before sessionTables.Close so a concurrent | |
| // openHandle racing Close either sees the empty map and takes the slow | |
| // path (which observes the closed cache and returns Unavailable) or has | |
| // already snapshotted a stale handle whose next RPC surfaces the | |
| // closed-pool error. Either outcome is honest; a lingering map entry | |
| // after Close would just keep the doomed handle reachable from Channel. | |
| c.handleFastPath.Range(func(k, _ any) bool { | |
| c.handleFastPath.Delete(k) | |
| return true | |
| }) | |
| c.sessionTables.Close() | |
| return c.sc.Close() | |
| } | |
| if c.sc == nil { | |
| return nil | |
| } | |
| c.closed.Store(true) | |
| // Drop fast-path references before sessionTables.Close so a concurrent | |
| // openHandle racing Close either sees the empty map and takes the slow | |
| // path (which observes the closed cache and returns Unavailable) or has | |
| // already snapshotted a stale handle whose next RPC surfaces the | |
| // closed-pool error. Either outcome is honest; a lingering map entry | |
| // after Close would just keep the doomed handle reachable from Channel. | |
| c.handleFastPath.Range(func(k, _ any) bool { | |
| c.handleFastPath.Delete(k) | |
| return true | |
| }) | |
| c.sessionTables.Close() | |
| return c.sc.Close() | |
| } |
References
- Ensure that Close() methods are idempotent, which can be achieved by wrapping the closing logic with sync.Once or using an atomic boolean.
…p handleFastPath Scope-reduce the PR to the correctness change (concurrent-miss collapse) and defer the per-RPC hot-path optimization for its own follow-up. Removed: - handleFastPath sync.Map + LoadOrStore publish inside openGroup.Do + Range/Delete drain in Close. Kept: - openGroup singleflight.Group + openHandle's slow-path body wrapped in openGroup.Do(res.Name, ...). This alone eliminates the loser-branch stress on sessionTables.GetOrOpen — the correctness property this PR is about. Behavioral effect: after the initial burst on a resource, singleflight releases the key and subsequent openHandle calls re-enter the Do body. Each pays the switch/parse/closure-alloc cost + a sessionTables.mu acquire, and hits sessionTables.GetOrOpen's fast path. No new sc.OpenTable invocations. Test unchanged in intent: 32 concurrent openHandle callers → 1 sc.OpenTable invocation. Post-burst follow-up assertion tightened to say "sessionTables cache hit, not re-open" rather than "fast path should not touch sc" (which conflated the two caches).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Channel.openHandleis called per RPC from the accelerator's dispatch path. On the first burst of traffic to a fresh resource (or after aTableCacheTTL eviction), N concurrent RPCs all reachsessionTables.GetOrOpenwith the same cold key. UnderGetOrOpen's miss handling, every racer builds its ownsessionTableviasc.OpenTable(...), N-1 lose the insertion race, and each loser'sClose()runs the sessionTable's release path. Against a shared-key session pool, that release tore down the pool the winner already memoized — every subsequent RPC through the winner's handle servedErrPoolClosed.Top-level
client.OpenTable(t)(seebigtable/open.go) never triggered the bug because it warms the cache exactly once at construction, single-threaded from the caller's setup path. The concurrent-miss race the tests inbigtable/internal/session/pr20366_repro_test.gopin was unreachable from that caller shape. But the accelerator's per-RPCopenHandlereached it on the very first QPS burst.Fix
Two additions to
*Channel:handleFastPath sync.Map— lock-free per-resource cache mirroringsessionTables. On hit, skips the closure allocation,res.Kindswitch, resource-name parse, andsessionTables.muacquire. Staleness is safe because*TableHandle.ReadRow/MutateRowself-heal viaTableCache's internaldispatch()on TTL eviction, so a pointer left behind here still routes to the live successor on next use.openGroup singleflight.Group— collapses the cold-start burst. All N concurrent missers for a resource funnel through oneopenFncall; the waiters block onsingleflight's internal chan and receive the same handle.sessionTables.GetOrOpen's loser branch fires zero times from this call site, regardless of the underlying session-layer state.Result:
sc.OpenTableis called exactly once per resource perChannellifetime, matching theclient.OpenTable(t)-then-share-the-shim shape that never triggered the bug in production.TableCachestill owns pool lifecycle, TTL, and cross-caller dedup; the accelerator just stops stressing its loser branch from its own call site.Composition with #20366 / #20368
Complementary, not overlapping:
opened()-gate +refsrefcount) and fix(bigtable): session client — single ownership + single-flight pool loading #20368 (per-ownerpoolCloser+ single-flightTableCache) fix the race at the session layer, so no caller shape can trigger it.Verified by running the ported repro tests on
pr-20368: bothTestPR20366_*PASS. On the current branch without any session-layer fix, this accelerator change alone makes the burst benign because only onesc.OpenTableper resource ever runs.Test plan
bigtable/internal/accelerator/accelerator_singleflight_test.goaddsTestOpenHandle_SingleFlightCollapsesColdStartBurst:mockSessionClientin aslowSessionClientthat blocksOpenTableon a barrier and counts real invocations.openHandlecallers for the same fresh resource.selectonsc.entered/time.After(500ms)that fires whichever comes first — with singleflight, only one goroutine reachesOpenTableand the barrier's arrival count never hits N; without singleflight, arrival hits N almost immediately.sc.openTableCalls == 1and asserts a follow-up call is a fast-path hit (still 1).sc.OpenTable was invoked 32 times; want 1.Full accelerator suite passes under
-race:Loadreturns asession.TableAPIinterface value whose underlying*TableHandleself-heals on TTL eviction; no need to proactively invalidatehandleFastPathon eviction events. (Duplicating that invariant would drift.)ClosedrainshandleFastPathbeforesessionTables.Closeso a concurrentopenHandleracingCloseeither sees the empty map or a stale handle whose next RPC surfaces the closed-pool error — no lingering strong reference fromChannel.