From bf895082704e016d17c2806f3af8dc6c50ebfe25 Mon Sep 17 00:00:00 2001 From: Andrea Marziali Date: Wed, 7 Oct 2026 10:06:55 +0200 Subject: [PATCH] Fix unfinished R2DBC query spans after filtering results --- .../r2dbc/R2dbcProxyFactory.java | 128 ++++++++++++++++++ .../r2dbc/R2dbcTracingSupport.java | 10 +- .../r2dbc/R2dbcFilteredResultTest.java | 66 +++++++++ 3 files changed, 200 insertions(+), 4 deletions(-) create mode 100644 dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcProxyFactory.java create mode 100644 dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/datadog/trace/instrumentation/r2dbc/R2dbcFilteredResultTest.java diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcProxyFactory.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcProxyFactory.java new file mode 100644 index 00000000000..679664a63a1 --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcProxyFactory.java @@ -0,0 +1,128 @@ +package datadog.trace.instrumentation.r2dbc; + +import datadog.trace.instrumentation.r2dbc.shaded.proxy.callback.ProxyConfig; +import datadog.trace.instrumentation.r2dbc.shaded.proxy.callback.ProxyFactory; +import datadog.trace.instrumentation.r2dbc.shaded.proxy.callback.ProxyFactoryFactory; +import datadog.trace.instrumentation.r2dbc.shaded.proxy.callback.QueriesExecutionContext; +import datadog.trace.instrumentation.r2dbc.shaded.proxy.core.ConnectionInfo; +import datadog.trace.instrumentation.r2dbc.shaded.proxy.core.QueryExecutionInfo; +import datadog.trace.instrumentation.r2dbc.shaded.proxy.core.StatementInfo; +import io.r2dbc.spi.Batch; +import io.r2dbc.spi.Connection; +import io.r2dbc.spi.ConnectionFactory; +import io.r2dbc.spi.Readable; +import io.r2dbc.spi.Result; +import io.r2dbc.spi.Row; +import io.r2dbc.spi.RowMetadata; +import io.r2dbc.spi.Statement; +import io.r2dbc.spi.Wrapped; +import java.util.function.BiFunction; +import java.util.function.Function; +import java.util.function.Predicate; +import org.reactivestreams.Publisher; + +/** Keeps results returned by filter attached to the proxy's query lifecycle. */ +public final class R2dbcProxyFactory implements ProxyFactory, ProxyFactoryFactory { + private final ProxyFactory delegate; + + public R2dbcProxyFactory(ProxyFactory delegate) { + this.delegate = delegate; + } + + @Override + public ProxyFactory create(ProxyConfig config) { + return this; + } + + @Override + public ConnectionFactory wrapConnectionFactory(ConnectionFactory factory) { + return delegate.wrapConnectionFactory(factory); + } + + @Override + public Connection wrapConnection(Connection connection, ConnectionInfo info) { + return delegate.wrapConnection(connection, info); + } + + @Override + public Batch wrapBatch(Batch batch, ConnectionInfo info) { + return delegate.wrapBatch(batch, info); + } + + @Override + public Statement wrapStatement( + Statement statement, StatementInfo statementInfo, ConnectionInfo info) { + return delegate.wrapStatement(statement, statementInfo, info); + } + + @Override + public Result wrapResult( + Result result, QueryExecutionInfo info, QueriesExecutionContext context) { + return new FilterAwareResult(delegate.wrapResult(result, info, context), this, info, context); + } + + @Override + public Row wrapRow(Row row, QueryExecutionInfo info) { + return delegate.wrapRow(row, info); + } + + @Override + public Result.RowSegment wrapRowSegment(Result.RowSegment segment, QueryExecutionInfo info) { + return delegate.wrapRowSegment(segment, info); + } + + public static final class FilterAwareResult implements Result, Wrapped { + private final Result delegate; + private final R2dbcProxyFactory factory; + private final QueryExecutionInfo info; + private final QueriesExecutionContext context; + + public FilterAwareResult( + Result delegate, + R2dbcProxyFactory factory, + QueryExecutionInfo info, + QueriesExecutionContext context) { + this.delegate = delegate; + this.factory = factory; + this.info = info; + this.context = context; + } + + @Override + public Publisher getRowsUpdated() { + return delegate.getRowsUpdated(); + } + + @Override + public Publisher map(BiFunction mapping) { + return delegate.map(mapping); + } + + @Override + public Publisher map(Function mapping) { + return delegate.map(mapping); + } + + @Override + public Result filter(Predicate predicate) { + // r2dbc-proxy 1.1.6 returns the driver's filtered Result without wrapping it. + // Reuse the execution context so consuming this result finishes the original query once. + return factory.wrapResult(delegate.filter(predicate), info, context); + } + + @Override + public Publisher flatMap(Function> mapping) { + return delegate.flatMap(mapping); + } + + @Override + public Result unwrap() { + return (Result) ((Wrapped) delegate).unwrap(); + } + + @Override + public E unwrap(Class type) { + return ((Wrapped) delegate).unwrap(type); + } + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcTracingSupport.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcTracingSupport.java index 4d5863093a4..ad46075a0e8 100644 --- a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcTracingSupport.java +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcTracingSupport.java @@ -2,6 +2,7 @@ import datadog.trace.bootstrap.ContextStore; import datadog.trace.instrumentation.r2dbc.shaded.proxy.ProxyConnectionFactory; +import datadog.trace.instrumentation.r2dbc.shaded.proxy.callback.ProxyConfig; import datadog.trace.instrumentation.r2dbc.shaded.proxy.core.ConnectionInfo; import datadog.trace.instrumentation.r2dbc.shaded.proxy.core.MethodExecutionInfo; import datadog.trace.instrumentation.r2dbc.shaded.proxy.listener.ProxyExecutionListener; @@ -41,10 +42,11 @@ public static ConnectionFactory wrapConnectionFactory( ConnectionMetadataListener metadataListener = new ConnectionMetadataListener(options, connectionOptionsStore); - return ProxyConnectionFactory.builder(factory) - .listener(queryListener) - .listener(metadataListener) - .build(); + ProxyConnectionFactory.Builder builder = + ProxyConnectionFactory.builder(factory).listener(queryListener).listener(metadataListener); + ProxyConfig config = builder.getProxyConfig(); + config.setProxyFactoryFactory(new R2dbcProxyFactory(config.getProxyFactory())); + return builder.build(); } /** diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/datadog/trace/instrumentation/r2dbc/R2dbcFilteredResultTest.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/datadog/trace/instrumentation/r2dbc/R2dbcFilteredResultTest.java new file mode 100644 index 00000000000..133f0b44c35 --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/datadog/trace/instrumentation/r2dbc/R2dbcFilteredResultTest.java @@ -0,0 +1,66 @@ +package datadog.trace.instrumentation.r2dbc; + +import static datadog.trace.agent.test.assertions.SpanMatcher.span; +import static datadog.trace.agent.test.assertions.TraceMatcher.SORT_BY_START_TIME; +import static datadog.trace.agent.test.assertions.TraceMatcher.trace; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.instrumentation.r2dbc.R2dbcInstrumentationTest.H2_QUERY; +import static java.util.Collections.singletonList; +import static org.junit.jupiter.api.Assertions.assertEquals; + +import datadog.context.ContextScope; +import datadog.trace.agent.test.AbstractInstrumentationTest; +import datadog.trace.api.DDSpanTypes; +import datadog.trace.core.DDSpan; +import io.r2dbc.spi.Connection; +import io.r2dbc.spi.ConnectionFactories; +import io.r2dbc.spi.Result; +import java.time.Duration; +import java.util.UUID; +import org.junit.jupiter.api.Test; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +class R2dbcFilteredResultTest extends AbstractInstrumentationTest { + private static final Duration TIMEOUT = Duration.ofSeconds(10); + + @Test + void filteredResultConsumptionFinishesQuerySpan() { + Connection connection = + Mono.from( + ConnectionFactories.get("r2dbc:h2:mem:///filtered_result_" + UUID.randomUUID()) + .create()) + .block(TIMEOUT); + DDSpan parent = (DDSpan) startSpan("test", "parent"); + try (ContextScope scope = activateSpan(parent)) { + assertEquals( + singletonList(42), + Flux.from(connection.createStatement("SELECT 42").execute()) + .flatMap( + result -> + result + .filter(segment -> segment instanceof Result.RowSegment) + .flatMap( + segment -> + Mono.just( + ((Result.RowSegment) segment).row().get(0, Integer.class)))) + .collectList() + .block(TIMEOUT)); + } finally { + parent.finish(); + Mono.from(connection.close()).block(TIMEOUT); + } + + assertTraces( + trace( + SORT_BY_START_TIME, + span().root().operationName("parent"), + span() + .childOf(parent.getSpanId()) + .operationName(H2_QUERY) + .type(DDSpanTypes.SQL) + .measured() + .error(false))); + } +}