From 90f95efb54fde58e49bb9587638719f0af249fd2 Mon Sep 17 00:00:00 2001 From: Yuri Golobokov Date: Sun, 19 Jul 2026 23:10:36 +0000 Subject: [PATCH 1/3] feat(bigtable): BigtableDataClientFactory session support MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 — all channels in ACTIVE or DRAINING state - routeObservations: Map — 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 — 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 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). --- .../data/v2/BigtableDataClientFactory.java | 6 +- .../bigtable/data/v2/internal/api/Client.java | 57 ++- .../data/v2/internal/api/TableBase.java | 10 +- .../internal/channels/ChannelPoolDpImpl.java | 399 ++++++++++-------- .../internal/channels/ChannelPoolOptions.java | 33 ++ .../data/v2/internal/channels/TenantKey.java | 56 +++ .../data/v2/internal/compat/ShimImpl.java | 58 ++- .../data/v2/stub/BigtableClientContext.java | 67 ++- .../channels/ChannelPoolDpImplTest.java | 33 +- 9 files changed, 477 insertions(+), 242 deletions(-) create mode 100644 java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/ChannelPoolOptions.java create mode 100644 java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/TenantKey.java diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactory.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactory.java index 61ab4cc12c12..5837d3956f4b 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactory.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactory.java @@ -74,12 +74,8 @@ public final class BigtableDataClientFactory implements AutoCloseable { */ public static BigtableDataClientFactory create(BigtableDataSettings defaultSettings) throws IOException { - BigtableDataSettings.Builder builder = defaultSettings.toBuilder(); - builder.stubSettings().setSessionsEnabled(false); - defaultSettings = builder.build(); - BigtableClientContext sharedClientContext = - BigtableClientContext.create(defaultSettings.getStubSettings()); + BigtableClientContext.createForFactory(defaultSettings.getStubSettings()); ClientOperationSettings perOpSettings = defaultSettings.getStubSettings().getPerOpSettings(); return new BigtableDataClientFactory(sharedClientContext, perOpSettings); } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java index 82ddff08c3cb..73c25747c071 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java @@ -73,7 +73,7 @@ public class Client implements AutoCloseable { private final Resource backgroundExecutor; private final CallOptions defaultCallOptions; - private final ChannelPool channelPool; + private final Resource channelPool; private final Resource metrics; private final Resource configManager; @@ -147,19 +147,23 @@ public static Client create(ClientSettings settings) throws IOException { return new Client( featureFlags, clientInfo, - settings.getChannelProvider(), Resource.createOwned(metrics, metrics::close), Resource.createOwned(configManager, configManager::close), - Resource.createOwned(backgroundExecutor, backgroundExecutor::shutdown)); + Resource.createOwned(backgroundExecutor, backgroundExecutor::shutdown), + settings.getChannelProvider()); } + /** + * Standard constructor used by non-factory clients. Builds an owned {@link SwitchingChannelPool} + * from the provided {@link ChannelProvider}. + */ public Client( FeatureFlags featureFlags, ClientInfo clientInfo, - ChannelProvider channelProvider, Resource metrics, Resource configManager, - Resource bgExecutor) + Resource bgExecutor, + ChannelProvider channelProvider) throws IOException { this.featureFlags = featureFlags; this.clientInfo = clientInfo; @@ -181,13 +185,34 @@ public Client( // TODO: consider localizing this for large reads .maxInboundMessageSize(256 * 1024 * 1024)); - channelPool = + SwitchingChannelPool switchingPool = new SwitchingChannelPool( configuredChannelProvider, configManager.get(), metrics.get(), backgroundExecutor.get()); - channelPool.start(); + switchingPool.start(); + this.channelPool = Resource.createOwned(switchingPool, switchingPool::close); + } + + /** + * 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. + */ + public Client( + FeatureFlags featureFlags, + ClientInfo clientInfo, + Resource metrics, + Resource configManager, + Resource bgExecutor, + Resource sharedChannelPool) { + this.featureFlags = featureFlags; + this.clientInfo = clientInfo; + this.metrics = metrics; + this.configManager = configManager; + this.backgroundExecutor = bgExecutor; + this.channelPool = sharedChannelPool; + defaultCallOptions = CallOptions.DEFAULT; } @Override @@ -200,7 +225,7 @@ public void close() { .setDescription("Client closing") .build())); metrics.close(); - channelPool.close(); + channelPool.close(); // no-op when Resource.createShared (factory child) configManager.close(); backgroundExecutor.close(); } @@ -211,7 +236,7 @@ public TableAsync openTableAsync(String tableId, Permission permission) { featureFlags, clientInfo, configManager.get(), - channelPool, + channelPool.get(), defaultCallOptions, tableId, permission, @@ -228,7 +253,7 @@ public AuthorizedViewAsync openAuthorizedViewAsync( featureFlags, clientInfo, configManager.get(), - channelPool, + channelPool.get(), defaultCallOptions, tableId, viewId, @@ -246,7 +271,7 @@ public MaterializedViewAsync openMaterializedViewAsync( featureFlags, clientInfo, configManager.get(), - channelPool, + channelPool.get(), defaultCallOptions, viewId, permission, @@ -256,6 +281,16 @@ public MaterializedViewAsync openMaterializedViewAsync( return viewAsync; } + /** Returns the underlying channel pool (e.g. for sharing with factory children). */ + public ChannelPool getChannelPool() { + return channelPool.get(); + } + + /** Returns the feature flags (e.g. for sharing with factory children). */ + public FeatureFlags getFeatureFlags() { + return featureFlags; + } + public static class Resource { private T value; private Runnable closer; diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableBase.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableBase.java index 8feef399d625..b32920c67cb4 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableBase.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableBase.java @@ -22,6 +22,8 @@ import com.google.bigtable.v2.SessionReadRowRequest; import com.google.bigtable.v2.SessionReadRowResponse; import com.google.cloud.bigtable.data.v2.internal.channels.ChannelPool; +import com.google.cloud.bigtable.data.v2.internal.channels.ChannelPoolOptions; +import com.google.cloud.bigtable.data.v2.internal.channels.TenantKey; import com.google.cloud.bigtable.data.v2.internal.csm.Metrics; import com.google.cloud.bigtable.data.v2.internal.csm.attributes.ClientInfo; import com.google.cloud.bigtable.data.v2.internal.csm.tracers.VRpcTracer; @@ -63,6 +65,12 @@ static TableBase createAndStart( Metrics metrics, ScheduledExecutorService executor) { + // Stamp the tenant key so ChannelPoolDpImpl can make tenant-aware placement decisions. + CallOptions stamped = + callOptions.withOption( + ChannelPoolOptions.TENANT_KEY_OPTION, + new TenantKey(clientInfo.getInstanceName(), clientInfo.getAppProfileId())); + SessionPool sessionPool = new SessionPoolImpl<>( metrics, @@ -70,7 +78,7 @@ static TableBase createAndStart( clientInfo, configManager, channelPool, - callOptions, + stamped, sessionDescriptor, sessionPoolName, executor); diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/ChannelPoolDpImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/ChannelPoolDpImpl.java index 7913b28ef14c..fccd60f56c38 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/ChannelPoolDpImpl.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/ChannelPoolDpImpl.java @@ -36,19 +36,18 @@ import java.time.Clock; import java.time.Duration; import java.time.Instant; -import java.util.ArrayDeque; import java.util.ArrayList; import java.util.Comparator; -import java.util.Deque; -import java.util.Iterator; +import java.util.HashMap; import java.util.List; import java.util.Locale; +import java.util.Map; import java.util.Objects; -import java.util.Optional; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; import java.util.function.Supplier; import java.util.logging.Level; import java.util.logging.Logger; @@ -57,9 +56,36 @@ import javax.annotation.concurrent.GuardedBy; /** - * Proof of concept channel pool that avoids parallel channels to the same afe. The pool is - * dynamically sized based on outstanding streams and tries to limit each channel to at most 10 - * streams. + * Dynamically-sized channel pool that spreads sessions across AFEs and supports multi-tenant + * factory usage. + * + *

Accounting is anchored on sessions rather than channels: + * + *

    + *
  • {@code sessionsPerAfeId} tracks how many active sessions each AFE holds globally. + *
  • {@code routeObservations} maps {@code (channelId, tenantKey) → lastObservedAfeId}, + * populated on every {@code onBeforeSessionStart}. Used to place new sessions on the channel + * most likely to route to the desired AFE for a given tenant. + *
+ * + *

Channels have a two-state lifecycle: ACTIVE (eligible for new sessions) and DRAINING (no new + * sessions; shut down when numOutstanding reaches 0). + * + *

Placement logic for each new stream: + * + *

    + *
  1. Capacity mode: find the AFE with the smallest session count below {@code softMaxPerGroup}. + *
      + *
    • 3a. Preferred: pick the ACTIVE channel whose last observed route for this tenant was + * that AFE, choosing the least-loaded. + *
    • 3b. Unobserved-tenant fallback: pick any ACTIVE channel with capacity (least-loaded). + * RLS may route to the target AFE or somewhere else; {@code routeObservations} will + * be updated on {@code onBeforeSessionStart} for future placements. + *
    + *
  2. Diversity mode (all known AFEs at cap, or none known): prefer an ACTIVE channel with fewer + * than {@code softMaxPerGroup / 2} streams; otherwise add a new channel and absorb whatever + * AFE it lands on. + *
*/ public class ChannelPoolDpImpl implements ChannelPool { private static final Logger LOGGER = Logger.getLogger(ChannelPoolDpImpl.class.getName()); @@ -70,6 +96,9 @@ public class ChannelPoolDpImpl implements ChannelPool { private static final int CONSECUTIVE_OPEN_SESSION_FAILURE_THRESHOLD = 5; private static final Duration INITIAL_RECYCLE_BACKOFF = Duration.ofMillis(1); private static final Duration MAX_RECYCLE_BACKOFF = Duration.ofMinutes(1); + // Preserves the existing desiredGroups = ceil(totalStreams / softMaxPerGroup * 2) factor, which + // keeps extra channels around at low load for AFE diversity. + private static final double HEADROOM = 2.0; private final String poolLogId; @@ -83,16 +112,23 @@ public class ChannelPoolDpImpl implements ChannelPool { private final DebugTagTracer debugTagTracer; + /** Flat list of all channels (ACTIVE + DRAINING). */ @GuardedBy("this") - private final List channelGroups; + private final List channels = new ArrayList<>(); + /** + * Per-(channel, tenant) observation of the last AFE seen on {@code onBeforeSessionStart}. + * Self-healing: updated on every session open, so stale entries from ~1hr AFE drift are + * corrected automatically. + */ @GuardedBy("this") - private final AfeChannelGroup startingGroup; + private final Map routeObservations = new HashMap<>(); + + private final AtomicLong nextChannelId = new AtomicLong(); @GuardedBy("this") private int totalStreams = 0; - // TODO: replace SocketAddress with AfeId @GuardedBy("this") private final Multiset sessionsPerAfeId = HashMultiset.create(); @@ -164,8 +200,6 @@ public ChannelPoolDpImpl( this.channelSupplier = channelSupplier; this.executor = executor; updateConfig(config); - channelGroups = new ArrayList<>(); - startingGroup = new AfeChannelGroup(null); this.debugTagTracer = debugTagTracer; } @@ -190,9 +224,7 @@ public synchronized void close() { serviceFuture.cancel(false); } - channelGroups.stream() - .flatMap(g -> g.channels.stream()) - .forEach(channelWrapper -> channelWrapper.channel.shutdown()); + channels.forEach(cw -> cw.channel.shutdown()); } @Override @@ -202,55 +234,72 @@ public synchronized SessionStream newStream( debugTagTracer.record(TelemetryConfiguration.Level.WARN, "channel_pool_new_stream_failed"); return new FailingSessionStream(Status.UNAVAILABLE.withDescription("ChannelPool is closed")); } - // Find the AFE with the fewest sessions - Optional maybeGroup = - channelGroups.stream() - .filter(g -> g.numStreams < softMaxPerGroup) - .min(Comparator.comparingInt(a -> a.numStreams)); - - // Find the first channel that has capacity - Optional maybeChannel = - maybeGroup.flatMap(g -> g.channels.stream().findFirst()); - - // try to find a channel that has capacity in the least loaded afe - final ChannelWrapper channelWrapper = - maybeChannel.orElseGet( - () -> { - log( - Level.FINE, - "Couldn't find an existing channel with capacity, num outstanding streams across" - + " all channel groups: %d, num groups: %d", - channelGroups.stream().mapToInt(g -> g.numStreams).sum(), - channelGroups.size()); - - // If there is no such channel, then try to add a channel that being resolved - // and if no channel is in the process of being added or its already nearing 50% - // utilization - // during initialization, then add a new channel - return startingGroup.channels.stream() - .filter(w -> w.numOutstanding < softMaxPerGroup / 2) - .min(Comparator.comparingInt(w -> w.numOutstanding)) - .map( - (w) -> { - log(Level.FINE, "Using a channel thats already connecting"); - return w; - }) - .orElseGet( - () -> { - log(Level.FINE, "Creating a new channel"); - return addChannel(); - }); - }); - - channelWrapper.numOutstanding++; - channelWrapper.group.numStreams++; + + TenantKey tenant = callOptions.getOption(ChannelPoolOptions.TENANT_KEY_OPTION); + + // --- Capacity mode: find the least-loaded AFE that still has room --- + AfeId targetAfe = + sessionsPerAfeId.entrySet().stream() + .filter(e -> e.getCount() < softMaxPerGroup) + .min(Comparator.comparingInt(Multiset.Entry::getCount)) + .map(Multiset.Entry::getElement) + .orElse(null); + + ChannelWrapper channelWrapper = null; + + if (targetAfe != null) { + // 3a: prefer the ACTIVE channel that already routes this tenant to targetAfe. + // For-loop avoids accessing the guarded routeObservations field inside a lambda. + 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) { + channelWrapper = c; + } + } + } + + // 3b: no observed route for this tenant → any ACTIVE channel with capacity (least-loaded). + if (channelWrapper == null) { + channelWrapper = + channels.stream() + .filter(c -> c.state == ChannelWrapper.State.ACTIVE) + .filter(c -> c.numOutstanding < softMaxPerGroup) + .min(Comparator.comparingInt(c -> c.numOutstanding)) + .orElse(null); + } + } + + // --- Diversity mode: all AFEs at cap, or pool is empty --- + if (channelWrapper == null) { + // Reuse an ACTIVE channel that has plenty of headroom before spawning a new one. + channelWrapper = + channels.stream() + .filter(c -> c.state == ChannelWrapper.State.ACTIVE) + .filter(c -> c.numOutstanding < softMaxPerGroup / 2) + .min(Comparator.comparingInt(c -> c.numOutstanding)) + .orElse(null); + + if (channelWrapper == null) { + log( + Level.FINE, + "Couldn't find an existing channel with capacity, num outstanding streams: %d," + + " num channels: %d", + totalStreams, + channels.size()); + channelWrapper = addChannel(); + } + } + + final ChannelWrapper fw = channelWrapper; + fw.numOutstanding++; totalStreams++; ClientCall innerCall = - channelWrapper.channel.newCall(desc, callOptions); + fw.channel.newCall(desc, callOptions); return new SessionStreamImpl(innerCall) { - // mark as null so that onClose can tell if onBeforeSessionStart was never called + // null until onBeforeSessionStart fires @Nullable AfeId afeId = null; @Override @@ -261,10 +310,11 @@ public void start(Listener responseListener, Metadata headers) { public void onBeforeSessionStart(PeerInfo peerInfo) { afeId = AfeId.extract(peerInfo); synchronized (ChannelPoolDpImpl.this) { - channelWrapper.consecutiveFailures = 0; + fw.consecutiveFailures = 0; recycleBackoff = INITIAL_RECYCLE_BACKOFF; - rehomeChannel(channelWrapper, afeId); sessionsPerAfeId.add(afeId); + routeObservations.put(new RouteKey(fw.id, tenant), afeId); + fw.lastUsedAt = Instant.now(clock); } super.onBeforeSessionStart(peerInfo); } @@ -275,9 +325,9 @@ public void onClose(Status status, Metadata trailers) { if (afeId != null) { sessionsPerAfeId.remove(afeId); } else if (!status.isOk() && status.getCode() != Code.CANCELLED) { - channelWrapper.consecutiveFailures++; + fw.consecutiveFailures++; } - releaseChannel(channelWrapper, status); + releaseChannel(fw, status); } super.onClose(status, trailers); } @@ -292,63 +342,35 @@ private ChannelWrapper addChannel() { debugTagTracer.checkPrecondition( !closed, "channel_pool_add_channel_failure", "Channel pool is closed"); log(Level.FINE, "Adding a new channel"); - ChannelWrapper wrapped = new ChannelWrapper(startingGroup, channelSupplier.get(), clock); - startingGroup.channels.addFirst(wrapped); + ChannelWrapper wrapped = + new ChannelWrapper(nextChannelId.getAndIncrement(), channelSupplier.get(), clock); + channels.add(wrapped); return wrapped; } + /** Removes a channel from the pool and purges its route observations. */ @GuardedBy("this") - private void removeGroup(AfeChannelGroup group) { - log(Level.FINE, "Removing group: %s", group.afeId); - group.channels.forEach(c -> c.channel.shutdown()); - channelGroups.remove(group); - // TODO: need to handle a group being added right after it was removed with lingering sessions - } - - @GuardedBy("this") - private void rehomeChannel(ChannelWrapper channelWrapper, AfeId afeId) { - // No need to rehome recycled channels. - if (channelWrapper.channel.isShutdown()) { - return; - } - - AfeChannelGroup origGroup = channelWrapper.group; - - if (Objects.equals(origGroup.afeId, afeId)) { - return; - } - - log(Level.FINE, "Rehoming channel from: %s to %s", origGroup.afeId, afeId); - - // Re-home the channel - AfeChannelGroup newGroup = - channelGroups.stream() - .filter(g -> Objects.equals(g.afeId, afeId)) - .findAny() - .orElseGet( - () -> { - AfeChannelGroup g = new AfeChannelGroup(afeId); - channelGroups.add(g); - return g; - }); - channelWrapper.group = newGroup; - origGroup.channels.remove(channelWrapper); - if (origGroup.channels.isEmpty()) { - channelGroups.remove(origGroup); - } - origGroup.numStreams -= channelWrapper.numOutstanding; - newGroup.channels.add(channelWrapper); - newGroup.numStreams += channelWrapper.numOutstanding; - - return; + private void removeChannel(ChannelWrapper cw) { + channels.remove(cw); + final long id = cw.id; + routeObservations.keySet().removeIf(k -> k.channelId == id); + cw.channel.shutdown(); } - // Update accounting when a stream is closed and releases its channel @GuardedBy("this") private void releaseChannel(ChannelWrapper channelWrapper, Status status) { totalStreams--; - channelWrapper.group.numStreams--; channelWrapper.numOutstanding--; + channelWrapper.lastUsedAt = Instant.now(clock); + + // Draining channels: don't recycle (that would spawn a replacement, undoing the drain). + // Just remove when they go quiet. + if (channelWrapper.state == ChannelWrapper.State.DRAINING) { + if (channelWrapper.numOutstanding == 0) { + removeChannel(channelWrapper); + } + return; + } if (shouldRecycleChannel(channelWrapper, status)) { recycleChannel(channelWrapper); @@ -391,12 +413,7 @@ private void recycleChannel(ChannelWrapper channelWrapper) { recycleBackoff = MAX_RECYCLE_BACKOFF; } - channelWrapper.group.channels.remove(channelWrapper); - channelWrapper.channel.shutdown(); - // Checking for starting group because we don't want to delete the stating group. - if (channelWrapper.group.channels.isEmpty() && channelWrapper.group != startingGroup) { - removeGroup(channelWrapper.group); - } + removeChannel(channelWrapper); addChannel(); } @@ -410,59 +427,31 @@ void serviceChannels() { } } - /* - - add new group if existing groups are overflowing - - shutdown older channels within a group - */ private synchronized void serviceChannelsSafe() { log(Level.FINE, "Servicing channels"); dumpState(); - // Thin out the channels in each group, so that each AFEGroup only has 1 channel - for (AfeChannelGroup group : channelGroups) { - if (LOGGER.isLoggable(Level.FINEST) && group.channels.size() > 1) { - log( - Level.FINEST, - "Shutting down %d parallel channels for %s", - group.channels.size() - 1, - group.afeId); - } - while (group.channels.size() > 1) { - // Clean up parallel channels. - // Recent channels added to the end, thus removing the oldest from the beginning. - group.channels.removeFirst().channel.shutdown(); - } + int target = (int) Math.ceil((totalStreams * HEADROOM) / softMaxPerGroup); + if (target > maxGroups) { + target = maxGroups; + } else if (target < minGroups) { + target = minGroups; } - // try to adjust the groups - int desiredGroups = (int) Math.ceil(((float) totalStreams / softMaxPerGroup) * 2); - if (desiredGroups > maxGroups) { - desiredGroups = maxGroups; - } else if (desiredGroups < minGroups) { - desiredGroups = minGroups; - } + int activeCount = (int) channels.stream().filter(c -> c.state == ChannelWrapper.State.ACTIVE).count(); - // Right size the groups - if (desiredGroups < channelGroups.size()) { - // Remove extra groups, oldest first - Iterator it = - channelGroups.stream() - .sorted(Comparator.comparing(g -> g.channels.peek().createdAt)) - .limit(channelGroups.size() - desiredGroups) - .iterator(); - - while (it.hasNext()) { - AfeChannelGroup group = it.next(); - log( - Level.FINE, - "Removing %d channel for %s due to lack of concurrency", - group.channels.size(), - group.afeId); - removeGroup(group); + if (activeCount > target) { + List candidates = pickDrainCandidates(activeCount - target); + for (ChannelWrapper cw : candidates) { + log(Level.FINE, "Draining channel (outstanding: %d)", cw.numOutstanding); + cw.state = ChannelWrapper.State.DRAINING; + if (cw.numOutstanding == 0) { + removeChannel(cw); + } } - } else if (desiredGroups > channelGroups.size() + startingGroup.channels.size()) { - log(Level.FINE, "Adding %d channels", desiredGroups - channelGroups.size()); - for (int i = channelGroups.size(); i < desiredGroups; i++) { + } else if (activeCount < target) { + log(Level.FINE, "Adding %d channels to reach target %d", target - activeCount, target); + for (int i = activeCount; i < target; i++) { addChannel(); } } @@ -471,6 +460,24 @@ private synchronized void serviceChannelsSafe() { dumpState(); } + /** + * Picks channels to drain, preferring idle channels (numOutstanding == 0) then by least-recently + * used, then by oldest createdAt. + */ + @GuardedBy("this") + private List pickDrainCandidates(int count) { + return channels.stream() + .filter(c -> c.state == ChannelWrapper.State.ACTIVE) + .sorted( + Comparator.comparingInt((ChannelWrapper c) -> c.numOutstanding) + .thenComparing( + c -> c.lastUsedAt == null ? Instant.MIN : c.lastUsedAt, + Comparator.naturalOrder()) + .thenComparing(c -> c.createdAt)) + .limit(count) + .collect(Collectors.toList()); + } + private void log(Level level, String msg, Throwable throwable) { LOGGER.log(level, String.format("[%s] %s", poolLogId, msg), throwable); } @@ -485,8 +492,8 @@ void dumpState() { return; } - int channels = - channelGroups.stream().mapToInt((AfeChannelGroup chg) -> chg.channels.size()).sum(); + long activeCount = channels.stream().filter(c -> c.state == ChannelWrapper.State.ACTIVE).count(); + long drainingCount = channels.stream().filter(c -> c.state == ChannelWrapper.State.DRAINING).count(); String s = sessionsPerAfeId.entrySet().stream() .sorted(Comparator.comparing(e -> e.getElement().toString())) @@ -495,12 +502,10 @@ void dumpState() { log( Level.FINE, - "ChannelPool channelGroups: " - + channelGroups.size() - + ", channels: " - + channels - + ", starting channels: " - + startingGroup.channels.size() + "ChannelPool active: " + + activeCount + + ", draining: " + + drainingCount + ", totalStreams: " + totalStreams + ", AFEs: " @@ -519,29 +524,51 @@ void dumpState() { } } - static class AfeChannelGroup { - private final AfeId afeId; - private final Deque channels; - private int numStreams; + /** + * Key for the route-observation map: the combination of a channel id and a tenant uniquely + * determines the last observed AFE (RLS is stable per channel + tenant). + */ + static final class RouteKey { + final long channelId; + final TenantKey tenantKey; + + RouteKey(long channelId, TenantKey tenantKey) { + this.channelId = channelId; + this.tenantKey = tenantKey; + } - public AfeChannelGroup(AfeId afeId) { - this.afeId = afeId; - channels = new ArrayDeque<>(); - numStreams = 0; + @Override + public boolean equals(Object o) { + if (!(o instanceof RouteKey)) return false; + RouteKey other = (RouteKey) o; + return channelId == other.channelId && Objects.equals(tenantKey, other.tenantKey); + } + + @Override + public int hashCode() { + return Objects.hash(channelId, tenantKey); } } static class ChannelWrapper { - private AfeChannelGroup group; - private final ManagedChannel channel; - private final Instant createdAt; - private int numOutstanding = 0; - private int consecutiveFailures = 0; - - public ChannelWrapper(AfeChannelGroup group, ManagedChannel channel, Clock clock) { - this.group = group; + enum State { + ACTIVE, + DRAINING + } + + final long id; + final ManagedChannel channel; + final Instant createdAt; + volatile State state = State.ACTIVE; + int numOutstanding = 0; + int consecutiveFailures = 0; + Instant lastUsedAt; + + ChannelWrapper(long id, ManagedChannel channel, Clock clock) { + this.id = id; this.channel = channel; - createdAt = Instant.now(clock); + this.createdAt = Instant.now(clock); + this.lastUsedAt = this.createdAt; } } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/ChannelPoolOptions.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/ChannelPoolOptions.java new file mode 100644 index 000000000000..da3dced8af8b --- /dev/null +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/ChannelPoolOptions.java @@ -0,0 +1,33 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigtable.data.v2.internal.channels; + +import io.grpc.CallOptions; + +/** CallOptions keys used by {@link ChannelPoolDpImpl} for tenant-aware stream placement. */ +public final class ChannelPoolOptions { + /** + * Identifies the tenant (instanceName + appProfileId) for a session stream request. Set by + * TableBase before constructing SessionPoolImpl so that the channel pool can route streams to the + * correct AFE when multiple tenants share a pool. Defaults to {@link TenantKey#UNKNOWN} for + * single-tenant clients. + */ + public static final CallOptions.Key TENANT_KEY_OPTION = + CallOptions.Key.createWithDefault("bigtable.tenantKey", TenantKey.UNKNOWN); + + private ChannelPoolOptions() {} +} diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/TenantKey.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/TenantKey.java new file mode 100644 index 000000000000..2465f715c17c --- /dev/null +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/TenantKey.java @@ -0,0 +1,56 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigtable.data.v2.internal.channels; + +import com.google.cloud.bigtable.data.v2.internal.api.InstanceName; +import java.util.Objects; +import javax.annotation.Nullable; + +/** + * Identifies a single (instanceName, appProfileId) tenant within a shared channel pool. Used as + * the tenant dimension of a {@link RouteKey} in {@link ChannelPoolDpImpl}. + */ +public final class TenantKey { + /** Sentinel used when no tenant is stamped (single-tenant / sessions-disabled clients). */ + public static final TenantKey UNKNOWN = new TenantKey(null, ""); + + @Nullable private final InstanceName instanceName; + private final String appProfileId; + + public TenantKey(@Nullable InstanceName instanceName, String appProfileId) { + this.instanceName = instanceName; + this.appProfileId = appProfileId == null ? "" : appProfileId; + } + + @Override + public boolean equals(Object o) { + if (!(o instanceof TenantKey)) return false; + TenantKey other = (TenantKey) o; + return Objects.equals(instanceName, other.instanceName) + && Objects.equals(appProfileId, other.appProfileId); + } + + @Override + public int hashCode() { + return Objects.hash(instanceName, appProfileId); + } + + @Override + public String toString() { + return instanceName + "/" + appProfileId; + } +} diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java index 353973dc8e58..57dfb7be2bde 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java @@ -32,6 +32,7 @@ import com.google.cloud.bigtable.data.v2.internal.api.ChannelProviders.ConfiguredChannelProvider; import com.google.cloud.bigtable.data.v2.internal.api.Client; import com.google.cloud.bigtable.data.v2.internal.api.Client.Resource; +import com.google.cloud.bigtable.data.v2.internal.channels.ChannelPool; import com.google.cloud.bigtable.data.v2.internal.compat.ops.DivertingUnaryCallable; import com.google.cloud.bigtable.data.v2.internal.compat.ops.MutateRowShim; import com.google.cloud.bigtable.data.v2.internal.compat.ops.ReadRowShim; @@ -72,6 +73,7 @@ public class ShimImpl implements Shim { private static final Duration DA_CHECK_TIMEOUT = Duration.ofSeconds(5); private final ClientConfigurationManager configManager; + private final Resource configManagerResource; private final Client client; private final ReadRowShimInner readRowShimInner; @@ -164,22 +166,66 @@ public static Shim create( new Client( clientChannelProvider.updateFeatureFlags(featureFlags), clientInfo, - clientChannelProvider, Resource.createShared(metrics), Resource.createShared(configManager), - Resource.createShared(bgExecutor)); + Resource.createShared(bgExecutor), + clientChannelProvider); - return new ShimImpl(configManager, client); + return new ShimImpl(Resource.createOwned(configManager, configManager::close), client); } - public ShimImpl(ClientConfigurationManager configManager, Client client) { - this.configManager = configManager; + public ShimImpl(Resource configManagerResource, Client client) { + this.configManagerResource = configManagerResource; + this.configManager = configManagerResource.get(); this.client = client; this.readRowShimInner = new ReadRowShimInner(client); this.mutateRowShim = new MutateRowShim(client); } + /** + * Creates a lightweight child shim for factory mode. Reuses the factory's already-started, + * shared channel pool and config manager. The pool and manager are not closed when this shim is + * closed — only the child's own session pools are cleaned up. + */ + public static Shim createForFactoryChild( + ClientInfo clientInfo, + Metrics metrics, + ScheduledExecutorService bgExecutor, + ChannelPool sharedChannelPool, + ClientConfigurationManager sharedConfigManager, + FeatureFlags parentFeatureFlags) { + // Inherit the parent's fully-computed feature flags (which include channel-provider flags like + // trafficDirectorEnabled/directAccessRequested in addition to sessionsRequired). + FeatureFlags featureFlags = parentFeatureFlags; + + Client client = + new Client( + featureFlags, + clientInfo, + Resource.createShared(metrics), + Resource.createShared(sharedConfigManager), + Resource.createShared(bgExecutor), + Resource.createShared(sharedChannelPool)); + + return new ShimImpl(Resource.createShared(sharedConfigManager), client); + } + + /** Returns the raw config manager (e.g. for sharing with factory children). */ + public ClientConfigurationManager getConfigManager() { + return configManager; + } + + /** Returns the channel pool (e.g. for sharing with factory children). */ + public ChannelPool getChannelPool() { + return client.getChannelPool(); + } + + /** Returns the feature flags (e.g. for sharing with factory children). */ + public FeatureFlags getFeatureFlags() { + return client.getFeatureFlags(); + } + private static ChannelProvider configureChannelProvider( ChannelProvider channelProvider, Metrics metrics) { return new ConfiguredChannelProvider( @@ -252,7 +298,7 @@ private static boolean checkDirectAccessAvailable( @Override public void close() { client.close(); - configManager.close(); + configManagerResource.close(); // no-op when Resource.createShared (factory child) } @Override diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java index a74c15613e0f..084122155f18 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java @@ -28,9 +28,11 @@ import com.google.auth.oauth2.ServiceAccountJwtAccessCredentials; import com.google.cloud.bigtable.data.v2.internal.JwtCredentialsWithAudience; import com.google.cloud.bigtable.data.v2.internal.api.InstanceName; +import com.google.cloud.bigtable.data.v2.internal.channels.ChannelPool; import com.google.cloud.bigtable.data.v2.internal.compat.DisabledShim; import com.google.cloud.bigtable.data.v2.internal.compat.Shim; import com.google.cloud.bigtable.data.v2.internal.compat.ShimImpl; +import com.google.cloud.bigtable.data.v2.internal.util.ClientConfigurationManager; import com.google.cloud.bigtable.data.v2.internal.csm.MetricRegistry; import com.google.cloud.bigtable.data.v2.internal.csm.Metrics; import com.google.cloud.bigtable.data.v2.internal.csm.MetricsImpl; @@ -42,7 +44,6 @@ import com.google.cloud.bigtable.data.v2.stub.metrics.CustomOpenTelemetryMetricsProvider; import com.google.cloud.bigtable.gaxx.grpc.BigtableTransportChannelProvider; import com.google.cloud.bigtable.gaxx.grpc.ChannelPrimer; -import com.google.common.base.Preconditions; import io.grpc.ManagedChannelBuilder; import io.opencensus.stats.Stats; import io.opencensus.stats.StatsRecorder; @@ -68,6 +69,8 @@ public class BigtableClientContext { private static final Logger logger = Logger.getLogger(BigtableClientContext.class.getName()); private final boolean isChild; + // true when this is the parent context of a BigtableDataClientFactory + private final boolean isFactory; private final ClientInfo clientInfo; private final Metrics metrics; private final ClientContext clientContext; @@ -78,12 +81,32 @@ public class BigtableClientContext { public static BigtableClientContext create(EnhancedBigtableStubSettings settings) throws IOException { - return create(settings, Tags.getTagger(), Stats.getStatsRecorder()); + return createInternal(settings, Tags.getTagger(), Stats.getStatsRecorder(), false); } public static BigtableClientContext create( EnhancedBigtableStubSettings settings, Tagger ocTagger, StatsRecorder ocRecorder) throws IOException { + return createInternal(settings, ocTagger, ocRecorder, false); + } + + /** + * Creates a context to be used as the shared parent of a {@link + * com.google.cloud.bigtable.data.v2.BigtableDataClientFactory}. The context eagerly builds the + * shared channel pool and config manager using the factory's default settings. Children are + * created via {@link #createChild} and share these resources. + */ + public static BigtableClientContext createForFactory(EnhancedBigtableStubSettings settings) + throws IOException { + return createInternal(settings, Tags.getTagger(), Stats.getStatsRecorder(), true); + } + + private static BigtableClientContext createInternal( + EnhancedBigtableStubSettings settings, + Tagger ocTagger, + StatsRecorder ocRecorder, + boolean isFactory) + throws IOException { ClientInfo clientInfo = ClientInfo.builder() .setInstanceName(InstanceName.of(settings.getProjectId(), settings.getInstanceId())) @@ -208,7 +231,7 @@ public static BigtableClientContext create( try { return new BigtableClientContext( - false, shim, clientInfo, clientContext, metrics, executorProvider); + false, isFactory, shim, clientInfo, clientContext, metrics, executorProvider); } catch (IOException | RuntimeException t) { metrics.close(); throw t; @@ -237,6 +260,7 @@ private static void configureGrpcOtel( private BigtableClientContext( boolean isChild, + boolean isFactory, Shim shim, ClientInfo clientInfo, ClientContext clientContext, @@ -244,6 +268,7 @@ private BigtableClientContext( ExecutorProvider backgroundExecutorProvider) throws IOException { this.isChild = isChild; + this.isFactory = isFactory; this.sessionShim = shim; this.clientInfo = clientInfo; @@ -268,27 +293,39 @@ public ClientContext getClientContext() { public BigtableClientContext createChild(InstanceName instanceName, String appProfileId) throws IOException { - // TODO: either mark BigtableDataClientFactory as deprecated or figure out how to make it - // work with Sessions - Preconditions.checkState( - sessionShim instanceof DisabledShim, "Sessions don't support BigtableDataClientFactory"); + ClientInfo childInfo = + clientInfo.toBuilder().setInstanceName(instanceName).setAppProfileId(appProfileId).build(); + + Shim childShim; + if (isFactory && sessionShim instanceof ShimImpl) { + // Factory mode: create a lightweight child shim that shares the parent's pool and manager. + ShimImpl parentShim = (ShimImpl) sessionShim; + childShim = + ShimImpl.createForFactoryChild( + childInfo, + metrics, + backgroundExecutorProvider.getExecutor(), + parentShim.getChannelPool(), + parentShim.getConfigManager(), + parentShim.getFeatureFlags()); + } else { + // Legacy mode (sessions disabled): share the parent's DisabledShim. + childShim = sessionShim; + } return new BigtableClientContext( - true, - sessionShim, - clientInfo.toBuilder().setInstanceName(instanceName).setAppProfileId(appProfileId).build(), - clientContext, - metrics, - backgroundExecutorProvider); + true, false, childShim, childInfo, clientContext, metrics, backgroundExecutorProvider); } public void close() throws Exception { + // Always close the shim: no-op for DisabledShim, closes session pools for factory children. + sessionShim.close(); + if (isChild) { + // Shared resources (metrics, executor, transport) are owned by the parent context. return; } - sessionShim.close(); - for (BackgroundResource resource : clientContext.getBackgroundResources()) { resource.close(); } diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/channels/ChannelPoolDpImplTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/channels/ChannelPoolDpImplTest.java index eff05e2fd2b0..4efca0494718 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/channels/ChannelPoolDpImplTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/channels/ChannelPoolDpImplTest.java @@ -26,6 +26,7 @@ import com.google.bigtable.v2.FakeSessionGrpc; import com.google.bigtable.v2.PeerInfo; +import com.google.bigtable.v2.PeerInfo.TransportType; import com.google.bigtable.v2.SessionClientConfiguration.ChannelPoolConfiguration; import com.google.cloud.bigtable.data.v2.internal.channels.SessionStream.Listener; import com.google.cloud.bigtable.data.v2.internal.csm.NoopMetrics; @@ -283,23 +284,18 @@ void testDownsizeToOptimal() { i++; } - // Now we have 25 sessions on 5 channel groups each of 1 channel. - // Let's close 6 sessions from different channels/AFEs. - i = 0; - for (ClientCall.Listener listener : listener.getAllValues()) { - if (i % 4 == 0 && i != 0) { - listener.onClose(Status.OK, new Metadata()); - } - i++; + // Close all sessions from channel 0 (the first softMaxPerGroup/2 listeners) so ch0 has + // 0 outstanding. serviceChannels() picks the idle channel as drain candidate and removes it + // immediately (DRAINING channels with 0 outstanding are removed at drain time, not later). + for (int j = 0; j < pool.softMaxPerGroup / 2; j++) { + listener.getAllValues().get(j).onClose(Status.OK, new Metadata()); } - // Now we should have 19 sessions on 5 channel groups each of 1 channel. - // I.e. dumpState - // FINE: ChannelPool channelGroups: 5, channels: 5, starting channels: 0, totalStreams: 19, - // AFEs: 5, distribution: [4, 4, 4, 4, 3] + // 20 sessions remain across 5 channels; ch0 has 0 outstanding, ch1-4 have 5 each. + // target = ceil(20 * 2.0 / 10) = 4; ch0 wins drain candidate sort (fewest outstanding). pool.serviceChannels(); - // Should scale down to 4 channels. 19 / 5 round up = 4. + // ch0 has 0 outstanding → removed immediately → shutdown() called once. verify(channel, times(numChannels - 4)).shutdown(); pool.close(); @@ -422,7 +418,6 @@ void testRecycledChannelDoesNotRejoinPool() throws InterruptedException { when(channelSupplier.get()).thenReturn(channel); when(channel.newCall(any(), any())).thenReturn(clientCall); doNothing().when(clientCall).start(listener.capture(), any()); - doReturn(Attributes.EMPTY).when(clientCall).getAttributes(); ChannelPoolDpImpl pool = new ChannelPoolDpImpl(channelSupplier, defaultConfig, debugTagTracer, bgExecutor); @@ -446,12 +441,14 @@ void testRecycledChannelDoesNotRejoinPool() throws InterruptedException { // 3. Recycle channel1 via stream2 listener2.onClose(Status.UNIMPLEMENTED, new Metadata()); verify(channel, times(1)).shutdown(); - // Now isShutdown for the channel1 returns true - when(channel.isShutdown()).thenReturn(true); + // channel1 is now removed from the pool's channel list by removeChannel(). // 4. stream1 (on recycled channel1) receives headers with AFE ID - // This triggers rehomeChannel - PeerInfo peerInfo = PeerInfo.newBuilder().setApplicationFrontendId(555).build(); + PeerInfo peerInfo = + PeerInfo.newBuilder() + .setApplicationFrontendId(555) + .setTransportType(TransportType.TRANSPORT_TYPE_SESSION_UNKNOWN) + .build(); Metadata headers = new Metadata(); headers.put( SessionStreamImpl.PEER_INFO_KEY, From 353c5ba4b97f3e8e25afc2f4943f1599aea573d0 Mon Sep 17 00:00:00 2001 From: Yuri Golobokov Date: Tue, 21 Jul 2026 21:16:37 +0000 Subject: [PATCH 2/3] address comments --- .../data/v2/BigtableDataClientFactory.java | 2 +- .../bigtable/data/v2/internal/api/Client.java | 5 ++- .../internal/channels/ChannelPoolDpImpl.java | 4 +- .../data/v2/internal/channels/TenantKey.java | 7 ++-- .../data/v2/internal/compat/ShimImpl.java | 2 +- .../data/v2/stub/BigtableClientContext.java | 38 +++---------------- 6 files changed, 17 insertions(+), 41 deletions(-) diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactory.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactory.java index 5837d3956f4b..f19726e2a315 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactory.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactory.java @@ -75,7 +75,7 @@ public final class BigtableDataClientFactory implements AutoCloseable { public static BigtableDataClientFactory create(BigtableDataSettings defaultSettings) throws IOException { BigtableClientContext sharedClientContext = - BigtableClientContext.createForFactory(defaultSettings.getStubSettings()); + BigtableClientContext.create(defaultSettings.getStubSettings()); ClientOperationSettings perOpSettings = defaultSettings.getStubSettings().getPerOpSettings(); return new BigtableDataClientFactory(sharedClientContext, perOpSettings); } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java index 73c25747c071..5348589814c8 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java @@ -197,7 +197,8 @@ public Client( /** * 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. + * ClientConfigurationManager}. The pool and config manager should be created with + * Resource.createShared() and closed when the factory closes. */ public Client( FeatureFlags featureFlags, @@ -225,7 +226,7 @@ public void close() { .setDescription("Client closing") .build())); metrics.close(); - channelPool.close(); // no-op when Resource.createShared (factory child) + channelPool.close(); configManager.close(); backgroundExecutor.close(); } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/ChannelPoolDpImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/ChannelPoolDpImpl.java index fccd60f56c38..c73f5c55e02c 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/ChannelPoolDpImpl.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/ChannelPoolDpImpl.java @@ -313,7 +313,9 @@ public void onBeforeSessionStart(PeerInfo peerInfo) { fw.consecutiveFailures = 0; recycleBackoff = INITIAL_RECYCLE_BACKOFF; sessionsPerAfeId.add(afeId); - routeObservations.put(new RouteKey(fw.id, tenant), afeId); + if (!fw.channel.isShutdown()) { + routeObservations.put(new RouteKey(fw.id, tenant), afeId); + } fw.lastUsedAt = Instant.now(clock); } super.onBeforeSessionStart(peerInfo); diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/TenantKey.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/TenantKey.java index 2465f715c17c..70ec9735cf84 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/TenantKey.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/TenantKey.java @@ -18,7 +18,6 @@ import com.google.cloud.bigtable.data.v2.internal.api.InstanceName; import java.util.Objects; -import javax.annotation.Nullable; /** * Identifies a single (instanceName, appProfileId) tenant within a shared channel pool. Used as @@ -26,12 +25,12 @@ */ public final class TenantKey { /** Sentinel used when no tenant is stamped (single-tenant / sessions-disabled clients). */ - public static final TenantKey UNKNOWN = new TenantKey(null, ""); + public static final TenantKey UNKNOWN = new TenantKey(InstanceName.of("", ""), ""); - @Nullable private final InstanceName instanceName; + private final InstanceName instanceName; private final String appProfileId; - public TenantKey(@Nullable InstanceName instanceName, String appProfileId) { + public TenantKey(InstanceName instanceName, String appProfileId) { this.instanceName = instanceName; this.appProfileId = appProfileId == null ? "" : appProfileId; } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java index 57dfb7be2bde..e3bc6ec90abf 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java @@ -298,7 +298,7 @@ private static boolean checkDirectAccessAvailable( @Override public void close() { client.close(); - configManagerResource.close(); // no-op when Resource.createShared (factory child) + configManagerResource.close(); } @Override diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java index 084122155f18..86bc18c8482e 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java @@ -69,8 +69,6 @@ public class BigtableClientContext { private static final Logger logger = Logger.getLogger(BigtableClientContext.class.getName()); private final boolean isChild; - // true when this is the parent context of a BigtableDataClientFactory - private final boolean isFactory; private final ClientInfo clientInfo; private final Metrics metrics; private final ClientContext clientContext; @@ -81,31 +79,13 @@ public class BigtableClientContext { public static BigtableClientContext create(EnhancedBigtableStubSettings settings) throws IOException { - return createInternal(settings, Tags.getTagger(), Stats.getStatsRecorder(), false); + return create(settings, Tags.getTagger(), Stats.getStatsRecorder()); } public static BigtableClientContext create( - EnhancedBigtableStubSettings settings, Tagger ocTagger, StatsRecorder ocRecorder) - throws IOException { - return createInternal(settings, ocTagger, ocRecorder, false); - } - - /** - * Creates a context to be used as the shared parent of a {@link - * com.google.cloud.bigtable.data.v2.BigtableDataClientFactory}. The context eagerly builds the - * shared channel pool and config manager using the factory's default settings. Children are - * created via {@link #createChild} and share these resources. - */ - public static BigtableClientContext createForFactory(EnhancedBigtableStubSettings settings) - throws IOException { - return createInternal(settings, Tags.getTagger(), Stats.getStatsRecorder(), true); - } - - private static BigtableClientContext createInternal( EnhancedBigtableStubSettings settings, Tagger ocTagger, - StatsRecorder ocRecorder, - boolean isFactory) + StatsRecorder ocRecorder) throws IOException { ClientInfo clientInfo = ClientInfo.builder() @@ -231,7 +211,7 @@ private static BigtableClientContext createInternal( try { return new BigtableClientContext( - false, isFactory, shim, clientInfo, clientContext, metrics, executorProvider); + false, shim, clientInfo, clientContext, metrics, executorProvider); } catch (IOException | RuntimeException t) { metrics.close(); throw t; @@ -260,7 +240,6 @@ private static void configureGrpcOtel( private BigtableClientContext( boolean isChild, - boolean isFactory, Shim shim, ClientInfo clientInfo, ClientContext clientContext, @@ -268,7 +247,6 @@ private BigtableClientContext( ExecutorProvider backgroundExecutorProvider) throws IOException { this.isChild = isChild; - this.isFactory = isFactory; this.sessionShim = shim; this.clientInfo = clientInfo; @@ -296,9 +274,8 @@ public BigtableClientContext createChild(InstanceName instanceName, String appPr ClientInfo childInfo = clientInfo.toBuilder().setInstanceName(instanceName).setAppProfileId(appProfileId).build(); - Shim childShim; - if (isFactory && sessionShim instanceof ShimImpl) { - // Factory mode: create a lightweight child shim that shares the parent's pool and manager. + Shim childShim = sessionShim; + if (sessionShim instanceof ShimImpl) { ShimImpl parentShim = (ShimImpl) sessionShim; childShim = ShimImpl.createForFactoryChild( @@ -308,13 +285,10 @@ public BigtableClientContext createChild(InstanceName instanceName, String appPr parentShim.getChannelPool(), parentShim.getConfigManager(), parentShim.getFeatureFlags()); - } else { - // Legacy mode (sessions disabled): share the parent's DisabledShim. - childShim = sessionShim; } return new BigtableClientContext( - true, false, childShim, childInfo, clientContext, metrics, backgroundExecutorProvider); + true, childShim, childInfo, clientContext, metrics, backgroundExecutorProvider); } public void close() throws Exception { From 744e3850d3c6e6f468be218cd3cc89193d050936 Mon Sep 17 00:00:00 2001 From: Yuri Golobokov Date: Wed, 22 Jul 2026 05:21:03 +0000 Subject: [PATCH 3/3] fix lint --- .../bigtable/data/v2/internal/api/Client.java | 2 +- .../v2/internal/channels/ChannelPoolDpImpl.java | 17 ++++++++++------- .../data/v2/internal/channels/TenantKey.java | 4 ++-- .../data/v2/internal/compat/ShimImpl.java | 6 +++--- .../data/v2/stub/BigtableClientContext.java | 4 +--- 5 files changed, 17 insertions(+), 16 deletions(-) diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java index 2d840e75a30b..b13786f4c4ad 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java @@ -189,7 +189,7 @@ public static Client create(ClientSettings settings) throws IOException { Resource.createOwned(configManager, configManager::close), Resource.createOwned(backgroundExecutor, backgroundExecutor::shutdown), Resource.createOwned( - userCallbackExecutor, () -> shutdownAndAwait(userCallbackExecutor)), + userCallbackExecutor, () -> shutdownAndAwait(userCallbackExecutor)), settings.getChannelProvider()); } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/ChannelPoolDpImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/ChannelPoolDpImpl.java index 0983350c33cc..ae2dacb2d1aa 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/ChannelPoolDpImpl.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/ChannelPoolDpImpl.java @@ -80,8 +80,8 @@ *
  • 3a. Preferred: pick the ACTIVE channel whose last observed route for this tenant was * that AFE, choosing the least-loaded. *
  • 3b. Unobserved-tenant fallback: pick any ACTIVE channel with capacity (least-loaded). - * RLS may route to the target AFE or somewhere else; {@code routeObservations} will - * be updated on {@code onBeforeSessionStart} for future placements. + * RLS may route to the target AFE or somewhere else; {@code routeObservations} will be + * updated on {@code onBeforeSessionStart} for future placements. * *
  • Diversity mode (all known AFEs at cap, or none known): prefer an ACTIVE channel with fewer * than {@code softMaxPerGroup / 2} streams; otherwise add a new channel and absorb whatever @@ -119,8 +119,8 @@ public class ChannelPoolDpImpl implements ChannelPool { /** * Per-(channel, tenant) observation of the last AFE seen on {@code onBeforeSessionStart}. - * Self-healing: updated on every session open, so stale entries from ~1hr AFE drift are - * corrected automatically. + * Self-healing: updated on every session open, so stale entries from ~1hr AFE drift are corrected + * automatically. */ @GuardedBy("this") private final Map routeObservations = new HashMap<>(); @@ -444,7 +444,8 @@ private synchronized void serviceChannelsSafe() { target = minGroups; } - int activeCount = (int) channels.stream().filter(c -> c.state == ChannelWrapper.State.ACTIVE).count(); + int activeCount = + (int) channels.stream().filter(c -> c.state == ChannelWrapper.State.ACTIVE).count(); if (activeCount > target) { List candidates = pickDrainCandidates(activeCount - target); @@ -498,8 +499,10 @@ void dumpState() { return; } - long activeCount = channels.stream().filter(c -> c.state == ChannelWrapper.State.ACTIVE).count(); - long drainingCount = channels.stream().filter(c -> c.state == ChannelWrapper.State.DRAINING).count(); + long activeCount = + channels.stream().filter(c -> c.state == ChannelWrapper.State.ACTIVE).count(); + long drainingCount = + channels.stream().filter(c -> c.state == ChannelWrapper.State.DRAINING).count(); String s = sessionsPerAfeId.entrySet().stream() .sorted(Comparator.comparing(e -> e.getElement().toString())) diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/TenantKey.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/TenantKey.java index 70ec9735cf84..b5a82195e7f1 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/TenantKey.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/TenantKey.java @@ -20,8 +20,8 @@ import java.util.Objects; /** - * Identifies a single (instanceName, appProfileId) tenant within a shared channel pool. Used as - * the tenant dimension of a {@link RouteKey} in {@link ChannelPoolDpImpl}. + * Identifies a single (instanceName, appProfileId) tenant within a shared channel pool. Used as the + * tenant dimension of a {@link RouteKey} in {@link ChannelPoolDpImpl}. */ public final class TenantKey { /** Sentinel used when no tenant is stamped (single-tenant / sessions-disabled clients). */ diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java index cfa8b03732b3..7fe645d76b80 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java @@ -195,9 +195,9 @@ public ShimImpl(Resource configManagerResource, Clie } /** - * Creates a lightweight child shim for factory mode. Reuses the factory's already-started, - * shared channel pool and config manager. The pool and manager are not closed when this shim is - * closed — only the child's own session pools are cleaned up. + * Creates a lightweight child shim for factory mode. Reuses the factory's already-started, shared + * channel pool and config manager. The pool and manager are not closed when this shim is closed — + * only the child's own session pools are cleaned up. */ public static Shim createForFactoryChild( ClientInfo clientInfo, diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java index 34be9f39264f..b8eb4c123001 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java @@ -81,9 +81,7 @@ public static BigtableClientContext create(EnhancedBigtableStubSettings settings } public static BigtableClientContext create( - EnhancedBigtableStubSettings settings, - Tagger ocTagger, - StatsRecorder ocRecorder) + EnhancedBigtableStubSettings settings, Tagger ocTagger, StatsRecorder ocRecorder) throws IOException { ClientInfo clientInfo = ClientInfo.builder()