From 990330b917d3bacf2fa2cf94db61c5aa3054dfa9 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Mon, 2 Dec 2024 09:57:10 -0500 Subject: [PATCH 01/13] fix: channel primer respect other client settings --- .../data/v2/stub/BigtableChannelPrimer.java | 22 +++++++------------ .../data/v2/stub/BigtableClientContext.java | 7 +----- .../v2/stub/BigtableChannelPrimerTest.java | 17 +++++++++----- 3 files changed, 20 insertions(+), 26 deletions(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java index ecbef85be5..3fd6614a3d 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java @@ -16,12 +16,11 @@ package com.google.cloud.bigtable.data.v2.stub; import com.google.api.core.BetaApi; -import com.google.api.gax.core.FixedCredentialsProvider; import com.google.api.gax.core.InstantiatingExecutorProvider; import com.google.api.gax.grpc.ChannelPrimer; import com.google.api.gax.grpc.GrpcTransportChannel; import com.google.api.gax.rpc.FixedTransportChannelProvider; -import com.google.auth.Credentials; +import com.google.api.gax.rpc.InstantiatingWatchdogProvider; import com.google.bigtable.v2.PingAndWarmRequest; import com.google.cloud.bigtable.data.v2.internal.NameUtil; import com.google.common.base.Preconditions; @@ -43,18 +42,13 @@ class BigtableChannelPrimer implements ChannelPrimer { private final EnhancedBigtableStubSettings settingsTemplate; - static BigtableChannelPrimer create( - Credentials credentials, String projectId, String instanceId, String appProfileId) { - EnhancedBigtableStubSettings.Builder builder = - EnhancedBigtableStubSettings.newBuilder() - .setProjectId(projectId) - .setInstanceId(instanceId) - .setAppProfileId(appProfileId) - .setCredentialsProvider(FixedCredentialsProvider.create(credentials)) - // Disable refreshing channel here to avoid creating settings in a loop - .setRefreshingChannel(false) - .setExecutorProvider( - InstantiatingExecutorProvider.newBuilder().setExecutorThreadCount(1).build()); + static BigtableChannelPrimer create(EnhancedBigtableStubSettings settings) { + EnhancedBigtableStubSettings.Builder builder = settings.toBuilder(); + builder + .setRefreshingChannel(false) + .setBackgroundExecutorProvider( + InstantiatingExecutorProvider.newBuilder().setExecutorThreadCount(1).build()) + .setStreamWatchdogProvider(InstantiatingWatchdogProvider.create()); return new BigtableChannelPrimer(builder.build()); } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java index d23b487caf..ebd4df0bcd 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java @@ -103,12 +103,7 @@ public static BigtableClientContext create(EnhancedBigtableStubSettings settings } // Inject channel priming if enabled if (builder.isRefreshingChannel()) { - transportProvider.setChannelPrimer( - BigtableChannelPrimer.create( - credentials, - settings.getProjectId(), - settings.getInstanceId(), - settings.getAppProfileId())); + transportProvider.setChannelPrimer(BigtableChannelPrimer.create(builder.build())); } builder.setTransportChannelProvider(transportProvider.build()); diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimerTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimerTest.java index e1f22bebbd..bfacd1d769 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimerTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimerTest.java @@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat; import com.google.api.core.ApiFunction; +import com.google.api.gax.core.FixedCredentialsProvider; import com.google.auth.oauth2.AccessToken; import com.google.auth.oauth2.OAuth2Credentials; import com.google.bigtable.v2.BigtableGrpc.BigtableImplBase; @@ -67,12 +68,16 @@ public void setup() throws IOException { server = FakeServiceBuilder.create(fakeService).intercept(metadataInterceptor).start(); - primer = - BigtableChannelPrimer.create( - OAuth2Credentials.create(new AccessToken(TOKEN_VALUE, null)), - "fake-project", - "fake-instance", - "fake-app-profile"); + EnhancedBigtableStubSettings settings = + EnhancedBigtableStubSettings.newBuilder() + .setProjectId("fake-project") + .setInstanceId("fake-instance") + .setAppProfileId("fake-app-profile") + .setCredentialsProvider( + FixedCredentialsProvider.create( + OAuth2Credentials.create(new AccessToken(TOKEN_VALUE, null)))) + .build(); + primer = BigtableChannelPrimer.create(settings); channel = ManagedChannelBuilder.forAddress("localhost", server.getPort()).usePlaintext().build(); From dd6b7ce3ed230b7666d1b8bba35ed66e88545f03 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Mon, 2 Dec 2024 13:53:41 -0500 Subject: [PATCH 02/13] do not use the stub --- .../data/v2/stub/BigtableChannelPrimer.java | 133 +++++++++++------- .../data/v2/stub/EnhancedBigtableStub.java | 28 ---- 2 files changed, 85 insertions(+), 76 deletions(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java index 3fd6614a3d..80e977f768 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java @@ -16,17 +16,22 @@ package com.google.cloud.bigtable.data.v2.stub; import com.google.api.core.BetaApi; -import com.google.api.gax.core.InstantiatingExecutorProvider; +import com.google.api.core.SettableApiFuture; import com.google.api.gax.grpc.ChannelPrimer; -import com.google.api.gax.grpc.GrpcTransportChannel; -import com.google.api.gax.rpc.FixedTransportChannelProvider; -import com.google.api.gax.rpc.InstantiatingWatchdogProvider; +import com.google.api.gax.grpc.GrpcCallContext; +import com.google.bigtable.v2.BigtableGrpc; +import com.google.bigtable.v2.InstanceName; import com.google.bigtable.v2.PingAndWarmRequest; -import com.google.cloud.bigtable.data.v2.internal.NameUtil; -import com.google.common.base.Preconditions; +import com.google.bigtable.v2.PingAndWarmResponse; +import com.google.common.net.PercentEscaper; +import io.grpc.CallCredentials; +import io.grpc.ClientCall; import io.grpc.ManagedChannel; +import io.grpc.Metadata; +import io.grpc.Status; +import io.grpc.auth.MoreCallCredentials; import java.io.IOException; -import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; import java.util.logging.Logger; /** @@ -40,22 +45,46 @@ class BigtableChannelPrimer implements ChannelPrimer { private static Logger LOG = Logger.getLogger(BigtableChannelPrimer.class.toString()); - private final EnhancedBigtableStubSettings settingsTemplate; + private final Metadata.Key requestParams = + Metadata.Key.of("x-goog-request-params", Metadata.ASCII_STRING_MARSHALLER); + private final PingAndWarmRequest request; + private final Metadata metadata = new Metadata(); + private final CallCredentials credentials; - static BigtableChannelPrimer create(EnhancedBigtableStubSettings settings) { - EnhancedBigtableStubSettings.Builder builder = settings.toBuilder(); - builder - .setRefreshingChannel(false) - .setBackgroundExecutorProvider( - InstantiatingExecutorProvider.newBuilder().setExecutorThreadCount(1).build()) - .setStreamWatchdogProvider(InstantiatingWatchdogProvider.create()); - - return new BigtableChannelPrimer(builder.build()); + static BigtableChannelPrimer create(EnhancedBigtableStubSettings settings) throws IOException { + return new BigtableChannelPrimer(settings); } - private BigtableChannelPrimer(EnhancedBigtableStubSettings settingsTemplate) { - Preconditions.checkNotNull(settingsTemplate, "settingsTemplate can't be null"); - this.settingsTemplate = settingsTemplate; + BigtableChannelPrimer(EnhancedBigtableStubSettings settings) throws IOException { + String projectId = settings.getProjectId(); + String instanceId = settings.getInstanceId(); + String appProfileId = settings.getAppProfileId(); + + if (settings.getCredentialsProvider().getCredentials() != null) { + credentials = MoreCallCredentials.from(settings.getCredentialsProvider().getCredentials()); + } else { + credentials = null; + } + + request = + PingAndWarmRequest.newBuilder() + .setName(InstanceName.format(projectId, instanceId)) + .setAppProfileId(appProfileId) + .build(); + + PercentEscaper escaper = new PercentEscaper("._-~", false); + + Metadata metadata = new Metadata(); + settings + .getHeaderProvider() + .getHeaders() + .forEach((k, v) -> metadata.put(Metadata.Key.of(k, Metadata.ASCII_STRING_MARSHALLER), v)); + + metadata.put( + requestParams, + escaper.escape( + String.format( + "name=%s&app_profile_id=%s", request.getName(), request.getAppProfileId()))); } @Override @@ -72,35 +101,43 @@ private void primeChannelUnsafe(ManagedChannel managedChannel) throws IOExceptio sendPrimeRequests(managedChannel); } - private void sendPrimeRequests(ManagedChannel managedChannel) throws IOException { - // Wrap the channel in a temporary stub - EnhancedBigtableStubSettings primingSettings = - settingsTemplate - .toBuilder() - .setTransportChannelProvider( - FixedTransportChannelProvider.create(GrpcTransportChannel.create(managedChannel))) - .build(); + private void sendPrimeRequests(ManagedChannel managedChannel) { + try { + ClientCall clientCall = + managedChannel.newCall( + BigtableGrpc.getPingAndWarmMethod(), + GrpcCallContext.createDefault().getCallOptions().withCallCredentials(credentials)); + + SettableApiFuture future = SettableApiFuture.create(); + clientCall.start( + new ClientCall.Listener() { + PingAndWarmResponse response; + + @Override + public void onMessage(PingAndWarmResponse message) { + response = message; + } + + @Override + public void onClose(Status status, Metadata trailers) { + if (status.isOk()) { + future.set(response); + } else { + future.setException(status.asException()); + } + } + }, + metadata); + clientCall.sendMessage(request); + clientCall.halfClose(); + clientCall.request(1); - try (EnhancedBigtableStub stub = EnhancedBigtableStub.create(primingSettings)) { - PingAndWarmRequest request = - PingAndWarmRequest.newBuilder() - .setName( - NameUtil.formatInstanceName( - primingSettings.getProjectId(), primingSettings.getInstanceId())) - .setAppProfileId(primingSettings.getAppProfileId()) - .build(); - - try { - stub.pingAndWarmCallable().call(request); - } catch (Throwable e) { - // TODO: Not sure if we should swallow the error here. We are pre-emptively swapping - // channels if the new - // channel is bad. - if (e instanceof ExecutionException) { - e = e.getCause(); - } - LOG.warning(String.format("Failed to prime channel: %s", e)); - } + future.get(1, TimeUnit.MINUTES); + } catch (Throwable e) { + // TODO: Not sure if we should swallow the error here. We are pre-emptively swapping + // channels if the new + // channel is bad. + LOG.warning(String.format("Failed to prime channel: %s", e)); } } } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java index 5cab91c92c..cc95243975 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java @@ -188,7 +188,6 @@ public class EnhancedBigtableStub implements AutoCloseable { private final UnaryCallable externalBulkMutateRowsCallable; private final UnaryCallable checkAndMutateRowCallable; private final UnaryCallable readModifyWriteRowCallable; - private final UnaryCallable pingAndWarmCallable; private final ServerStreamingCallable generateInitialChangeStreamPartitionsCallable; @@ -321,7 +320,6 @@ public EnhancedBigtableStub( createGenerateInitialChangeStreamPartitionsCallable(); readChangeStreamCallable = createReadChangeStreamCallable(new DefaultChangeStreamRecordAdapter()); - pingAndWarmCallable = createPingAndWarmCallable(); executeQueryCallable = createExecuteQueryCallable(); } @@ -1252,28 +1250,6 @@ ServerStreamingCallSettings convertUnaryToServerStreamingSettings( .build(); } - private UnaryCallable createPingAndWarmCallable() { - UnaryCallable pingAndWarm = - GrpcRawCallableFactory.createUnaryCallable( - GrpcCallSettings.newBuilder() - .setMethodDescriptor(BigtableGrpc.getPingAndWarmMethod()) - .setParamsExtractor( - new RequestParamsExtractor() { - @Override - public Map extract(PingAndWarmRequest request) { - return ImmutableMap.of( - "name", request.getName(), - "app_profile_id", request.getAppProfileId()); - } - }) - .build(), - Collections.emptySet()); - return pingAndWarm.withDefaultCallContext( - clientContext - .getDefaultCallContext() - .withRetrySettings(settings.pingAndWarmSettings().getRetrySettings())); - } - private UnaryCallable withRetries( UnaryCallable innerCallable, UnaryCallSettings unaryCallSettings) { UnaryCallable retrying; @@ -1381,10 +1357,6 @@ public ExecuteQueryCallable executeQueryCallable() { return executeQueryCallable; } - UnaryCallable pingAndWarmCallable() { - return pingAndWarmCallable; - } - // private SpanName getSpanName(String methodName) { From 793d4dadcd87fde0155190d30834c07343d28cc5 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Mon, 2 Dec 2024 14:02:36 -0500 Subject: [PATCH 03/13] format --- .../cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java | 2 -- 1 file changed, 2 deletions(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java index cc95243975..46377fbc41 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java @@ -61,8 +61,6 @@ import com.google.bigtable.v2.GenerateInitialChangeStreamPartitionsResponse; import com.google.bigtable.v2.MutateRowsRequest; import com.google.bigtable.v2.MutateRowsResponse; -import com.google.bigtable.v2.PingAndWarmRequest; -import com.google.bigtable.v2.PingAndWarmResponse; import com.google.bigtable.v2.ReadChangeStreamRequest; import com.google.bigtable.v2.ReadChangeStreamResponse; import com.google.bigtable.v2.ReadRowsRequest; From 4859af07ce9eba4c5d453c79809d6f177cf25771 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Mon, 2 Dec 2024 14:32:18 -0500 Subject: [PATCH 04/13] address comments --- google-cloud-bigtable/pom.xml | 1 - .../data/v2/stub/BigtableChannelPrimer.java | 14 +++++++++----- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/google-cloud-bigtable/pom.xml b/google-cloud-bigtable/pom.xml index 7d027c995c..4717a24deb 100644 --- a/google-cloud-bigtable/pom.xml +++ b/google-cloud-bigtable/pom.xml @@ -709,7 +709,6 @@ grpc-auth is not directly used transitively, but is pulled to align with other grpc parts opencensus-impl-core is brought in transitively through opencensus-impl --> - io.grpc:grpc-auth io.opencensus:opencensus-impl-core diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java index 80e977f768..7c3b59efca 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java @@ -26,6 +26,7 @@ import com.google.common.net.PercentEscaper; import io.grpc.CallCredentials; import io.grpc.ClientCall; +import io.grpc.Deadline; import io.grpc.ManagedChannel; import io.grpc.Metadata; import io.grpc.Status; @@ -80,11 +81,11 @@ static BigtableChannelPrimer create(EnhancedBigtableStubSettings settings) throw .getHeaders() .forEach((k, v) -> metadata.put(Metadata.Key.of(k, Metadata.ASCII_STRING_MARSHALLER), v)); + String escapedName = escaper.escape(request.getName()); + String escapedAppProfile = escaper.escape(request.getAppProfileId()); metadata.put( requestParams, - escaper.escape( - String.format( - "name=%s&app_profile_id=%s", request.getName(), request.getAppProfileId()))); + escaper.escape(String.format("name=%s&app_profile_id=%s", escapedName, escapedAppProfile))); } @Override @@ -106,7 +107,10 @@ private void sendPrimeRequests(ManagedChannel managedChannel) { ClientCall clientCall = managedChannel.newCall( BigtableGrpc.getPingAndWarmMethod(), - GrpcCallContext.createDefault().getCallOptions().withCallCredentials(credentials)); + GrpcCallContext.createDefault() + .getCallOptions() + .withCallCredentials(credentials) + .withDeadline(Deadline.after(1, TimeUnit.MINUTES))); SettableApiFuture future = SettableApiFuture.create(); clientCall.start( @@ -130,7 +134,7 @@ public void onClose(Status status, Metadata trailers) { metadata); clientCall.sendMessage(request); clientCall.halfClose(); - clientCall.request(1); + clientCall.request(Integer.MAX_VALUE); future.get(1, TimeUnit.MINUTES); } catch (Throwable e) { From 87613134b6da0d5e33f76c7b8cd40b4172ad87a8 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Mon, 2 Dec 2024 14:34:12 -0500 Subject: [PATCH 05/13] update --- .../cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java index 7c3b59efca..faa5485996 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java @@ -81,11 +81,12 @@ static BigtableChannelPrimer create(EnhancedBigtableStubSettings settings) throw .getHeaders() .forEach((k, v) -> metadata.put(Metadata.Key.of(k, Metadata.ASCII_STRING_MARSHALLER), v)); - String escapedName = escaper.escape(request.getName()); - String escapedAppProfile = escaper.escape(request.getAppProfileId()); metadata.put( requestParams, - escaper.escape(String.format("name=%s&app_profile_id=%s", escapedName, escapedAppProfile))); + escaper.escape( + String.format( + "name=%s&app_profile_id=%s", + escaper.escape(request.getName()), escaper.escape(request.getAppProfileId())))); } @Override From 954828a11b7de18fa9b768947498edf88996b12a Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Mon, 2 Dec 2024 15:02:42 -0500 Subject: [PATCH 06/13] update --- .../data/v2/stub/BigtableChannelPrimer.java | 42 +++++++++++-------- .../data/v2/stub/BigtableClientContext.java | 8 +++- .../v2/stub/BigtableChannelPrimerTest.java | 19 ++++----- 3 files changed, 40 insertions(+), 29 deletions(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java index faa5485996..fe47c2798e 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java @@ -19,6 +19,7 @@ import com.google.api.core.SettableApiFuture; import com.google.api.gax.grpc.ChannelPrimer; import com.google.api.gax.grpc.GrpcCallContext; +import com.google.auth.Credentials; import com.google.bigtable.v2.BigtableGrpc; import com.google.bigtable.v2.InstanceName; import com.google.bigtable.v2.PingAndWarmRequest; @@ -32,6 +33,7 @@ import io.grpc.Status; import io.grpc.auth.MoreCallCredentials; import java.io.IOException; +import java.util.Map; import java.util.concurrent.TimeUnit; import java.util.logging.Logger; @@ -50,21 +52,29 @@ class BigtableChannelPrimer implements ChannelPrimer { Metadata.Key.of("x-goog-request-params", Metadata.ASCII_STRING_MARSHALLER); private final PingAndWarmRequest request; private final Metadata metadata = new Metadata(); - private final CallCredentials credentials; - - static BigtableChannelPrimer create(EnhancedBigtableStubSettings settings) throws IOException { - return new BigtableChannelPrimer(settings); + private final CallCredentials callCredentials; + + static BigtableChannelPrimer create( + String projectId, + String instanceId, + String appProfileId, + Credentials credentials, + Map headers) + throws IOException { + return new BigtableChannelPrimer(projectId, instanceId, appProfileId, credentials, headers); } - BigtableChannelPrimer(EnhancedBigtableStubSettings settings) throws IOException { - String projectId = settings.getProjectId(); - String instanceId = settings.getInstanceId(); - String appProfileId = settings.getAppProfileId(); - - if (settings.getCredentialsProvider().getCredentials() != null) { - credentials = MoreCallCredentials.from(settings.getCredentialsProvider().getCredentials()); + BigtableChannelPrimer( + String projectId, + String instanceId, + String appProfileId, + Credentials credentials, + Map headers) + throws IOException { + if (credentials != null) { + callCredentials = MoreCallCredentials.from(credentials); } else { - credentials = null; + callCredentials = null; } request = @@ -76,10 +86,8 @@ static BigtableChannelPrimer create(EnhancedBigtableStubSettings settings) throw PercentEscaper escaper = new PercentEscaper("._-~", false); Metadata metadata = new Metadata(); - settings - .getHeaderProvider() - .getHeaders() - .forEach((k, v) -> metadata.put(Metadata.Key.of(k, Metadata.ASCII_STRING_MARSHALLER), v)); + headers.forEach( + (k, v) -> metadata.put(Metadata.Key.of(k, Metadata.ASCII_STRING_MARSHALLER), v)); metadata.put( requestParams, @@ -110,7 +118,7 @@ private void sendPrimeRequests(ManagedChannel managedChannel) { BigtableGrpc.getPingAndWarmMethod(), GrpcCallContext.createDefault() .getCallOptions() - .withCallCredentials(credentials) + .withCallCredentials(callCredentials) .withDeadline(Deadline.after(1, TimeUnit.MINUTES))); SettableApiFuture future = SettableApiFuture.create(); diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java index ebd4df0bcd..170c5f24e8 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java @@ -103,7 +103,13 @@ public static BigtableClientContext create(EnhancedBigtableStubSettings settings } // Inject channel priming if enabled if (builder.isRefreshingChannel()) { - transportProvider.setChannelPrimer(BigtableChannelPrimer.create(builder.build())); + transportProvider.setChannelPrimer( + BigtableChannelPrimer.create( + builder.getProjectId(), + builder.getInstanceId(), + builder.getAppProfileId(), + credentials, + builder.getHeaderProvider().getHeaders())); } builder.setTransportChannelProvider(transportProvider.build()); diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimerTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimerTest.java index bfacd1d769..b1d3a9b11e 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimerTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimerTest.java @@ -18,13 +18,13 @@ import static com.google.common.truth.Truth.assertThat; import com.google.api.core.ApiFunction; -import com.google.api.gax.core.FixedCredentialsProvider; import com.google.auth.oauth2.AccessToken; import com.google.auth.oauth2.OAuth2Credentials; import com.google.bigtable.v2.BigtableGrpc.BigtableImplBase; import com.google.bigtable.v2.PingAndWarmRequest; import com.google.bigtable.v2.PingAndWarmResponse; import com.google.cloud.bigtable.data.v2.FakeServiceBuilder; +import com.google.common.collect.ImmutableMap; import io.grpc.ManagedChannel; import io.grpc.ManagedChannelBuilder; import io.grpc.Metadata; @@ -68,16 +68,13 @@ public void setup() throws IOException { server = FakeServiceBuilder.create(fakeService).intercept(metadataInterceptor).start(); - EnhancedBigtableStubSettings settings = - EnhancedBigtableStubSettings.newBuilder() - .setProjectId("fake-project") - .setInstanceId("fake-instance") - .setAppProfileId("fake-app-profile") - .setCredentialsProvider( - FixedCredentialsProvider.create( - OAuth2Credentials.create(new AccessToken(TOKEN_VALUE, null)))) - .build(); - primer = BigtableChannelPrimer.create(settings); + primer = + BigtableChannelPrimer.create( + "fake-project", + "fake-instance", + "fake-app-profile", + OAuth2Credentials.create(new AccessToken(TOKEN_VALUE, null)), + ImmutableMap.of()); channel = ManagedChannelBuilder.forAddress("localhost", server.getPort()).usePlaintext().build(); From cea1ffe815515e465cc5b277140cbf4ee12fbc1a Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Mon, 2 Dec 2024 15:07:10 -0500 Subject: [PATCH 07/13] fix escape --- .../cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java index fe47c2798e..be57e61db2 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java @@ -91,10 +91,9 @@ static BigtableChannelPrimer create( metadata.put( requestParams, - escaper.escape( - String.format( - "name=%s&app_profile_id=%s", - escaper.escape(request.getName()), escaper.escape(request.getAppProfileId())))); + String.format( + "name=%s&app_profile_id=%s", + escaper.escape(request.getName()), escaper.escape(request.getAppProfileId()))); } @Override From 802a197d91335c787213c77c906f1683af964cf3 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Mon, 2 Dec 2024 15:32:16 -0500 Subject: [PATCH 08/13] ues a different encoder --- .../data/v2/stub/BigtableChannelPrimer.java | 26 ++++++++++--------- 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java index be57e61db2..c04b5746e8 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java @@ -24,7 +24,6 @@ import com.google.bigtable.v2.InstanceName; import com.google.bigtable.v2.PingAndWarmRequest; import com.google.bigtable.v2.PingAndWarmResponse; -import com.google.common.net.PercentEscaper; import io.grpc.CallCredentials; import io.grpc.ClientCall; import io.grpc.Deadline; @@ -33,6 +32,8 @@ import io.grpc.Status; import io.grpc.auth.MoreCallCredentials; import java.io.IOException; +import java.io.UnsupportedEncodingException; +import java.net.URLEncoder; import java.util.Map; import java.util.concurrent.TimeUnit; import java.util.logging.Logger; @@ -59,8 +60,7 @@ static BigtableChannelPrimer create( String instanceId, String appProfileId, Credentials credentials, - Map headers) - throws IOException { + Map headers) { return new BigtableChannelPrimer(projectId, instanceId, appProfileId, credentials, headers); } @@ -69,8 +69,7 @@ static BigtableChannelPrimer create( String instanceId, String appProfileId, Credentials credentials, - Map headers) - throws IOException { + Map headers) { if (credentials != null) { callCredentials = MoreCallCredentials.from(credentials); } else { @@ -83,17 +82,20 @@ static BigtableChannelPrimer create( .setAppProfileId(appProfileId) .build(); - PercentEscaper escaper = new PercentEscaper("._-~", false); - Metadata metadata = new Metadata(); headers.forEach( (k, v) -> metadata.put(Metadata.Key.of(k, Metadata.ASCII_STRING_MARSHALLER), v)); - metadata.put( - requestParams, - String.format( - "name=%s&app_profile_id=%s", - escaper.escape(request.getName()), escaper.escape(request.getAppProfileId()))); + try { + metadata.put( + requestParams, + String.format( + "name=%s&app_profile_id=%s", + URLEncoder.encode(request.getName(), "UTF-8"), + URLEncoder.encode(request.getAppProfileId(), "UTF-8"))); + } catch (UnsupportedEncodingException e) { + LOG.warning(String.format("failed to encode request params %s", e)); + } } @Override From d113ba4e38bd55265d5a70351c3f74092a87404c Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Mon, 2 Dec 2024 17:53:14 -0500 Subject: [PATCH 09/13] fix logging --- .../cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java index c04b5746e8..adcc5d77cc 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java @@ -36,6 +36,7 @@ import java.net.URLEncoder; import java.util.Map; import java.util.concurrent.TimeUnit; +import java.util.logging.Level; import java.util.logging.Logger; /** @@ -94,7 +95,7 @@ static BigtableChannelPrimer create( URLEncoder.encode(request.getName(), "UTF-8"), URLEncoder.encode(request.getAppProfileId(), "UTF-8"))); } catch (UnsupportedEncodingException e) { - LOG.warning(String.format("failed to encode request params %s", e)); + LOG.log(Level.WARNING, "Failed to encode request params", e); } } From 415ef20f90b93676ba79696894ac6d54af53f044 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Mon, 2 Dec 2024 20:16:02 -0500 Subject: [PATCH 10/13] update --- .../data/v2/stub/BigtableChannelPrimer.java | 43 +++++++++++-------- .../v2/stub/BigtableChannelPrimerTest.java | 16 ++++++- 2 files changed, 39 insertions(+), 20 deletions(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java index adcc5d77cc..5a4778e890 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java @@ -50,11 +50,11 @@ class BigtableChannelPrimer implements ChannelPrimer { private static Logger LOG = Logger.getLogger(BigtableChannelPrimer.class.toString()); - private final Metadata.Key requestParams = + static final Metadata.Key REQUEST_PARAMS = Metadata.Key.of("x-goog-request-params", Metadata.ASCII_STRING_MARSHALLER); private final PingAndWarmRequest request; - private final Metadata metadata = new Metadata(); private final CallCredentials callCredentials; + private final Map headers; static BigtableChannelPrimer create( String projectId, @@ -83,20 +83,7 @@ static BigtableChannelPrimer create( .setAppProfileId(appProfileId) .build(); - Metadata metadata = new Metadata(); - headers.forEach( - (k, v) -> metadata.put(Metadata.Key.of(k, Metadata.ASCII_STRING_MARSHALLER), v)); - - try { - metadata.put( - requestParams, - String.format( - "name=%s&app_profile_id=%s", - URLEncoder.encode(request.getName(), "UTF-8"), - URLEncoder.encode(request.getAppProfileId(), "UTF-8"))); - } catch (UnsupportedEncodingException e) { - LOG.log(Level.WARNING, "Failed to encode request params", e); - } + this.headers = headers; } @Override @@ -104,8 +91,7 @@ public void primeChannel(ManagedChannel managedChannel) { try { primeChannelUnsafe(managedChannel); } catch (IOException | RuntimeException e) { - LOG.warning( - String.format("Unexpected error while trying to prime a channel: %s", e.getMessage())); + LOG.log(Level.WARNING, "Unexpected error while trying to prime a channel", e); } } @@ -142,7 +128,7 @@ public void onClose(Status status, Metadata trailers) { } } }, - metadata); + createMetadata(headers, request)); clientCall.sendMessage(request); clientCall.halfClose(); clientCall.request(Integer.MAX_VALUE); @@ -155,4 +141,23 @@ public void onClose(Status status, Metadata trailers) { LOG.warning(String.format("Failed to prime channel: %s", e)); } } + + private Metadata createMetadata(Map headers, PingAndWarmRequest request) { + Metadata metadata = new Metadata(); + + headers.forEach( + (k, v) -> metadata.put(Metadata.Key.of(k, Metadata.ASCII_STRING_MARSHALLER), v)); + try { + metadata.put( + REQUEST_PARAMS, + String.format( + "name=%s&app_profile_id=%s", + URLEncoder.encode(request.getName(), "UTF-8"), + URLEncoder.encode(request.getAppProfileId(), "UTF-8"))); + } catch (UnsupportedEncodingException e) { + LOG.log(Level.WARNING, "Failed to encode request params", e); + } + + return metadata; + } } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimerTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimerTest.java index b1d3a9b11e..03d290ad24 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimerTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimerTest.java @@ -74,7 +74,7 @@ public void setup() throws IOException { "fake-instance", "fake-app-profile", OAuth2Credentials.create(new AccessToken(TOKEN_VALUE, null)), - ImmutableMap.of()); + ImmutableMap.of("bigtable-feature", "fake-feature")); channel = ManagedChannelBuilder.forAddress("localhost", server.getPort()).usePlaintext().build(); @@ -152,6 +152,20 @@ public void testChannelErrorsAreLogged() { } } + @Test + public void testHeadersAreSent() { + primer.primeChannel(channel); + + for (Metadata metadata : metadataInterceptor.metadataList) { + assertThat(metadata.get(BigtableChannelPrimer.REQUEST_PARAMS)) + .isEqualTo( + "name=projects%2Ffake-project%2Finstances%2Ffake-instance&app_profile_id=fake-app-profile"); + assertThat( + metadata.get(Metadata.Key.of("bigtable-feature", Metadata.ASCII_STRING_MARSHALLER))) + .isEqualTo("fake-feature"); + } + } + private static class MetadataInterceptor implements ServerInterceptor { ConcurrentLinkedQueue metadataList = new ConcurrentLinkedQueue<>(); From 5f916238d2391ccfce4075fde1a2002a6b6880ac Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 3 Dec 2024 10:32:53 -0500 Subject: [PATCH 11/13] address comments` --- .../bigtable/data/v2/stub/BigtableChannelPrimer.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java index 5a4778e890..99600f4485 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java @@ -25,6 +25,7 @@ import com.google.bigtable.v2.PingAndWarmRequest; import com.google.bigtable.v2.PingAndWarmResponse; import io.grpc.CallCredentials; +import io.grpc.CallOptions; import io.grpc.ClientCall; import io.grpc.Deadline; import io.grpc.ManagedChannel; @@ -104,8 +105,7 @@ private void sendPrimeRequests(ManagedChannel managedChannel) { ClientCall clientCall = managedChannel.newCall( BigtableGrpc.getPingAndWarmMethod(), - GrpcCallContext.createDefault() - .getCallOptions() + CallOptions.DEFAULT .withCallCredentials(callCredentials) .withDeadline(Deadline.after(1, TimeUnit.MINUTES))); @@ -138,11 +138,11 @@ public void onClose(Status status, Metadata trailers) { // TODO: Not sure if we should swallow the error here. We are pre-emptively swapping // channels if the new // channel is bad. - LOG.warning(String.format("Failed to prime channel: %s", e)); + LOG.log(Level.WARNING, "failed to prime channel", e); } } - private Metadata createMetadata(Map headers, PingAndWarmRequest request) { + private static Metadata createMetadata(Map headers, PingAndWarmRequest request) { Metadata metadata = new Metadata(); headers.forEach( From 9a31cc39e0a7dba52374d3b51a90129f98f498bd Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 3 Dec 2024 10:37:51 -0500 Subject: [PATCH 12/13] remove import --- .../cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java | 1 - 1 file changed, 1 deletion(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java index 99600f4485..74f2b28384 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java @@ -18,7 +18,6 @@ import com.google.api.core.BetaApi; import com.google.api.core.SettableApiFuture; import com.google.api.gax.grpc.ChannelPrimer; -import com.google.api.gax.grpc.GrpcCallContext; import com.google.auth.Credentials; import com.google.bigtable.v2.BigtableGrpc; import com.google.bigtable.v2.InstanceName; From 8e5af00d7706e252dd33eb1d54b00d6323d94345 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 3 Dec 2024 11:38:45 -0500 Subject: [PATCH 13/13] fix test --- .../cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java | 2 +- .../bigtable/data/v2/stub/BigtableChannelPrimerTest.java | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java index 74f2b28384..7495ca6ceb 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java @@ -137,7 +137,7 @@ public void onClose(Status status, Metadata trailers) { // TODO: Not sure if we should swallow the error here. We are pre-emptively swapping // channels if the new // channel is bad. - LOG.log(Level.WARNING, "failed to prime channel", e); + LOG.log(Level.WARNING, "Failed to prime channel", e); } } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimerTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimerTest.java index 03d290ad24..709b482477 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimerTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimerTest.java @@ -135,7 +135,7 @@ public PingAndWarmResponse apply(PingAndWarmRequest pingAndWarmRequest) { assertThat(logHandler.logs).hasSize(1); for (LogRecord log : logHandler.logs) { - assertThat(log.getMessage()).contains("FAILED_PRECONDITION"); + assertThat(log.getThrown().getMessage()).contains("FAILED_PRECONDITION"); } } @@ -148,7 +148,7 @@ public void testChannelErrorsAreLogged() { assertThat(logHandler.logs).hasSize(1); for (LogRecord log : logHandler.logs) { - assertThat(log.getMessage()).contains("UnsupportedOperationException"); + assertThat(log.getThrown()).isInstanceOf(UnsupportedOperationException.class); } }