Fix unfinished R2DBC query spans after filtering results

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<Result> {
+    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<Long> getRowsUpdated() {
+      return delegate.getRowsUpdated();
+    }
+
+    @Override
+    public <T> Publisher<T> map(BiFunction<Row, RowMetadata, ? extends T> mapping) {
+      return delegate.map(mapping);
+    }
+
+    @Override
+    public <T> Publisher<T> map(Function<? super Readable, ? extends T> mapping) {
+      return delegate.map(mapping);
+    }
+
+    @Override
+    public Result filter(Predicate<Segment> 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 <T> Publisher<T> flatMap(Function<Segment, ? extends Publisher<? extends T>> mapping) {
+      return delegate.flatMap(mapping);
+    }
+
+    @Override
+    public Result unwrap() {
+      return (Result) ((Wrapped<?>) delegate).unwrap();
+    }
+
+    @Override
+    public <E> E unwrap(Class<E> type) {
+      return ((Wrapped<?>) delegate).unwrap(type);
+    }
+  }
+}

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();
   }
 
   /**

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)));
+  }
+}