added check that connection might be closed at a time when statement …

…is executed (fixes #35)
gavlyukovskiy · Sep 6, 2019 · 713fe14 · 713fe14
1 parent ca6f2cf
commit 713fe14
Show file tree

Hide file tree

Showing 4 changed files with 192 additions and 1 deletion.
diff --git a/...onfigure/src/main/java/com/github/gavlyukovskiy/cloud/sleuth/TracingListenerStrategy.java b/...onfigure/src/main/java/com/github/gavlyukovskiy/cloud/sleuth/TracingListenerStrategy.java
@@ -53,7 +53,12 @@ void beforeQuery(CON connectionKey, STMT statementKey, String dataSourceName) {
             spanWithScope = new SpanWithScope(statementSpan, tracer.withSpanInScope(statementSpan));
         }
         StatementInfo statementInfo = new StatementInfo(spanWithScope);
-        openConnections.get(connectionKey).getNestedStatements().put(statementKey, statementInfo);
+        ConnectionInfo connectionInfo = openConnections.get(connectionKey);
+        if (connectionInfo == null) {
+            // Connection may be closed after statement preparation, but before statement execution.
+            return;
+        }
+        connectionInfo.getNestedStatements().put(statementKey, statementInfo);
     }
 
     void addQueryRowCount(CON connectionKey, STMT statementKey, int rowCount) {
@@ -66,6 +71,10 @@ void addQueryRowCount(CON connectionKey, STMT statementKey, int rowCount) {
 
     void afterQuery(CON connectionKey, STMT statementKey, String sql, Throwable t) {
         ConnectionInfo connectionInfo = openConnections.get(connectionKey);
+        if (connectionInfo == null) {
+            // Connection may be closed after statement preparation, but before statement execution.
+            return;
+        }
         StatementInfo statementInfo = connectionInfo.getNestedStatements().get(statementKey);
         statementInfo.getSpan().ifPresent(statementSpan -> {
             statementSpan.getSpan().tag(SleuthListenerAutoConfiguration.SPAN_SQL_QUERY_TAG_NAME, sql);

diff --git a/...test/java/com/github/gavlyukovskiy/boot/jdbc/decorator/p6spy/P6SpyConfigurationTests.java b/...test/java/com/github/gavlyukovskiy/boot/jdbc/decorator/p6spy/P6SpyConfigurationTests.java
@@ -38,6 +38,7 @@
 import javax.sql.DataSource;
 
 import java.sql.Connection;
+import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.util.concurrent.ThreadLocalRandom;
 
@@ -116,6 +117,31 @@ void testCanSetCustomLoggingFormat() {
         });
     }
 
+    @Test
+    void name() {
+        System.setProperty("p6spy.config.logMessageFormat", "com.p6spy.engine.spy.appender.CustomLineFormat");
+        System.setProperty("p6spy.config.excludecategories", "debug");
+        ApplicationContextRunner contextRunner = this.contextRunner.withPropertyValues("decorator.datasource.p6spy.multiline:true");
+
+        contextRunner.run(context -> {
+            JdbcEventListenerFactory jdbcEventListenerFactory = context.getBean(JdbcEventListenerFactory.class);
+            CompoundJdbcEventListener jdbcEventListener = (CompoundJdbcEventListener) jdbcEventListenerFactory.createJdbcEventListener();
+
+            DataSource dataSource = context.getBean(DataSource.class);
+
+            Connection connection = dataSource.getConnection();
+            ResultSet resultSet = connection.createStatement().executeQuery("SELECT 1, 2 UNION ALL SELECT 2, 3");
+            while (resultSet.next()) {
+                resultSet.getString(1);
+                resultSet.getString(2);
+            }
+            resultSet.close();
+            connection.close();
+
+            System.out.println();
+        });
+    }
+
     @Configuration
     static class CustomListenerConfiguration {
 

diff --git a/...re/src/test/java/com/github/gavlyukovskiy/cloud/sleuth/TracingJdbcEventListenerTests.java b/...re/src/test/java/com/github/gavlyukovskiy/cloud/sleuth/TracingJdbcEventListenerTests.java
@@ -31,6 +31,7 @@
 import javax.sql.DataSource;
 
 import java.sql.Connection;
+import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
@@ -509,4 +510,81 @@ void testShouldIncludeOnlyQueryAndFetchTraces() {
             assertThat(resultSetSpan.name()).isEqualTo("jdbc:/test/fetch");
         });
     }
+
+    @Test
+    void testShouldNotOverrideExceptionWhenConnectionWasClosedBeforeExecutingQuery() {
+        contextRunner.run(context -> {
+            DataSource dataSource = context.getBean(DataSource.class);
+            ArrayListSpanReporter spanReporter = context.getBean(ArrayListSpanReporter.class);
+
+            Connection connection = dataSource.getConnection();
+            PreparedStatement statement = connection.prepareStatement("SELECT NOW()");
+            connection.close();
+            try {
+                statement.executeQuery();
+                fail("should throw SQLException");
+            }
+            catch (SQLException expected) {
+            }
+
+            assertThat(spanReporter.getSpans()).hasSize(1);
+            Span connectionSpan = spanReporter.getSpans().get(0);
+            assertThat(connectionSpan.name()).isEqualTo("jdbc:/test/connection");
+        });
+    }
+
+    @Test
+    void testShouldNotOverrideExceptionWhenStatementWasClosedBeforeExecutingQuery() {
+        contextRunner.run(context -> {
+            DataSource dataSource = context.getBean(DataSource.class);
+            ArrayListSpanReporter spanReporter = context.getBean(ArrayListSpanReporter.class);
+
+            Connection connection = dataSource.getConnection();
+            PreparedStatement statement = connection.prepareStatement("SELECT NOW()");
+            statement.close();
+            try {
+                statement.executeQuery();
+                fail("should throw SQLException");
+            }
+            catch (SQLException expected) {
+            }
+            connection.close();
+
+            assertThat(spanReporter.getSpans()).hasSize(2);
+            Span connectionSpan = spanReporter.getSpans().get(1);
+            Span statementSpan = spanReporter.getSpans().get(0);
+            assertThat(connectionSpan.name()).isEqualTo("jdbc:/test/connection");
+            assertThat(statementSpan.name()).isEqualTo("jdbc:/test/query");
+        });
+    }
+
+    @Test
+    void testShouldNotOverrideExceptionWhenResultSetWasClosedBeforeNext() {
+        contextRunner.run(context -> {
+            DataSource dataSource = context.getBean(DataSource.class);
+            ArrayListSpanReporter spanReporter = context.getBean(ArrayListSpanReporter.class);
+
+            Connection connection = dataSource.getConnection();
+            Statement statement = connection.createStatement();
+            ResultSet resultSet = statement.executeQuery("SELECT NOW()");
+            resultSet.close();
+            try {
+                resultSet.next();
+                fail("should throw SQLException");
+            }
+            catch (SQLException expected) {
+            }
+            statement.close();
+            connection.close();
+
+            assertThat(spanReporter.getSpans()).hasSize(3);
+            Span connectionSpan = spanReporter.getSpans().get(2);
+            Span resultSetSpan = spanReporter.getSpans().get(1);
+            Span statementSpan = spanReporter.getSpans().get(0);
+            assertThat(connectionSpan.name()).isEqualTo("jdbc:/test/connection");
+            assertThat(statementSpan.name()).isEqualTo("jdbc:/test/query");
+            assertThat(resultSetSpan.name()).isEqualTo("jdbc:/test/fetch");
+            assertThat(statementSpan.tags()).containsEntry(SleuthListenerAutoConfiguration.SPAN_SQL_QUERY_TAG_NAME, "SELECT NOW()");
+        });
+    }
 }
diff --git a/...c/test/java/com/github/gavlyukovskiy/cloud/sleuth/TracingQueryExecutionListenerTests.java b/...c/test/java/com/github/gavlyukovskiy/cloud/sleuth/TracingQueryExecutionListenerTests.java
@@ -32,6 +32,7 @@
 import javax.sql.DataSource;
 
 import java.sql.Connection;
+import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
@@ -532,4 +533,81 @@ void testShouldIncludeOnlyQueryAndFetchTraces() {
             assertThat(resultSetSpan.name()).isEqualTo("jdbc:/test/fetch");
         });
     }
+
+    @Test
+    void testShouldNotOverrideExceptionWhenConnectionWasClosedBeforeExecutingQuery() {
+        contextRunner.run(context -> {
+            DataSource dataSource = context.getBean(DataSource.class);
+            ArrayListSpanReporter spanReporter = context.getBean(ArrayListSpanReporter.class);
+
+            Connection connection = dataSource.getConnection();
+            PreparedStatement statement = connection.prepareStatement("SELECT NOW()");
+            connection.close();
+            try {
+                statement.executeQuery();
+                fail("should throw SQLException");
+            }
+            catch (SQLException expected) {
+            }
+
+            assertThat(spanReporter.getSpans()).hasSize(1);
+            Span connectionSpan = spanReporter.getSpans().get(0);
+            assertThat(connectionSpan.name()).isEqualTo("jdbc:/test/connection");
+        });
+    }
+
+    @Test
+    void testShouldNotOverrideExceptionWhenStatementWasClosedBeforeExecutingQuery() {
+        contextRunner.run(context -> {
+            DataSource dataSource = context.getBean(DataSource.class);
+            ArrayListSpanReporter spanReporter = context.getBean(ArrayListSpanReporter.class);
+
+            Connection connection = dataSource.getConnection();
+            PreparedStatement statement = connection.prepareStatement("SELECT NOW()");
+            statement.close();
+            try {
+                statement.executeQuery();
+                fail("should throw SQLException");
+            }
+            catch (SQLException expected) {
+            }
+            connection.close();
+
+            assertThat(spanReporter.getSpans()).hasSize(2);
+            Span connectionSpan = spanReporter.getSpans().get(1);
+            Span statementSpan = spanReporter.getSpans().get(0);
+            assertThat(connectionSpan.name()).isEqualTo("jdbc:/test/connection");
+            assertThat(statementSpan.name()).isEqualTo("jdbc:/test/query");
+        });
+    }
+
+    @Test
+    void testShouldNotOverrideExceptionWhenResultSetWasClosedBeforeNext() {
+        contextRunner.run(context -> {
+            DataSource dataSource = context.getBean(DataSource.class);
+            ArrayListSpanReporter spanReporter = context.getBean(ArrayListSpanReporter.class);
+
+            Connection connection = dataSource.getConnection();
+            Statement statement = connection.createStatement();
+            ResultSet resultSet = statement.executeQuery("SELECT NOW()");
+            resultSet.close();
+            try {
+                resultSet.next();
+                fail("should throw SQLException");
+            }
+            catch (SQLException expected) {
+            }
+            statement.close();
+            connection.close();
+
+            assertThat(spanReporter.getSpans()).hasSize(3);
+            Span connectionSpan = spanReporter.getSpans().get(2);
+            Span resultSetSpan = spanReporter.getSpans().get(1);
+            Span statementSpan = spanReporter.getSpans().get(0);
+            assertThat(connectionSpan.name()).isEqualTo("jdbc:/test/connection");
+            assertThat(statementSpan.name()).isEqualTo("jdbc:/test/query");
+            assertThat(resultSetSpan.name()).isEqualTo("jdbc:/test/fetch");
+            assertThat(statementSpan.tags()).containsEntry(SleuthListenerAutoConfiguration.SPAN_SQL_QUERY_TAG_NAME, "SELECT NOW()");
+        });
+    }
 }