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..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 @@ -74,10 +74,6 @@ 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()); ClientOperationSettings perOpSettings = defaultSettings.getStubSettings().getPerOpSettings(); 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 7f3c50fa68e6..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 @@ -100,7 +100,7 @@ public class Client implements AutoCloseable { private final BigtableTimer sessionTimer; private final CallOptions defaultCallOptions; - private final ChannelPool channelPool; + private final Resource channelPool; private final Resource metrics; private final Resource configManager; @@ -185,22 +185,26 @@ 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( - userCallbackExecutor, () -> shutdownAndAwait(userCallbackExecutor))); + userCallbackExecutor, () -> shutdownAndAwait(userCallbackExecutor)), + 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 userCallbackExecutor) + Resource userCallbackExecutor, + ChannelProvider channelProvider) throws IOException { this.featureFlags = featureFlags; this.clientInfo = clientInfo; @@ -224,13 +228,38 @@ 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 and other resources should be created with + * Resource.createShared() and closed when the factory closes. + */ + public Client( + FeatureFlags featureFlags, + ClientInfo clientInfo, + Resource metrics, + Resource configManager, + Resource bgExecutor, + Resource userCallbackExecutor, + Resource sharedChannelPool) { + this.featureFlags = featureFlags; + this.clientInfo = clientInfo; + this.metrics = metrics; + this.configManager = configManager; + this.backgroundExecutor = bgExecutor; + this.userCallbackExecutor = userCallbackExecutor; + this.channelPool = sharedChannelPool; + this.sessionTimer = new HashedWheelTimer("bigtable-session-timer"); + defaultCallOptions = CallOptions.DEFAULT; } @Override @@ -329,7 +358,7 @@ public TableAsync openTableAsync(String tableId, Permission permission) { featureFlags, clientInfo, configManager.get(), - channelPool, + channelPool.get(), defaultCallOptions, tableId, permission, @@ -353,7 +382,7 @@ public AuthorizedViewAsync openAuthorizedViewAsync( featureFlags, clientInfo, configManager.get(), - channelPool, + channelPool.get(), defaultCallOptions, tableId, viewId, @@ -378,7 +407,7 @@ public MaterializedViewAsync openMaterializedViewAsync( featureFlags, clientInfo, configManager.get(), - channelPool, + channelPool.get(), defaultCallOptions, viewId, permission, @@ -391,6 +420,21 @@ public MaterializedViewAsync openMaterializedViewAsync( } } + /** 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; + } + + /** Returns the underlying user callback executor (e.g. for sharing with factory children). */ + public Executor getUserCallbackExecutor() { + return userCallbackExecutor.get(); + } + public static class Resource { private final T value; private final 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 5565e9687a3d..4892a0cfb738 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; @@ -66,6 +68,12 @@ static TableBase createAndStart( Executor backgroundExecutor, Executor userCallbackExecutor) { + // 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, @@ -73,7 +81,7 @@ static TableBase createAndStart( clientInfo, configManager, channelPool, - callOptions, + stamped, sessionDescriptor, sessionPoolName, timer, 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 402151676c7c..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 @@ -37,19 +37,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; @@ -58,9 +57,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()); @@ -71,6 +97,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; @@ -84,16 +113,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(); @@ -165,8 +201,6 @@ public ChannelPoolDpImpl( this.channelSupplier = channelSupplier; this.executor = executor; updateConfig(config); - channelGroups = new ArrayList<>(); - startingGroup = new AfeChannelGroup(null); this.debugTagTracer = debugTagTracer; } @@ -191,9 +225,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 @@ -203,48 +235,65 @@ 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++; // DirectExecutor: gRPC/Netty delivers SessionStream.Listener callbacks directly on the @@ -254,7 +303,7 @@ public synchronized SessionStream newStream( desc, callOptions.withExecutor(MoreExecutors.directExecutor())); 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 @@ -265,10 +314,13 @@ 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); + if (!fw.channel.isShutdown()) { + routeObservations.put(new RouteKey(fw.id, tenant), afeId); + } + fw.lastUsedAt = Instant.now(clock); } super.onBeforeSessionStart(peerInfo); } @@ -279,9 +331,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); } @@ -296,63 +348,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); @@ -395,12 +419,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(); } @@ -414,59 +433,32 @@ 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(); - } - } - - // 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 target = (int) Math.ceil((totalStreams * HEADROOM) / softMaxPerGroup); + if (target > maxGroups) { + target = maxGroups; + } else if (target < minGroups) { + target = minGroups; } - // 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); + int activeCount = + (int) channels.stream().filter(c -> c.state == ChannelWrapper.State.ACTIVE).count(); + + 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(); } } @@ -475,6 +467,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); } @@ -489,8 +499,10 @@ 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())) @@ -499,12 +511,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: " @@ -523,29 +533,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; + } + + @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); + } - public AfeChannelGroup(AfeId afeId) { - this.afeId = afeId; - channels = new ArrayDeque<>(); - numStreams = 0; + @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..b5a82195e7f1 --- /dev/null +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/TenantKey.java @@ -0,0 +1,55 @@ +/* + * 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; + +/** + * 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(InstanceName.of("", ""), ""); + + private final InstanceName instanceName; + private final String appProfileId; + + public TenantKey(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 e51f65b828a9..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 @@ -33,6 +33,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; @@ -78,6 +79,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; @@ -174,23 +176,73 @@ public static Shim create( new Client( clientChannelProvider.updateFeatureFlags(featureFlags), clientInfo, - clientChannelProvider, Resource.createShared(metrics), Resource.createShared(configManager), Resource.createShared(bgExecutor), - userCallbackExecutor); + userCallbackExecutor, + 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, + Executor userCallbackExecutor, + 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(userCallbackExecutor), + 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(); + } + + public Executor getUserCallbackExecutor() { + return client.getUserCallbackExecutor(); + } + // If the user configured an executor — either via InstantiatingGrpcChannelProvider#setExecutor // or the legacy StubSettings#setExecutorProvider — reuse it for user-callback dispatch so the // transport pool and callback pool are the same. Ownership: transport-set is borrowed by @@ -291,7 +343,7 @@ private static boolean checkDirectAccessAvailable( @Override public void close() { client.close(); - configManager.close(); + 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 a74c15613e0f..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 @@ -42,7 +42,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; @@ -268,27 +267,36 @@ 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 = sessionShim; + if (sessionShim instanceof ShimImpl) { + ShimImpl parentShim = (ShimImpl) sessionShim; + childShim = + ShimImpl.createForFactoryChild( + childInfo, + metrics, + backgroundExecutorProvider.getExecutor(), + parentShim.getUserCallbackExecutor(), + parentShim.getChannelPool(), + parentShim.getConfigManager(), + parentShim.getFeatureFlags()); + } return new BigtableClientContext( - true, - sessionShim, - clientInfo.toBuilder().setInstanceName(instanceName).setAppProfileId(appProfileId).build(), - clientContext, - metrics, - backgroundExecutorProvider); + true, 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,