Skip to content

Commit 5e47749

Browse files
authored
fix(gax): Implement lazy resource name evaluation in ApiTracerContext (#12618)
This PR - Move the resource name extraction logic from `TracedUnaryCallable` to `ApiTracerContext`. This ensures that the extraction is done lazily only when resource name is needed in tracing. Hence it does not affect regular customers at all if tracing is not enabled. - Prepends urlDomain to resource name.
1 parent f5101d9 commit 5e47749

File tree

4 files changed

+168
-49
lines changed

4 files changed

+168
-49
lines changed

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/tracing/ApiTracerContext.java

Lines changed: 44 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
import com.google.common.base.Strings;
3838
import java.util.HashMap;
3939
import java.util.Map;
40+
import java.util.function.Supplier;
4041
import javax.annotation.Nullable;
4142

4243
/**
@@ -164,9 +165,46 @@ String rpcSystemName() {
164165
@Nullable
165166
public abstract String urlDomain();
166167

167-
/** The destination resource id of the request (e.g. projects/p/locations/l/topics/t). */
168168
@Nullable
169-
public abstract String destinationResourceId();
169+
protected abstract Supplier<String> destinationResourceIdSupplier();
170+
171+
/**
172+
* The destination resource id of the request (e.g.
173+
* //pubsub.googleapis.com/projects/p/locations/l/topics/t).
174+
*/
175+
@Nullable
176+
public String destinationResourceId() {
177+
Supplier<String> supplier = destinationResourceIdSupplier();
178+
if (supplier == null) {
179+
return null;
180+
}
181+
String resourceId = supplier.get();
182+
if (Strings.isNullOrEmpty(resourceId)) {
183+
return null;
184+
}
185+
if (Strings.isNullOrEmpty(urlDomain())) {
186+
return resourceId;
187+
}
188+
return "//" + urlDomain() + "/" + resourceId;
189+
}
190+
191+
<RequestT> ApiTracerContext withResourceNameExtractor(
192+
@Nullable RequestT request,
193+
@Nullable com.google.api.gax.rpc.ResourceNameExtractor<RequestT> extractor) {
194+
if (extractor == null || request == null) {
195+
return this;
196+
}
197+
return toBuilder()
198+
.setDestinationResourceIdSupplier(
199+
() -> {
200+
try {
201+
return extractor.extract(request);
202+
} catch (Exception e) {
203+
return null;
204+
}
205+
})
206+
.build();
207+
}
170208

171209
/**
172210
* @return a map of attributes to be included in attempt-level spans
@@ -284,8 +322,8 @@ ApiTracerContext merge(ApiTracerContext other) {
284322
if (!Strings.isNullOrEmpty(other.urlDomain())) {
285323
builder.setUrlDomain(other.urlDomain());
286324
}
287-
if (other.destinationResourceId() != null) {
288-
builder.setDestinationResourceId(other.destinationResourceId());
325+
if (other.destinationResourceIdSupplier() != null) {
326+
builder.setDestinationResourceIdSupplier(other.destinationResourceIdSupplier());
289327
}
290328
return builder.build();
291329
}
@@ -322,7 +360,8 @@ public abstract static class Builder {
322360

323361
public abstract Builder setUrlDomain(@Nullable String urlDomain);
324362

325-
public abstract Builder setDestinationResourceId(@Nullable String destinationResourceId);
363+
public abstract Builder setDestinationResourceIdSupplier(
364+
@Nullable Supplier<String> destinationResourceIdSupplier);
326365

327366
public abstract ApiTracerContext build();
328367
}

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedUnaryCallable.java

Lines changed: 4 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -37,8 +37,6 @@
3737
import com.google.api.gax.rpc.ResourceNameExtractor;
3838
import com.google.api.gax.rpc.UnaryCallable;
3939
import com.google.api.gax.tracing.ApiTracerFactory.OperationType;
40-
import com.google.common.annotations.VisibleForTesting;
41-
import com.google.common.base.Strings;
4240
import com.google.common.util.concurrent.MoreExecutors;
4341
import javax.annotation.Nullable;
4442

@@ -90,8 +88,10 @@ public TracedUnaryCallable(
9088
public ApiFuture<ResponseT> futureCall(RequestT request, ApiCallContext context) {
9189
ApiTracer tracer;
9290
if (apiTracerContext != null) {
93-
ApiTracerContext finalContext = extractResourceNameToApiTracerContext(request);
94-
tracer = tracerFactory.newTracer(context.getTracer(), finalContext);
91+
tracer =
92+
tracerFactory.newTracer(
93+
context.getTracer(),
94+
apiTracerContext.withResourceNameExtractor(request, resourceNameExtractor));
9595
} else {
9696
tracer = tracerFactory.newTracer(context.getTracer(), spanName, OperationType.Unary);
9797
}
@@ -108,15 +108,4 @@ public ApiFuture<ResponseT> futureCall(RequestT request, ApiCallContext context)
108108
throw e;
109109
}
110110
}
111-
112-
@VisibleForTesting
113-
ApiTracerContext extractResourceNameToApiTracerContext(RequestT request) {
114-
ApiTracerContext finalContext = apiTracerContext;
115-
String resourceName =
116-
resourceNameExtractor != null ? resourceNameExtractor.extract(request) : null;
117-
if (!Strings.isNullOrEmpty(resourceName)) {
118-
finalContext = finalContext.toBuilder().setDestinationResourceId(resourceName).build();
119-
}
120-
return finalContext;
121-
}
122111
}

sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/tracing/ApiTracerContextTest.java

Lines changed: 120 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -310,7 +310,7 @@ void testGetAttemptAttributes_destinationResourceId() {
310310
ApiTracerContext context =
311311
ApiTracerContext.newBuilder()
312312
.setLibraryMetadata(LibraryMetadata.empty())
313-
.setDestinationResourceId("projects/123/instances/abc")
313+
.setDestinationResourceIdSupplier(() -> "projects/123/instances/abc")
314314
.build();
315315
Map<String, Object> attributes = context.getAttemptAttributes();
316316

@@ -383,13 +383,13 @@ void testMerge_destinationResourceId() {
383383
ApiTracerContext context1 =
384384
ApiTracerContext.newBuilder()
385385
.setLibraryMetadata(LibraryMetadata.empty())
386-
.setDestinationResourceId("name1")
386+
.setDestinationResourceIdSupplier(() -> "name1")
387387
.build();
388388

389389
ApiTracerContext context2 =
390390
ApiTracerContext.newBuilder()
391391
.setLibraryMetadata(LibraryMetadata.empty())
392-
.setDestinationResourceId("name2")
392+
.setDestinationResourceIdSupplier(() -> "name2")
393393
.build();
394394

395395
ApiTracerContext merged = context1.merge(context2);
@@ -434,4 +434,121 @@ void testMerge_httpFields() {
434434
assertThat(merged.httpMethod()).isEqualTo("GET");
435435
assertThat(merged.httpPathTemplate()).isEqualTo("v1/projects/{project}");
436436
}
437+
438+
@Test
439+
void testDestinationResourceId_withUrlDomain() {
440+
ApiTracerContext context =
441+
ApiTracerContext.newBuilder()
442+
.setLibraryMetadata(LibraryMetadata.empty())
443+
.setUrlDomain("compute.googleapis.com")
444+
.setDestinationResourceIdSupplier(() -> "projects/my-project/locations/us-central1")
445+
.build();
446+
447+
assertThat(context.destinationResourceId())
448+
.isEqualTo("//compute.googleapis.com/projects/my-project/locations/us-central1");
449+
}
450+
451+
@Test
452+
void testDestinationResourceId_withoutUrlDomain() {
453+
ApiTracerContext context =
454+
ApiTracerContext.newBuilder()
455+
.setLibraryMetadata(LibraryMetadata.empty())
456+
.setDestinationResourceIdSupplier(() -> "projects/my-project/locations/us-central1")
457+
.build();
458+
459+
assertThat(context.destinationResourceId())
460+
.isEqualTo("projects/my-project/locations/us-central1");
461+
}
462+
463+
@Test
464+
void testWithResourceNameExtractor_nullExtractor() {
465+
ApiTracerContext context = ApiTracerContext.empty();
466+
ApiTracerContext result = context.withResourceNameExtractor("request", null);
467+
assertThat(result).isSameInstanceAs(context);
468+
}
469+
470+
@Test
471+
void testWithResourceNameExtractor_extractorReturnsNull() {
472+
ApiTracerContext context = ApiTracerContext.empty();
473+
ApiTracerContext result = context.withResourceNameExtractor("request", req -> null);
474+
assertThat(result.destinationResourceId()).isNull();
475+
}
476+
477+
@Test
478+
void testWithResourceNameExtraction_lazyExtractor() {
479+
ApiTracerContext context = ApiTracerContext.empty();
480+
boolean[] extracted = {false};
481+
ApiTracerContext result =
482+
context.withResourceNameExtractor(
483+
"request",
484+
req -> {
485+
extracted[0] = true;
486+
return "extracted-id";
487+
});
488+
489+
assertThat(extracted[0]).isFalse(); // Should be lazily evaluated
490+
assertThat(result.destinationResourceId()).isEqualTo("extracted-id");
491+
assertThat(extracted[0]).isTrue();
492+
}
493+
494+
@Test
495+
void testWithResourceNameExtractor_extractorThrowsException() {
496+
ApiTracerContext context = ApiTracerContext.empty();
497+
ApiTracerContext result =
498+
context.withResourceNameExtractor(
499+
"request",
500+
req -> {
501+
throw new RuntimeException("Intentional mock extraction failure");
502+
});
503+
504+
assertThat(result.destinationResourceId()).isNull();
505+
}
506+
507+
@Test
508+
void testDestinationResourceId_emptyIdWithUrlDomain() {
509+
ApiTracerContext context =
510+
ApiTracerContext.newBuilder()
511+
.setLibraryMetadata(LibraryMetadata.empty())
512+
.setUrlDomain("compute.googleapis.com")
513+
.setDestinationResourceIdSupplier(() -> "")
514+
.build();
515+
516+
assertThat(context.destinationResourceId()).isNull();
517+
}
518+
519+
@Test
520+
void testMerge_preservesLaziness() {
521+
ApiTracerContext context1 = ApiTracerContext.empty();
522+
boolean[] extracted = {false};
523+
ApiTracerContext context2 =
524+
ApiTracerContext.empty()
525+
.withResourceNameExtractor(
526+
"request",
527+
req -> {
528+
extracted[0] = true;
529+
return "lazy-id";
530+
});
531+
532+
ApiTracerContext merged = context1.merge(context2);
533+
assertThat(extracted[0]).isFalse(); // Should not be evaluated during merge
534+
assertThat(merged.destinationResourceId()).isEqualTo("lazy-id");
535+
assertThat(extracted[0]).isTrue(); // Evaluated upon calling getter
536+
}
537+
538+
@Test
539+
void testDestinationResourceId_evaluatedEveryTime() {
540+
ApiTracerContext context = ApiTracerContext.empty();
541+
int[] counter = {0};
542+
ApiTracerContext result =
543+
context.withResourceNameExtractor(
544+
"request",
545+
req -> {
546+
counter[0]++;
547+
return "extracted-id-" + counter[0];
548+
});
549+
550+
assertThat(result.destinationResourceId()).isEqualTo("extracted-id-1");
551+
assertThat(result.destinationResourceId()).isEqualTo("extracted-id-2");
552+
assertThat(counter[0]).isEqualTo(2);
553+
}
437554
}

sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/tracing/TracedUnaryCallableTest.java

Lines changed: 0 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -208,30 +208,4 @@ void testResourceNameExtractorUsed() {
208208
assertThat(contextCaptor.getValue().destinationResourceId())
209209
.isEqualTo("extracted-resource-name");
210210
}
211-
212-
@Test
213-
void testExtractResourceNameToApiTracerContext_nullExtractor() {
214-
tracedUnaryCallable =
215-
new TracedUnaryCallable<>(innerCallable, tracerFactory, TRACER_CONTEXT, null);
216-
ApiTracerContext context = tracedUnaryCallable.extractResourceNameToApiTracerContext("request");
217-
assertThat(context).isEqualTo(TRACER_CONTEXT);
218-
}
219-
220-
@Test
221-
void testExtractResourceNameToApiTracerContext_extractorReturnsNull() {
222-
tracedUnaryCallable =
223-
new TracedUnaryCallable<>(innerCallable, tracerFactory, TRACER_CONTEXT, request -> null);
224-
ApiTracerContext context = tracedUnaryCallable.extractResourceNameToApiTracerContext("request");
225-
assertThat(context).isEqualTo(TRACER_CONTEXT);
226-
}
227-
228-
@Test
229-
void testExtractResourceNameToApiTracerContext_extractorReturnsResourceId() {
230-
tracedUnaryCallable =
231-
new TracedUnaryCallable<>(
232-
innerCallable, tracerFactory, TRACER_CONTEXT, request -> "extracted-id");
233-
ApiTracerContext context = tracedUnaryCallable.extractResourceNameToApiTracerContext("request");
234-
assertThat(context).isNotSameInstanceAs(TRACER_CONTEXT);
235-
assertThat(context.destinationResourceId()).isEqualTo("extracted-id");
236-
}
237211
}

0 commit comments

Comments
 (0)