Skip to content

Commit 914f0cc

Browse files
igorbernstein2igorbernstein2
authored
fix: client.close() should wait until the channels are terminated before shutting down the executor (#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
1 parent e3ab43f commit 914f0cc

File tree

3 files changed

+248
-1
lines changed

3 files changed

+248
-1
lines changed

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

+5-1
Original file line numberDiff line numberDiff line change
@@ -803,7 +803,11 @@ private SpanName getSpanName(String methodName) {
803803
@Override
804804
public void close() {
805805
for (BackgroundResource backgroundResource : clientContext.getBackgroundResources()) {
806-
backgroundResource.shutdown();
806+
try {
807+
backgroundResource.close();
808+
} catch (Exception e) {
809+
throw new IllegalStateException("Failed to close resource", e);
810+
}
807811
}
808812
}
809813
}

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

+120
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
/*
2+
* Copyright 2021 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.bigtable.data.v2.stub;
17+
18+
import static com.google.common.truth.Truth.assertThat;
19+
20+
import com.google.api.core.ApiFuture;
21+
import com.google.api.gax.core.NoCredentialsProvider;
22+
import com.google.bigtable.v2.BigtableGrpc;
23+
import com.google.bigtable.v2.ReadRowsRequest;
24+
import com.google.bigtable.v2.ReadRowsResponse;
25+
import com.google.cloud.bigtable.data.v2.BigtableDataSettings;
26+
import com.google.cloud.bigtable.data.v2.FakeServiceHelper;
27+
import com.google.cloud.bigtable.data.v2.models.Query;
28+
import com.google.cloud.bigtable.data.v2.models.Row;
29+
import io.grpc.Status;
30+
import io.grpc.stub.StreamObserver;
31+
import java.util.List;
32+
import java.util.concurrent.ArrayBlockingQueue;
33+
import java.util.concurrent.BlockingQueue;
34+
import java.util.concurrent.ExecutionException;
35+
import java.util.concurrent.ExecutorService;
36+
import java.util.concurrent.Executors;
37+
import java.util.concurrent.RejectedExecutionException;
38+
import java.util.concurrent.atomic.AtomicInteger;
39+
import org.junit.After;
40+
import org.junit.Assert;
41+
import org.junit.Before;
42+
import org.junit.Test;
43+
import org.junit.function.ThrowingRunnable;
44+
import org.junit.runner.RunWith;
45+
import org.junit.runners.JUnit4;
46+
47+
/** Ensures that closing a client during exponential retry will not hang any requests. */
48+
@RunWith(JUnit4.class)
49+
public class EnhancedBigtableStubCloseRetryTest {
50+
private static final String PROJECT_ID = "fake-project";
51+
private static final String INSTANCE_ID = "fake-instance";
52+
53+
private ExecutorService testExecutor;
54+
private BlockingQueue<ReadRowsRequest> requests;
55+
private AtomicInteger numRequests;
56+
57+
private FakeServiceHelper serviceHelper;
58+
private EnhancedBigtableStub stub;
59+
60+
@Before
61+
public void setUp() throws Exception {
62+
testExecutor = Executors.newCachedThreadPool();
63+
requests = new ArrayBlockingQueue<>(10);
64+
numRequests = new AtomicInteger();
65+
66+
serviceHelper = new FakeServiceHelper(new FakeBigtable());
67+
serviceHelper.start();
68+
69+
BigtableDataSettings.Builder settingBuilder =
70+
BigtableDataSettings.newBuilderForEmulator(serviceHelper.getPort())
71+
.setProjectId(PROJECT_ID)
72+
.setInstanceId(INSTANCE_ID)
73+
.setCredentialsProvider(NoCredentialsProvider.create())
74+
.setRefreshingChannel(false);
75+
76+
stub = EnhancedBigtableStub.create(settingBuilder.build().getStubSettings());
77+
}
78+
79+
@After
80+
public void tearDown() throws Exception {
81+
testExecutor.shutdown();
82+
stub.close();
83+
serviceHelper.shutdown();
84+
}
85+
86+
@Test
87+
public void outstandingRequestsFinishAfterClose() throws Exception {
88+
final ApiFuture<List<Row>> resultFuture =
89+
stub.readRowsCallable().all().futureCall(Query.create("table1"));
90+
91+
// wait for the first request to hit the server
92+
requests.take();
93+
// wait enough time for a retry attempt to be scheduled before closing the client
94+
Thread.sleep(100);
95+
stub.close();
96+
97+
ExecutionException error =
98+
Assert.assertThrows(
99+
ExecutionException.class,
100+
new ThrowingRunnable() {
101+
@Override
102+
public void run() throws Throwable {
103+
resultFuture.get();
104+
}
105+
});
106+
107+
assertThat(error.getCause()).isInstanceOf(RejectedExecutionException.class);
108+
}
109+
110+
class FakeBigtable extends BigtableGrpc.BigtableImplBase {
111+
@Override
112+
public void readRows(
113+
ReadRowsRequest request, StreamObserver<ReadRowsResponse> responseObserver) {
114+
115+
requests.add(request);
116+
// Keep returning a retriable error
117+
responseObserver.onError(Status.UNAVAILABLE.asRuntimeException());
118+
}
119+
}
120+
}

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

+123
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
/*
2+
* Copyright 2021 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.bigtable.data.v2.stub;
17+
18+
import static com.google.common.truth.Truth.assertThat;
19+
import static java.util.concurrent.TimeUnit.MINUTES;
20+
21+
import com.google.api.core.ApiFuture;
22+
import com.google.api.gax.core.NoCredentialsProvider;
23+
import com.google.bigtable.v2.BigtableGrpc;
24+
import com.google.bigtable.v2.ReadRowsRequest;
25+
import com.google.bigtable.v2.ReadRowsResponse;
26+
import com.google.cloud.bigtable.data.v2.BigtableDataSettings;
27+
import com.google.cloud.bigtable.data.v2.FakeServiceHelper;
28+
import com.google.cloud.bigtable.data.v2.models.Query;
29+
import com.google.cloud.bigtable.data.v2.models.Row;
30+
import com.google.common.util.concurrent.SettableFuture;
31+
import io.grpc.stub.StreamObserver;
32+
import java.util.List;
33+
import java.util.concurrent.ExecutorService;
34+
import java.util.concurrent.Executors;
35+
import org.junit.After;
36+
import org.junit.Before;
37+
import org.junit.Test;
38+
import org.junit.runner.RunWith;
39+
import org.junit.runners.JUnit4;
40+
41+
/** Ensure that an outstanding RPC will finish during a close */
42+
@RunWith(JUnit4.class)
43+
public class EnhancedBigtableStubCloseTest {
44+
private static final String PROJECT_ID = "fake-project";
45+
private static final String INSTANCE_ID = "fake-instance";
46+
47+
private ExecutorService testExecutor;
48+
private SettableFuture<Void> requestReceivedBarrier = SettableFuture.create();
49+
private SettableFuture<Void> clientClosedBarrier = SettableFuture.create();
50+
51+
private FakeServiceHelper serviceHelper;
52+
private EnhancedBigtableStub stub;
53+
54+
@Before
55+
public void setUp() throws Exception {
56+
testExecutor = Executors.newCachedThreadPool();
57+
requestReceivedBarrier = SettableFuture.create();
58+
clientClosedBarrier = SettableFuture.create();
59+
60+
serviceHelper = new FakeServiceHelper(new FakeBigtable());
61+
serviceHelper.start();
62+
63+
EnhancedBigtableStubSettings stubSettings =
64+
BigtableDataSettings.newBuilderForEmulator(serviceHelper.getPort())
65+
.setProjectId(PROJECT_ID)
66+
.setInstanceId(INSTANCE_ID)
67+
.setCredentialsProvider(NoCredentialsProvider.create())
68+
.setRefreshingChannel(false)
69+
.build()
70+
.getStubSettings();
71+
72+
stub = EnhancedBigtableStub.create(stubSettings);
73+
}
74+
75+
@After
76+
public void tearDown() throws Exception {
77+
testExecutor.shutdown();
78+
stub.close();
79+
serviceHelper.shutdown();
80+
}
81+
82+
@Test
83+
public void outstandingRequestsFinishAfterClose() throws Exception {
84+
ApiFuture<List<Row>> resultFuture =
85+
stub.readRowsCallable().all().futureCall(Query.create("table1"));
86+
87+
// Wait for the server to receive the request
88+
requestReceivedBarrier.get(1, MINUTES);
89+
// Close the client - must happen in a separate thread because close will block until all
90+
// requests have completed, which can't happen until the clientClosedBarrier is released.
91+
testExecutor.submit(
92+
new Runnable() {
93+
@Override
94+
public void run() {
95+
stub.close();
96+
}
97+
});
98+
Thread.sleep(200); // give the closer a chance to run
99+
clientClosedBarrier.set(null);
100+
101+
assertThat(resultFuture.get(1, MINUTES)).isEmpty();
102+
}
103+
104+
class FakeBigtable extends BigtableGrpc.BigtableImplBase {
105+
@Override
106+
public void readRows(
107+
ReadRowsRequest request, StreamObserver<ReadRowsResponse> responseObserver) {
108+
109+
// signal that server received the request
110+
requestReceivedBarrier.set(null);
111+
// wait until the main thread closes the client
112+
try {
113+
clientClosedBarrier.get();
114+
} catch (Exception e) {
115+
// Shouldn't happen
116+
responseObserver.onError(e);
117+
}
118+
119+
// send the response
120+
responseObserver.onCompleted();
121+
}
122+
}
123+
}

0 commit comments

Comments
 (0)