Skip to content

feat(bigtable): BigtableDataClientFactory session support#13829

Open
nimf wants to merge 3 commits into
mainfrom
dataclientfactorysession
Open

feat(bigtable): BigtableDataClientFactory session support#13829
nimf wants to merge 3 commits into
mainfrom
dataclientfactorysession

Conversation

@nimf

@nimf nimf commented Jul 19, 2026

Copy link
Copy Markdown
Contributor

This PR 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).

@nimf
nimf requested review from a team as code owners July 19, 2026 23:23
@nimf
nimf requested a review from mutianf July 19, 2026 23:23

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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).
@nimf
nimf force-pushed the dataclientfactorysession branch from 78841fa to 90f95ef Compare July 19, 2026 23:28
@nimf nimf changed the title Rewrite ChannelPoolDpImpl to support BigtableDataClientFactory feat(bigtable): BigtableDataClientFactory session support Jul 19, 2026

/**
* 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.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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"

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

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;

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

should we also check c.numOutstanding < softmax?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

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.

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