From 7592d7a8b88ebeb5c25c7c17038169b68a4153aa Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Fri, 3 Jul 2026 06:54:56 -0700 Subject: [PATCH] refactor: use RxJava's `using` operator instead of the `defer` + `doFinally` pattern ### Key Changes: * **Improved Resource Management**: Migrates `traceFlowable` and `TracingTransformer` in `Tracing.java` to use `Flowable.using`, `Single.using`, etc. This ensures that the OpenTelemetry `Scope` is strictly bound to the stream's lifecycle, providing better guarantees for closure during cancellation or errors compared to `doFinally`. * **Fixes Lazy Context Activation**: Replaced manual `try-with-resources` blocks in `BaseLlmFlow.java` with deferred execution (`Maybe.defer`) and context re-activation (`Tracing.withContext`). This ensures the context is active when the RxJava pipeline actually executes, rather than when it is assembled. * **Refactored TracingObserver**: Removed static factory methods in `TracingObserver` in favor of direct constructor calls within the `lift` operators. * **Added Testing**: Introduced `testTraceFlowableLazy` in `ContextPropagationTest` to verify that tracing is applied lazily at subscription time and that the scope is correctly managed. * **Cleanup**: Updated documentation and added `@SuppressWarnings("MustBeClosedChecker")` where the scope lifecycle is intentionally managed by RxJava's resource management operators. PiperOrigin-RevId: 942123391 --- .../adk/flows/llmflows/BaseLlmFlow.java | 60 ++++++------ .../com/google/adk/telemetry/Tracing.java | 94 +++++++++---------- .../adk/telemetry/ContextPropagationTest.java | 31 +++++- 3 files changed, 104 insertions(+), 81 deletions(-) diff --git a/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java b/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java index 7d62c8505..7e59ec0d6 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java @@ -427,6 +427,7 @@ private Single handleAfterModelCallback( * @throws LlmCallsLimitExceededException if the agent exceeds allowed LLM invocations. * @throws IllegalStateException if a transfer agent is specified but not found. */ + @SuppressWarnings("MustBeClosedChecker") // Scope lifecycle managed by RxJava using private Flowable runOneStep(Context spanContext, InvocationContext context) { AtomicReference llmRequestRef = new AtomicReference<>(LlmRequest.builder().build()); @@ -560,11 +561,12 @@ private Flowable run( * @return A {@link Flowable} of {@link Event}s streamed in real-time. */ @Override + @SuppressWarnings("MustBeClosedChecker") // Scope lifecycle managed by RxJava using public Flowable runLive(InvocationContext invocationContext) { AtomicReference llmRequestRef = new AtomicReference<>(LlmRequest.builder().build()); Flowable preprocessEvents = preprocess(invocationContext, llmRequestRef); // Capture agent context at assembly time to use as parent for agent transfer at subscription - // time. See Flowable.defer() usages below. + // time. See Flowable.using() usages below. Context spanContext = Context.current(); return preprocessEvents.concatWith( @@ -687,12 +689,10 @@ public void onError(Throwable e) { "Agent not found: " + event.actions().transferToAgent().get()); } Flowable nextAgentEvents = - Flowable.defer( - () -> { - try (Scope scope = spanContext.makeCurrent()) { - return nextAgent.get().runLive(invocationContext); - } - }); + Flowable.using( + spanContext::makeCurrent, + scope -> nextAgent.get().runLive(invocationContext), + Scope::close); events = Flowable.concat(events, nextAgentEvents); } return events; @@ -757,28 +757,30 @@ private Flowable buildPostprocessingEvents( return processorEvents.concatWith(Flowable.just(modelResponseEvent)); } - Flowable functionEvents; - try (Scope scope = parentContext.makeCurrent()) { - Maybe maybeFunctionResponseEvent = - context.runConfig().streamingMode() == StreamingMode.BIDI - ? Functions.handleFunctionCallsLive(context, modelResponseEvent, llmRequest.tools()) - : Functions.handleFunctionCalls(context, modelResponseEvent, llmRequest.tools()); - functionEvents = - maybeFunctionResponseEvent.flatMapPublisher( - functionResponseEvent -> { - Optional toolConfirmationEvent = - Functions.generateRequestConfirmationEvent( - context, modelResponseEvent, functionResponseEvent); - List events = new ArrayList<>(); - toolConfirmationEvent.ifPresent(events::add); - events.add(functionResponseEvent); - OutputSchema.getStructuredModelResponse(functionResponseEvent) - .ifPresent( - json -> - events.add(OutputSchema.createFinalModelResponseEvent(context, json))); - return Flowable.fromIterable(events); - }); - } + Flowable functionEvents = + Maybe.defer( + () -> + context.runConfig().streamingMode() == StreamingMode.BIDI + ? Functions.handleFunctionCallsLive( + context, modelResponseEvent, llmRequest.tools()) + : Functions.handleFunctionCalls( + context, modelResponseEvent, llmRequest.tools())) + .compose(Tracing.withContext(parentContext)) + .flatMapPublisher( + functionResponseEvent -> { + Optional toolConfirmationEvent = + Functions.generateRequestConfirmationEvent( + context, modelResponseEvent, functionResponseEvent); + List events = new ArrayList<>(); + toolConfirmationEvent.ifPresent(events::add); + events.add(functionResponseEvent); + OutputSchema.getStructuredModelResponse(functionResponseEvent) + .ifPresent( + json -> + events.add( + OutputSchema.createFinalModelResponseEvent(context, json))); + return Flowable.fromIterable(events); + }); return processorEvents.concatWith(Flowable.just(modelResponseEvent)).concatWith(functionEvents); } diff --git a/core/src/main/java/com/google/adk/telemetry/Tracing.java b/core/src/main/java/com/google/adk/telemetry/Tracing.java index d57e46c54..c0b22f7a0 100644 --- a/core/src/main/java/com/google/adk/telemetry/Tracing.java +++ b/core/src/main/java/com/google/adk/telemetry/Tracing.java @@ -434,8 +434,8 @@ public static Tracer getTracer() { * scope before the Flowable subscribes, causing Context.current() to return ROOT in nested * operations and breaking parent-child span relationships (fragmenting traces). * - *

The scope is properly closed via doFinally when the stream terminates, ensuring no resource - * leaks regardless of completion mode (success, error, or cancellation). + *

The scope is properly closed via RxJava using when the stream terminates, ensuring no + * resource leaks regardless of completion mode (success, error, or cancellation). * * @param spanContext The context containing the span to activate * @param span The span to end when the stream completes @@ -443,17 +443,16 @@ public static Tracer getTracer() { * @param The type of items emitted by the Flowable * @return Flowable with OpenTelemetry scope lifecycle management */ - @SuppressWarnings("MustBeClosedChecker") // Scope lifecycle managed by RxJava doFinally + @SuppressWarnings("MustBeClosedChecker") // Scope lifecycle managed by RxJava using public static Flowable traceFlowable( Context spanContext, Span span, Supplier> flowableSupplier) { - Scope scope = spanContext.makeCurrent(); - return flowableSupplier - .get() - .doFinally( - () -> { - scope.close(); - span.end(); - }); + return Flowable.using( + spanContext::makeCurrent, + scope -> flowableSupplier.get(), + scope -> { + scope.close(); + span.end(); + }); } /** @@ -536,11 +535,11 @@ private Context getParentContext() { } private final class TracingLifecycle { - private Span span; - private Scope scope; + private final Span span; + private final Scope scope; @SuppressWarnings("MustBeClosedChecker") - void start() { + TracingLifecycle() { span = tracer.spanBuilder(spanName).setParent(getParentContext()).startSpan(); spanConfigurers.forEach(c -> c.accept(span)); scope = span.makeCurrent(); @@ -564,16 +563,16 @@ void end() { */ @Override public Publisher apply(Flowable upstream) { - return Flowable.defer( - () -> { - TracingLifecycle lifecycle = new TracingLifecycle(); - lifecycle.start(); + return Flowable.using( + TracingLifecycle::new, + lifecycle -> { Flowable pipeline = upstream; if (onSuccessConsumer != null) { pipeline = pipeline.doOnNext(t -> onSuccessConsumer.accept(lifecycle.span, t)); } - return pipeline.doFinally(lifecycle::end); - }); + return pipeline; + }, + TracingLifecycle::end); } /** @@ -584,16 +583,16 @@ public Publisher apply(Flowable upstream) { */ @Override public SingleSource apply(Single upstream) { - return Single.defer( - () -> { - TracingLifecycle lifecycle = new TracingLifecycle(); - lifecycle.start(); + return Single.using( + TracingLifecycle::new, + lifecycle -> { Single pipeline = upstream; if (onSuccessConsumer != null) { pipeline = pipeline.doOnSuccess(t -> onSuccessConsumer.accept(lifecycle.span, t)); } - return pipeline.doFinally(lifecycle::end); - }); + return pipeline; + }, + TracingLifecycle::end); } /** @@ -604,16 +603,16 @@ public SingleSource apply(Single upstream) { */ @Override public MaybeSource apply(Maybe upstream) { - return Maybe.defer( - () -> { - TracingLifecycle lifecycle = new TracingLifecycle(); - lifecycle.start(); + return Maybe.using( + TracingLifecycle::new, + lifecycle -> { Maybe pipeline = upstream; if (onSuccessConsumer != null) { pipeline = pipeline.doOnSuccess(t -> onSuccessConsumer.accept(lifecycle.span, t)); } - return pipeline.doFinally(lifecycle::end); - }); + return pipeline; + }, + TracingLifecycle::end); } /** @@ -624,12 +623,7 @@ public MaybeSource apply(Maybe upstream) { */ @Override public CompletableSource apply(Completable upstream) { - return Completable.defer( - () -> { - TracingLifecycle lifecycle = new TracingLifecycle(); - lifecycle.start(); - return upstream.doFinally(lifecycle::end); - }); + return Completable.using(TracingLifecycle::new, lifecycle -> upstream, TracingLifecycle::end); } } @@ -669,7 +663,7 @@ private ContextTransformer(Context context) { */ @Override public Publisher apply(Flowable upstream) { - return upstream.lift(subscriber -> TracingObserver.wrap(context, subscriber)); + return upstream.lift(subscriber -> new TracingObserver<>(context, subscriber)); } /** @@ -680,7 +674,7 @@ public Publisher apply(Flowable upstream) { */ @Override public SingleSource apply(Single upstream) { - return upstream.lift(observer -> TracingObserver.wrap(context, observer)); + return upstream.lift(observer -> new TracingObserver<>(context, observer)); } /** @@ -691,7 +685,7 @@ public SingleSource apply(Single upstream) { */ @Override public MaybeSource apply(Maybe upstream) { - return upstream.lift(observer -> TracingObserver.wrap(context, observer)); + return upstream.lift(observer -> new TracingObserver<>(context, observer)); } /** @@ -702,7 +696,7 @@ public MaybeSource apply(Maybe upstream) { */ @Override public CompletableSource apply(Completable upstream) { - return upstream.lift(observer -> TracingObserver.wrap(context, observer)); + return upstream.lift(observer -> new TracingObserver<>(context, observer)); } } @@ -739,20 +733,20 @@ private TracingObserver( this.completableObserver = completableObserver; } - static TracingObserver wrap(Context context, Subscriber subscriber) { - return new TracingObserver<>(context, subscriber, null, null, null); + TracingObserver(Context context, Subscriber subscriber) { + this(context, subscriber, null, null, null); } - static TracingObserver wrap(Context context, SingleObserver observer) { - return new TracingObserver<>(context, null, observer, null, null); + TracingObserver(Context context, SingleObserver observer) { + this(context, null, observer, null, null); } - static TracingObserver wrap(Context context, MaybeObserver observer) { - return new TracingObserver<>(context, null, null, observer, null); + TracingObserver(Context context, MaybeObserver observer) { + this(context, null, null, observer, null); } - static TracingObserver wrap(Context context, CompletableObserver observer) { - return new TracingObserver<>(context, null, null, null, observer); + TracingObserver(Context context, CompletableObserver observer) { + this(context, null, null, null, observer); } private void runInContext(Runnable action) { diff --git a/core/src/test/java/com/google/adk/telemetry/ContextPropagationTest.java b/core/src/test/java/com/google/adk/telemetry/ContextPropagationTest.java index ecfba4995..1bac85265 100644 --- a/core/src/test/java/com/google/adk/telemetry/ContextPropagationTest.java +++ b/core/src/test/java/com/google/adk/telemetry/ContextPropagationTest.java @@ -18,6 +18,8 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotEquals; import static org.junit.Assert.assertTrue; import com.google.adk.agents.BaseAgent; @@ -62,6 +64,7 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import org.junit.After; import org.junit.Before; @@ -129,6 +132,30 @@ public void testTraceFlowable() throws InterruptedException { assertTrue(flowableSpanData.hasEnded()); } + @Test + public void testTraceFlowableLazy() throws InterruptedException { + Span flowableSpan = tracer.spanBuilder("flowable-lazy").startSpan(); + AtomicBoolean supplierCalled = new AtomicBoolean(false); + Flowable flowable = + Tracing.traceFlowable( + Context.current().with(flowableSpan), + flowableSpan, + () -> { + supplierCalled.set(true); + assertEquals( + flowableSpan.getSpanContext().getSpanId(), + Span.current().getSpanContext().getSpanId()); + return Flowable.just(1); + }); + assertFalse(supplierCalled.get()); + assertNotEquals( + flowableSpan.getSpanContext().getSpanId(), Span.current().getSpanContext().getSpanId()); + + flowable.test().await().assertComplete(); + assertTrue(supplierCalled.get()); + assertTrue(findSpanByName("flowable-lazy").hasEnded()); + } + @Test public void testWithContextFlowable() throws InterruptedException { ContextKey testKey = ContextKey.named("test-key"); @@ -640,7 +667,7 @@ public void testAgentWithToolCallTraceHierarchy() throws InterruptedException { List callLlmSpans = openTelemetryRule.getSpans().stream() .filter(s -> s.getName().equals("call_llm")) - .sorted(Comparator.comparing(SpanData::getStartEpochNanos)) + .sorted(Comparator.comparingLong(SpanData::getStartEpochNanos)) .toList(); assertThat(callLlmSpans).hasSize(2); SpanData callLlm1 = callLlmSpans.get(0); @@ -716,7 +743,7 @@ public void testNestedAgentTraceHierarchy() throws InterruptedException { List callLlmSpans = openTelemetryRule.getSpans().stream() .filter(s -> s.getName().equals("call_llm")) - .sorted(Comparator.comparing(SpanData::getStartEpochNanos)) + .sorted(Comparator.comparingLong(SpanData::getStartEpochNanos)) .toList(); assertThat(callLlmSpans).hasSize(2);