Skip to content
This repository has been archived by the owner on Sep 26, 2023. It is now read-only.

Commit

Permalink
feat: add custom options to ApiCallContext (#1435)
Browse files Browse the repository at this point in the history
  • Loading branch information
@mutianf
mutianf committed Aug 16, 2021
1 parent fc2f997 commit 0fe20f3
Show file tree
Hide file tree
Showing 8 changed files with 455 additions and 4 deletions.
46 changes: 44 additions & 2 deletions gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcCallContext.java
View file
Expand Up @@ -34,6 +34,7 @@
import com.google.api.gax.rpc.ApiCallContext;
import com.google.api.gax.rpc.StatusCode;
import com.google.api.gax.rpc.TransportChannel;
import com.google.api.gax.rpc.internal.ApiCallContextOptions;
import com.google.api.gax.rpc.internal.Headers;
import com.google.api.gax.tracing.ApiTracer;
import com.google.api.gax.tracing.BaseApiTracer;
Expand All @@ -43,7 +44,6 @@
import com.google.common.collect.ImmutableSet;
import io.grpc.CallCredentials;
import io.grpc.CallOptions;
import io.grpc.CallOptions.Key;
import io.grpc.Channel;
import io.grpc.Deadline;
import io.grpc.Metadata;
Expand All @@ -66,7 +66,7 @@
*/
@BetaApi("Reference ApiCallContext instead - this class is likely to experience breaking changes")
public final class GrpcCallContext implements ApiCallContext {
static final CallOptions.Key<ApiTracer> TRACER_KEY = Key.create("gax.tracer");
static final CallOptions.Key<ApiTracer> TRACER_KEY = CallOptions.Key.create("gax.tracer");

private final Channel channel;
private final CallOptions callOptions;
Expand All @@ -77,6 +77,7 @@ public final class GrpcCallContext implements ApiCallContext {
@Nullable private final RetrySettings retrySettings;
@Nullable private final ImmutableSet<StatusCode.Code> retryableCodes;
private final ImmutableMap<String, List<String>> extraHeaders;
private final ApiCallContextOptions options;

/** Returns an empty instance with a null channel and default {@link CallOptions}. */
public static GrpcCallContext createDefault() {
Expand All @@ -88,6 +89,7 @@ public static GrpcCallContext createDefault() {
null,
null,
ImmutableMap.<String, List<String>>of(),
ApiCallContextOptions.getDefaultOptions(),
null,
null);
}
Expand All @@ -102,6 +104,7 @@ public static GrpcCallContext of(Channel channel, CallOptions callOptions) {
null,
null,
ImmutableMap.<String, List<String>>of(),
ApiCallContextOptions.getDefaultOptions(),
null,
null);
}
Expand All @@ -114,6 +117,7 @@ private GrpcCallContext(
@Nullable Duration streamIdleTimeout,
@Nullable Integer channelAffinity,
ImmutableMap<String, List<String>> extraHeaders,
ApiCallContextOptions options,
@Nullable RetrySettings retrySettings,
@Nullable Set<StatusCode.Code> retryableCodes) {
this.channel = channel;
Expand All @@ -123,6 +127,7 @@ private GrpcCallContext(
this.streamIdleTimeout = streamIdleTimeout;
this.channelAffinity = channelAffinity;
this.extraHeaders = Preconditions.checkNotNull(extraHeaders);
this.options = Preconditions.checkNotNull(options);
this.retrySettings = retrySettings;
this.retryableCodes = retryableCodes == null ? null : ImmutableSet.copyOf(retryableCodes);
}
Expand Down Expand Up @@ -187,6 +192,7 @@ public GrpcCallContext withTimeout(@Nullable Duration timeout) {
this.streamIdleTimeout,
this.channelAffinity,
this.extraHeaders,
this.options,
this.retrySettings,
this.retryableCodes);
}
Expand All @@ -212,6 +218,7 @@ public GrpcCallContext withStreamWaitTimeout(@Nullable Duration streamWaitTimeou
this.streamIdleTimeout,
this.channelAffinity,
this.extraHeaders,
this.options,
this.retrySettings,
this.retryableCodes);
}
Expand All @@ -231,6 +238,7 @@ public GrpcCallContext withStreamIdleTimeout(@Nullable Duration streamIdleTimeou
streamIdleTimeout,
this.channelAffinity,
this.extraHeaders,
this.options,
this.retrySettings,
this.retryableCodes);
}
Expand All @@ -245,6 +253,7 @@ public GrpcCallContext withChannelAffinity(@Nullable Integer affinity) {
this.streamIdleTimeout,
affinity,
this.extraHeaders,
this.options,
this.retrySettings,
this.retryableCodes);
}
Expand All @@ -263,6 +272,7 @@ public GrpcCallContext withExtraHeaders(Map<String, List<String>> extraHeaders)
this.streamIdleTimeout,
this.channelAffinity,
newExtraHeaders,
this.options,
this.retrySettings,
this.retryableCodes);
}
Expand All @@ -282,6 +292,7 @@ public GrpcCallContext withRetrySettings(RetrySettings retrySettings) {
this.streamIdleTimeout,
this.channelAffinity,
this.extraHeaders,
this.options,
retrySettings,
this.retryableCodes);
}
Expand All @@ -301,6 +312,7 @@ public GrpcCallContext withRetryableCodes(Set<StatusCode.Code> retryableCodes) {
this.streamIdleTimeout,
this.channelAffinity,
this.extraHeaders,
this.options,
this.retrySettings,
retryableCodes);
}
Expand Down Expand Up @@ -370,6 +382,8 @@ public ApiCallContext merge(ApiCallContext inputCallContext) {
ImmutableMap<String, List<String>> newExtraHeaders =
Headers.mergeHeaders(this.extraHeaders, grpcCallContext.extraHeaders);

ApiCallContextOptions newOptions = options.merge(grpcCallContext.options);

CallOptions newCallOptions =
grpcCallContext
.callOptions
Expand All @@ -388,6 +402,7 @@ public ApiCallContext merge(ApiCallContext inputCallContext) {
newStreamIdleTimeout,
newChannelAffinity,
newExtraHeaders,
newOptions,
newRetrySettings,
newRetryableCodes);
}
Expand Down Expand Up @@ -448,6 +463,7 @@ public GrpcCallContext withChannel(Channel newChannel) {
this.streamIdleTimeout,
this.channelAffinity,
this.extraHeaders,
this.options,
this.retrySettings,
this.retryableCodes);
}
Expand All @@ -462,6 +478,7 @@ public GrpcCallContext withCallOptions(CallOptions newCallOptions) {
this.streamIdleTimeout,
this.channelAffinity,
this.extraHeaders,
this.options,
this.retrySettings,
this.retryableCodes);
}
Expand Down Expand Up @@ -491,6 +508,29 @@ public GrpcCallContext withTracer(@Nonnull ApiTracer tracer) {
return withCallOptions(callOptions.withOption(TRACER_KEY, tracer));
}

/** {@inheritDoc} */
@Override
public <T> GrpcCallContext withOption(Key<T> key, T value) {
ApiCallContextOptions newOptions = options.withOption(key, value);
return new GrpcCallContext(
this.channel,
this.callOptions,
this.timeout,
this.streamWaitTimeout,
this.streamIdleTimeout,
this.channelAffinity,
this.extraHeaders,
newOptions,
this.retrySettings,
this.retryableCodes);
}

/** {@inheritDoc} */
@Override
public <T> T getOption(Key<T> key) {
return options.getOption(key);
}

@Override
public int hashCode() {
return Objects.hash(
Expand All @@ -501,6 +541,7 @@ public int hashCode() {
streamIdleTimeout,
channelAffinity,
extraHeaders,
options,
retrySettings,
retryableCodes);
}
Expand All @@ -522,6 +563,7 @@ public boolean equals(Object o) {
&& Objects.equals(this.streamIdleTimeout, that.streamIdleTimeout)
&& Objects.equals(this.channelAffinity, that.channelAffinity)
&& Objects.equals(this.extraHeaders, that.extraHeaders)
&& Objects.equals(this.options, that.options)
&& Objects.equals(this.retrySettings, that.retrySettings)
&& Objects.equals(this.retryableCodes, that.retryableCodes);
}
Expand Down
42 changes: 42 additions & 0 deletions gax-grpc/src/test/java/com/google/api/gax/grpc/GrpcCallContextTest.java
View file
Expand Up @@ -354,6 +354,48 @@ public void testWithRetryableCodes() {
assertNotNull(context.getRetryableCodes());
}

@Test
public void testWithOptions() {
GrpcCallContext emptyCallContext = GrpcCallContext.createDefault();
ApiCallContext.Key<String> contextKey1 = ApiCallContext.Key.create("testKey1");
ApiCallContext.Key<String> contextKey2 = ApiCallContext.Key.create("testKey2");
String testContext1 = "test1";
String testContext2 = "test2";
String testContextOverwrite = "test1Overwrite";
GrpcCallContext context =
emptyCallContext
.withOption(contextKey1, testContext1)
.withOption(contextKey2, testContext2);
assertEquals(testContext1, context.getOption(contextKey1));
assertEquals(testContext2, context.getOption(contextKey2));
GrpcCallContext newContext = context.withOption(contextKey1, testContextOverwrite);
assertEquals(testContextOverwrite, newContext.getOption(contextKey1));
}

@Test
public void testMergeOptions() {
GrpcCallContext emptyCallContext = GrpcCallContext.createDefault();
ApiCallContext.Key<String> contextKey1 = ApiCallContext.Key.create("testKey1");
ApiCallContext.Key<String> contextKey2 = ApiCallContext.Key.create("testKey2");
ApiCallContext.Key<String> contextKey3 = ApiCallContext.Key.create("testKey3");
String testContext1 = "test1";
String testContext2 = "test2";
String testContext3 = "test3";
String testContextOverwrite = "test1Overwrite";
GrpcCallContext context1 =
emptyCallContext
.withOption(contextKey1, testContext1)
.withOption(contextKey2, testContext2);
GrpcCallContext context2 =
emptyCallContext
.withOption(contextKey1, testContextOverwrite)
.withOption(contextKey3, testContext3);
ApiCallContext mergedContext = context1.merge(context2);
assertEquals(testContextOverwrite, mergedContext.getOption(contextKey1));
assertEquals(testContext2, mergedContext.getOption(contextKey2));
assertEquals(testContext3, mergedContext.getOption(contextKey3));
}

private static Map<String, List<String>> createTestExtraHeaders(String... keyValues) {
Map<String, List<String>> extraHeaders = new HashMap<>();
for (int i = 0; i < keyValues.length; i += 2) {
Expand Down

0 comments on commit 0fe20f3

Please sign in to comment.