Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Appearance settings
This repository was archived by the owner on Sep 26, 2023. It is now read-only.

Commit 0fe20f3

Browse filesBrowse files
mutianfmutianf
authored
feat: add custom options to ApiCallContext (#1435)
1 parent fc2f997 commit 0fe20f3
Copy full SHA for 0fe20f3

File tree

8 files changed

+455
-4
lines changed
Filter options

8 files changed

+455
-4
lines changed

‎gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcCallContext.java

Copy file name to clipboardExpand all lines: gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcCallContext.java
+44-2Lines changed: 44 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
import com.google.api.gax.rpc.ApiCallContext;
3535
import com.google.api.gax.rpc.StatusCode;
3636
import com.google.api.gax.rpc.TransportChannel;
37+
import com.google.api.gax.rpc.internal.ApiCallContextOptions;
3738
import com.google.api.gax.rpc.internal.Headers;
3839
import com.google.api.gax.tracing.ApiTracer;
3940
import com.google.api.gax.tracing.BaseApiTracer;
@@ -43,7 +44,6 @@
4344
import com.google.common.collect.ImmutableSet;
4445
import io.grpc.CallCredentials;
4546
import io.grpc.CallOptions;
46-
import io.grpc.CallOptions.Key;
4747
import io.grpc.Channel;
4848
import io.grpc.Deadline;
4949
import io.grpc.Metadata;
@@ -66,7 +66,7 @@
6666
*/
6767
@BetaApi("Reference ApiCallContext instead - this class is likely to experience breaking changes")
6868
public final class GrpcCallContext implements ApiCallContext {
69-
static final CallOptions.Key<ApiTracer> TRACER_KEY = Key.create("gax.tracer");
69+
static final CallOptions.Key<ApiTracer> TRACER_KEY = CallOptions.Key.create("gax.tracer");
7070

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

8182
/** Returns an empty instance with a null channel and default {@link CallOptions}. */
8283
public static GrpcCallContext createDefault() {
@@ -88,6 +89,7 @@ public static GrpcCallContext createDefault() {
8889
null,
8990
null,
9091
ImmutableMap.<String, List<String>>of(),
92+
ApiCallContextOptions.getDefaultOptions(),
9193
null,
9294
null);
9395
}
@@ -102,6 +104,7 @@ public static GrpcCallContext of(Channel channel, CallOptions callOptions) {
102104
null,
103105
null,
104106
ImmutableMap.<String, List<String>>of(),
107+
ApiCallContextOptions.getDefaultOptions(),
105108
null,
106109
null);
107110
}
@@ -114,6 +117,7 @@ private GrpcCallContext(
114117
@Nullable Duration streamIdleTimeout,
115118
@Nullable Integer channelAffinity,
116119
ImmutableMap<String, List<String>> extraHeaders,
120+
ApiCallContextOptions options,
117121
@Nullable RetrySettings retrySettings,
118122
@Nullable Set<StatusCode.Code> retryableCodes) {
119123
this.channel = channel;
@@ -123,6 +127,7 @@ private GrpcCallContext(
123127
this.streamIdleTimeout = streamIdleTimeout;
124128
this.channelAffinity = channelAffinity;
125129
this.extraHeaders = Preconditions.checkNotNull(extraHeaders);
130+
this.options = Preconditions.checkNotNull(options);
126131
this.retrySettings = retrySettings;
127132
this.retryableCodes = retryableCodes == null ? null : ImmutableSet.copyOf(retryableCodes);
128133
}
@@ -187,6 +192,7 @@ public GrpcCallContext withTimeout(@Nullable Duration timeout) {
187192
this.streamIdleTimeout,
188193
this.channelAffinity,
189194
this.extraHeaders,
195+
this.options,
190196
this.retrySettings,
191197
this.retryableCodes);
192198
}
@@ -212,6 +218,7 @@ public GrpcCallContext withStreamWaitTimeout(@Nullable Duration streamWaitTimeou
212218
this.streamIdleTimeout,
213219
this.channelAffinity,
214220
this.extraHeaders,
221+
this.options,
215222
this.retrySettings,
216223
this.retryableCodes);
217224
}
@@ -231,6 +238,7 @@ public GrpcCallContext withStreamIdleTimeout(@Nullable Duration streamIdleTimeou
231238
streamIdleTimeout,
232239
this.channelAffinity,
233240
this.extraHeaders,
241+
this.options,
234242
this.retrySettings,
235243
this.retryableCodes);
236244
}
@@ -245,6 +253,7 @@ public GrpcCallContext withChannelAffinity(@Nullable Integer affinity) {
245253
this.streamIdleTimeout,
246254
affinity,
247255
this.extraHeaders,
256+
this.options,
248257
this.retrySettings,
249258
this.retryableCodes);
250259
}
@@ -263,6 +272,7 @@ public GrpcCallContext withExtraHeaders(Map<String, List<String>> extraHeaders)
263272
this.streamIdleTimeout,
264273
this.channelAffinity,
265274
newExtraHeaders,
275+
this.options,
266276
this.retrySettings,
267277
this.retryableCodes);
268278
}
@@ -282,6 +292,7 @@ public GrpcCallContext withRetrySettings(RetrySettings retrySettings) {
282292
this.streamIdleTimeout,
283293
this.channelAffinity,
284294
this.extraHeaders,
295+
this.options,
285296
retrySettings,
286297
this.retryableCodes);
287298
}
@@ -301,6 +312,7 @@ public GrpcCallContext withRetryableCodes(Set<StatusCode.Code> retryableCodes) {
301312
this.streamIdleTimeout,
302313
this.channelAffinity,
303314
this.extraHeaders,
315+
this.options,
304316
this.retrySettings,
305317
retryableCodes);
306318
}
@@ -370,6 +382,8 @@ public ApiCallContext merge(ApiCallContext inputCallContext) {
370382
ImmutableMap<String, List<String>> newExtraHeaders =
371383
Headers.mergeHeaders(this.extraHeaders, grpcCallContext.extraHeaders);
372384

385+
ApiCallContextOptions newOptions = options.merge(grpcCallContext.options);
386+
373387
CallOptions newCallOptions =
374388
grpcCallContext
375389
.callOptions
@@ -388,6 +402,7 @@ public ApiCallContext merge(ApiCallContext inputCallContext) {
388402
newStreamIdleTimeout,
389403
newChannelAffinity,
390404
newExtraHeaders,
405+
newOptions,
391406
newRetrySettings,
392407
newRetryableCodes);
393408
}
@@ -448,6 +463,7 @@ public GrpcCallContext withChannel(Channel newChannel) {
448463
this.streamIdleTimeout,
449464
this.channelAffinity,
450465
this.extraHeaders,
466+
this.options,
451467
this.retrySettings,
452468
this.retryableCodes);
453469
}
@@ -462,6 +478,7 @@ public GrpcCallContext withCallOptions(CallOptions newCallOptions) {
462478
this.streamIdleTimeout,
463479
this.channelAffinity,
464480
this.extraHeaders,
481+
this.options,
465482
this.retrySettings,
466483
this.retryableCodes);
467484
}
@@ -491,6 +508,29 @@ public GrpcCallContext withTracer(@Nonnull ApiTracer tracer) {
491508
return withCallOptions(callOptions.withOption(TRACER_KEY, tracer));
492509
}
493510

511+
/** {@inheritDoc} */
512+
@Override
513+
public <T> GrpcCallContext withOption(Key<T> key, T value) {
514+
ApiCallContextOptions newOptions = options.withOption(key, value);
515+
return new GrpcCallContext(
516+
this.channel,
517+
this.callOptions,
518+
this.timeout,
519+
this.streamWaitTimeout,
520+
this.streamIdleTimeout,
521+
this.channelAffinity,
522+
this.extraHeaders,
523+
newOptions,
524+
this.retrySettings,
525+
this.retryableCodes);
526+
}
527+
528+
/** {@inheritDoc} */
529+
@Override
530+
public <T> T getOption(Key<T> key) {
531+
return options.getOption(key);
532+
}
533+
494534
@Override
495535
public int hashCode() {
496536
return Objects.hash(
@@ -501,6 +541,7 @@ public int hashCode() {
501541
streamIdleTimeout,
502542
channelAffinity,
503543
extraHeaders,
544+
options,
504545
retrySettings,
505546
retryableCodes);
506547
}
@@ -522,6 +563,7 @@ public boolean equals(Object o) {
522563
&& Objects.equals(this.streamIdleTimeout, that.streamIdleTimeout)
523564
&& Objects.equals(this.channelAffinity, that.channelAffinity)
524565
&& Objects.equals(this.extraHeaders, that.extraHeaders)
566+
&& Objects.equals(this.options, that.options)
525567
&& Objects.equals(this.retrySettings, that.retrySettings)
526568
&& Objects.equals(this.retryableCodes, that.retryableCodes);
527569
}

‎gax-grpc/src/test/java/com/google/api/gax/grpc/GrpcCallContextTest.java

Copy file name to clipboardExpand all lines: gax-grpc/src/test/java/com/google/api/gax/grpc/GrpcCallContextTest.java
+42Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -354,6 +354,48 @@ public void testWithRetryableCodes() {
354354
assertNotNull(context.getRetryableCodes());
355355
}
356356

357+
@Test
358+
public void testWithOptions() {
359+
GrpcCallContext emptyCallContext = GrpcCallContext.createDefault();
360+
ApiCallContext.Key<String> contextKey1 = ApiCallContext.Key.create("testKey1");
361+
ApiCallContext.Key<String> contextKey2 = ApiCallContext.Key.create("testKey2");
362+
String testContext1 = "test1";
363+
String testContext2 = "test2";
364+
String testContextOverwrite = "test1Overwrite";
365+
GrpcCallContext context =
366+
emptyCallContext
367+
.withOption(contextKey1, testContext1)
368+
.withOption(contextKey2, testContext2);
369+
assertEquals(testContext1, context.getOption(contextKey1));
370+
assertEquals(testContext2, context.getOption(contextKey2));
371+
GrpcCallContext newContext = context.withOption(contextKey1, testContextOverwrite);
372+
assertEquals(testContextOverwrite, newContext.getOption(contextKey1));
373+
}
374+
375+
@Test
376+
public void testMergeOptions() {
377+
GrpcCallContext emptyCallContext = GrpcCallContext.createDefault();
378+
ApiCallContext.Key<String> contextKey1 = ApiCallContext.Key.create("testKey1");
379+
ApiCallContext.Key<String> contextKey2 = ApiCallContext.Key.create("testKey2");
380+
ApiCallContext.Key<String> contextKey3 = ApiCallContext.Key.create("testKey3");
381+
String testContext1 = "test1";
382+
String testContext2 = "test2";
383+
String testContext3 = "test3";
384+
String testContextOverwrite = "test1Overwrite";
385+
GrpcCallContext context1 =
386+
emptyCallContext
387+
.withOption(contextKey1, testContext1)
388+
.withOption(contextKey2, testContext2);
389+
GrpcCallContext context2 =
390+
emptyCallContext
391+
.withOption(contextKey1, testContextOverwrite)
392+
.withOption(contextKey3, testContext3);
393+
ApiCallContext mergedContext = context1.merge(context2);
394+
assertEquals(testContextOverwrite, mergedContext.getOption(contextKey1));
395+
assertEquals(testContext2, mergedContext.getOption(contextKey2));
396+
assertEquals(testContext3, mergedContext.getOption(contextKey3));
397+
}
398+
357399
private static Map<String, List<String>> createTestExtraHeaders(String... keyValues) {
358400
Map<String, List<String>> extraHeaders = new HashMap<>();
359401
for (int i = 0; i < keyValues.length; i += 2) {

0 commit comments

Comments
0 (0)
Morty Proxy This is a proxified and sanitized view of the page, visit original site.