() {
+ @Override
+ public void onMessage(RespT msg) {
+ syncContext.execute(
+ () -> {
+ if (currentState != Active.this) {
+ LOG.log(
+ Level.FINE,
+ "Discarding response {0} because the attempt is no longer active.",
+ msg);
+ return;
+ }
+ tracer.onResponseReceived();
+ Stopwatch appTimer = Stopwatch.createStarted();
+ try {
+ listener.onMessage(msg);
+ } finally {
+ tracer.recordApplicationBlockingLatencies(appTimer.elapsed());
+ }
+ });
+ }
+
+ @Override
+ public void onClose(VRpcResult result) {
+ syncContext.execute(
+ () -> {
+ tracer.onAttemptFinish(result);
+ if (currentState != Active.this) {
+ LOG.log(
+ Level.FINE,
+ "Discarding server close with result {0} because the the attempt is no longer active.",
+ result);
+ return;
+ }
+ if (shouldRetry(result)) {
+ context = context.createForNextAttempt();
+ Duration retryDelay =
+ Optional.ofNullable(result.getRetryInfo())
+ .map(RetryInfo::getRetryDelay)
+ .orElse(Durations.ZERO);
+ if (Durations.compare(retryDelay, Durations.ZERO) > 0) {
+ Scheduled scheduled = new Scheduled(retryDelay);
+ onStateChange(scheduled);
+ } else {
+ onStateChange(new Idle());
+ }
+ return;
+ }
+
+ onStateChange(new Done(result));
+ });
+ }
+ });
+ }
+
+ @Override
+ public void onCancel(String reason, Throwable throwable) {
+ // attempt could be null if attemptFactory.get() throws an exception. In which case sync
+ // context uncaught exception handler will be called, which calls cancel on the current
+ // state before transition into done state.
+ if (attempt != null) {
+ attempt.cancel(reason, throwable);
+ }
+ }
+
+ boolean shouldRetry(VRpcResult result) {
+ // If the error has RetryInfo, it means it comes from the server and should
+ // be retried.
+ if (!result.getStatus().isOk()
+ && result.getRetryInfo() != null
+ && result.getRetryInfo().hasRetryDelay()) {
+ long retryDelay = Durations.toMillis(result.getRetryInfo().getRetryDelay());
+ // only schedule retry if there's still a chance for the request to succeed
+ return context.getOperationInfo().getDeadline().timeRemaining(TimeUnit.MILLISECONDS)
+ - retryDelay
+ > 1;
+ }
+ // Do not retry result that is explicitly rejected
+ if (result.getRejected()) {
+ return false;
+ }
+ // If the error didn't leave the client or failed in transport and is idempotent, we
+ // can retry up to 3 times.
+ boolean isRetryable =
+ (result.getState() == VRpcResult.State.UNCOMMITED)
+ || (context.getOperationInfo().isIdempotent()
+ && result.getState() == VRpcResult.State.TRANSPORT_FAILURE);
+ if (isRetryable && context.getOperationInfo().getAttemptNumber() < 3) {
+ return true;
+ }
+ return false;
+ }
+ }
+
+ class Scheduled extends State {
+ private final Duration retryDelay;
+ private SynchronizationContext.ScheduledHandle future;
+
+ Scheduled(Duration retryDelay) {
+ this.retryDelay = retryDelay;
+ }
+
+ @Override
+ public void onStart() {
+ try {
+ future =
+ syncContext.schedule(
+ () -> grpcContext.wrap(() -> onStateChange(new Idle())).run(),
+ Durations.toMillis(retryDelay),
+ TimeUnit.MILLISECONDS,
+ executor);
+ } catch (RejectedExecutionException e) {
+ onStateChange(
+ new Done(
+ VRpcResult.createRejectedError(
+ Status.CANCELLED
+ .withDescription(
+ "Executor shutting down, can't schedule operation for retry.")
+ .withCause(e))));
+ }
+ }
+
+ @Override
+ public void onCancel(String reason, Throwable throwable) {
+ // future can be null if schedule throws an exception that's not RejectedExecutionException.
+ // In which case sync context uncaught exception handler will be called, which calls cancel on
+ // the current
+ // state before transition into done state.
+ if (future != null && future.isPending()) {
+ future.cancel();
+ }
+ }
+ }
+
+ class Done extends State {
+
+ private final VRpcResult result;
+
+ Done(VRpcResult result) {
+ this.result = result;
+ }
+
+ @Override
+ public void onStart() {
+ if (!started) {
+ LOG.fine("operation is not started yet.");
+ return;
+ }
+ Stopwatch appTimer = Stopwatch.createStarted();
+ try {
+ listener.onClose(result);
+ } finally {
+ tracer.recordApplicationBlockingLatencies(appTimer.elapsed());
+ tracer.onOperationFinish(result);
+ }
+ }
+
+ @Override
+ public boolean isDone() {
+ return true;
+ }
+ }
+}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/VRpc.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/VRpc.java
new file mode 100644
index 000000000000..a29a51fd72c6
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/VRpc.java
@@ -0,0 +1,328 @@
+/*
+ * 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.middleware;
+
+import com.google.auto.value.AutoValue;
+import com.google.bigtable.v2.ClusterInformation;
+import com.google.bigtable.v2.ErrorResponse;
+import com.google.bigtable.v2.VirtualRpcResponse;
+import com.google.cloud.bigtable.data.v2.internal.csm.tracers.VRpcTracer;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Ticker;
+import com.google.common.collect.ImmutableList;
+import com.google.protobuf.Any;
+import com.google.rpc.RetryInfo;
+import io.grpc.Context;
+import io.grpc.Deadline;
+import io.grpc.Metadata;
+import io.grpc.Status;
+import io.grpc.protobuf.StatusProto;
+import java.time.Duration;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+import javax.annotation.Nullable;
+
+/**
+ * Internal Bigtable representation of an RPC.
+ *
+ * The primary intent is to model the RPC as a chain of composable middleware. This abstraction
+ * models Unary and ServerStreaming style RPCs by generalizing them as ServerStreaming.
+ */
+public interface VRpc {
+
+ /**
+ * Start the RPC and send the request message.
+ *
+ * This must be called before any other method on this object. The {@code listener} will be
+ * notified of the results. Please note, that the first response of the response stream will be
+ * delivered by default.
+ */
+ void start(ReqT req, VRpcCallContext ctx, VRpcListener listener);
+
+ /** Cancel a started RPC. This will be done by best effort. */
+ void cancel(@Nullable String message, @Nullable Throwable cause);
+
+ /**
+ * TBD - server streaming rpcs. This will be used to request more data. Unlike gRPC's request(n),
+ * starting a call will implicitly request the first message.
+ */
+ void requestNext();
+
+ /**
+ * Bigtable specific version of ClientCall.Listener - simplified to only cater to client unary
+ * RPCs. Methods on this class are guaranteed to be called sequentially.
+ */
+ interface VRpcListener {
+
+ /** Called when the next response message is the received. */
+ void onMessage(RespT msg);
+
+ /** Called when the vRPC completes. */
+ void onClose(VRpcResult result);
+ }
+
+ /** vRPC side band data. */
+ // TODO: set grpc deadline on callOptions or context
+ @AutoValue
+ abstract class VRpcCallContext {
+
+ private static final Logger logger = Logger.getLogger(VRpcCallContext.class.getName());
+
+ /** Retry related metadata. */
+ public abstract OperationInfo getOperationInfo();
+
+ /** The TraceId of the caller. */
+ public abstract String getTraceParent();
+
+ public abstract VRpcTracer getTracer();
+
+ // TODO: csm
+ // Clientside metrics instrument
+ // public abstract BigtableTracer getTracer();
+
+ public static VRpcCallContext create(
+ Deadline deadline, boolean isIdempotent, VRpcTracer tracer) {
+
+ Deadline grpcContextDeadline = Context.current().getDeadline();
+
+ Duration operationTimeout;
+ if (grpcContextDeadline != null && grpcContextDeadline.isBefore(deadline)) {
+ logger.log(
+ Level.FINE,
+ "grpc Context deadline {} is shorter than VrpcCallContext deadline {}",
+ new Object[] {grpcContextDeadline, deadline});
+ operationTimeout =
+ Duration.ofNanos(grpcContextDeadline.timeRemaining(TimeUnit.NANOSECONDS));
+ } else {
+ operationTimeout = Duration.ofNanos(deadline.timeRemaining(TimeUnit.NANOSECONDS));
+ }
+
+ return new AutoValue_VRpc_VRpcCallContext(
+ OperationInfo.create(operationTimeout, isIdempotent), "TODO", tracer);
+ }
+
+ public VRpcCallContext createForNextAttempt() {
+ return new AutoValue_VRpc_VRpcCallContext(
+ getOperationInfo().createForNextAttempt(), getTraceParent(), getTracer());
+ }
+ }
+
+ /** Sideband data required for retries and/or hedging. */
+ @AutoValue
+ abstract class OperationInfo {
+ /**
+ * Monotonically increasing number of retry attempt per operation. Starts with 0 for the
+ * original attempt.
+ */
+ public abstract int getAttemptNumber();
+
+ /**
+ * When the caller started the operation (ie. when the first attempt was sent). This is the
+ * machine time.
+ */
+ abstract long getOperationStartTickNs();
+
+ /**
+ * Original timeout of the operation (OperationStart + OperationTimeout = original deadline).
+ */
+ public abstract Duration getOperationTimeout();
+
+ /** If it's safe to retry the vRPC after its been commited. */
+ public abstract boolean isIdempotent();
+
+ abstract Ticker getTicker();
+
+ public Deadline getDeadline() {
+ return Deadline.after(
+ getOperationTimeout().toNanos() + getOperationStartTickNs() - getTicker().read(),
+ TimeUnit.NANOSECONDS);
+ }
+
+ /** Create a new copy of the {@link OperationInfo} for the next retry/heding attempt. */
+ public OperationInfo createForNextAttempt() {
+ return new AutoValue_VRpc_OperationInfo(
+ getAttemptNumber() + 1,
+ getOperationStartTickNs(),
+ getOperationTimeout(),
+ isIdempotent(),
+ getTicker());
+ }
+
+ /** Create new {@link OperationInfo} for the first attempt. */
+ public static OperationInfo create(Duration operationTimeout, boolean isIdempotent) {
+ return create(Ticker.systemTicker(), operationTimeout, isIdempotent);
+ }
+
+ @VisibleForTesting
+ static OperationInfo create(Ticker ticker, Duration operationTimeout, boolean isIdempotent) {
+ return new AutoValue_VRpc_OperationInfo(
+ 0, ticker.read(), operationTimeout, isIdempotent, ticker);
+ }
+ }
+
+ /** Represents the final state of a vRPC. */
+ @AutoValue
+ abstract class VRpcResult {
+ /**
+ * Describes how far the vRPC progressed prior to failure:
+ *
+ *
+ * - {@code UNCOMMITED}
+ *
- The vRPC never left the client
+ *
- {@code TRANSPORT_FAILURE}
+ *
- The vRPC failed due to transport
+ *
- {@code SERVER_RESULT}
+ *
- The vRPC result was explicitly communicated by the server
+ *
- {@code USER_FAILURE}
+ *
- The vRPC failed due to errors in users callback
+ *
+ */
+ public enum State {
+ UNCOMMITED,
+ TRANSPORT_FAILURE,
+ SERVER_RESULT,
+ USER_FAILURE
+ }
+
+ /**
+ * How far the vRPC progressed before reaching terminal state.
+ *
+ * @see State
+ */
+ public abstract State getState();
+
+ public abstract Status getStatus();
+
+ /** The status details. */
+ @Nullable
+ protected abstract List getDetails();
+
+ /**
+ * Side channel metadata for client side metrics - describes the cluster that handled the
+ * request.
+ */
+ @Nullable
+ public abstract ClusterInformation getClusterInfo();
+
+ /** Latency returned in SessionRequestStats. */
+ public abstract Duration getBackendLatency();
+
+ /** Server directed retries. */
+ // Server directed retries
+ @Nullable
+ public abstract RetryInfo getRetryInfo();
+
+ /** If the vrpc should be rejected for retry. */
+ public abstract boolean getRejected();
+
+ public static VRpcResult createUncommitedError(Status status) {
+ return new AutoValue_VRpc_VRpcResult(
+ State.UNCOMMITED, status, ImmutableList.of(), null, Duration.ZERO, null, false);
+ }
+
+ public static VRpcResult createRejectedError(Status status) {
+ return new AutoValue_VRpc_VRpcResult(
+ State.UNCOMMITED, status, ImmutableList.of(), null, Duration.ZERO, null, true);
+ }
+
+ /**
+ * vRPC failed because the underlying transport failed. We don't know if the vRPC made it to the
+ * nodes, so we must assume complete uncertainty.
+ */
+ public static VRpcResult createRemoteTransportError(Status transportStatus, Metadata trailers) {
+ Status status = Status.UNAVAILABLE.withDescription("vRPC failed due to transport error");
+
+ if (transportStatus.getCause() != null) {
+ status = status.withCause(transportStatus.getCause());
+ } else {
+ status =
+ status.augmentDescription(
+ String.format(
+ "Transport error: %s: %s",
+ transportStatus.getCode(), transportStatus.getDescription()));
+ }
+
+ List details =
+ StatusProto.fromStatusAndTrailers(transportStatus, trailers).getDetailsList();
+
+ return new AutoValue_VRpc_VRpcResult(
+ // TODO: need clusterInfo
+ State.TRANSPORT_FAILURE, status, details, null, Duration.ZERO, null, false);
+ }
+
+ /**
+ * vRPC failed because the underlying transport failed. We don't know if the the vRPC made it to
+ * the nodes, so we must assume complete uncertainty.
+ */
+ public static VRpcResult createLocalTransportError(Status status) {
+ return new AutoValue_VRpc_VRpcResult(
+ State.TRANSPORT_FAILURE, status, null, null, Duration.ZERO, null, false);
+ }
+
+ public static VRpcResult createUserError(Throwable throwable) {
+ return new AutoValue_VRpc_VRpcResult(
+ State.USER_FAILURE,
+ Status.CANCELLED
+ .withCause(throwable)
+ .withDescription("Cancelling RPC due to exception thrown by user callback"),
+ ImmutableList.of(),
+ null,
+ Duration.ZERO,
+ // TODO: use server retry delay if available
+ null,
+ true);
+ }
+
+ /** Wrap an OK from the server. */
+ public static VRpcResult createServerOk(VirtualRpcResponse r) {
+ return new AutoValue_VRpc_VRpcResult(
+ State.SERVER_RESULT,
+ Status.OK,
+ ImmutableList.of(),
+ normalizeClusterInfo(r.getClusterInfo()),
+ Duration.ofSeconds(
+ r.getStats().getBackendLatency().getSeconds(),
+ r.getStats().getBackendLatency().getNanos()),
+ null,
+ false);
+ }
+
+ /** Wrap the error response from the server. */
+ public static VRpcResult createServerError(ErrorResponse r) {
+ Status grpcStatus =
+ Status.fromCodeValue(r.getStatus().getCode()).withDescription(r.getStatus().getMessage());
+ return new AutoValue_VRpc_VRpcResult(
+ State.SERVER_RESULT,
+ grpcStatus,
+ r.getStatus().getDetailsList(),
+ normalizeClusterInfo(r.getClusterInfo()),
+ Duration.ZERO,
+ r.getRetryInfo(),
+ false);
+ }
+
+ @Nullable
+ private static ClusterInformation normalizeClusterInfo(ClusterInformation clusterInformation) {
+ if (ClusterInformation.getDefaultInstance().equals(clusterInformation)) {
+ return null;
+ }
+ return clusterInformation;
+ }
+ }
+}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/DynamicPicker.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/DynamicPicker.java
new file mode 100644
index 000000000000..664f8e64657a
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/DynamicPicker.java
@@ -0,0 +1,70 @@
+/*
+ * 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
+ *
+ * http://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.session;
+
+import com.google.bigtable.v2.LoadBalancingOptions;
+import com.google.bigtable.v2.SessionClientConfiguration;
+import com.google.cloud.bigtable.data.v2.internal.session.SessionList.SessionHandle;
+import java.util.Optional;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+/**
+ * A Picker that delegates to a concrete Picker implementation based on the current configuration.
+ */
+class DynamicPicker extends Picker {
+ private static final Logger LOGGER = Logger.getLogger(DynamicPicker.class.getName());
+ private final SessionList sessions;
+
+ private volatile Picker delegate;
+ private LoadBalancingOptions.LoadBalancingStrategyCase currentStrategy;
+
+ public DynamicPicker(
+ SessionList sessions, LoadBalancingOptions.LoadBalancingStrategyCase initialStrategy) {
+ this.sessions = sessions;
+ this.currentStrategy = initialStrategy;
+ this.delegate = createPicker(initialStrategy);
+ }
+
+ @Override
+ public Optional pickSession() {
+ return delegate.pickSession();
+ }
+
+ public void updateConfig(SessionClientConfiguration.SessionPoolConfiguration config) {
+ LoadBalancingOptions.LoadBalancingStrategyCase newStrategy =
+ config.getLoadBalancingOptions().getLoadBalancingStrategyCase();
+ if (newStrategy != currentStrategy) {
+ delegate = createPicker(newStrategy);
+ currentStrategy = newStrategy;
+ }
+ }
+
+ private Picker createPicker(LoadBalancingOptions.LoadBalancingStrategyCase strategy) {
+ switch (strategy) {
+ case RANDOM:
+ return new SimplePicker(sessions);
+ case LEAST_IN_FLIGHT:
+ return new LeastInFlightPicker(sessions);
+ default:
+ LOGGER.log(
+ Level.FINE, "got load balancing strategy {0} which was not implemented", strategy);
+ // TODO: implement PeakEwma
+ return new LeastInFlightPicker(sessions);
+ }
+ }
+}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/LeastInFlightPicker.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/LeastInFlightPicker.java
new file mode 100644
index 000000000000..8a3e195b1fdb
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/LeastInFlightPicker.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.session;
+
+import com.google.cloud.bigtable.data.v2.internal.session.SessionList.AfeHandle;
+import com.google.cloud.bigtable.data.v2.internal.session.SessionList.SessionHandle;
+import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.ThreadLocalRandom;
+
+/** Pick the AFE with the fewest in-flight requests. Experimental for now. */
+class LeastInFlightPicker extends Picker {
+ private final SessionList sessionList;
+
+ public LeastInFlightPicker(SessionList sessionList) {
+ this.sessionList = sessionList;
+ }
+
+ @Override
+ Optional pickSession() {
+ List readyAfes = sessionList.getAfesWithReadySessions();
+ int size = readyAfes.size();
+
+ if (size == 0) {
+ return Optional.empty();
+ }
+
+ ThreadLocalRandom random = ThreadLocalRandom.current();
+ AfeHandle selected = readyAfes.get(random.nextInt(size));
+
+ // If we have options, pick a second candidate and keep the better one
+ if (size > 1) {
+ AfeHandle candidate2 = readyAfes.get(random.nextInt(size));
+ if (candidate2.getNumOutstanding() < selected.getNumOutstanding()) {
+ selected = candidate2;
+ }
+ }
+
+ return sessionList.checkoutSession(selected);
+ }
+}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/Picker.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/Picker.java
new file mode 100644
index 000000000000..ce4bc5a68c98
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/Picker.java
@@ -0,0 +1,24 @@
+/*
+ * 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.session;
+
+import com.google.cloud.bigtable.data.v2.internal.session.SessionList.SessionHandle;
+import java.util.Optional;
+
+abstract class Picker {
+ abstract Optional pickSession();
+}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/PoolSizer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/PoolSizer.java
new file mode 100644
index 000000000000..01a83a661486
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/PoolSizer.java
@@ -0,0 +1,108 @@
+/*
+ * 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.session;
+
+import com.google.bigtable.v2.GoAwayResponse;
+import com.google.bigtable.v2.SessionClientConfiguration;
+import com.google.cloud.bigtable.data.v2.internal.session.SessionList.PoolStats;
+import com.google.rpc.Status;
+import javax.annotation.CheckReturnValue;
+
+class PoolSizer {
+
+ // Fraction of idle sessions to keep in order to manage an increase in requests-in-flight. For
+ // example, a headroom of 50% will keep enough sessions to deal with a 50% increase in QPS or a
+ // 50% increase in latency (since that also increases requests-in-flight).
+ private volatile float idlesSessionHeadRoom;
+ private volatile int minIdleSessions;
+ private volatile int maxIdleSessions;
+
+ private final PoolStats stats;
+ private final Sized pendingRpcs;
+
+ // This resolves cold start overprovisioning issue.
+ // Session creation sometimes involves creating new connection and warming the backend.
+ // So it takes time, but accumulated vRPCs may be processed rather quickly as simple
+ // point reads are expected to complete within 1-4ms.
+ // We assume that it is okay to collect 10 pending calls per starting session so that after
+ // the session is ready it takes 10-40ms to process all pending calls.
+ private volatile int pendingVRpcsPerSession;
+
+ interface Sized {
+ int getSize();
+ }
+
+ PoolSizer(
+ PoolStats stats,
+ Sized pendingRpcs,
+ SessionClientConfiguration.SessionPoolConfiguration poolConfig) {
+ this.stats = stats;
+ this.pendingRpcs = pendingRpcs;
+
+ this.idlesSessionHeadRoom = poolConfig.getHeadroom();
+ this.minIdleSessions = poolConfig.getMinSessionCount();
+ this.maxIdleSessions = poolConfig.getMaxSessionCount();
+ this.pendingVRpcsPerSession = poolConfig.getNewSessionQueueLength();
+ }
+
+ void updateConfig(SessionClientConfiguration.SessionPoolConfiguration poolConfiguration) {
+ this.idlesSessionHeadRoom = poolConfiguration.getHeadroom();
+ this.minIdleSessions = poolConfiguration.getMinSessionCount();
+ this.maxIdleSessions = poolConfiguration.getMaxSessionCount();
+ this.pendingVRpcsPerSession = poolConfiguration.getNewSessionQueueLength();
+ }
+
+ public int getScaleDelta() {
+ // Assume each session handles 1 RPC at a time. This should be revisited if sessions get
+ // support for multiplexing.
+ int effectivePending = (int) Math.ceil((float) pendingRpcs.getSize() / pendingVRpcsPerSession);
+ int sessionsInUse = effectivePending + stats.getInUseCount();
+ int unboundedDesiredIdleSessions = (int) Math.ceil(sessionsInUse * idlesSessionHeadRoom);
+ int desiredIdleSessions =
+ Math.max(Math.min(unboundedDesiredIdleSessions, maxIdleSessions), minIdleSessions);
+
+ int desiredCapacity = sessionsInUse + desiredIdleSessions;
+ int eventualCapacity = stats.getExpectedCapacity();
+ int immediateCapacity = eventualCapacity - stats.getStartingCount();
+
+ if (desiredCapacity < immediateCapacity) {
+ return desiredCapacity - immediateCapacity;
+ } else if (desiredCapacity > eventualCapacity) {
+ return desiredCapacity - eventualCapacity;
+ } else {
+ return 0;
+ }
+ }
+
+ /** Returns true if the session should be replaced */
+ @CheckReturnValue
+ public boolean handleGoAway(GoAwayResponse msg) {
+ return getScaleDelta() >= 0;
+ }
+
+ /** Returns true if the closed session should be replaced */
+ @CheckReturnValue
+ boolean handleSessionClose(Status statusProto) {
+ return getScaleDelta() >= 0;
+ }
+
+ /** Returns true if a new session should be added. */
+ @CheckReturnValue
+ boolean handleNewCall() {
+ return getScaleDelta() > 0;
+ }
+}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/Session.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/Session.java
new file mode 100644
index 000000000000..35f8f98bc571
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/Session.java
@@ -0,0 +1,124 @@
+/*
+ * 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.session;
+
+import com.google.auto.value.AutoValue;
+import com.google.bigtable.v2.CloseSessionRequest;
+import com.google.bigtable.v2.GoAwayResponse;
+import com.google.bigtable.v2.OpenSessionRequest;
+import com.google.bigtable.v2.OpenSessionResponse;
+import com.google.bigtable.v2.PeerInfo;
+import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc;
+import com.google.protobuf.Message;
+import io.grpc.Metadata;
+import io.grpc.Status;
+import java.time.Instant;
+
+/**
+ * A Bigtable abstraction over a bidirectional {@link io.grpc.ClientCall} to treat it as a stateful
+ * vRPC transport.
+ */
+public interface Session {
+ /** State transitions of Session, will happen in defined order */
+ enum SessionState {
+ NEW(0),
+ STARTING(1),
+ READY(2),
+ CLOSING(3),
+ WAIT_SERVER_CLOSE(4),
+ CLOSED(5);
+
+ final int phase;
+
+ SessionState(int phase) {
+ this.phase = phase;
+ }
+ }
+
+ @AutoValue
+ abstract class OpenParams {
+ abstract Metadata metadata();
+
+ abstract OpenSessionRequest request();
+
+ static OpenParams create(Metadata metadata, OpenSessionRequest request) {
+ return new AutoValue_Session_OpenParams(metadata, request);
+ }
+
+ OpenParams withResetConsecutiveAttempt() {
+ return new AutoValue_Session_OpenParams(
+ metadata(), request().toBuilder().setConsecutiveFailedConnectionAttempts(0).build());
+ }
+
+ OpenParams withIncrementedAttempts() {
+ return new AutoValue_Session_OpenParams(
+ metadata(),
+ request().toBuilder()
+ .setConsecutiveFailedConnectionAttempts(
+ request().getConsecutiveFailedConnectionAttempts() + 1)
+ .build());
+ }
+ }
+
+ SessionState getState();
+
+ Instant getLastStateChange();
+
+ OpenParams getOpenParams();
+
+ boolean isOpenParamsUpdated();
+
+ PeerInfo getPeerInfo();
+
+ String getLogName();
+
+ Instant getNextHeartbeat();
+
+ /**
+ * Start the session by sending the opening message.
+ *
+ * This must be the first method called on the {@link Session}. It will try setup a session
+ * using the req (for the server) and headers (for rls).
+ */
+ void start(OpenSessionRequest req, Metadata headers, Listener sessionListener);
+
+ /**
+ * Caller instructed close of the Session. Will immediately close the session and abort all
+ * outstanding vRPCs.
+ */
+ void close(CloseSessionRequest req);
+
+ /** Force close a session. */
+ void forceClose(CloseSessionRequest reason);
+
+ /**
+ * Start a new vRPC. This method can only be called after {@link
+ * Listener#onReady(OpenSessionResponse)} has been notified.
+ */
+ VRpc newCall(
+ VRpcDescriptor descriptor) throws IllegalStateException;
+
+ /** Callback for Session lifecycyle transitions. Methods are called sequentially. */
+ interface Listener {
+
+ void onReady(OpenSessionResponse msg);
+
+ void onGoAway(GoAwayResponse msg);
+
+ void onClose(SessionState prevState, Status status, Metadata trailers);
+ }
+}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionCreationBudget.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionCreationBudget.java
new file mode 100644
index 000000000000..bfa8dd2213a9
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionCreationBudget.java
@@ -0,0 +1,138 @@
+/*
+ * 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.session;
+
+import com.google.bigtable.v2.SessionClientConfiguration;
+import com.google.common.annotations.VisibleForTesting;
+import java.time.Clock;
+import java.time.Duration;
+import java.time.Instant;
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+import javax.annotation.concurrent.NotThreadSafe;
+
+@NotThreadSafe
+class SessionCreationBudget {
+ private static final Logger DEFAULT_LOGGER =
+ Logger.getLogger(SessionCreationBudget.class.getName());
+
+ // TODO: add a metric for session can't be created because there's no budget
+ private volatile int maxConcurrentRequest;
+ private volatile Duration penalty;
+ private final Clock clock;
+
+ private int concurrentRequests = 0;
+
+ // Whenever a session creation failed, add the time to the list
+ private final List delayedCreationTokens = new ArrayList<>();
+
+ static SessionCreationBudget create(int max, Duration penalty) {
+ return new SessionCreationBudget(max, penalty, Clock.systemUTC());
+ }
+
+ @VisibleForTesting
+ SessionCreationBudget(int max, Duration penalty, Clock clock) {
+ this.maxConcurrentRequest = max;
+ this.penalty = penalty;
+ this.clock = clock;
+ }
+
+ Instant getNextAvailableBudget() {
+ if (concurrentRequests < maxConcurrentRequest) {
+ return Instant.now();
+ }
+
+ if (delayedCreationTokens.isEmpty()) {
+ return Instant.now();
+ }
+
+ return delayedCreationTokens.get(0);
+ }
+
+ boolean tryReserveSession() {
+ sanityCheck();
+
+ if (concurrentRequests == maxConcurrentRequest) {
+ drainCreationFailures();
+ }
+
+ if (concurrentRequests == maxConcurrentRequest) {
+ return false;
+ }
+
+ concurrentRequests++;
+ return true;
+ }
+
+ void onSessionCreationFailure() {
+ delayedCreationTokens.add(Instant.now(clock).plus(penalty));
+ }
+
+ void onSessionCreationSuccess() {
+ concurrentRequests--;
+ }
+
+ private void drainCreationFailures() {
+ Instant now = Instant.now(clock);
+ Iterator iter = delayedCreationTokens.listIterator();
+ while (iter.hasNext()) {
+ if (iter.next().isBefore(now)) {
+ concurrentRequests--;
+ iter.remove();
+ } else {
+ // The list should be roughly sorted. Exit early when we encounter
+ // something expires later.
+ break;
+ }
+ }
+ }
+
+ private void sanityCheck() {
+ // This could happen if the budget is updated
+ if (concurrentRequests < 0) {
+ DEFAULT_LOGGER.log(
+ Level.FINE,
+ "concurrent request can't be negative: {0}. Resetting it to 0.",
+ concurrentRequests);
+ concurrentRequests = 0;
+ }
+ if (concurrentRequests > maxConcurrentRequest) {
+ DEFAULT_LOGGER.log(
+ Level.FINE,
+ "Concurrent requests out of range: {0}. Resetting it to max.",
+ concurrentRequests);
+ concurrentRequests = maxConcurrentRequest;
+ }
+ }
+
+ public int getMaxConcurrentRequest() {
+ return maxConcurrentRequest;
+ }
+
+ void updateConfig(SessionClientConfiguration.SessionPoolConfiguration config) {
+ int oldBudget = this.maxConcurrentRequest;
+ this.maxConcurrentRequest = config.getNewSessionCreationBudget();
+ this.penalty = SessionUtil.toJavaDuration(config.getNewSessionCreationPenalty());
+ DEFAULT_LOGGER.log(
+ Level.FINE,
+ "updated session creation budget from {0} to {1}.",
+ new Object[] {oldBudget, config.getNewSessionCreationBudget()});
+ }
+}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionFactory.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionFactory.java
new file mode 100644
index 000000000000..20246775b6cb
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionFactory.java
@@ -0,0 +1,44 @@
+/*
+ * 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.session;
+
+import com.google.bigtable.v2.SessionRequest;
+import com.google.bigtable.v2.SessionResponse;
+import com.google.cloud.bigtable.data.v2.internal.channels.ChannelPool;
+import com.google.cloud.bigtable.data.v2.internal.channels.SessionStream;
+import io.grpc.CallOptions;
+import io.grpc.MethodDescriptor;
+
+/** Wrapper around a channel to centralize per call session customizations. */
+public final class SessionFactory {
+ private final ChannelPool channelPool;
+ private final MethodDescriptor methodDescriptor;
+ private final CallOptions callOptions;
+
+ public SessionFactory(
+ ChannelPool channelPool,
+ MethodDescriptor methodDescriptor,
+ CallOptions callOptions) {
+ this.channelPool = channelPool;
+ this.methodDescriptor = methodDescriptor;
+ this.callOptions = callOptions;
+ }
+
+ public SessionStream createNew() {
+ return channelPool.newStream(methodDescriptor, callOptions);
+ }
+}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImpl.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImpl.java
new file mode 100644
index 000000000000..9cd58da876d2
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImpl.java
@@ -0,0 +1,689 @@
+/*
+ * 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.session;
+
+import com.google.bigtable.v2.CloseSessionRequest;
+import com.google.bigtable.v2.CloseSessionRequest.CloseSessionReason;
+import com.google.bigtable.v2.ErrorResponse;
+import com.google.bigtable.v2.GoAwayResponse;
+import com.google.bigtable.v2.HeartbeatResponse;
+import com.google.bigtable.v2.OpenSessionRequest;
+import com.google.bigtable.v2.OpenSessionResponse;
+import com.google.bigtable.v2.PeerInfo;
+import com.google.bigtable.v2.SessionParametersResponse;
+import com.google.bigtable.v2.SessionRefreshConfig;
+import com.google.bigtable.v2.SessionRequest;
+import com.google.bigtable.v2.SessionResponse;
+import com.google.bigtable.v2.TelemetryConfiguration;
+import com.google.bigtable.v2.VirtualRpcRequest;
+import com.google.bigtable.v2.VirtualRpcResponse;
+import com.google.cloud.bigtable.data.v2.internal.channels.SessionStream;
+import com.google.cloud.bigtable.data.v2.internal.csm.Metrics;
+import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DebugTagTracer;
+import com.google.cloud.bigtable.data.v2.internal.csm.tracers.SessionTracer;
+import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc;
+import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc.VRpcResult;
+import com.google.cloud.bigtable.data.v2.internal.session.VRpcImpl.VRpcSessionApi;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.protobuf.Message;
+import com.google.protobuf.TextFormat;
+import com.google.protobuf.util.Durations;
+import io.grpc.Metadata;
+import io.grpc.Status;
+import java.time.Clock;
+import java.time.Duration;
+import java.time.Instant;
+import java.util.Locale;
+import java.util.Optional;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+import javax.annotation.Nullable;
+import javax.annotation.concurrent.GuardedBy;
+
+/** Wraps a Bidi ClientCall and layers session semantics on top. */
+@VisibleForTesting
+public class SessionImpl implements Session, VRpcSessionApi {
+ private static final Logger DEFAULT_LOGGER = Logger.getLogger(SessionImpl.class.getName());
+ private Logger logger = DEFAULT_LOGGER;
+
+ private static final SessionParametersResponse DEFAULT_SESSION_PARAMS =
+ SessionParametersResponse.newBuilder().setKeepAlive(Durations.fromMillis(100)).build();
+
+ static final Duration HEARTBEAT_CHECK_INTERVAL =
+ Duration.ofMillis(Durations.toMillis(DEFAULT_SESSION_PARAMS.getKeepAlive()));
+
+ @VisibleForTesting
+ // A time in the future to skip heartbeat checks when there's no active vRPCs on the session
+ static final Duration FUTURE_TIME = Duration.ofMinutes(30);
+
+ /*
+ * This lock should be mostly uncontended - all access should be naturally interleaved. Contention
+ * can only really happen when an unsolicited gRPC control message (ie GOAWAY) arrives at the same
+ * time as newCall or cancel.
+ * TODO: Contention will increase when multiplexing is implemented.
+ */
+ private final Object lock = new Object();
+
+ private final Clock clock;
+
+ private final SessionTracer tracer;
+ private final DebugTagTracer debugTagTracer;
+
+ private final SessionInfo info;
+
+ @GuardedBy("lock")
+ private final SessionStream stream;
+
+ @GuardedBy("lock")
+ private SessionState state = SessionState.NEW;
+
+ @GuardedBy("lock")
+ private Instant lastStateChangedAt;
+
+ private Listener sessionListener;
+
+ private volatile OpenParams openParams;
+
+ private volatile boolean openParamsUpdated;
+
+ @Nullable private CloseSessionRequest closeReason = null;
+
+ @GuardedBy("lock")
+ private long nextRpcId = 1;
+
+ // TODO: replace with a map when implementing multiplexing
+ @GuardedBy("lock")
+ private VRpcImpl currentRpc = null;
+
+ @GuardedBy("lock")
+ private VRpcResult currentCancel = null;
+
+ private SessionParametersResponse sessionParameters = DEFAULT_SESSION_PARAMS;
+ private Duration heartbeatInterval =
+ Duration.ofMillis(Durations.toMillis(sessionParameters.getKeepAlive()));
+
+ private volatile Instant nextHeartbeat;
+
+ public SessionImpl(
+ Metrics metrics, SessionPoolInfo poolInfo, long sessionNum, SessionStream stream) {
+ this(metrics, Clock.systemUTC(), poolInfo, sessionNum, stream);
+ }
+
+ SessionImpl(
+ Metrics metrics,
+ Clock clock,
+ SessionPoolInfo poolInfo,
+ long sessionNum,
+ SessionStream stream) {
+ this.clock = clock;
+ this.info = SessionInfo.create(poolInfo, sessionNum);
+ this.stream = stream;
+ this.tracer = metrics.newSessionTracer(poolInfo);
+ this.debugTagTracer = metrics.getDebugTagTracer();
+ this.nextHeartbeat = clock.instant().plus(FUTURE_TIME);
+ this.openParamsUpdated = false;
+ }
+
+ @Override
+ public SessionState getState() {
+ synchronized (lock) {
+ return state;
+ }
+ }
+
+ @Override
+ public Instant getLastStateChange() {
+ synchronized (lock) {
+ return lastStateChangedAt;
+ }
+ }
+
+ @Override
+ public OpenParams getOpenParams() {
+ return openParams;
+ }
+
+ @Override
+ public boolean isOpenParamsUpdated() {
+ return openParamsUpdated;
+ }
+
+ @Override
+ public Instant getNextHeartbeat() {
+ return nextHeartbeat;
+ }
+
+ @Override
+ public PeerInfo getPeerInfo() {
+ // This lock might not be necessary, its populated once on a gRPC callback which should
+ // establish a happens before relationship. However access to the underlying stream is guarded
+ // with errorprone, so sync block is required to get around the lint.
+ // TODO: consider removing the sync block
+ synchronized (lock) {
+ return stream.getPeerInfo();
+ }
+ }
+
+ @Override
+ public String getLogName() {
+ return info.getLogName();
+ }
+
+ @Override
+ public void forceClose(CloseSessionRequest closeReason) {
+ synchronized (lock) {
+ debugTagTracer.checkPrecondition(
+ state != SessionState.NEW,
+ "session_force_close_wrong_state",
+ "Tried to forceClose an unstarted session %s in state %s",
+ info.getLogName(),
+ state);
+
+ if (state == SessionState.CLOSED) {
+ return;
+ }
+
+ updateState(SessionState.WAIT_SERVER_CLOSE);
+ this.closeReason = closeReason;
+
+ // Not sending the CloseSessionRequest because cancel() will just drop it
+ stream.forceClose(closeReason.getDescription(), null);
+ // Listeners will be notified by dispatchStreamClosed
+ }
+ }
+
+ @Override
+ public void start(OpenSessionRequest req, Metadata headers, Listener sessionListener) {
+ synchronized (lock) {
+ debugTagTracer.checkPrecondition(
+ state == SessionState.NEW,
+ "session_start_wrong_state",
+ "Tried to start a started session, current state: %s",
+ state);
+
+ logger.fine(String.format("Starting session %s", info.getLogName()));
+ tracer.onStart();
+
+ updateState(SessionState.STARTING);
+ openParams = OpenParams.create(headers, req);
+ this.sessionListener = sessionListener;
+
+ SessionRequest wrappedReq = SessionRequest.newBuilder().setOpenSession(req).build();
+ stream.start(
+ new SessionStream.Listener() {
+ @Override
+ public void onBeforeSessionStart(PeerInfo peerInfo) {}
+
+ @Override
+ public void onMessage(SessionResponse message) {
+ dispatchResponseMessage(message);
+ }
+
+ @Override
+ public void onClose(Status status, Metadata trailers) {
+ dispatchStreamClosed(status, trailers);
+ }
+ },
+ headers);
+
+ stream.sendMessage(wrappedReq);
+ }
+ }
+
+ @Override
+ public void close(CloseSessionRequest req) {
+ logger.fine(String.format("Closing session %s for reason: %s", info.getLogName(), req));
+
+ synchronized (lock) {
+ // Throw an exception because this is a bug and we dont have a listener
+ debugTagTracer.checkPrecondition(
+ state != SessionState.NEW,
+ "session_close_wrong_state",
+ "Session error: Caller tried to close session %s before starting it with the reason: %s",
+ info.getLogName(),
+ req);
+
+ // Multiple close is a no-op
+ if (state.phase >= SessionState.CLOSING.phase) {
+ logger.fine(
+ String.format(
+ "Session error: Caller tried to close a session %s that is %s for reason: %s",
+ info.getLogName(), state, req));
+ return;
+ }
+
+ closeReason = req;
+ updateState(SessionState.CLOSING);
+
+ if (currentRpc == null) {
+ startGracefulClose();
+ }
+ }
+ }
+
+ /** Wraps the flow of closing a session. */
+ @GuardedBy("lock")
+ private void startGracefulClose() {
+ debugTagTracer.checkPrecondition(
+ state == SessionState.CLOSING,
+ "session_graceful_close_wrong_state",
+ "Session error: %s tried to actuate session closing when not in the correct state. State: %s",
+ info.getLogName(),
+ state);
+
+ // TODO: send metrics
+ updateState(SessionState.WAIT_SERVER_CLOSE);
+
+ // Should never happen
+ if (closeReason == null) {
+ debugTagTracer.record(TelemetryConfiguration.Level.WARN, "session_close_no_reason");
+ logger.log(
+ Level.WARNING,
+ String.format("%s graceful shutdown started without a reason", info.getLogName()),
+ new IllegalStateException("Tried to close a session without a reason"));
+ // Synthesize a reason so that we let the server know of the problem instead
+ closeReason =
+ CloseSessionRequest.newBuilder()
+ .setReason(CloseSessionReason.CLOSE_SESSION_REASON_ERROR)
+ .setDescription("Started graceful shutdown close without a reason set")
+ .build();
+ }
+ stream.sendMessage(SessionRequest.newBuilder().setCloseSession(closeReason).build());
+ // TODO: remove this after the server is updated
+ stream.halfClose();
+ }
+
+ @Override
+ public
+ VRpc newCall(VRpcDescriptor