Skip to content

Commit c2713e1

Browse files
lauritlaurit
authored
Fix tracing CoroutineCrudRepository.findById (#12131)
1 parent 34740ee commit c2713e1

File tree

7 files changed

+202
-2
lines changed

7 files changed

+202
-2
lines changed

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

+18-2
Original file line numberDiff line numberDiff line change
@@ -89,18 +89,27 @@ public void postProcess(ProxyFactory factory, RepositoryInformation repositoryIn
8989
}
9090

9191
static final class RepositoryInterceptor implements MethodInterceptor {
92+
private static final Class<?> MONO_CLASS = loadClass("reactor.core.publisher.Mono");
9293
private final Class<?> repositoryInterface;
9394

9495
RepositoryInterceptor(Class<?> repositoryInterface) {
9596
this.repositoryInterface = repositoryInterface;
9697
}
9798

99+
private static Class<?> loadClass(String name) {
100+
try {
101+
return Class.forName(name);
102+
} catch (ClassNotFoundException exception) {
103+
return null;
104+
}
105+
}
106+
98107
@Override
99108
public Object invoke(MethodInvocation methodInvocation) throws Throwable {
100109
Context parentContext = currentContext();
101110
Method method = methodInvocation.getMethod();
102111
// Since this interceptor is the outermost interceptor, non-Repository methods
103-
// including Object methods will also flow through here. Don't create spans for those.
112+
// including Object methods will also flow through here. Don't create spans for those.
104113
boolean isRepositoryOp = !Object.class.equals(method.getDeclaringClass());
105114
ClassAndMethod classAndMethod = ClassAndMethod.create(repositoryInterface, method.getName());
106115
if (!isRepositoryOp || !instrumenter().shouldStart(parentContext, classAndMethod)) {
@@ -110,7 +119,14 @@ public Object invoke(MethodInvocation methodInvocation) throws Throwable {
110119
Context context = instrumenter().start(parentContext, classAndMethod);
111120
try (Scope ignored = context.makeCurrent()) {
112121
Object result = methodInvocation.proceed();
113-
return AsyncOperationEndSupport.create(instrumenter(), Void.class, method.getReturnType())
122+
Class<?> type = method.getReturnType();
123+
// the return type for
124+
// org.springframework.data.repository.kotlin.CoroutineCrudRepository#findById
125+
// is Object but the method may actually return a Mono
126+
if (Object.class == type && MONO_CLASS != null && MONO_CLASS.isInstance(result)) {
127+
type = MONO_CLASS;
128+
}
129+
return AsyncOperationEndSupport.create(instrumenter(), Void.class, type)
114130
.asyncEnd(context, classAndMethod, result, null);
115131
} catch (Throwable t) {
116132
instrumenter().end(context, classAndMethod, null, t);

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

+32
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
import org.jetbrains.kotlin.gradle.dsl.JvmTarget
2+
3+
plugins {
4+
id("otel.javaagent-testing")
5+
id("org.jetbrains.kotlin.jvm")
6+
}
7+
8+
dependencies {
9+
testInstrumentation(project(":instrumentation:r2dbc-1.0:javaagent"))
10+
testInstrumentation(project(":instrumentation:reactor:reactor-3.1:javaagent"))
11+
testInstrumentation(project(":instrumentation:spring:spring-core-2.0:javaagent"))
12+
testInstrumentation(project(":instrumentation:spring:spring-data:spring-data-1.8:javaagent"))
13+
14+
testLibrary("org.springframework.data:spring-data-r2dbc:3.0.0")
15+
16+
testImplementation("org.jetbrains.kotlinx:kotlinx-coroutines-reactor:1.8.1")
17+
testImplementation("org.jetbrains.kotlin:kotlin-reflect")
18+
19+
testImplementation("org.testcontainers:testcontainers")
20+
testImplementation("io.r2dbc:r2dbc-h2:1.0.0.RELEASE")
21+
testImplementation("com.h2database:h2:1.4.197")
22+
}
23+
24+
otelJava {
25+
minJavaVersionSupported.set(JavaVersion.VERSION_17)
26+
}
27+
28+
kotlin {
29+
compilerOptions {
30+
jvmTarget.set(JvmTarget.JVM_17)
31+
}
32+
}

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

+60
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.javaagent.instrumentation.spring.data.v3_0
7+
8+
import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension
9+
import io.opentelemetry.javaagent.instrumentation.spring.data.v3_0.repository.CustomerRepository
10+
import io.opentelemetry.javaagent.instrumentation.spring.data.v3_0.repository.PersistenceConfig
11+
import kotlinx.coroutines.runBlocking
12+
import org.assertj.core.api.Assertions
13+
import org.junit.jupiter.api.AfterAll
14+
import org.junit.jupiter.api.BeforeAll
15+
import org.junit.jupiter.api.Test
16+
import org.junit.jupiter.api.TestInstance
17+
import org.junit.jupiter.api.extension.RegisterExtension
18+
import org.springframework.context.ConfigurableApplicationContext
19+
import org.springframework.context.annotation.AnnotationConfigApplicationContext
20+
21+
@TestInstance(TestInstance.Lifecycle.PER_CLASS)
22+
class KotlinSpringDataTest {
23+
24+
companion object {
25+
@JvmStatic
26+
@RegisterExtension
27+
val testing = AgentInstrumentationExtension.create()
28+
}
29+
30+
private var applicationContext: ConfigurableApplicationContext? = null
31+
private var customerRepository: CustomerRepository? = null
32+
33+
@BeforeAll
34+
fun setUp() {
35+
applicationContext = AnnotationConfigApplicationContext(PersistenceConfig::class.java)
36+
customerRepository = applicationContext!!.getBean(CustomerRepository::class.java)
37+
}
38+
39+
@AfterAll
40+
fun cleanUp() {
41+
applicationContext!!.close()
42+
}
43+
44+
@Test
45+
fun `trace findById`() {
46+
runBlocking {
47+
val customer = customerRepository?.findById(1)
48+
Assertions.assertThat(customer?.name).isEqualTo("Name")
49+
}
50+
51+
testing.waitAndAssertTraces({
52+
trace ->
53+
trace.hasSpansSatisfyingExactly({
54+
it.hasName("CustomerRepository.findById").hasNoParent()
55+
}, {
56+
it.hasName("SELECT db.customer").hasParent(trace.getSpan(0))
57+
})
58+
})
59+
}
60+
}

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

+16
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.javaagent.instrumentation.spring.data.v3_0.repository
7+
8+
import org.springframework.data.annotation.Id
9+
import org.springframework.data.relational.core.mapping.Column
10+
import org.springframework.data.relational.core.mapping.Table
11+
12+
@Table("customer")
13+
data class Customer(
14+
@Id @Column("id") val id: Long,
15+
@Column("name") val name: String,
16+
)

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

+12
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.javaagent.instrumentation.spring.data.v3_0.repository
7+
8+
import org.springframework.data.repository.kotlin.CoroutineCrudRepository
9+
import org.springframework.stereotype.Repository
10+
11+
@Repository
12+
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

+63
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.javaagent.instrumentation.spring.data.v3_0.repository
7+
8+
import io.r2dbc.spi.ConnectionFactories
9+
import io.r2dbc.spi.ConnectionFactory
10+
import io.r2dbc.spi.ConnectionFactoryOptions
11+
import io.r2dbc.spi.Option
12+
import org.springframework.context.annotation.Bean
13+
import org.springframework.core.io.ByteArrayResource
14+
import org.springframework.data.r2dbc.core.R2dbcEntityTemplate
15+
import org.springframework.data.r2dbc.dialect.H2Dialect
16+
import org.springframework.data.r2dbc.repository.config.EnableR2dbcRepositories
17+
import org.springframework.r2dbc.connection.init.ConnectionFactoryInitializer
18+
import org.springframework.r2dbc.connection.init.ResourceDatabasePopulator
19+
import org.springframework.r2dbc.core.DatabaseClient
20+
import java.nio.charset.StandardCharsets
21+
22+
@EnableR2dbcRepositories(basePackages = ["io.opentelemetry.javaagent.instrumentation.spring.data.v3_0.repository"])
23+
class PersistenceConfig {
24+
25+
@Bean
26+
fun connectionFactory(): ConnectionFactory? {
27+
return ConnectionFactories.find(
28+
ConnectionFactoryOptions.builder()
29+
.option(ConnectionFactoryOptions.DRIVER, "h2")
30+
.option(ConnectionFactoryOptions.PROTOCOL, "mem")
31+
.option(ConnectionFactoryOptions.HOST, "localhost")
32+
.option(ConnectionFactoryOptions.USER, "sa")
33+
.option(ConnectionFactoryOptions.PASSWORD, "")
34+
.option(ConnectionFactoryOptions.DATABASE, "db")
35+
.option(Option.valueOf("DB_CLOSE_DELAY"), "-1")
36+
.build()
37+
)
38+
}
39+
40+
@Bean
41+
fun initializer(connectionFactory: ConnectionFactory): ConnectionFactoryInitializer {
42+
val initializer = ConnectionFactoryInitializer()
43+
initializer.setConnectionFactory(connectionFactory)
44+
initializer.setDatabasePopulator(
45+
ResourceDatabasePopulator(
46+
ByteArrayResource(
47+
("CREATE TABLE customer (id INT PRIMARY KEY, name VARCHAR(100) NOT NULL);" +
48+
"INSERT INTO customer (id, name) VALUES ('1', 'Name');")
49+
.toByteArray(StandardCharsets.UTF_8)
50+
)
51+
)
52+
)
53+
54+
return initializer
55+
}
56+
57+
@Bean
58+
fun r2dbcEntityTemplate(connectionFactory: ConnectionFactory): R2dbcEntityTemplate {
59+
val databaseClient = DatabaseClient.create(connectionFactory)
60+
61+
return R2dbcEntityTemplate(databaseClient, H2Dialect.INSTANCE)
62+
}
63+
}

‎settings.gradle.kts

+1
Original file line numberDiff line numberDiff line change
@@ -563,6 +563,7 @@ include(":instrumentation:spring:spring-cloud-gateway:spring-cloud-gateway-commo
563563
include(":instrumentation:spring:spring-core-2.0:javaagent")
564564
include(":instrumentation:spring:spring-data:spring-data-1.8:javaagent")
565565
include(":instrumentation:spring:spring-data:spring-data-3.0:testing")
566+
include(":instrumentation:spring:spring-data:spring-data-3.0:kotlin-testing")
566567
include(":instrumentation:spring:spring-data:spring-data-common:testing")
567568
include(":instrumentation:spring:spring-integration-4.1:javaagent")
568569
include(":instrumentation:spring:spring-integration-4.1:library")

0 commit comments

Comments
 (0)