From acf1fec2c255f134b9ca8273211502708fb3c3a4 Mon Sep 17 00:00:00 2001 From: Weiran Fang <8175562+WeiranFang@users.noreply.github.com> Date: Thu, 28 May 2020 09:19:28 -0700 Subject: [PATCH] feat: add setAttemptDirectPath(boolean) for InstantiatingGrpcChannelProvider.Builder (#1015) * Use boolean flag instead of env var to determine whether to attempt DirectPath. * Remove EnvironmentProvider because we are not using DirectPath env var any more * Keep DP env var checker until setAttemptDirectPath is adapted * fix logic OR * Check env var only when attemptDirectPath is not set --- .../InstantiatingGrpcChannelProvider.java | 28 ++++++++++++------- .../InstantiatingGrpcChannelProviderTest.java | 27 ++++-------------- 2 files changed, 24 insertions(+), 31 deletions(-) diff --git a/gax-grpc/src/main/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProvider.java b/gax-grpc/src/main/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProvider.java index deaa59176..77a933354 100644 --- a/gax-grpc/src/main/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProvider.java +++ b/gax-grpc/src/main/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProvider.java @@ -89,6 +89,7 @@ public final class InstantiatingGrpcChannelProvider implements TransportChannelP @Nullable private final Integer poolSize; @Nullable private final Credentials credentials; @Nullable private final ChannelPrimer channelPrimer; + @Nullable private final Boolean attemptDirectPath; @Nullable private final ApiFunction channelConfigurator; @@ -109,6 +110,7 @@ private InstantiatingGrpcChannelProvider(Builder builder) { this.channelConfigurator = builder.channelConfigurator; this.credentials = builder.credentials; this.channelPrimer = builder.channelPrimer; + this.attemptDirectPath = builder.attemptDirectPath; } @Override @@ -211,9 +213,13 @@ public ManagedChannel createSingleChannel() throws IOException { return GrpcTransportChannel.create(outerChannel); } - // The environment variable is used during the rollout phase for directpath. - // This checker function will be removed once directpath is stable. + // TODO(weiranf): Use attemptDirectPath as the only indicator once setAttemptDirectPath is adapted + // and the env var is removed from client environment. private boolean isDirectPathEnabled(String serviceAddress) { + if (attemptDirectPath != null) { + return attemptDirectPath; + } + // Only check DIRECT_PATH_ENV_VAR when attemptDirectPath is not set. String whiteList = envProvider.getenv(DIRECT_PATH_ENV_VAR); if (whiteList == null) return false; for (String service : whiteList.split(",")) { @@ -238,8 +244,7 @@ private ManagedChannel createSingleChannel() throws IOException { ManagedChannelBuilder builder; - // TODO(weiranf): Add a new API in ComputeEngineCredentials to check whether it's using default - // service account. + // TODO(weiranf): Add API in ComputeEngineCredentials to check default service account. if (isDirectPathEnabled(serviceAddress) && credentials instanceof ComputeEngineCredentials) { builder = ComputeEngineChannelBuilder.forAddress(serviceAddress, port); // Set default keepAliveTime and keepAliveTimeout when directpath environment is enabled. @@ -363,6 +368,7 @@ public static final class Builder { @Nullable private ApiFunction channelConfigurator; @Nullable private Credentials credentials; @Nullable private ChannelPrimer channelPrimer; + @Nullable private Boolean attemptDirectPath; private Builder() { processorCount = Runtime.getRuntime().availableProcessors(); @@ -385,6 +391,7 @@ private Builder(InstantiatingGrpcChannelProvider provider) { this.channelConfigurator = provider.channelConfigurator; this.credentials = provider.credentials; this.channelPrimer = provider.channelPrimer; + this.attemptDirectPath = provider.attemptDirectPath; } /** Sets the number of available CPUs, used internally for testing. */ @@ -393,12 +400,6 @@ Builder setProcessorCount(int processorCount) { return this; } - /** Sets the environment variable provider used for testing. */ - Builder setEnvironmentProvider(EnvironmentProvider envProvider) { - this.envProvider = envProvider; - return this; - } - /** * Sets the ExecutorProvider for this TransportChannelProvider. * @@ -565,6 +566,13 @@ public Builder setChannelPrimer(ChannelPrimer channelPrimer) { return this; } + /** Whether attempt DirectPath. */ + @InternalApi("For internal use by google-cloud-java clients only") + public Builder setAttemptDirectPath(boolean attemptDirectPath) { + this.attemptDirectPath = attemptDirectPath; + return this; + } + public InstantiatingGrpcChannelProvider build() { return new InstantiatingGrpcChannelProvider(this); } diff --git a/gax-grpc/src/test/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProviderTest.java b/gax-grpc/src/test/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProviderTest.java index 9d2df5756..9d1a4fd68 100644 --- a/gax-grpc/src/test/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProviderTest.java +++ b/gax-grpc/src/test/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProviderTest.java @@ -29,7 +29,6 @@ */ package com.google.api.gax.grpc; -import static com.google.api.gax.grpc.InstantiatingGrpcChannelProvider.DIRECT_PATH_ENV_VAR; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertEquals; import static org.junit.Assert.fail; @@ -37,7 +36,6 @@ import com.google.api.core.ApiFunction; import com.google.api.gax.core.ExecutorProvider; import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider.Builder; -import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider.EnvironmentProvider; import com.google.api.gax.rpc.HeaderProvider; import com.google.api.gax.rpc.TransportChannelProvider; import com.google.auth.oauth2.CloudShellCredentials; @@ -215,9 +213,6 @@ public void testChannelConfigurator() throws IOException { @Test public void testWithGCECredentials() throws IOException { - EnvironmentProvider mockEnvProvider = Mockito.mock(EnvironmentProvider.class); - Mockito.when(mockEnvProvider.getenv(DIRECT_PATH_ENV_VAR)).thenReturn("localhost"); - ScheduledExecutorService executor = new ScheduledThreadPoolExecutor(1); executor.shutdown(); @@ -231,7 +226,7 @@ public ManagedChannelBuilder apply(ManagedChannelBuilder channelBuilder) { TransportChannelProvider provider = InstantiatingGrpcChannelProvider.newBuilder() - .setEnvironmentProvider(mockEnvProvider) + .setAttemptDirectPath(true) .setChannelConfigurator(channelConfigurator) .build() .withExecutor(executor) @@ -247,9 +242,6 @@ public ManagedChannelBuilder apply(ManagedChannelBuilder channelBuilder) { @Test public void testWithNonGCECredentials() throws IOException { - EnvironmentProvider mockEnvProvider = Mockito.mock(EnvironmentProvider.class); - Mockito.when(mockEnvProvider.getenv(DIRECT_PATH_ENV_VAR)).thenReturn("localhost"); - ScheduledExecutorService executor = new ScheduledThreadPoolExecutor(1); executor.shutdown(); @@ -264,7 +256,7 @@ public ManagedChannelBuilder apply(ManagedChannelBuilder channelBuilder) { TransportChannelProvider provider = InstantiatingGrpcChannelProvider.newBuilder() - .setEnvironmentProvider(mockEnvProvider) + .setAttemptDirectPath(true) .setChannelConfigurator(channelConfigurator) .build() .withExecutor(executor) @@ -280,16 +272,13 @@ public ManagedChannelBuilder apply(ManagedChannelBuilder channelBuilder) { @Test public void testWithDirectPathDisabled() throws IOException { - EnvironmentProvider mockEnvProvider = Mockito.mock(EnvironmentProvider.class); - Mockito.when(mockEnvProvider.getenv(DIRECT_PATH_ENV_VAR)).thenReturn("otherhost"); - ScheduledExecutorService executor = new ScheduledThreadPoolExecutor(1); executor.shutdown(); ApiFunction channelConfigurator = new ApiFunction() { public ManagedChannelBuilder apply(ManagedChannelBuilder channelBuilder) { - // Clients without DirectPath environment variable will not attempt DirectPath + // Clients without setting attemptDirectPath flag will not attempt DirectPath assertThat(channelBuilder instanceof ComputeEngineChannelBuilder).isFalse(); return channelBuilder; } @@ -297,7 +286,7 @@ public ManagedChannelBuilder apply(ManagedChannelBuilder channelBuilder) { TransportChannelProvider provider = InstantiatingGrpcChannelProvider.newBuilder() - .setEnvironmentProvider(mockEnvProvider) + .setAttemptDirectPath(false) .setChannelConfigurator(channelConfigurator) .build() .withExecutor(executor) @@ -312,17 +301,14 @@ public ManagedChannelBuilder apply(ManagedChannelBuilder channelBuilder) { } @Test - public void testWithNoDirectPathEnvironment() throws IOException { - EnvironmentProvider mockEnvProvider = Mockito.mock(EnvironmentProvider.class); - Mockito.when(mockEnvProvider.getenv(DIRECT_PATH_ENV_VAR)).thenReturn(null); - + public void testWithNoDirectPathFlagSet() throws IOException { ScheduledExecutorService executor = new ScheduledThreadPoolExecutor(1); executor.shutdown(); ApiFunction channelConfigurator = new ApiFunction() { public ManagedChannelBuilder apply(ManagedChannelBuilder channelBuilder) { - // Clients without DirectPath environment variable will not attempt DirectPath + // Clients without setting attemptDirectPath flag will not attempt DirectPath assertThat(channelBuilder instanceof ComputeEngineChannelBuilder).isFalse(); return channelBuilder; } @@ -330,7 +316,6 @@ public ManagedChannelBuilder apply(ManagedChannelBuilder channelBuilder) { TransportChannelProvider provider = InstantiatingGrpcChannelProvider.newBuilder() - .setEnvironmentProvider(mockEnvProvider) .setChannelConfigurator(channelConfigurator) .build() .withExecutor(executor)