Skip to content

Commit fc74164

Browse files
mutianfmutianf
authored
fix: Make refreshing channel compatible with BigtableDataClientFactory (#474)
Set up channel priming in EnhancedBigtableStubSettings when refreshingChannel is true. BigtableDataClientFactory will set up refreshing channels in the shared client context. Derivative clients cannot reconfigure this setting and will all have refreshing channels.
1 parent f38a8ec commit fc74164

File tree

5 files changed

+216
-6
lines changed

5 files changed

+216
-6
lines changed

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

+5-2
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
import com.google.api.gax.rpc.FixedHeaderProvider;
2424
import com.google.api.gax.rpc.FixedTransportChannelProvider;
2525
import com.google.api.gax.rpc.FixedWatchdogProvider;
26-
import com.google.api.gax.rpc.StubSettings;
26+
import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStubSettings;
2727
import java.io.IOException;
2828
import javax.annotation.Nonnull;
2929

@@ -189,8 +189,11 @@ public BigtableDataClient createForInstance(
189189
}
190190

191191
// Update stub settings to use shared resources in this factory
192-
private void patchStubSettings(StubSettings.Builder stubSettings) {
192+
private void patchStubSettings(EnhancedBigtableStubSettings.Builder stubSettings) {
193193
stubSettings
194+
// Channel refreshing will be configured in the shared ClientContext. Derivative clients
195+
// won't be able to reconfigure the refreshing logic
196+
.setRefreshingChannel(false)
194197
.setTransportChannelProvider(
195198
FixedTransportChannelProvider.create(sharedClientContext.getTransportChannel()))
196199
.setCredentialsProvider(

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

+19
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import com.google.api.gax.batching.BatchingSettings;
2121
import com.google.api.gax.batching.FlowControlSettings;
2222
import com.google.api.gax.batching.FlowController.LimitExceededBehavior;
23+
import com.google.api.gax.core.FixedCredentialsProvider;
2324
import com.google.api.gax.core.GoogleCredentialsProvider;
2425
import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
2526
import com.google.api.gax.retrying.RetrySettings;
@@ -29,6 +30,7 @@
2930
import com.google.api.gax.rpc.StubSettings;
3031
import com.google.api.gax.rpc.TransportChannelProvider;
3132
import com.google.api.gax.rpc.UnaryCallSettings;
33+
import com.google.auth.Credentials;
3234
import com.google.cloud.bigtable.Version;
3335
import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation;
3436
import com.google.cloud.bigtable.data.v2.models.KeyOffset;
@@ -42,6 +44,7 @@
4244
import com.google.common.collect.ImmutableList;
4345
import com.google.common.collect.ImmutableMap;
4446
import com.google.common.collect.ImmutableSet;
47+
import java.io.IOException;
4548
import java.util.List;
4649
import java.util.Map;
4750
import java.util.Set;
@@ -787,6 +790,22 @@ public EnhancedBigtableStubSettings build() {
787790
Preconditions.checkArgument(
788791
getTransportChannelProvider() instanceof InstantiatingGrpcChannelProvider,
789792
"refreshingChannel only works with InstantiatingGrpcChannelProviders");
793+
InstantiatingGrpcChannelProvider.Builder channelProviderBuilder =
794+
((InstantiatingGrpcChannelProvider) getTransportChannelProvider()).toBuilder();
795+
Credentials credentials = null;
796+
if (getCredentialsProvider() != null) {
797+
try {
798+
credentials = getCredentialsProvider().getCredentials();
799+
} catch (IOException e) {
800+
throw new RuntimeException(e);
801+
}
802+
}
803+
// Use shared credentials
804+
this.setCredentialsProvider(FixedCredentialsProvider.create(credentials));
805+
channelProviderBuilder.setChannelPrimer(
806+
BigtableChannelPrimer.create(
807+
credentials, projectId, instanceId, appProfileId, primedTableIds));
808+
this.setTransportChannelProvider(channelProviderBuilder.build());
790809
}
791810
return new EnhancedBigtableStubSettings(this);
792811
}

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

+120-3
Original file line numberDiff line numberDiff line change
@@ -18,19 +18,32 @@
1818
import static com.google.common.truth.Truth.assertThat;
1919

2020
import com.google.api.core.ApiClock;
21+
import com.google.api.core.ApiFunction;
2122
import com.google.api.gax.core.CredentialsProvider;
2223
import com.google.api.gax.core.ExecutorProvider;
24+
import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
2325
import com.google.api.gax.rpc.TransportChannelProvider;
2426
import com.google.api.gax.rpc.WatchdogProvider;
2527
import com.google.bigtable.v2.BigtableGrpc;
2628
import com.google.bigtable.v2.MutateRowRequest;
2729
import com.google.bigtable.v2.MutateRowResponse;
30+
import com.google.bigtable.v2.ReadRowsRequest;
31+
import com.google.bigtable.v2.ReadRowsResponse;
32+
import com.google.bigtable.v2.RowFilter;
33+
import com.google.bigtable.v2.RowSet;
2834
import com.google.cloud.bigtable.data.v2.internal.NameUtil;
2935
import com.google.cloud.bigtable.data.v2.models.RowMutation;
3036
import com.google.common.base.Preconditions;
37+
import com.google.protobuf.ByteString;
38+
import io.grpc.Attributes;
39+
import io.grpc.ServerTransportFilter;
3140
import io.grpc.stub.StreamObserver;
3241
import java.io.IOException;
3342
import java.lang.reflect.Method;
43+
import java.util.LinkedList;
44+
import java.util.List;
45+
import java.util.concurrent.BlockingQueue;
46+
import java.util.concurrent.LinkedBlockingDeque;
3447
import org.junit.After;
3548
import org.junit.Before;
3649
import org.junit.Rule;
@@ -60,16 +73,33 @@ public class BigtableDataClientFactoryTest {
6073
private WatchdogProvider watchdogProvider;
6174
private ApiClock apiClock;
6275
private BigtableDataSettings defaultSettings;
76+
private int port;
77+
78+
private final BlockingQueue<Attributes> setUpAttributes = new LinkedBlockingDeque<>();
79+
private final BlockingQueue<Attributes> terminateAttributes = new LinkedBlockingDeque<>();
6380

6481
@Before
6582
public void setUp() throws IOException {
6683
service = new FakeBigtableService();
67-
68-
serviceHelper = new FakeServiceHelper(service);
84+
ServerTransportFilter transportFilter =
85+
new ServerTransportFilter() {
86+
@Override
87+
public Attributes transportReady(Attributes transportAttrs) {
88+
setUpAttributes.add(transportAttrs);
89+
return super.transportReady(transportAttrs);
90+
}
91+
92+
@Override
93+
public void transportTerminated(Attributes transportAttrs) {
94+
terminateAttributes.add(transportAttrs);
95+
}
96+
};
97+
serviceHelper = new FakeServiceHelper(null, transportFilter, service);
98+
port = serviceHelper.getPort();
6999
serviceHelper.start();
70100

71101
BigtableDataSettings.Builder builder =
72-
BigtableDataSettings.newBuilderForEmulator(serviceHelper.getPort())
102+
BigtableDataSettings.newBuilderForEmulator(port)
73103
.setProjectId(DEFAULT_PROJECT_ID)
74104
.setInstanceId(DEFAULT_INSTANCE_ID)
75105
.setAppProfileId(DEFAULT_APP_PROFILE_ID);
@@ -191,8 +221,94 @@ public void testCreateForInstanceWithAppProfileHasCorrectSettings() throws Excep
191221
assertThat(service.lastRequest.getAppProfileId()).isEqualTo("other-app-profile");
192222
}
193223

224+
@Test
225+
public void testCreateWithRefreshingChannel() throws Exception {
226+
String[] tableIds = {"fake-table1", "fake-table2"};
227+
int poolSize = 3;
228+
BigtableDataSettings.Builder builder =
229+
BigtableDataSettings.newBuilderForEmulator(port)
230+
.setProjectId(DEFAULT_PROJECT_ID)
231+
.setInstanceId(DEFAULT_INSTANCE_ID)
232+
.setAppProfileId(DEFAULT_APP_PROFILE_ID)
233+
.setPrimingTableIds(tableIds)
234+
.setRefreshingChannel(true);
235+
builder
236+
.stubSettings()
237+
.setCredentialsProvider(credentialsProvider)
238+
.setStreamWatchdogProvider(watchdogProvider)
239+
.setExecutorProvider(executorProvider);
240+
InstantiatingGrpcChannelProvider channelProvider =
241+
(InstantiatingGrpcChannelProvider) builder.stubSettings().getTransportChannelProvider();
242+
InstantiatingGrpcChannelProvider.Builder channelProviderBuilder = channelProvider.toBuilder();
243+
channelProviderBuilder.setPoolSize(poolSize);
244+
builder.stubSettings().setTransportChannelProvider(channelProviderBuilder.build());
245+
246+
BigtableDataClientFactory factory = BigtableDataClientFactory.create(builder.build());
247+
factory.createDefault();
248+
factory.createForAppProfile("other-appprofile");
249+
factory.createForInstance("other-project", "other-instance");
250+
251+
// Make sure that only 1 instance is created for all clients
252+
Mockito.verify(credentialsProvider, Mockito.times(1)).getCredentials();
253+
Mockito.verify(executorProvider, Mockito.times(1)).getExecutor();
254+
Mockito.verify(watchdogProvider, Mockito.times(1)).getWatchdog();
255+
256+
// Make sure that the clients are sharing the same ChannelPool
257+
assertThat(setUpAttributes).hasSize(poolSize);
258+
259+
// Make sure that prime requests were sent only once per table per connection
260+
assertThat(service.readRowsRequests).hasSize(poolSize * tableIds.length);
261+
List<ReadRowsRequest> expectedRequests = new LinkedList<>();
262+
for (String tableId : tableIds) {
263+
for (int i = 0; i < poolSize; i++) {
264+
expectedRequests.add(
265+
ReadRowsRequest.newBuilder()
266+
.setTableName(
267+
String.format(
268+
"projects/%s/instances/%s/tables/%s",
269+
DEFAULT_PROJECT_ID, DEFAULT_INSTANCE_ID, tableId))
270+
.setAppProfileId(DEFAULT_APP_PROFILE_ID)
271+
.setRows(
272+
RowSet.newBuilder()
273+
.addRowKeys(ByteString.copyFromUtf8("nonexistent-priming-row")))
274+
.setFilter(RowFilter.newBuilder().setBlockAllFilter(true).build())
275+
.setRowsLimit(1)
276+
.build());
277+
}
278+
}
279+
assertThat(service.readRowsRequests).containsExactly(expectedRequests.toArray());
280+
281+
// Wait for all the connections to close asynchronously
282+
factory.close();
283+
long sleepTimeMs = 1000;
284+
Thread.sleep(sleepTimeMs);
285+
// Verify that all the channels are closed
286+
assertThat(terminateAttributes).hasSize(poolSize);
287+
}
288+
194289
private static class FakeBigtableService extends BigtableGrpc.BigtableImplBase {
290+
195291
volatile MutateRowRequest lastRequest;
292+
BlockingQueue<ReadRowsRequest> readRowsRequests = new LinkedBlockingDeque<>();
293+
private ApiFunction<ReadRowsRequest, ReadRowsResponse> readRowsCallback =
294+
new ApiFunction<ReadRowsRequest, ReadRowsResponse>() {
295+
@Override
296+
public ReadRowsResponse apply(ReadRowsRequest readRowsRequest) {
297+
return ReadRowsResponse.getDefaultInstance();
298+
}
299+
};
300+
301+
@Override
302+
public void readRows(
303+
ReadRowsRequest request, StreamObserver<ReadRowsResponse> responseObserver) {
304+
try {
305+
readRowsRequests.add(request);
306+
responseObserver.onNext(readRowsCallback.apply(request));
307+
responseObserver.onCompleted();
308+
} catch (RuntimeException e) {
309+
responseObserver.onError(e);
310+
}
311+
}
196312

197313
@Override
198314
public void mutateRow(
@@ -204,6 +320,7 @@ public void mutateRow(
204320
}
205321

206322
private static class BuilderAnswer<T> implements Answer {
323+
207324
private final Class<T> targetClass;
208325
private T targetInstance;
209326

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

+12
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import io.grpc.Server;
2020
import io.grpc.ServerBuilder;
2121
import io.grpc.ServerInterceptor;
22+
import io.grpc.ServerTransportFilter;
2223
import java.io.IOException;
2324
import java.net.ServerSocket;
2425

@@ -33,13 +34,24 @@ public FakeServiceHelper(BindableService... services) throws IOException {
3334

3435
public FakeServiceHelper(ServerInterceptor interceptor, BindableService... services)
3536
throws IOException {
37+
this(interceptor, null, services);
38+
}
39+
40+
public FakeServiceHelper(
41+
ServerInterceptor interceptor,
42+
ServerTransportFilter transportFilter,
43+
BindableService... services)
44+
throws IOException {
3645
try (ServerSocket ss = new ServerSocket(0)) {
3746
port = ss.getLocalPort();
3847
}
3948
ServerBuilder builder = ServerBuilder.forPort(port);
4049
if (interceptor != null) {
4150
builder = builder.intercept(interceptor);
4251
}
52+
if (transportFilter != null) {
53+
builder = builder.addTransportFilter(transportFilter);
54+
}
4355
for (BindableService service : services) {
4456
builder = builder.addService(service);
4557
}

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

+60-1
Original file line numberDiff line numberDiff line change
@@ -19,20 +19,27 @@
1919

2020
import com.google.api.gax.batching.BatchingSettings;
2121
import com.google.api.gax.core.CredentialsProvider;
22+
import com.google.api.gax.core.FixedCredentialsProvider;
2223
import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
2324
import com.google.api.gax.retrying.RetrySettings;
2425
import com.google.api.gax.rpc.ServerStreamingCallSettings;
2526
import com.google.api.gax.rpc.StatusCode.Code;
2627
import com.google.api.gax.rpc.UnaryCallSettings;
2728
import com.google.api.gax.rpc.WatchdogProvider;
29+
import com.google.auth.Credentials;
2830
import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation;
2931
import com.google.cloud.bigtable.data.v2.models.KeyOffset;
3032
import com.google.cloud.bigtable.data.v2.models.Query;
3133
import com.google.cloud.bigtable.data.v2.models.Row;
3234
import com.google.cloud.bigtable.data.v2.models.RowMutation;
35+
import com.google.common.collect.ImmutableMap;
3336
import com.google.common.collect.ImmutableSet;
3437
import com.google.common.collect.Range;
38+
import java.io.IOException;
39+
import java.net.URI;
40+
import java.util.Arrays;
3541
import java.util.List;
42+
import java.util.Map;
3643
import java.util.Set;
3744
import org.junit.Test;
3845
import org.junit.runner.RunWith;
@@ -61,7 +68,7 @@ public void settingsAreNotLostTest() {
6168
String projectId = "my-project";
6269
String instanceId = "my-instance";
6370
String appProfileId = "my-app-profile-id";
64-
boolean isRefreshingChannel = true;
71+
boolean isRefreshingChannel = false;
6572
String endpoint = "some.other.host:123";
6673
CredentialsProvider credentialsProvider = Mockito.mock(CredentialsProvider.class);
6774
WatchdogProvider watchdogProvider = Mockito.mock(WatchdogProvider.class);
@@ -612,4 +619,56 @@ public void isRefreshingChannelFalseValueTest() {
612619
assertThat(builder.build().isRefreshingChannel()).isFalse();
613620
assertThat(builder.build().toBuilder().isRefreshingChannel()).isFalse();
614621
}
622+
623+
@Test
624+
public void refreshingChannelSetFixedCredentialProvider() throws Exception {
625+
String dummyProjectId = "my-project";
626+
String dummyInstanceId = "my-instance";
627+
CredentialsProvider credentialsProvider = Mockito.mock(CredentialsProvider.class);
628+
FakeCredentials expectedCredentials = new FakeCredentials();
629+
Mockito.when(credentialsProvider.getCredentials())
630+
.thenReturn(expectedCredentials, new FakeCredentials(), new FakeCredentials());
631+
EnhancedBigtableStubSettings.Builder builder =
632+
EnhancedBigtableStubSettings.newBuilder()
633+
.setProjectId(dummyProjectId)
634+
.setInstanceId(dummyInstanceId)
635+
.setRefreshingChannel(true)
636+
.setCredentialsProvider(credentialsProvider);
637+
assertThat(builder.isRefreshingChannel()).isTrue();
638+
// Verify that isRefreshing setting is not lost and stubSettings will always return the same
639+
// credential
640+
EnhancedBigtableStubSettings stubSettings = builder.build();
641+
assertThat(stubSettings.isRefreshingChannel()).isTrue();
642+
assertThat(stubSettings.getCredentialsProvider()).isInstanceOf(FixedCredentialsProvider.class);
643+
assertThat(stubSettings.getCredentialsProvider().getCredentials())
644+
.isEqualTo(expectedCredentials);
645+
assertThat(stubSettings.toBuilder().isRefreshingChannel()).isTrue();
646+
assertThat(stubSettings.toBuilder().getCredentialsProvider().getCredentials())
647+
.isEqualTo(expectedCredentials);
648+
}
649+
650+
private static class FakeCredentials extends Credentials {
651+
@Override
652+
public String getAuthenticationType() {
653+
return "fake";
654+
}
655+
656+
@Override
657+
public Map<String, List<String>> getRequestMetadata(URI uri) throws IOException {
658+
return ImmutableMap.of("my-header", Arrays.asList("fake-credential"));
659+
}
660+
661+
@Override
662+
public boolean hasRequestMetadata() {
663+
return true;
664+
}
665+
666+
@Override
667+
public boolean hasRequestMetadataOnly() {
668+
return true;
669+
}
670+
671+
@Override
672+
public void refresh() throws IOException {}
673+
}
615674
}

0 commit comments

Comments
 (0)