fix: client.close() should wait until the channels are terminated bef…

…ore shutting down the executor (googleapis#916)

* fix: client.close() should wait until the channels are terminated before shutting down the executor

Previously the client.close() would simply shutdown the channel pool and the executor immediately after. Unfortunately this leads to RPCs that would hang forever because an outstanding RPC didnt have an executor to notify of its completion. This PR ensures that the channels are drained before shutting down the executor

* copyright

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java

@@ -803,7 +803,11 @@ private SpanName getSpanName(String methodName) {
   @Override
   public void close() {
     for (BackgroundResource backgroundResource : clientContext.getBackgroundResources()) {
-      backgroundResource.shutdown();
+      try {
+        backgroundResource.close();
+      } catch (Exception e) {
+        throw new IllegalStateException("Failed to close resource", e);
+      }
     }
   }
 }

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubCloseRetryTest.java

@@ -0,0 +1,120 @@
+/*
+ * Copyright 2021 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.stub;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.api.core.ApiFuture;
+import com.google.api.gax.core.NoCredentialsProvider;
+import com.google.bigtable.v2.BigtableGrpc;
+import com.google.bigtable.v2.ReadRowsRequest;
+import com.google.bigtable.v2.ReadRowsResponse;
+import com.google.cloud.bigtable.data.v2.BigtableDataSettings;
+import com.google.cloud.bigtable.data.v2.FakeServiceHelper;
+import com.google.cloud.bigtable.data.v2.models.Query;
+import com.google.cloud.bigtable.data.v2.models.Row;
+import io.grpc.Status;
+import io.grpc.stub.StreamObserver;
+import java.util.List;
+import java.util.concurrent.ArrayBlockingQueue;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.atomic.AtomicInteger;
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.function.ThrowingRunnable;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Ensures that closing a client during exponential retry will not hang any requests. */
+@RunWith(JUnit4.class)
+public class EnhancedBigtableStubCloseRetryTest {
+  private static final String PROJECT_ID = "fake-project";
+  private static final String INSTANCE_ID = "fake-instance";
+
+  private ExecutorService testExecutor;
+  private BlockingQueue<ReadRowsRequest> requests;
+  private AtomicInteger numRequests;
+
+  private FakeServiceHelper serviceHelper;
+  private EnhancedBigtableStub stub;
+
+  @Before
+  public void setUp() throws Exception {
+    testExecutor = Executors.newCachedThreadPool();
+    requests = new ArrayBlockingQueue<>(10);
+    numRequests = new AtomicInteger();
+
+    serviceHelper = new FakeServiceHelper(new FakeBigtable());
+    serviceHelper.start();
+
+    BigtableDataSettings.Builder settingBuilder =
+        BigtableDataSettings.newBuilderForEmulator(serviceHelper.getPort())
+            .setProjectId(PROJECT_ID)
+            .setInstanceId(INSTANCE_ID)
+            .setCredentialsProvider(NoCredentialsProvider.create())
+            .setRefreshingChannel(false);
+
+    stub = EnhancedBigtableStub.create(settingBuilder.build().getStubSettings());
+  }
+
+  @After
+  public void tearDown() throws Exception {
+    testExecutor.shutdown();
+    stub.close();
+    serviceHelper.shutdown();
+  }
+
+  @Test
+  public void outstandingRequestsFinishAfterClose() throws Exception {
+    final ApiFuture<List<Row>> resultFuture =
+        stub.readRowsCallable().all().futureCall(Query.create("table1"));
+
+    // wait for the first request to hit the server
+    requests.take();
+    // wait enough time for a retry attempt to be scheduled before closing the client
+    Thread.sleep(100);
+    stub.close();
+
+    ExecutionException error =
+        Assert.assertThrows(
+            ExecutionException.class,
+            new ThrowingRunnable() {
+              @Override
+              public void run() throws Throwable {
+                resultFuture.get();
+              }
+            });
+
+    assertThat(error.getCause()).isInstanceOf(RejectedExecutionException.class);
+  }
+
+  class FakeBigtable extends BigtableGrpc.BigtableImplBase {
+    @Override
+    public void readRows(
+        ReadRowsRequest request, StreamObserver<ReadRowsResponse> responseObserver) {
+
+      requests.add(request);
+      // Keep returning a retriable error
+      responseObserver.onError(Status.UNAVAILABLE.asRuntimeException());
+    }
+  }
+}

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubCloseTest.java

@@ -0,0 +1,123 @@
+/*
+ * Copyright 2021 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.stub;
+
+import static com.google.common.truth.Truth.assertThat;
+import static java.util.concurrent.TimeUnit.MINUTES;
+
+import com.google.api.core.ApiFuture;
+import com.google.api.gax.core.NoCredentialsProvider;
+import com.google.bigtable.v2.BigtableGrpc;
+import com.google.bigtable.v2.ReadRowsRequest;
+import com.google.bigtable.v2.ReadRowsResponse;
+import com.google.cloud.bigtable.data.v2.BigtableDataSettings;
+import com.google.cloud.bigtable.data.v2.FakeServiceHelper;
+import com.google.cloud.bigtable.data.v2.models.Query;
+import com.google.cloud.bigtable.data.v2.models.Row;
+import com.google.common.util.concurrent.SettableFuture;
+import io.grpc.stub.StreamObserver;
+import java.util.List;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Ensure that an outstanding RPC will finish during a close */
+@RunWith(JUnit4.class)
+public class EnhancedBigtableStubCloseTest {
+  private static final String PROJECT_ID = "fake-project";
+  private static final String INSTANCE_ID = "fake-instance";
+
+  private ExecutorService testExecutor;
+  private SettableFuture<Void> requestReceivedBarrier = SettableFuture.create();
+  private SettableFuture<Void> clientClosedBarrier = SettableFuture.create();
+
+  private FakeServiceHelper serviceHelper;
+  private EnhancedBigtableStub stub;
+
+  @Before
+  public void setUp() throws Exception {
+    testExecutor = Executors.newCachedThreadPool();
+    requestReceivedBarrier = SettableFuture.create();
+    clientClosedBarrier = SettableFuture.create();
+
+    serviceHelper = new FakeServiceHelper(new FakeBigtable());
+    serviceHelper.start();
+
+    EnhancedBigtableStubSettings stubSettings =
+        BigtableDataSettings.newBuilderForEmulator(serviceHelper.getPort())
+            .setProjectId(PROJECT_ID)
+            .setInstanceId(INSTANCE_ID)
+            .setCredentialsProvider(NoCredentialsProvider.create())
+            .setRefreshingChannel(false)
+            .build()
+            .getStubSettings();
+
+    stub = EnhancedBigtableStub.create(stubSettings);
+  }
+
+  @After
+  public void tearDown() throws Exception {
+    testExecutor.shutdown();
+    stub.close();
+    serviceHelper.shutdown();
+  }
+
+  @Test
+  public void outstandingRequestsFinishAfterClose() throws Exception {
+    ApiFuture<List<Row>> resultFuture =
+        stub.readRowsCallable().all().futureCall(Query.create("table1"));
+
+    // Wait for the server to receive the request
+    requestReceivedBarrier.get(1, MINUTES);
+    // Close the client - must happen in a separate thread because close will block until all
+    // requests have completed, which can't happen until the clientClosedBarrier is released.
+    testExecutor.submit(
+        new Runnable() {
+          @Override
+          public void run() {
+            stub.close();
+          }
+        });
+    Thread.sleep(200); // give the closer a chance to run
+    clientClosedBarrier.set(null);
+
+    assertThat(resultFuture.get(1, MINUTES)).isEmpty();
+  }
+
+  class FakeBigtable extends BigtableGrpc.BigtableImplBase {
+    @Override
+    public void readRows(
+        ReadRowsRequest request, StreamObserver<ReadRowsResponse> responseObserver) {
+
+      // signal that server received the request
+      requestReceivedBarrier.set(null);
+      // wait until the main thread closes the client
+      try {
+        clientClosedBarrier.get();
+      } catch (Exception e) {
+        // Shouldn't happen
+        responseObserver.onError(e);
+      }
+
+      // send the response
+      responseObserver.onCompleted();
+    }
+  }
+}