Fix tracing CoroutineCrudRepository.findById

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

@@ -89,18 +89,27 @@ public void postProcess(ProxyFactory factory, RepositoryInformation repositoryIn
   }
 
   static final class RepositoryInterceptor implements MethodInterceptor {
+    private static final Class<?> MONO_CLASS = loadClass("reactor.core.publisher.Mono");
     private final Class<?> repositoryInterface;
 
     RepositoryInterceptor(Class<?> repositoryInterface) {
       this.repositoryInterface = repositoryInterface;
     }
 
+    private static Class<?> loadClass(String name) {
+      try {
+        return Class.forName(name);
+      } catch (ClassNotFoundException exception) {
+        return null;
+      }
+    }
+
     @Override
     public Object invoke(MethodInvocation methodInvocation) throws Throwable {
       Context parentContext = currentContext();
       Method method = methodInvocation.getMethod();
       // Since this interceptor is the outermost interceptor, non-Repository methods
-      // including Object methods will also flow through here.  Don't create spans for those.
+      // including Object methods will also flow through here. Don't create spans for those.
       boolean isRepositoryOp = !Object.class.equals(method.getDeclaringClass());
       ClassAndMethod classAndMethod = ClassAndMethod.create(repositoryInterface, method.getName());
       if (!isRepositoryOp || !instrumenter().shouldStart(parentContext, classAndMethod)) {
@@ -110,7 +119,14 @@ public Object invoke(MethodInvocation methodInvocation) throws Throwable {
       Context context = instrumenter().start(parentContext, classAndMethod);
       try (Scope ignored = context.makeCurrent()) {
         Object result = methodInvocation.proceed();
-        return AsyncOperationEndSupport.create(instrumenter(), Void.class, method.getReturnType())
+        Class<?> type = method.getReturnType();
+        // the return type for
+        // org.springframework.data.repository.kotlin.CoroutineCrudRepository#findById
+        // is Object but the method may actually return a Mono
+        if (Object.class == type && MONO_CLASS != null && MONO_CLASS.isInstance(result)) {
+          type = MONO_CLASS;
+        }
+        return AsyncOperationEndSupport.create(instrumenter(), Void.class, type)
             .asyncEnd(context, classAndMethod, result, null);
       } catch (Throwable t) {
         instrumenter().end(context, classAndMethod, null, t);

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

@@ -0,0 +1,32 @@
+import org.jetbrains.kotlin.gradle.dsl.JvmTarget
+
+plugins {
+  id("otel.javaagent-testing")
+  id("org.jetbrains.kotlin.jvm")
+}
+
+dependencies {
+  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"))
+
+  testLibrary("org.springframework.data:spring-data-r2dbc:3.0.0")
+
+  testImplementation("org.jetbrains.kotlinx:kotlinx-coroutines-reactor:1.8.1")
+  testImplementation("org.jetbrains.kotlin:kotlin-reflect")
+
+  testImplementation("org.testcontainers:testcontainers")
+  testImplementation("io.r2dbc:r2dbc-h2:1.0.0.RELEASE")
+  testImplementation("com.h2database:h2:1.4.197")
+}
+
+otelJava {
+  minJavaVersionSupported.set(JavaVersion.VERSION_17)
+}
+
+kotlin {
+  compilerOptions {
+    jvmTarget.set(JvmTarget.JVM_17)
+  }
+}

instrumentation/spring/spring-data/spring-data-3.0/kotlin-testing/src/test/kotlin/io/opentelemetry/javaagent/instrumentation/spring/data/v3_0/KotlinSpringDataTest.kt

@@ -0,0 +1,60 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.spring.data.v3_0
+
+import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension
+import io.opentelemetry.javaagent.instrumentation.spring.data.v3_0.repository.CustomerRepository
+import io.opentelemetry.javaagent.instrumentation.spring.data.v3_0.repository.PersistenceConfig
+import kotlinx.coroutines.runBlocking
+import org.assertj.core.api.Assertions
+import org.junit.jupiter.api.AfterAll
+import org.junit.jupiter.api.BeforeAll
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.TestInstance
+import org.junit.jupiter.api.extension.RegisterExtension
+import org.springframework.context.ConfigurableApplicationContext
+import org.springframework.context.annotation.AnnotationConfigApplicationContext
+
+@TestInstance(TestInstance.Lifecycle.PER_CLASS)
+class KotlinSpringDataTest {
+
+  companion object {
+    @JvmStatic
+    @RegisterExtension
+    val testing = AgentInstrumentationExtension.create()
+  }
+
+  private var applicationContext: ConfigurableApplicationContext? = null
+  private var customerRepository: CustomerRepository? = null
+
+  @BeforeAll
+  fun setUp() {
+    applicationContext = AnnotationConfigApplicationContext(PersistenceConfig::class.java)
+    customerRepository = applicationContext!!.getBean(CustomerRepository::class.java)
+  }
+
+  @AfterAll
+  fun cleanUp() {
+    applicationContext!!.close()
+  }
+
+  @Test
+  fun `trace findById`() {
+    runBlocking {
+      val customer = customerRepository?.findById(1)
+      Assertions.assertThat(customer?.name).isEqualTo("Name")
+    }
+
+    testing.waitAndAssertTraces({
+        trace ->
+      trace.hasSpansSatisfyingExactly({
+        it.hasName("CustomerRepository.findById").hasNoParent()
+      }, {
+        it.hasName("SELECT db.customer").hasParent(trace.getSpan(0))
+      })
+    })
+  }
+}

instrumentation/spring/spring-data/spring-data-3.0/kotlin-testing/src/test/kotlin/io/opentelemetry/javaagent/instrumentation/spring/data/v3_0/repository/Customer.kt

@@ -0,0 +1,16 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.spring.data.v3_0.repository
+
+import org.springframework.data.annotation.Id
+import org.springframework.data.relational.core.mapping.Column
+import org.springframework.data.relational.core.mapping.Table
+
+@Table("customer")
+data class Customer(
+  @Id @Column("id") val id: Long,
+  @Column("name") val name: String,
+)

instrumentation/spring/spring-data/spring-data-3.0/kotlin-testing/src/test/kotlin/io/opentelemetry/javaagent/instrumentation/spring/data/v3_0/repository/CustomerRepository.kt

@@ -0,0 +1,12 @@
+/*
+ * 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.kotlin.CoroutineCrudRepository
+import org.springframework.stereotype.Repository
+
+@Repository
+interface CustomerRepository : CoroutineCrudRepository<Customer, Long>

instrumentation/spring/spring-data/spring-data-3.0/kotlin-testing/src/test/kotlin/io/opentelemetry/javaagent/instrumentation/spring/data/v3_0/repository/PersistenceConfig.kt

@@ -0,0 +1,63 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.spring.data.v3_0.repository
+
+import io.r2dbc.spi.ConnectionFactories
+import io.r2dbc.spi.ConnectionFactory
+import io.r2dbc.spi.ConnectionFactoryOptions
+import io.r2dbc.spi.Option
+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
+import java.nio.charset.StandardCharsets
+
+@EnableR2dbcRepositories(basePackages = ["io.opentelemetry.javaagent.instrumentation.spring.data.v3_0.repository"])
+class PersistenceConfig {
+
+  @Bean
+  fun connectionFactory(): ConnectionFactory? {
+    return ConnectionFactories.find(
+      ConnectionFactoryOptions.builder()
+        .option(ConnectionFactoryOptions.DRIVER, "h2")
+        .option(ConnectionFactoryOptions.PROTOCOL, "mem")
+        .option(ConnectionFactoryOptions.HOST, "localhost")
+        .option(ConnectionFactoryOptions.USER, "sa")
+        .option(ConnectionFactoryOptions.PASSWORD, "")
+        .option(ConnectionFactoryOptions.DATABASE, "db")
+        .option(Option.valueOf("DB_CLOSE_DELAY"), "-1")
+        .build()
+    )
+  }
+
+  @Bean
+  fun initializer(connectionFactory: ConnectionFactory): ConnectionFactoryInitializer {
+    val initializer = ConnectionFactoryInitializer()
+    initializer.setConnectionFactory(connectionFactory)
+    initializer.setDatabasePopulator(
+      ResourceDatabasePopulator(
+        ByteArrayResource(
+          ("CREATE TABLE customer (id INT PRIMARY KEY, name VARCHAR(100) NOT NULL);" +
+            "INSERT INTO customer (id, name) VALUES ('1', 'Name');")
+            .toByteArray(StandardCharsets.UTF_8)
+        )
+      )
+    )
+
+    return initializer
+  }
+
+  @Bean
+  fun r2dbcEntityTemplate(connectionFactory: ConnectionFactory): R2dbcEntityTemplate {
+    val databaseClient = DatabaseClient.create(connectionFactory)
+
+    return R2dbcEntityTemplate(databaseClient, H2Dialect.INSTANCE)
+  }
+}

settings.gradle.kts

@@ -560,6 +560,7 @@ include(":instrumentation:spring:spring-cloud-gateway:spring-cloud-gateway-commo
 include(":instrumentation:spring:spring-core-2.0:javaagent")
 include(":instrumentation:spring:spring-data:spring-data-1.8:javaagent")
 include(":instrumentation:spring:spring-data:spring-data-3.0:testing")
+include(":instrumentation:spring:spring-data:spring-data-3.0:kotlin-testing")
 include(":instrumentation:spring:spring-data:spring-data-common:testing")
 include(":instrumentation:spring:spring-integration-4.1:javaagent")
 include(":instrumentation:spring:spring-integration-4.1:library")