Threading refactor (5/7): Close ordering & shutdown hardening

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java

@@ -43,15 +43,27 @@
 import io.opencensus.tags.Tags;
 import io.opentelemetry.sdk.OpenTelemetrySdk;
 import java.io.IOException;
+import java.time.Duration;
+import java.util.ArrayList;
 import java.util.Collections;
+import java.util.List;
 import java.util.Set;
 import java.util.WeakHashMap;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 
 public class Client implements AutoCloseable {
+  private static final Logger logger = Logger.getLogger(Client.class.getName());
+
+  // Per-pool drain budget during close. One full watchdog tick (5 min) plus 1 min buffer; if a
+  // pool can't drain in that window, something is genuinely wrong on the server side and we give
+  // up on it so close() returns. The watchdog interval is what makes the worst case finite.
+  private static final Duration POOL_DRAIN_TIMEOUT = Duration.ofMinutes(6);
+
   public static final FeatureFlags BASE_FEATURE_FLAGS =
       FeatureFlags.newBuilder()
           .setReverseScans(false)
@@ -91,6 +103,10 @@ public class Client implements AutoCloseable {
   private final Resource<ClientConfigurationManager> configManager;
 
   private final Set<SessionPool<?>> sessionPools = Collections.newSetFromMap(new WeakHashMap<>());
+  // Guarded by sessionPools' monitor: close() sets it before snapshotting the pool set, and the
+  // open* methods check it before adding a new pool, so a racing open cannot insert a pool that
+  // close() has already missed in its snapshot.
+  private boolean closed = false;
 
   public static Client create(ClientSettings settings) throws IOException {
     FeatureFlags featureFlags =
@@ -218,78 +234,134 @@ public Client(
 
   @Override
   public void close() {
-    sessionPools.forEach(
-        pool ->
-            pool.close(
-                CloseSessionRequest.newBuilder()
-                    .setReason(CloseSessionReason.CLOSE_SESSION_REASON_USER)
-                    .setDescription("Client closing")
-                    .build()));
-    // Drain user-callback first so pool.close's cancelWithResult listener notifications complete
-    // before we tear down the surrounding executors and timer. Without this, the late onClose
-    // submissions race the shutdown and get RejectedExecutionException, silently dropping the
-    // user's terminal onClose.
+    List<SessionPool<?>> toClose;
+    synchronized (sessionPools) {
+      if (closed) {
+        return; // idempotent
+      }
+      closed = true;
+      toClose = new ArrayList<>(sessionPools);
+    }
+
+    CloseSessionRequest closeReq =
+        CloseSessionRequest.newBuilder()
+            .setReason(CloseSessionReason.CLOSE_SESSION_REASON_USER)
+            .setDescription("Client closing")
+            .build();
+
+    // Phase 1: initiate graceful close on each pool. Returns immediately; sessions transition
+    // CLOSING → graceful CloseSessionRequest → WAIT_SERVER_CLOSE → CLOSED asynchronously.
+    toClose.forEach(p -> p.close(closeReq));
+
+    // Phase 2: wait for sessions to drain. The pool's watchdog stays alive during this wait and
+    // escalates anything stuck in WAIT_SERVER_CLOSE longer than its tick interval (5 min). Once
+    // a pool's last session reaches CLOSED, drainedFuture completes and awaitTerminated returns.
+    // Sequential: worst case is POOL_DRAIN_TIMEOUT * N pools, but the happy path drains in << 1s.
+    for (SessionPool<?> pool : toClose) {
+      try {
+        if (!pool.awaitTerminated(POOL_DRAIN_TIMEOUT)) {
+          logger.warning(
+              "SessionPool did not drain within "
+                  + POOL_DRAIN_TIMEOUT
+                  + "; abandoning and continuing shutdown");
+        }
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+        logger.log(Level.WARNING, "Interrupted while draining SessionPool", e);
+        break;
+      }
+    }
+
+    // Phase 3: tear down infrastructure.
+    //
+    // sessionTimer.stop() runs FIRST so its onStop hooks can drive any pending Scheduled retries
+    // to a terminal Done — that delivery hops through op executor → userCallbackExecutor, both
+    // of which must still be alive at this moment.
+    //
+    // userCallbackExecutor.close() next, with a 5s drain to catch the listener.onClose tasks
+    // queued by both the session drain (Phase 2) and the just-fired retry shutdowns.
+    //
+    // backgroundExecutor must close last because it's the timer's dispatcher and the op
+    // executor's chain ultimately runs ScheduledExecutorService tasks here.
+    sessionTimer.stop();
     userCallbackExecutor.close();
     metrics.close();
     channelPool.close();
     configManager.close();
-    // Stop the timer before tearing down backgroundExecutor (the timer's dispatcher).
-    sessionTimer.stop();
     backgroundExecutor.close();
   }
 
+  // The closed check and pool insertion run under sessionPools' monitor so close() (which flips
+  // closed under the same monitor) cannot snapshot the pool set between our check and our insert.
+  // Opens are infrequent (typically once per table at app startup), so holding the monitor across
+  // createAndStart is acceptable.
   public TableAsync openTableAsync(String tableId, Permission permission) {
-    TableAsync tableAsync =
-        TableAsync.createAndStart(
-            featureFlags,
-            clientInfo,
-            configManager.get(),
-            channelPool,
-            defaultCallOptions,
-            tableId,
-            permission,
-            metrics.get(),
-            sessionTimer,
-            userCallbackExecutor.get());
-    sessionPools.add(tableAsync.getSessionPool());
-    return tableAsync;
+    synchronized (sessionPools) {
+      if (closed) {
+        throw new IllegalStateException("Client is closed");
+      }
+      TableAsync tableAsync =
+          TableAsync.createAndStart(
+              featureFlags,
+              clientInfo,
+              configManager.get(),
+              channelPool,
+              defaultCallOptions,
+              tableId,
+              permission,
+              metrics.get(),
+              sessionTimer,
+              userCallbackExecutor.get());
+      sessionPools.add(tableAsync.getSessionPool());
+      return tableAsync;
+    }
   }
 
   public AuthorizedViewAsync openAuthorizedViewAsync(
       String tableId, String viewId, OpenAuthorizedViewRequest.Permission permission) {
-    AuthorizedViewAsync viewAsync =
-        AuthorizedViewAsync.createAndStart(
-            featureFlags,
-            clientInfo,
-            configManager.get(),
-            channelPool,
-            defaultCallOptions,
-            tableId,
-            viewId,
-            permission,
-            metrics.get(),
-            sessionTimer,
-            userCallbackExecutor.get());
-    sessionPools.add(viewAsync.getSessionPool());
-    return viewAsync;
+    synchronized (sessionPools) {
+      if (closed) {
+        throw new IllegalStateException("Client is closed");
+      }
+      AuthorizedViewAsync viewAsync =
+          AuthorizedViewAsync.createAndStart(
+              featureFlags,
+              clientInfo,
+              configManager.get(),
+              channelPool,
+              defaultCallOptions,
+              tableId,
+              viewId,
+              permission,
+              metrics.get(),
+              sessionTimer,
+              userCallbackExecutor.get());
+      sessionPools.add(viewAsync.getSessionPool());
+      return viewAsync;
+    }
   }
 
   public MaterializedViewAsync openMaterializedViewAsync(
       String viewId, OpenMaterializedViewRequest.Permission permission) {
-    MaterializedViewAsync viewAsync =
-        MaterializedViewAsync.createAndStart(
-            featureFlags,
-            clientInfo,
-            configManager.get(),
-            channelPool,
-            defaultCallOptions,
-            viewId,
-            permission,
-            metrics.get(),
-            sessionTimer,
-            userCallbackExecutor.get());
-    sessionPools.add(viewAsync.getSessionPool());
-    return viewAsync;
+    synchronized (sessionPools) {
+      if (closed) {
+        throw new IllegalStateException("Client is closed");
+      }
+      MaterializedViewAsync viewAsync =
+          MaterializedViewAsync.createAndStart(
+              featureFlags,
+              clientInfo,
+              configManager.get(),
+              channelPool,
+              defaultCallOptions,
+              viewId,
+              permission,
+              metrics.get(),
+              sessionTimer,
+              userCallbackExecutor.get());
+      sessionPools.add(viewAsync.getSessionPool());
+      return viewAsync;
+    }
   }
 
   public static class Resource<T> {
@@ -324,8 +396,10 @@ public T get() {
   }
 
   // Drain in-flight listener.onClose tasks before the executor is shut down; bound the wait at 5s
-  // so close() doesn't hang the caller on a pathological listener.
-  private static void shutdownAndAwait(ExecutorService exec) {
+  // so close() doesn't hang the caller on a pathological listener. Public so the compat
+  // ShimImpl (different package) can reuse the same shutdown semantics for the user-callback
+  // executor it owns.
+  public static void shutdownAndAwait(ExecutorService exec) {
     exec.shutdown();
     try {
       if (!exec.awaitTermination(5, TimeUnit.SECONDS)) {

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

@@ -178,7 +178,8 @@ public static Shim create(
             Resource.createShared(metrics),
             Resource.createShared(configManager),
             Resource.createShared(bgExecutor),
-            Resource.createOwned(userCallbackExecutor, userCallbackExecutor::shutdown));
+            Resource.createOwned(
+                userCallbackExecutor, () -> Client.shutdownAndAwait(userCallbackExecutor)));
 
     return new ShimImpl(configManager, client);
   }

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<>();