From d3477cd37da5239ab09ecf8d160733262e96b752 Mon Sep 17 00:00:00 2001 From: cleverchuk Date: Thu, 17 Oct 2024 17:12:51 -0400 Subject: [PATCH] NH-93561: use DeferredResultSetAccess to add `sw` span --- .../hibernate/v6_0/Commenter.java | 2 +- .../hibernate/v6_0/DrsaInstrumentation.java | 106 ++++++++++++++++++ .../v6_0/HibernateInstrumentationModule.java | 2 +- .../hibernate/v6_0/HibernateInstrumenter.java | 1 + 4 files changed, 109 insertions(+), 2 deletions(-) create mode 100644 instrumentation/hibernate-6.0/src/main/java/com/solarwinds/opentelemetry/instrumentation/hibernate/v6_0/DrsaInstrumentation.java diff --git a/instrumentation/hibernate-6.0/src/main/java/com/solarwinds/opentelemetry/instrumentation/hibernate/v6_0/Commenter.java b/instrumentation/hibernate-6.0/src/main/java/com/solarwinds/opentelemetry/instrumentation/hibernate/v6_0/Commenter.java index 55ef496e..3b996bc9 100644 --- a/instrumentation/hibernate-6.0/src/main/java/com/solarwinds/opentelemetry/instrumentation/hibernate/v6_0/Commenter.java +++ b/instrumentation/hibernate-6.0/src/main/java/com/solarwinds/opentelemetry/instrumentation/hibernate/v6_0/Commenter.java @@ -32,7 +32,7 @@ public static String generateComment(Context context) { String traceContext = "00-" + spanContext.getTraceId() + "-" + spanContext.getSpanId() + "-" + flags; - String tag = String.format("traceparent='%s'", traceContext); + String tag = String.format("/*traceparent='%s'*/", traceContext); span.setAttribute("QueryTag", tag); return tag; } diff --git a/instrumentation/hibernate-6.0/src/main/java/com/solarwinds/opentelemetry/instrumentation/hibernate/v6_0/DrsaInstrumentation.java b/instrumentation/hibernate-6.0/src/main/java/com/solarwinds/opentelemetry/instrumentation/hibernate/v6_0/DrsaInstrumentation.java new file mode 100644 index 00000000..b08b2e26 --- /dev/null +++ b/instrumentation/hibernate-6.0/src/main/java/com/solarwinds/opentelemetry/instrumentation/hibernate/v6_0/DrsaInstrumentation.java @@ -0,0 +1,106 @@ +/* + * Copyright The OpenTelemetry Authors + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.solarwinds.opentelemetry.instrumentation.hibernate.v6_0; + +import static com.solarwinds.opentelemetry.instrumentation.hibernate.v6_0.Commenter.generateComment; +import static io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge.currentContext; +import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed; +import static net.bytebuddy.matcher.ElementMatchers.isMethod; +import static net.bytebuddy.matcher.ElementMatchers.named; + +import io.opentelemetry.context.Context; +import io.opentelemetry.context.Scope; +import io.opentelemetry.javaagent.bootstrap.CallDepth; +import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation; +import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer; +import java.lang.reflect.Field; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; +import org.hibernate.query.CommonQueryContract; +import org.hibernate.sql.results.jdbc.internal.DeferredResultSetAccess; + +public class DrsaInstrumentation implements TypeInstrumentation { + + @Override + public ElementMatcher classLoaderOptimization() { + return hasClassesNamed("org.hibernate.sql.results.jdbc.internal.DeferredResultSetAccess"); + } + + @Override + public ElementMatcher typeMatcher() { + return named("org.hibernate.sql.results.jdbc.internal.DeferredResultSetAccess"); + } + + @Override + public void transform(TypeTransformer transformer) { + transformer.applyAdviceToMethod( + isMethod().and(named("getResultSet")), + DrsaInstrumentation.class.getName() + "$QueryMethodAdvice"); + } + + @SuppressWarnings("unused") + public static class QueryMethodAdvice { + + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void startMethod( + @Advice.This DeferredResultSetAccess drsa, + @Advice.Local("swoCallDepth") CallDepth callDepth, + @Advice.Local("swoSqlContext") String swoSql, + @Advice.Local("swoContext") Context context, + @Advice.Local("swoScope") Scope scope) { + callDepth = CallDepth.forClass(CommonQueryContract.class); + if (callDepth.getAndIncrement() > 0) { + return; + } + try { + Class clazz = drsa.getClass(); + Field privateField = clazz.getDeclaredField("finalSql"); + privateField.setAccessible(true); + + String queryString = (String) privateField.get(drsa); + int lastIndexOf = queryString.lastIndexOf("*/"); + if (lastIndexOf != -1) { + queryString = queryString.substring(lastIndexOf); + } + + Context parentContext = currentContext(); + if (queryString == null + || !HibernateInstrumenter.getInstance().shouldStart(parentContext, queryString)) { + return; + } + + context = HibernateInstrumenter.getInstance().start(parentContext, queryString); + scope = context.makeCurrent(); + String comment = generateComment(currentContext()); + + privateField.set(drsa, String.format("%s %s", comment, queryString)); + swoSql = queryString; + + } catch (NoSuchFieldException | IllegalAccessException e) { + e.printStackTrace(); + } + } + + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) + public static void endMethod( + @Advice.Thrown Throwable throwable, + @Advice.Local("swoCallDepth") CallDepth callDepth, + @Advice.Local("swoSqlContext") String swoSql, + @Advice.Local("swoContext") Context context, + @Advice.Local("swoScope") Scope scope) { + + if (callDepth == null || callDepth.decrementAndGet() > 0) { + return; + } + + if (scope != null) { + scope.close(); + HibernateInstrumenter.getInstance().end(context, swoSql, null, throwable); + } + } + } +} diff --git a/instrumentation/hibernate-6.0/src/main/java/com/solarwinds/opentelemetry/instrumentation/hibernate/v6_0/HibernateInstrumentationModule.java b/instrumentation/hibernate-6.0/src/main/java/com/solarwinds/opentelemetry/instrumentation/hibernate/v6_0/HibernateInstrumentationModule.java index fa5d38c5..338fee87 100644 --- a/instrumentation/hibernate-6.0/src/main/java/com/solarwinds/opentelemetry/instrumentation/hibernate/v6_0/HibernateInstrumentationModule.java +++ b/instrumentation/hibernate-6.0/src/main/java/com/solarwinds/opentelemetry/instrumentation/hibernate/v6_0/HibernateInstrumentationModule.java @@ -34,7 +34,7 @@ public boolean isHelperClass(String className) { @Override public List typeInstrumentations() { - return List.of(new QueryInstrumentation()); + return List.of(new DrsaInstrumentation()); } @Override diff --git a/instrumentation/hibernate-6.0/src/main/java/com/solarwinds/opentelemetry/instrumentation/hibernate/v6_0/HibernateInstrumenter.java b/instrumentation/hibernate-6.0/src/main/java/com/solarwinds/opentelemetry/instrumentation/hibernate/v6_0/HibernateInstrumenter.java index afcb54c5..f37d4343 100644 --- a/instrumentation/hibernate-6.0/src/main/java/com/solarwinds/opentelemetry/instrumentation/hibernate/v6_0/HibernateInstrumenter.java +++ b/instrumentation/hibernate-6.0/src/main/java/com/solarwinds/opentelemetry/instrumentation/hibernate/v6_0/HibernateInstrumenter.java @@ -23,6 +23,7 @@ public final class HibernateInstrumenter { private static final Instrumenter INSTANCE = Instrumenter.builder( GlobalOpenTelemetry.get(), "com.solarwinds.jdbc", (sql) -> "sw.jdbc.context") + .addAttributesExtractor(new SqlAttributeExtractor()) .buildInstrumenter(); public static Instrumenter getInstance() {