Improve spring data reactive instrumentation (#9561)

instrumentation/spring/spring-data/spring-data-1.8/javaagent/build.gradle.kts

@@ -26,6 +26,7 @@ muzzle {
 dependencies {
   library("org.springframework.data:spring-data-commons:1.8.0.RELEASE")
   compileOnly("org.springframework:spring-aop:1.2")
+  compileOnly(project(":instrumentation-annotations-support"))
 
   testInstrumentation(project(":instrumentation:jdbc:javaagent"))
 

instrumentation/spring/spring-data/spring-data-1.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/data/v1_8/SpringDataInstrumentationModule.java

@@ -14,6 +14,7 @@ import static net.bytebuddy.matcher.ElementMatchers.named;
 import com.google.auto.service.AutoService;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
+import io.opentelemetry.instrumentation.api.annotation.support.async.AsyncOperationEndSupport;
 import io.opentelemetry.instrumentation.api.instrumenter.util.ClassAndMethod;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
@@ -109,8 +110,8 @@ public class SpringDataInstrumentationModule extends InstrumentationModule {
       Context context = instrumenter().start(parentContext, classAndMethod);
       try (Scope ignored = context.makeCurrent()) {
         Object result = methodInvocation.proceed();
-        instrumenter().end(context, classAndMethod, null, null);
-        return result;
+        return AsyncOperationEndSupport.create(instrumenter(), Void.class, method.getReturnType())
+            .asyncEnd(context, classAndMethod, result, null);
       } catch (Throwable t) {
         instrumenter().end(context, classAndMethod, null, t);
         throw t;

instrumentation/spring/spring-data/spring-data-3.0/testing/build.gradle.kts

@@ -4,6 +4,8 @@ plugins {
 
 dependencies {
   testInstrumentation(project(":instrumentation:jdbc:javaagent"))
+  testInstrumentation(project(":instrumentation:r2dbc-1.0:javaagent"))
+  testInstrumentation(project(":instrumentation:reactor:reactor-3.1:javaagent"))
   testInstrumentation(project(":instrumentation:spring:spring-core-2.0:javaagent"))
   testInstrumentation(project(":instrumentation:spring:spring-data:spring-data-1.8:javaagent"))
 
@@ -12,9 +14,12 @@ dependencies {
   testLibrary("org.hibernate.orm:hibernate-core:6.0.0.Final")
   testLibrary("org.springframework.data:spring-data-commons:3.0.0")
   testLibrary("org.springframework.data:spring-data-jpa:3.0.0")
+  testLibrary("org.springframework.data:spring-data-r2dbc:3.0.0")
   testLibrary("org.springframework:spring-test:6.0.0")
 
   testImplementation("org.hsqldb:hsqldb:2.0.0")
+  testImplementation("com.h2database:h2:1.4.197")
+  testImplementation("io.r2dbc:r2dbc-h2:1.0.0.RELEASE")
 
   latestDepTestLibrary("org.hibernate.orm:hibernate-core:6.2.+")
 }
@@ -23,10 +28,27 @@ otelJava {
   minJavaVersionSupported.set(JavaVersion.VERSION_17)
 }
 
+testing {
+  suites {
+    val reactiveTest by registering(JvmTestSuite::class) {
+      dependencies {
+        implementation("org.springframework.data:spring-data-r2dbc:3.0.0")
+        implementation("org.testcontainers:testcontainers")
+        implementation("io.r2dbc:r2dbc-h2:1.0.0.RELEASE")
+        implementation("com.h2database:h2:1.4.197")
+      }
+    }
+  }
+}
+
 tasks {
   test {
     jvmArgs("--add-opens=java.base/java.lang.invoke=ALL-UNNAMED")
     jvmArgs("--add-opens=java.base/java.lang=ALL-UNNAMED")
     jvmArgs("-XX:+IgnoreUnrecognizedVMOptions")
   }
+
+  check {
+    dependsOn(testing.suites)
+  }
 }

instrumentation/spring/spring-data/spring-data-3.0/testing/src/reactiveTest/java/io/opentelemetry/javaagent/instrumentation/spring/data/v3_0/ReactiveSpringDataTest.java

@@ -0,0 +1,93 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.spring.data.v3_0;
+
+import static io.opentelemetry.sdk.testing.assertj.OpenTelemetryAssertions.equalTo;
+import static io.opentelemetry.semconv.SemanticAttributes.DB_CONNECTION_STRING;
+import static io.opentelemetry.semconv.SemanticAttributes.DB_NAME;
+import static io.opentelemetry.semconv.SemanticAttributes.DB_OPERATION;
+import static io.opentelemetry.semconv.SemanticAttributes.DB_SQL_TABLE;
+import static io.opentelemetry.semconv.SemanticAttributes.DB_STATEMENT;
+import static io.opentelemetry.semconv.SemanticAttributes.DB_SYSTEM;
+import static io.opentelemetry.semconv.SemanticAttributes.DB_USER;
+import static io.opentelemetry.semconv.SemanticAttributes.NET_PEER_NAME;
+import static org.assertj.core.api.Assertions.assertThat;
+
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
+import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
+import io.opentelemetry.javaagent.instrumentation.spring.data.v3_0.repository.CustomerRepository;
+import io.opentelemetry.javaagent.instrumentation.spring.data.v3_0.repository.PersistenceConfig;
+import io.opentelemetry.semconv.SemanticAttributes;
+import java.time.Duration;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.springframework.context.ConfigurableApplicationContext;
+import org.springframework.context.annotation.AnnotationConfigApplicationContext;
+
+@SuppressWarnings("deprecation") // until old http semconv are dropped in 2.0
+class ReactiveSpringDataTest {
+
+  @RegisterExtension
+  private static final InstrumentationExtension testing = AgentInstrumentationExtension.create();
+
+  private static ConfigurableApplicationContext applicationContext;
+  private static CustomerRepository customerRepository;
+
+  @BeforeAll
+  static void setUp() {
+    applicationContext = new AnnotationConfigApplicationContext(PersistenceConfig.class);
+    customerRepository = applicationContext.getBean(CustomerRepository.class);
+  }
+
+  @AfterAll
+  static void cleanUp() {
+    applicationContext.close();
+  }
+
+  @Test
+  void testFindAll() {
+    long count =
+        testing
+            .runWithSpan("parent", () -> customerRepository.findAll())
+            .count()
+            .block(Duration.ofSeconds(30));
+    assertThat(count).isEqualTo(1);
+
+    testing.waitAndAssertTraces(
+        trace ->
+            trace.hasSpansSatisfyingExactly(
+                span -> span.hasName("parent").hasKind(SpanKind.INTERNAL),
+                span ->
+                    span.hasName("CustomerRepository.findAll")
+                        .hasKind(SpanKind.INTERNAL)
+                        .hasAttributesSatisfyingExactly(
+                            equalTo(
+                                SemanticAttributes.CODE_NAMESPACE,
+                                CustomerRepository.class.getName()),
+                            equalTo(SemanticAttributes.CODE_FUNCTION, "findAll")),
+                span ->
+                    span.hasName("SELECT db.CUSTOMER")
+                        .hasKind(SpanKind.CLIENT)
+                        .hasParent(trace.getSpan(1))
+                        // assert that this span ends before its parent span
+                        .satisfies(
+                            spanData ->
+                                assertThat(spanData.getEndEpochNanos())
+                                    .isLessThanOrEqualTo(trace.getSpan(1).getEndEpochNanos()))
+                        .hasAttributesSatisfyingExactly(
+                            equalTo(DB_SYSTEM, "h2"),
+                            equalTo(DB_NAME, "db"),
+                            equalTo(DB_USER, "sa"),
+                            equalTo(DB_STATEMENT, "SELECT CUSTOMER.* FROM CUSTOMER"),
+                            equalTo(DB_OPERATION, "SELECT"),
+                            equalTo(DB_SQL_TABLE, "CUSTOMER"),
+                            equalTo(DB_CONNECTION_STRING, "h2:mem://localhost"),
+                            equalTo(NET_PEER_NAME, "localhost"))));
+  }
+}

instrumentation/spring/spring-data/spring-data-3.0/testing/src/reactiveTest/java/io/opentelemetry/javaagent/instrumentation/spring/data/v3_0/repository/Customer.java

@@ -0,0 +1,73 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.spring.data.v3_0.repository;
+
+import java.util.Objects;
+import javax.annotation.Nullable;
+import org.springframework.data.annotation.Id;
+
+public class Customer {
+
+  @Id private Long id;
+
+  private String firstName;
+  private String lastName;
+
+  protected Customer() {}
+
+  public Customer(String firstName, String lastName) {
+    this.firstName = firstName;
+    this.lastName = lastName;
+  }
+
+  public Long getId() {
+    return id;
+  }
+
+  public void setId(Long id) {
+    this.id = id;
+  }
+
+  public String getFirstName() {
+    return firstName;
+  }
+
+  public void setFirstName(String firstName) {
+    this.firstName = firstName;
+  }
+
+  public String getLastName() {
+    return lastName;
+  }
+
+  public void setLastName(String lastName) {
+    this.lastName = lastName;
+  }
+
+  @Override
+  public String toString() {
+    return String.format("Customer[id=%d, firstName='%s', lastName='%s']", id, firstName, lastName);
+  }
+
+  @Override
+  public boolean equals(@Nullable Object obj) {
+    if (obj == this) {
+      return true;
+    }
+    if (!(obj instanceof Customer)) {
+      return false;
+    }
+    Customer other = (Customer) obj;
+    return Objects.equals(id, other.id)
+        && Objects.equals(firstName, other.firstName)
+        && Objects.equals(lastName, other.lastName);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(id, firstName, lastName);
+  }
+}

instrumentation/spring/spring-data/spring-data-3.0/testing/src/reactiveTest/java/io/opentelemetry/javaagent/instrumentation/spring/data/v3_0/repository/CustomerRepository.java

@@ -0,0 +1,10 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.spring.data.v3_0.repository;
+
+import org.springframework.data.repository.reactive.ReactiveCrudRepository;
+
+public interface CustomerRepository extends ReactiveCrudRepository<Customer, Long> {}

instrumentation/spring/spring-data/spring-data-3.0/testing/src/reactiveTest/java/io/opentelemetry/javaagent/instrumentation/spring/data/v3_0/repository/PersistenceConfig.java

@@ -0,0 +1,67 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.spring.data.v3_0.repository;
+
+import static io.r2dbc.spi.ConnectionFactoryOptions.DATABASE;
+import static io.r2dbc.spi.ConnectionFactoryOptions.DRIVER;
+import static io.r2dbc.spi.ConnectionFactoryOptions.HOST;
+import static io.r2dbc.spi.ConnectionFactoryOptions.PASSWORD;
+import static io.r2dbc.spi.ConnectionFactoryOptions.PROTOCOL;
+import static io.r2dbc.spi.ConnectionFactoryOptions.USER;
+
+import io.r2dbc.spi.ConnectionFactories;
+import io.r2dbc.spi.ConnectionFactory;
+import io.r2dbc.spi.ConnectionFactoryOptions;
+import io.r2dbc.spi.Option;
+import java.nio.charset.StandardCharsets;
+import org.springframework.context.annotation.Bean;
+import org.springframework.core.io.ByteArrayResource;
+import org.springframework.data.r2dbc.core.R2dbcEntityTemplate;
+import org.springframework.data.r2dbc.dialect.H2Dialect;
+import org.springframework.data.r2dbc.repository.config.EnableR2dbcRepositories;
+import org.springframework.r2dbc.connection.init.ConnectionFactoryInitializer;
+import org.springframework.r2dbc.connection.init.ResourceDatabasePopulator;
+import org.springframework.r2dbc.core.DatabaseClient;
+
+@EnableR2dbcRepositories(
+    basePackages = "io.opentelemetry.javaagent.instrumentation.spring.data.v3_0.repository")
+public class PersistenceConfig {
+
+  @Bean
+  ConnectionFactory connectionFactory() {
+    return ConnectionFactories.find(
+        ConnectionFactoryOptions.builder()
+            .option(DRIVER, "h2")
+            .option(PROTOCOL, "mem")
+            .option(HOST, "localhost")
+            .option(USER, "sa")
+            .option(PASSWORD, "")
+            .option(DATABASE, "db")
+            .option(Option.valueOf("DB_CLOSE_DELAY"), "-1")
+            .build());
+  }
+
+  @Bean
+  ConnectionFactoryInitializer initializer(ConnectionFactory connectionFactory) {
+    ConnectionFactoryInitializer initializer = new ConnectionFactoryInitializer();
+    initializer.setConnectionFactory(connectionFactory);
+    initializer.setDatabasePopulator(
+        new ResourceDatabasePopulator(
+            new ByteArrayResource(
+                ("CREATE TABLE customer (id INT PRIMARY KEY, firstname VARCHAR(100) NOT NULL, lastname VARCHAR(100) NOT NULL);"
+                        + "INSERT INTO customer (id, firstname, lastname) VALUES ('1', 'First', 'Last');")
+                    .getBytes(StandardCharsets.UTF_8))));
+
+    return initializer;
+  }
+
+  @Bean
+  public R2dbcEntityTemplate r2dbcEntityTemplate(ConnectionFactory connectionFactory) {
+    DatabaseClient databaseClient = DatabaseClient.create(connectionFactory);
+
+    return new R2dbcEntityTemplate(databaseClient, H2Dialect.INSTANCE);
+  }
+}

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/ignore/AdditionalLibraryIgnoredTypesConfigurer.java

@@ -154,7 +154,8 @@ public class AdditionalLibraryIgnoredTypesConfigurer implements IgnoredTypesConf
         .allowClass("org.springframework.core.task.")
         .allowClass("org.springframework.core.DecoratingClassLoader")
         .allowClass("org.springframework.core.OverridingClassLoader")
-        .allowClass("org.springframework.core.ReactiveAdapterRegistry$EmptyCompletableFuture");
+        .allowClass("org.springframework.core.ReactiveAdapterRegistry$EmptyCompletableFuture")
+        .allowClass("org.springframework.core.io.buffer.DataBufferUtils$");
 
     builder
         .ignoreClass("org.springframework.instrument.")