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

Commit

Permalink
refactor
Browse files Browse the repository at this point in the history
  • Loading branch information
@mutianf
mutianf committed Jul 28, 2021
1 parent c57c195 commit 9939551
Show file tree
Hide file tree
Showing 8 changed files with 203 additions and 192 deletions.
56 changes: 24 additions & 32 deletions gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcCallContext.java
View file
Expand Up @@ -34,7 +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.ApiCallContextUtil;
import com.google.api.gax.rpc.internal.ApiCallCustomContexts;
import com.google.api.gax.rpc.internal.Headers;
import com.google.api.gax.tracing.ApiTracer;
import com.google.api.gax.tracing.BaseApiTracer;
Expand Down Expand Up @@ -77,7 +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 ImmutableMap<Key, Object> extraContexts;
private final ApiCallCustomContexts customContexts;

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

ImmutableMap<Key, Object> newExtraContexts =
ApiCallContextUtil.mergeExtraContexts(this.extraContexts, grpcCallContext.extraContexts);
ApiCallCustomContexts newCustomContexts = customContexts.merge(grpcCallContext.customContexts);

CallOptions newCallOptions =
grpcCallContext
Expand All @@ -403,7 +402,7 @@ public ApiCallContext merge(ApiCallContext inputCallContext) {
newStreamIdleTimeout,
newChannelAffinity,
newExtraHeaders,
newExtraContexts,
newCustomContexts,
newRetrySettings,
newRetryableCodes);
}
Expand Down Expand Up @@ -464,7 +463,7 @@ public GrpcCallContext withChannel(Channel newChannel) {
this.streamIdleTimeout,
this.channelAffinity,
this.extraHeaders,
this.extraContexts,
this.customContexts,
this.retrySettings,
this.retryableCodes);
}
Expand All @@ -479,7 +478,7 @@ public GrpcCallContext withCallOptions(CallOptions newCallOptions) {
this.streamIdleTimeout,
this.channelAffinity,
this.extraHeaders,
this.extraContexts,
this.customContexts,
this.retrySettings,
this.retryableCodes);
}
Expand Down Expand Up @@ -511,11 +510,8 @@ public GrpcCallContext withTracer(@Nonnull ApiTracer tracer) {

/** {@inheritDoc} */
@Override
public <T> GrpcCallContext withExtraContext(Key<T> key, T extraContext) {
Preconditions.checkNotNull(key);

ImmutableMap<Key, Object> newExtraContexts =
ApiCallContextUtil.addExtraContext(extraContexts, key, extraContext);
public <T> GrpcCallContext withCustomContext(Key<T> key, T value) {
ApiCallCustomContexts newCustomContexts = customContexts.withCustomContext(key, value);
return new GrpcCallContext(
this.channel,
this.callOptions,
Expand All @@ -524,19 +520,15 @@ public <T> GrpcCallContext withExtraContext(Key<T> key, T extraContext) {
this.streamIdleTimeout,
this.channelAffinity,
this.extraHeaders,
newExtraContexts,
newCustomContexts,
this.retrySettings,
this.retryableCodes);
}

/** {@inheritDoc} */
@Override
public <T> T getExtraContext(Key<T> key) {
Preconditions.checkNotNull(key);
if (extraContexts.containsKey(key)) {
return (T) extraContexts.get(key);
}
return key.getDefault();
public <T> T getCustomContext(Key<T> key) {
return customContexts.getCustomContext(key);
}

@Override
Expand All @@ -549,7 +541,7 @@ public int hashCode() {
streamIdleTimeout,
channelAffinity,
extraHeaders,
extraContexts,
customContexts,
retrySettings,
retryableCodes);
}
Expand All @@ -571,7 +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.extraContexts, that.extraContexts)
&& Objects.equals(this.customContexts, that.customContexts)
&& Objects.equals(this.retrySettings, that.retrySettings)
&& Objects.equals(this.retryableCodes, that.retryableCodes);
}
Expand Down
30 changes: 15 additions & 15 deletions gax-grpc/src/test/java/com/google/api/gax/grpc/GrpcCallContextTest.java
View file
Expand Up @@ -355,7 +355,7 @@ public void testWithRetryableCodes() {
}

@Test
public void testWithExtraContext() {
public void testWithCustomContexts() {
GrpcCallContext emptyCallContext = GrpcCallContext.createDefault();
ApiCallContext.Key<String> contextKey1 = ApiCallContext.Key.create("testKey1");
ApiCallContext.Key<String> contextKey2 = ApiCallContext.Key.create("testKey2");
Expand All @@ -364,16 +364,16 @@ public void testWithExtraContext() {
String testContextOverwrite = "test1Overwrite";
GrpcCallContext context =
emptyCallContext
.withExtraContext(contextKey1, testContext1)
.withExtraContext(contextKey2, testContext2);
assertEquals(testContext1, context.getExtraContext(contextKey1));
assertEquals(testContext2, context.getExtraContext(contextKey2));
GrpcCallContext newContext = context.withExtraContext(contextKey1, testContextOverwrite);
assertEquals(testContextOverwrite, newContext.getExtraContext(contextKey1));
.withCustomContext(contextKey1, testContext1)
.withCustomContext(contextKey2, testContext2);
assertEquals(testContext1, context.getCustomContext(contextKey1));
assertEquals(testContext2, context.getCustomContext(contextKey2));
GrpcCallContext newContext = context.withCustomContext(contextKey1, testContextOverwrite);
assertEquals(testContextOverwrite, newContext.getCustomContext(contextKey1));
}

@Test
public void testMergeExtraContext() {
public void testMergeCustomContexts() {
GrpcCallContext emptyCallContext = GrpcCallContext.createDefault();
ApiCallContext.Key<String> contextKey1 = ApiCallContext.Key.create("testKey1");
ApiCallContext.Key<String> contextKey2 = ApiCallContext.Key.create("testKey2");
Expand All @@ -384,16 +384,16 @@ public void testMergeExtraContext() {
String testContextOverwrite = "test1Overwrite";
GrpcCallContext context1 =
emptyCallContext
.withExtraContext(contextKey1, testContext1)
.withExtraContext(contextKey2, testContext2);
.withCustomContext(contextKey1, testContext1)
.withCustomContext(contextKey2, testContext2);
GrpcCallContext context2 =
emptyCallContext
.withExtraContext(contextKey1, testContextOverwrite)
.withExtraContext(contextKey3, testContext3);
.withCustomContext(contextKey1, testContextOverwrite)
.withCustomContext(contextKey3, testContext3);
ApiCallContext mergedContext = context1.merge(context2);
assertEquals(testContextOverwrite, mergedContext.getExtraContext(contextKey1));
assertEquals(testContext2, mergedContext.getExtraContext(contextKey2));
assertEquals(testContext3, mergedContext.getExtraContext(contextKey3));
assertEquals(testContextOverwrite, mergedContext.getCustomContext(contextKey1));
assertEquals(testContext2, mergedContext.getCustomContext(contextKey2));
assertEquals(testContext3, mergedContext.getCustomContext(contextKey3));
}

private static Map<String, List<String>> createTestExtraHeaders(String... keyValues) {
Expand Down

0 comments on commit 9939551

Please sign in to comment.