feat(bigtable): BigtableDataClientFactory session support#13829
Conversation
There was a problem hiding this comment.
Code Review
This pull request refactors the Bigtable client and channel pool implementation to support multi-tenant factory usage and dynamic channel pool sizing. It introduces tenant-aware routing via ChannelPoolOptions and TenantKey, and updates ChannelPoolDpImpl to manage channel lifecycles using ACTIVE and DRAINING states. The review feedback highlights a compilation error in ChannelPoolDpImpl.java due to a missing import for Collectors, and suggests an optimization in diversity mode to prevent excessive channel creation under burst loads. Additionally, a thread-safety issue was identified in dumpState(), where unsynchronized access to guarded fields could cause concurrent modification exceptions, with a recommendation to use explicit locks instead of the synchronized keyword.
This commit extends the session-based API path to support BigtableDataClientFactory.
Now the factory creates a single shared ChannelPool and
ClientConfigurationManager once and hands lightweight child clients a reference to both.
The bulk of the implementation work is a rewrite of ChannelPoolDpImpl to correctly handle multiple tenants (different instances/app profiles) sharing a single channel pool.
---
ChannelPoolDpImpl: New Placement Model
Old design
The old pool organized channels into AfeChannelGroup objects — one deque of channels per observed AFE. When a channel's first session revealed its AFE (via onBeforeSessionStart), the channel was rehomed from
the temporary startingGroup into the matching AfeChannelGroup, creating one if it didn't exist. New streams were then placed by picking the group with the fewest sessions.
This design had a problem for the multi-tenant factory case:
AFE affinity was tracked per channel, but for factory children different tenants on the same channel can be routed to different AFEs by RLS. A channel in group "AFE-7" might route tenant A to AFE-7 but
tenant B somewhere else entirely.
New design
The new pool replaces the nested group structure with a flat channels list plus a routeObservations map keyed on (channelId, tenantKey) → lastObservedAfeId.
Data structures:
- channels: List<ChannelWrapper> — all channels in ACTIVE or DRAINING state
- routeObservations: Map<RouteKey, AfeId> — the last AFE seen for a given (channel, tenant) pair, populated on every onBeforeSessionStart and self-healing as RLS routes drift over time
- sessionsPerAfeId: Multiset<AfeId> — global count of live sessions per AFE, used to find underloaded targets
- TenantKey (new value type) — stamped onto CallOptions by TableBase so the pool can distinguish tenants within the same shared pool
Placement (newStream):
1. Capacity mode — find the AFE with the smallest session count that is still below softMaxPerGroup:
- 3a — Preferred: pick the ACTIVE channel whose last observed route for this tenant was that AFE, choosing the least-loaded among candidates. This is the steady-state path: RLS is stable per (channel,
tenant), so repeat sessions for the same tenant land on the same AFE.
- 3b — Unobserved-tenant fallback: if no route observation exists for this tenant on any channel, pick any ACTIVE channel with remaining capacity (least-loaded). The actual AFE is unknown until
onBeforeSessionStart fires; routeObservations is updated then for future placements.
2. Diversity mode — entered when all known AFEs are at cap, or the pool is empty: prefer an ACTIVE channel with fewer than softMaxPerGroup / 2 outstanding streams. If none exists, open a new channel and
absorb whichever AFE it lands on.
Scale-in (DRAINING state):
The old "thinning" logic would call channel.shutdown() during serviceChannels(). The new approach introduces a two-state channel lifecycle
(ACTIVE / DRAINING). When serviceChannels() decides to shrink the pool it marks excess channels DRAINING — they stop accepting new streams but continue serving existing ones. A DRAINING channel is removed
only when its numOutstanding reaches zero (in releaseChannel). Drain candidates are chosen by pickDrainCandidates, which prefers idle channels first, then least-recently-used, then oldest.
Route observation cleanup:
removeChannel purges all routeObservations entries for the removed channel's ID, preventing the map from growing unboundedly as channels cycle.
---
Other changes
BigtableDataClientFactory / BigtableClientContext:
- BigtableDataClientFactory.create() now calls BigtableClientContext.createForFactory() (a new entry point) instead of manually disabling sessions. The factory context initializes a shared ShimImpl (channel
pool + config manager) that is reused by all children.
- BigtableClientContext.createChild() now creates a lightweight ShimImpl child via ShimImpl.createForFactoryChild() rather than sharing a DisabledShim.
ShimImpl / Client:
- Client.channelPool changed from a bare ChannelPool reference to Resource<ChannelPool> so the factory-child constructor can hold a shared (non-closing) reference and the standard constructor holds an owned
one.
- ShimImpl.configManager similarly wrapped in Resource<> so close() is a no-op for factory children (the parent owns the manager's lifecycle).
78841fa to
90f95ef
Compare
|
|
||
| /** | ||
| * Factory-child constructor. Uses a pre-built, shared {@link ChannelPool} and {@link | ||
| * ClientConfigurationManager}. The pool is already started and must not be closed by this client. |
There was a problem hiding this comment.
nit: maybe reword "The pool is already started and must not be closed by this client." to something like "The pool should be created with Resource.createShared() and closed when the factory closes"
There was a problem hiding this comment.
The two public Client(...) constructors are distinguished only by the final parameter (ChannelProvider vs Resource). This is fragile — a factory (a static factory method or clearly named builders) would be safer than type-only overload resolution, especially since one throws IOException and the other doesn't.
There was a problem hiding this comment.
Sorry, I can wrap it in static factory, but looks like we'll still need two constructors like these. Leaving it as is for now.
|
|
||
| public TenantKey(@Nullable InstanceName instanceName, String appProfileId) { | ||
| this.instanceName = instanceName; | ||
| this.appProfileId = appProfileId == null ? "" : appProfileId; |
There was a problem hiding this comment.
nit: is this check necessary? I think we can just do this.appProfileId = appProfileId
| for (ChannelWrapper c : channels) { | ||
| if (c.state == ChannelWrapper.State.ACTIVE | ||
| && targetAfe.equals(routeObservations.get(new RouteKey(c.id, tenant)))) { | ||
| if (channelWrapper == null || c.numOutstanding < channelWrapper.numOutstanding) { |
There was a problem hiding this comment.
should we also check c.numOutstanding < softmax?
There was a problem hiding this comment.
no, we already checked that above when we picked target AFE. If target AFE has < softMaxPerGroup across all of its channels we should have at least one channel here with < softMaxPerGroup sessions.
This PR extends the session-based API path to support
BigtableDataClientFactory.Now the factory creates a single shared
ChannelPoolandClientConfigurationManageronce and hands lightweight child clients a reference to both.The bulk of the implementation work is a rewrite of
ChannelPoolDpImplto correctly handle multiple tenants (different instances/app profiles) sharing a single channel pool.ChannelPoolDpImpl: New Placement ModelOld design
The old pool organized channels into
AfeChannelGroupobjects — one deque of channels per observed AFE. When a channel's first session revealed its AFE (viaonBeforeSessionStart), the channel was rehomed from the temporarystartingGroupinto the matchingAfeChannelGroup, creating one if it didn't exist. New streams were then placed by picking the group with the fewest sessions.This design had a problem for the multi-tenant factory case:
AFE affinity was tracked per channel, but for factory children different tenants on the same channel can be routed to different AFEs by RLS. A channel in group "AFE-7" might route tenant A to AFE-7 but tenant B somewhere else entirely.
New design
The new pool replaces the nested group structure with a flat
channelslist plus arouteObservationsmap keyed on(channelId, tenantKey) → lastObservedAfeId.Data structures:
channels: List<ChannelWrapper>— all channels in ACTIVE or DRAINING staterouteObservations: Map<RouteKey, AfeId>— the last AFE seen for a given (channel, tenant) pair, populated on everyonBeforeSessionStartand self-healing as RLS routes drift over timesessionsPerAfeId: Multiset<AfeId>— global count of live sessions per AFE, used to find underloaded targetsTenantKey(new value type) — stamped ontoCallOptionsbyTableBaseso the pool can distinguish tenants within the same shared poolPlacement (
newStream):softMaxPerGroup:- 3a — Preferred: pick the ACTIVE channel whose last observed route for this tenant was that AFE, choosing the least-loaded among candidates. This is the steady-state path: RLS is stable per (channel, tenant), so repeat sessions for the same tenant land on the same AFE.
- 3b — Unobserved-tenant fallback: if no route observation exists for this tenant on any channel, pick any ACTIVE channel with remaining capacity (least-loaded). The actual AFE is unknown until
onBeforeSessionStartfires;routeObservationsis updated then for future placements.softMaxPerGroup / 2outstanding streams. If none exists, open a new channel and absorb whichever AFE it lands on.Scale-in (DRAINING state):
The old "thinning" logic would call
channel.shutdown()during serviceChannels(). The new approach introduces a two-state channel lifecycle (ACTIVE/DRAINING). WhenserviceChannels()decides to shrink the pool it marks excess channels DRAINING — they stop accepting new streams but continue serving existing ones. A DRAINING channel is removed only when itsnumOutstandingreaches zero (inreleaseChannel). Drain candidates are chosen bypickDrainCandidates, which prefers idle channels first, then least-recently-used, then oldest.Route observation cleanup:
removeChannelpurges allrouteObservationsentries for the removed channel's ID, preventing the map from growing unboundedly as channels cycle.Other changes
BigtableDataClientFactory/BigtableClientContext:BigtableDataClientFactory.create()now callsBigtableClientContext.createForFactory()(a new entry point) instead of manually disabling sessions. The factory context initializes a sharedShimImpl(channel pool + config manager) that is reused by all children.BigtableClientContext.createChild()now creates a lightweightShimImplchild viaShimImpl.createForFactoryChild()rather than sharing aDisabledShim.ShimImpl/Client:Client.channelPoolchanged from a bareChannelPoolreference toResource<ChannelPool>so the factory-child constructor can hold a shared (non-closing) reference and the standard constructor holds an owned one.ShimImpl.configManagersimilarly wrapped inResource<>soclose()is a no-op for factory children (the parent owns the manager's lifecycle).