fix: lift shim loader throws to failed futures via SessionPoolMap

ReadRowShim and MutateRowShim cache per-target handles in a Guava
LoadingCache whose loader calls Client.openTableAsync (and friends).
After the Client-close hardening, that loader throws
IllegalStateException post-close — and getUnchecked wraps it in
UncheckedExecutionException, so callers see an unchecked exception
thrown out of readRow / mutateRow instead of a failed CompletableFuture
as the surface contract promises.

Introduce SessionPoolMap, a small wrapper over the existing
Util.createSessionMap cache that owns the conversion: get() unwraps
UncheckedExecutionException to surface the original cause, and apply()
converts a loader throw into a failed CompletableFuture for the async
call paths.

Replace the inline LoadingCache fields in both shim ops files.
Covered by a new focused unit test.

Author: igorbernstein2

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowShim.java

@@ -22,28 +22,26 @@
 import com.google.cloud.bigtable.data.v2.internal.api.Client;
 import com.google.cloud.bigtable.data.v2.internal.api.TableAsync;
 import com.google.cloud.bigtable.data.v2.internal.compat.ShimImpl;
-import com.google.cloud.bigtable.data.v2.internal.compat.Util;
 import com.google.cloud.bigtable.data.v2.internal.session.SessionPool;
 import com.google.cloud.bigtable.data.v2.models.AuthorizedViewId;
 import com.google.cloud.bigtable.data.v2.models.RowMutation;
 import com.google.cloud.bigtable.data.v2.models.TableId;
 import com.google.cloud.bigtable.data.v2.models.TargetId;
-import com.google.common.cache.LoadingCache;
 import io.grpc.Deadline;
 import java.io.IOException;
 import java.util.concurrent.CompletableFuture;
 
 public class MutateRowShim implements UnaryShim<RowMutation, Void> {
 
-  private final LoadingCache<TableId, TableAsync> tables;
-  private final LoadingCache<AuthorizedViewId, AuthorizedViewAsync> authViews;
+  private final SessionPoolMap<TableId, TableAsync> tables;
+  private final SessionPoolMap<AuthorizedViewId, AuthorizedViewAsync> authViews;
 
   public MutateRowShim(Client client) {
     tables =
-        Util.createSessionMap(
+        new SessionPoolMap<>(
             k -> client.openTableAsync(k.getTableId(), Permission.PERMISSION_WRITE));
     authViews =
-        Util.createSessionMap(
+        new SessionPoolMap<>(
             k ->
                 client.openAuthorizedViewAsync(
                     k.getTableId(),
@@ -63,9 +61,9 @@ public boolean supports(RowMutation request) {
     SessionPool<?> pool;
     // TODO: avoid double lookup
     if (targetId instanceof TableId) {
-      pool = tables.getUnchecked((TableId) targetId).getSessionPool();
+      pool = tables.get((TableId) targetId).getSessionPool();
     } else if (targetId instanceof AuthorizedViewId) {
-      pool = authViews.getUnchecked((AuthorizedViewId) targetId).getSessionPool();
+      pool = authViews.get((AuthorizedViewId) targetId).getSessionPool();
     } else {
       return false;
     }
@@ -83,16 +81,12 @@ public CompletableFuture<Void> call(RowMutation request, Deadline deadline) {
     SessionMutateRowRequest innerReq = request.toSessionProto();
 
     if (targetId instanceof TableId) {
-      return tables
-          .getUnchecked((TableId) targetId)
-          .mutateRow(innerReq, deadline)
-          .thenApply(r -> null);
+      return tables.apply(
+          (TableId) targetId, t -> t.mutateRow(innerReq, deadline).thenApply(r -> null));
     }
     if (targetId instanceof AuthorizedViewId) {
-      return authViews
-          .getUnchecked((AuthorizedViewId) targetId)
-          .mutateRow(innerReq, deadline)
-          .thenApply(r -> null);
+      return authViews.apply(
+          (AuthorizedViewId) targetId, v -> v.mutateRow(innerReq, deadline).thenApply(r -> null));
     }
 
     CompletableFuture<Void> f = new CompletableFuture<>();

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/ReadRowShimInner.java

@@ -25,35 +25,33 @@
 import com.google.cloud.bigtable.data.v2.internal.api.MaterializedViewAsync;
 import com.google.cloud.bigtable.data.v2.internal.api.TableAsync;
 import com.google.cloud.bigtable.data.v2.internal.compat.ShimImpl;
-import com.google.cloud.bigtable.data.v2.internal.compat.Util;
 import com.google.cloud.bigtable.data.v2.internal.session.SessionPool;
 import com.google.cloud.bigtable.data.v2.models.AuthorizedViewId;
 import com.google.cloud.bigtable.data.v2.models.MaterializedViewId;
 import com.google.cloud.bigtable.data.v2.models.Query;
 import com.google.cloud.bigtable.data.v2.models.TableId;
 import com.google.cloud.bigtable.data.v2.models.TargetId;
-import com.google.common.cache.LoadingCache;
 import io.grpc.Deadline;
 import java.util.concurrent.CompletableFuture;
 
 public class ReadRowShimInner implements UnaryShim<Query, SessionReadRowResponse> {
-  private final LoadingCache<TableId, TableAsync> tables;
-  private final LoadingCache<AuthorizedViewId, AuthorizedViewAsync> authViews;
-  private final LoadingCache<MaterializedViewId, MaterializedViewAsync> matViews;
+  private final SessionPoolMap<TableId, TableAsync> tables;
+  private final SessionPoolMap<AuthorizedViewId, AuthorizedViewAsync> authViews;
+  private final SessionPoolMap<MaterializedViewId, MaterializedViewAsync> matViews;
 
   public ReadRowShimInner(Client client) {
     tables =
-        Util.createSessionMap(
+        new SessionPoolMap<>(
             k -> client.openTableAsync(k.getTableId(), Permission.PERMISSION_READ));
     authViews =
-        Util.createSessionMap(
+        new SessionPoolMap<>(
             k ->
                 client.openAuthorizedViewAsync(
                     k.getTableId(),
                     k.getAuthorizedViewId(),
                     OpenAuthorizedViewRequest.Permission.PERMISSION_READ));
     matViews =
-        Util.createSessionMap(
+        new SessionPoolMap<>(
             k ->
                 client.openMaterializedViewAsync(
                     k.getMaterializedViewId(),
@@ -73,11 +71,11 @@ public boolean supports(Query request) {
     SessionPool<?> pool;
     // TODO avoid double lookup
     if (targetId instanceof TableId) {
-      pool = tables.getUnchecked((TableId) targetId).getSessionPool();
+      pool = tables.get((TableId) targetId).getSessionPool();
     } else if (targetId instanceof AuthorizedViewId) {
-      pool = authViews.getUnchecked((AuthorizedViewId) targetId).getSessionPool();
+      pool = authViews.get((AuthorizedViewId) targetId).getSessionPool();
     } else if (targetId instanceof MaterializedViewId) {
-      pool = matViews.getUnchecked((MaterializedViewId) targetId).getSessionPool();
+      pool = matViews.get((MaterializedViewId) targetId).getSessionPool();
     } else {
       return false;
     }
@@ -95,13 +93,13 @@ public CompletableFuture<SessionReadRowResponse> call(Query query, Deadline dead
     SessionReadRowRequest innerReq = query.toSessionPointProto();
 
     if (targetId instanceof TableId) {
-      return tables.getUnchecked((TableId) targetId).readRow(innerReq, deadline);
+      return tables.apply((TableId) targetId, t -> t.readRow(innerReq, deadline));
     }
     if (targetId instanceof AuthorizedViewId) {
-      return authViews.getUnchecked((AuthorizedViewId) targetId).readRow(innerReq, deadline);
+      return authViews.apply((AuthorizedViewId) targetId, v -> v.readRow(innerReq, deadline));
     }
     if (targetId instanceof MaterializedViewId) {
-      return matViews.getUnchecked((MaterializedViewId) targetId).readRow(innerReq, deadline);
+      return matViews.apply((MaterializedViewId) targetId, v -> v.readRow(innerReq, deadline));
     }
 
     CompletableFuture<SessionReadRowResponse> f = new CompletableFuture<>();

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionPoolMap.java

@@ -0,0 +1,77 @@
+/*
+ * 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.compat.ops;
+
+import com.google.cloud.bigtable.data.v2.internal.compat.Util;
+import com.google.common.cache.LoadingCache;
+import com.google.common.util.concurrent.UncheckedExecutionException;
+import java.io.Closeable;
+import java.util.concurrent.CompletableFuture;
+import java.util.function.Function;
+
+/**
+ * Lazily-loaded map from target IDs to per-target handles (TableAsync / AuthorizedViewAsync /
+ * MaterializedViewAsync, each holding a SessionPool), backed by a Guava {@link LoadingCache}.
+ * Centralizes the conversion from Guava's {@link UncheckedExecutionException} (which wraps any
+ * loader throw) into the appropriate shape for each call site: a raw {@link RuntimeException} for
+ * sync inspections, or a failed {@link CompletableFuture} for async ops.
+ *
+ * <p>The most common loader failure is {@link IllegalStateException} from {@code
+ * Client.openTableAsync} after {@code Client.close}; without unwrapping, callers see {@code
+ * UncheckedExecutionException} instead of the documented exception type and async callers see a
+ * thrown exception instead of a failed future.
+ */
+public class SessionPoolMap<K, V extends Closeable> {
+  private final LoadingCache<K, V> cache;
+
+  public SessionPoolMap(Function<K, V> loader) {
+    this.cache = Util.createSessionMap(loader);
+  }
+
+  /** Returns the cached value, loading on miss. Loader throws surface as the original cause. */
+  public V get(K key) {
+    try {
+      return cache.getUnchecked(key);
+    } catch (UncheckedExecutionException e) {
+      Throwable cause = e.getCause() != null ? e.getCause() : e;
+      if (cause instanceof RuntimeException) throw (RuntimeException) cause;
+      if (cause instanceof Error) throw (Error) cause;
+      throw new RuntimeException(cause);
+    }
+  }
+
+  /**
+   * Looks up the cached value and applies {@code op}. If the lookup throws (e.g. {@link
+   * IllegalStateException} from a closed Client), the throw is converted to a failed future so
+   * callers consistently observe failures through the future surface.
+   */
+  public <R> CompletableFuture<R> apply(K key, Function<V, CompletableFuture<R>> op) {
+    V v;
+    try {
+      v = get(key);
+    } catch (Exception e) {
+      CompletableFuture<R> f = new CompletableFuture<>();
+      f.completeExceptionally(e);
+      return f;
+    }
+    return op.apply(v);
+  }
+
+  /** Evicts all entries, triggering Closeable.close on each via the cache's removal listener. */
+  public void invalidateAll() {
+    cache.invalidateAll();
+  }
+}

java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionPoolMapTest.java

@@ -0,0 +1,98 @@
+/*
+ * 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.compat.ops;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import java.io.Closeable;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import org.junit.jupiter.api.Test;
+
+class SessionPoolMapTest {
+
+  // Minimal Closeable value type — tracks how many times close() was called so invalidateAll's
+  // removal-listener wiring is observable from tests.
+  private static final class CountingHandle implements Closeable {
+    final AtomicInteger closeCount = new AtomicInteger();
+
+    @Override
+    public void close() {
+      closeCount.incrementAndGet();
+    }
+  }
+
+  @Test
+  void get_unwrapsLoaderRuntimeException() {
+    // Production trigger: Client.openTableAsync throws IllegalStateException after Client.close.
+    // The Guava LoadingCache wraps it in UncheckedExecutionException; SessionPoolMap.get must
+    // surface the original IllegalStateException so callers can pattern-match on it.
+    SessionPoolMap<String, CountingHandle> map =
+        new SessionPoolMap<>(
+            key -> {
+              throw new IllegalStateException("Client is closed");
+            });
+
+    IllegalStateException thrown = assertThrows(IllegalStateException.class, () -> map.get("k"));
+    assertThat(thrown).hasMessageThat().isEqualTo("Client is closed");
+  }
+
+  @Test
+  void apply_convertsLoaderThrowToFailedFuture() throws Exception {
+    // The async surface must not throw — callers wire onFailure handlers on the returned future.
+    SessionPoolMap<String, CountingHandle> map =
+        new SessionPoolMap<>(
+            key -> {
+              throw new IllegalStateException("Client is closed");
+            });
+
+    CompletableFuture<String> result =
+        map.apply("k", v -> CompletableFuture.completedFuture("unreached"));
+
+    assertThat(result.isCompletedExceptionally()).isTrue();
+    ExecutionException ee =
+        assertThrows(ExecutionException.class, () -> result.get(1, TimeUnit.SECONDS));
+    assertThat(ee).hasCauseThat().isInstanceOf(IllegalStateException.class);
+    assertThat(ee).hasCauseThat().hasMessageThat().isEqualTo("Client is closed");
+  }
+
+  @Test
+  void apply_happyPathInvokesOp() throws Exception {
+    CountingHandle handle = new CountingHandle();
+    SessionPoolMap<String, CountingHandle> map = new SessionPoolMap<>(key -> handle);
+
+    CompletableFuture<String> result =
+        map.apply("k", v -> CompletableFuture.completedFuture("ok"));
+
+    assertThat(result.get(1, TimeUnit.SECONDS)).isEqualTo("ok");
+  }
+
+  @Test
+  void invalidateAll_closesCachedValues() {
+    CountingHandle handle = new CountingHandle();
+    SessionPoolMap<String, CountingHandle> map = new SessionPoolMap<>(key -> handle);
+
+    // Populate the cache so invalidateAll has something to evict.
+    map.get("k");
+    assertThat(handle.closeCount.get()).isEqualTo(0);
+
+    map.invalidateAll();
+    assertThat(handle.closeCount.get()).isEqualTo(1);
+  }
+}