Fix JDBC instrumentation: recursive statements inside Connection#getMetaData

instrumentation/jdbc/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jdbc/PreparedStatementInstrumentation.java

@@ -17,8 +17,10 @@
 
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
+import io.opentelemetry.javaagent.instrumentation.api.CallDepthThreadLocalMap;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
 import java.sql.PreparedStatement;
+import java.sql.Statement;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
@@ -51,6 +53,17 @@ public static void onEnter(
         @Advice.This PreparedStatement statement,
         @Advice.Local("otelContext") Context context,
         @Advice.Local("otelScope") Scope scope) {
+      // Connection#getMetaData() may execute a Statement or PreparedStatement to retrieve DB info
+      // this happens before the DB CLIENT span is started (and put in the current context), so this
+      // instrumentation runs again and the shouldStartSpan() check always returns true - and so on
+      // until we get a StackOverflowError
+      // using CallDepth prevents this, because this check happens before Connection#getMetadata()
+      // is called - the first recursive Statement call is just skipped and we do not create a span
+      // for it
+      if (CallDepthThreadLocalMap.getCallDepth(Statement.class).getAndIncrement() > 0) {
+        return;
+      }
+
       Context parentContext = currentContext();
       if (!tracer().shouldStartSpan(parentContext)) {
         return;
@@ -68,6 +81,7 @@ public static void stopSpan(
       if (scope == null) {
         return;
       }
+      CallDepthThreadLocalMap.reset(Statement.class);
 
       scope.close();
       if (throwable == null) {

instrumentation/jdbc/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jdbc/StatementInstrumentation.java

@@ -17,6 +17,7 @@
 
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
+import io.opentelemetry.javaagent.instrumentation.api.CallDepthThreadLocalMap;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
 import java.sql.Statement;
 import java.util.Map;
@@ -52,6 +53,17 @@ public static void onEnter(
         @Advice.This Statement statement,
         @Advice.Local("otelContext") Context context,
         @Advice.Local("otelScope") Scope scope) {
+      // Connection#getMetaData() may execute a Statement or PreparedStatement to retrieve DB info
+      // this happens before the DB CLIENT span is started (and put in the current context), so this
+      // instrumentation runs again and the shouldStartSpan() check always returns true - and so on
+      // until we get a StackOverflowError
+      // using CallDepth prevents this, because this check happens before Connection#getMetadata()
+      // is called - the first recursive Statement call is just skipped and we do not create a span
+      // for it
+      if (CallDepthThreadLocalMap.getCallDepth(Statement.class).getAndIncrement() > 0) {
+        return;
+      }
+
       Context parentContext = currentContext();
       if (!tracer().shouldStartSpan(parentContext)) {
         return;
@@ -69,6 +81,7 @@ public static void stopSpan(
       if (scope == null) {
         return;
       }
+      CallDepthThreadLocalMap.reset(Statement.class);
 
       scope.close();
       if (throwable == null) {

instrumentation/jdbc/javaagent/src/test/groovy/JdbcInstrumentationTest.groovy

@@ -14,8 +14,10 @@ import io.opentelemetry.instrumentation.test.AgentInstrumentationSpecification
 import io.opentelemetry.semconv.trace.attributes.SemanticAttributes
 import java.sql.CallableStatement
 import java.sql.Connection
+import java.sql.DatabaseMetaData
 import java.sql.PreparedStatement
 import java.sql.ResultSet
+import java.sql.SQLException
 import java.sql.Statement
 import javax.sql.DataSource
 import org.apache.derby.jdbc.EmbeddedDataSource
@@ -724,4 +726,63 @@ class JdbcInstrumentationTest extends AgentInstrumentationSpecification {
     "tomcat"           | _
     "c3p0"             | _
   }
+
+  // regression test for https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/2644
+  @Unroll
+  def "should handle recursive Statements inside Connection.getMetaData(): #desc"() {
+    given:
+    def connection = new DbCallingConnection(usePreparedStatementInConnection)
+
+    when:
+    runUnderTrace("parent") {
+      executeQueryFunction(connection, "SELECT * FROM table")
+    }
+
+    then:
+    assertTraces(1) {
+      trace(0, 2) {
+        basicSpan(it, 0, "parent")
+        span(1) {
+          name "SELECT table"
+          kind CLIENT
+          childOf span(0)
+          errored false
+          attributes {
+            "$SemanticAttributes.DB_SYSTEM.key" "testdb"
+            "$SemanticAttributes.DB_CONNECTION_STRING.key" "testdb://localhost"
+            "$SemanticAttributes.DB_STATEMENT.key" "SELECT * FROM table"
+            "$SemanticAttributes.DB_OPERATION.key" "SELECT"
+            "$SemanticAttributes.DB_SQL_TABLE.key" "table"
+          }
+        }
+      }
+    }
+
+    where:
+    desc                                                           | usePreparedStatementInConnection | executeQueryFunction
+    "getMetaData() uses Statement, test Statement"                 | false                            | { con, query -> con.createStatement().executeQuery(query) }
+    "getMetaData() uses PreparedStatement, test Statement"         | true                             | { con, query -> con.createStatement().executeQuery(query) }
+    "getMetaData() uses Statement, test PreparedStatement"         | false                            | { con, query -> con.prepareStatement(query).executeQuery() }
+    "getMetaData() uses PreparedStatement, test PreparedStatement" | true                             | { con, query -> con.prepareStatement(query).executeQuery() }
+  }
+
+  class DbCallingConnection extends TestConnection {
+    final boolean usePreparedStatement
+
+    DbCallingConnection(boolean usePreparedStatement) {
+      super(false)
+      this.usePreparedStatement = usePreparedStatement
+    }
+
+    @Override
+    DatabaseMetaData getMetaData() throws SQLException {
+      // simulate retrieving DB metadata from the DB itself
+      if (usePreparedStatement) {
+        prepareStatement("SELECT * from DB_METADATA").executeQuery()
+      } else {
+        createStatement().executeQuery("SELECT * from DB_METADATA")
+      }
+      return super.getMetaData()
+    }
+  }
 }

instrumentation/jdbc/javaagent/src/test/groovy/test/TestConnection.groovy

@@ -41,7 +41,7 @@ class TestConnection implements Connection {
 
   @Override
   PreparedStatement prepareStatement(String sql) throws SQLException {
-    return null
+    return new TestPreparedStatement(this)
   }
 
   @Override