Skip to content

refactor: use RxJava's using operator instead of the defer + doFinally pattern #1327

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
+104 −81
Open

refactor: use RxJava's using operator instead of the defer + doFinally pattern #1327

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 31 additions & 29 deletions core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -427,6 +427,7 @@ private Single<LlmResponse> 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<Event> runOneStep(Context spanContext, InvocationContext context) {
AtomicReference<LlmRequest> llmRequestRef = new AtomicReference<>(LlmRequest.builder().build());

Expand Down Expand Up @@ -560,11 +561,12 @@ private Flowable<Event> run(
* @return A {@link Flowable} of {@link Event}s streamed in real-time.
*/
@Override
@SuppressWarnings("MustBeClosedChecker") // Scope lifecycle managed by RxJava using
public Flowable<Event> runLive(InvocationContext invocationContext) {
AtomicReference<LlmRequest> llmRequestRef = new AtomicReference<>(LlmRequest.builder().build());
Flowable<Event> 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(
Expand Down Expand Up @@ -687,12 +689,10 @@ public void onError(Throwable e) {
"Agent not found: " + event.actions().transferToAgent().get());
}
Flowable<Event> 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;
Expand Down Expand Up @@ -757,28 +757,30 @@ private Flowable<Event> buildPostprocessingEvents(
return processorEvents.concatWith(Flowable.just(modelResponseEvent));
}

Flowable<Event> functionEvents;
try (Scope scope = parentContext.makeCurrent()) {
Maybe<Event> maybeFunctionResponseEvent =
context.runConfig().streamingMode() == StreamingMode.BIDI
? Functions.handleFunctionCallsLive(context, modelResponseEvent, llmRequest.tools())
: Functions.handleFunctionCalls(context, modelResponseEvent, llmRequest.tools());
functionEvents =
maybeFunctionResponseEvent.flatMapPublisher(
functionResponseEvent -> {
Optional<Event> toolConfirmationEvent =
Functions.generateRequestConfirmationEvent(
context, modelResponseEvent, functionResponseEvent);
List<Event> 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<Event> 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<Event> toolConfirmationEvent =
Functions.generateRequestConfirmationEvent(
context, modelResponseEvent, functionResponseEvent);
List<Event> 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);
}
Expand Down
94 changes: 44 additions & 50 deletions core/src/main/java/com/google/adk/telemetry/Tracing.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -434,26 +434,25 @@ 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).
*
* <p>The scope is properly closed via doFinally when the stream terminates, ensuring no resource
* leaks regardless of completion mode (success, error, or cancellation).
* <p>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
* @param flowableSupplier Supplier that creates the Flowable to execute with active scope
* @param <T> 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 <T> Flowable<T> traceFlowable(
Context spanContext, Span span, Supplier<Flowable<T>> 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();
});
}

/**
Expand Down Expand Up @@ -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();
Expand All @@ -564,16 +563,16 @@ void end() {
*/
@Override
public Publisher<T> apply(Flowable<T> upstream) {
return Flowable.defer(
() -> {
TracingLifecycle lifecycle = new TracingLifecycle();
lifecycle.start();
return Flowable.using(
TracingLifecycle::new,
lifecycle -> {
Flowable<T> pipeline = upstream;
if (onSuccessConsumer != null) {
pipeline = pipeline.doOnNext(t -> onSuccessConsumer.accept(lifecycle.span, t));
}
return pipeline.doFinally(lifecycle::end);
});
return pipeline;
},
TracingLifecycle::end);
}

/**
Expand All @@ -584,16 +583,16 @@ public Publisher<T> apply(Flowable<T> upstream) {
*/
@Override
public SingleSource<T> apply(Single<T> upstream) {
return Single.defer(
() -> {
TracingLifecycle lifecycle = new TracingLifecycle();
lifecycle.start();
return Single.using(
TracingLifecycle::new,
lifecycle -> {
Single<T> pipeline = upstream;
if (onSuccessConsumer != null) {
pipeline = pipeline.doOnSuccess(t -> onSuccessConsumer.accept(lifecycle.span, t));
}
return pipeline.doFinally(lifecycle::end);
});
return pipeline;
},
TracingLifecycle::end);
}

/**
Expand All @@ -604,16 +603,16 @@ public SingleSource<T> apply(Single<T> upstream) {
*/
@Override
public MaybeSource<T> apply(Maybe<T> upstream) {
return Maybe.defer(
() -> {
TracingLifecycle lifecycle = new TracingLifecycle();
lifecycle.start();
return Maybe.using(
TracingLifecycle::new,
lifecycle -> {
Maybe<T> pipeline = upstream;
if (onSuccessConsumer != null) {
pipeline = pipeline.doOnSuccess(t -> onSuccessConsumer.accept(lifecycle.span, t));
}
return pipeline.doFinally(lifecycle::end);
});
return pipeline;
},
TracingLifecycle::end);
}

/**
Expand All @@ -624,12 +623,7 @@ public MaybeSource<T> apply(Maybe<T> 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);
}
}

Expand Down Expand Up @@ -669,7 +663,7 @@ private ContextTransformer(Context context) {
*/
@Override
public Publisher<T> apply(Flowable<T> upstream) {
return upstream.lift(subscriber -> TracingObserver.wrap(context, subscriber));
return upstream.lift(subscriber -> new TracingObserver<>(context, subscriber));
}

/**
Expand All @@ -680,7 +674,7 @@ public Publisher<T> apply(Flowable<T> upstream) {
*/
@Override
public SingleSource<T> apply(Single<T> upstream) {
return upstream.lift(observer -> TracingObserver.wrap(context, observer));
return upstream.lift(observer -> new TracingObserver<>(context, observer));
}

/**
Expand All @@ -691,7 +685,7 @@ public SingleSource<T> apply(Single<T> upstream) {
*/
@Override
public MaybeSource<T> apply(Maybe<T> upstream) {
return upstream.lift(observer -> TracingObserver.wrap(context, observer));
return upstream.lift(observer -> new TracingObserver<>(context, observer));
}

/**
Expand All @@ -702,7 +696,7 @@ public MaybeSource<T> apply(Maybe<T> upstream) {
*/
@Override
public CompletableSource apply(Completable upstream) {
return upstream.lift(observer -> TracingObserver.wrap(context, observer));
return upstream.lift(observer -> new TracingObserver<>(context, observer));
}
}

Expand Down Expand Up @@ -739,20 +733,20 @@ private TracingObserver(
this.completableObserver = completableObserver;
}

static <T> TracingObserver<T> wrap(Context context, Subscriber<? super T> subscriber) {
return new TracingObserver<>(context, subscriber, null, null, null);
TracingObserver(Context context, Subscriber<? super T> subscriber) {
this(context, subscriber, null, null, null);
}

static <T> TracingObserver<T> wrap(Context context, SingleObserver<? super T> observer) {
return new TracingObserver<>(context, null, observer, null, null);
TracingObserver(Context context, SingleObserver<? super T> observer) {
this(context, null, observer, null, null);
}

static <T> TracingObserver<T> wrap(Context context, MaybeObserver<? super T> observer) {
return new TracingObserver<>(context, null, null, observer, null);
TracingObserver(Context context, MaybeObserver<? super T> observer) {
this(context, null, null, observer, null);
}

static <T> TracingObserver<T> 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) {
Expand Down
31 changes: 29 additions & 2 deletions core/src/test/java/com/google/adk/telemetry/ContextPropagationTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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<Integer> 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<String> testKey = ContextKey.named("test-key");
Expand Down Expand Up @@ -640,7 +667,7 @@ public void testAgentWithToolCallTraceHierarchy() throws InterruptedException {
List<SpanData> 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);
Expand Down Expand Up @@ -716,7 +743,7 @@ public void testNestedAgentTraceHierarchy() throws InterruptedException {
List<SpanData> 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);

Expand Down