diff --git a/datasource-micrometer/src/main/java/net/ttddyy/observation/tracing/DataSourceObservationListener.java b/datasource-micrometer/src/main/java/net/ttddyy/observation/tracing/DataSourceObservationListener.java index 9865152..ef1ee19 100644 --- a/datasource-micrometer/src/main/java/net/ttddyy/observation/tracing/DataSourceObservationListener.java +++ b/datasource-micrometer/src/main/java/net/ttddyy/observation/tracing/DataSourceObservationListener.java @@ -269,42 +269,43 @@ private void handleGetConnectionBefore(MethodExecutionContext executionContext) private void handleGetConnectionAfter(MethodExecutionContext executionContext) { DataSource dataSource = (DataSource) executionContext.getTarget(); Connection connection = (Connection) executionContext.getResult(); - URI connectionUrl = getConnectionUrl(connection); - Observation.Scope scopeToUse = executionContext.getCustomValue(Observation.Scope.class.getName(), Observation.Scope.class); - ConnectionInfo connectionInfo = executionContext.getConnectionInfo(); - - ConnectionAttributes connectionAttributes = new ConnectionAttributes(); - connectionAttributes.connectionInfo = connectionInfo; - connectionAttributes.scope = scopeToUse; - connectionAttributes.dataSource = dataSource; - if (connectionUrl != null) { - connectionAttributes.host = connectionUrl.getHost(); - connectionAttributes.port = connectionUrl.getPort(); - } - - String connectionId = connectionInfo.getConnectionId(); - this.connectionAttributesManager.put(connectionId, connectionAttributes); - - ConnectionContext connectionContext = executionContext.getCustomValue(ConnectionContext.class.getName(), - ConnectionContext.class); - populateFromConnectionAttributes(connectionContext, connectionId); - Throwable throwable = executionContext.getThrown(); - if (throwable != null && scopeToUse != null) { + if (connection == null) { // Handle closing the observation due to an error from getConnection(). // For normal case, observation is stopped when connection is closed. // see "handleConnectionClose()". - try (Observation.Scope scope = scopeToUse) { - this.connectionAttributesManager.remove(connectionId); - stopObservation(scope.getCurrentObservation(), throwable); + if (scopeToUse != null) { + Throwable throwable = executionContext.getThrown(); + try (Observation.Scope scope = scopeToUse) { + stopObservation(scope.getCurrentObservation(), throwable); + } } } + else { + ConnectionInfo connectionInfo = executionContext.getConnectionInfo(); + ConnectionAttributes connectionAttributes = new ConnectionAttributes(); + connectionAttributes.connectionInfo = connectionInfo; + connectionAttributes.scope = scopeToUse; + connectionAttributes.dataSource = dataSource; + + URI connectionUrl = getConnectionUrl(connection); + if (connectionUrl != null) { + connectionAttributes.host = connectionUrl.getHost(); + connectionAttributes.port = connectionUrl.getPort(); + } - // When "getConnection" was successful, a connection is acquired. - if (scopeToUse != null) { - scopeToUse.getCurrentObservation().event(JdbcEvents.CONNECTION_ACQUIRED); + String connectionId = connectionInfo.getConnectionId(); + this.connectionAttributesManager.put(connectionId, connectionAttributes); + ConnectionContext connectionContext = executionContext.getCustomValue(ConnectionContext.class.getName(), + ConnectionContext.class); + populateFromConnectionAttributes(connectionContext, connectionId); + + // When "getConnection" was successful, a connection is acquired. + if (scopeToUse != null) { + scopeToUse.getCurrentObservation().event(JdbcEvents.CONNECTION_ACQUIRED); + } } } diff --git a/datasource-micrometer/src/test/java/net/ttddyy/observation/tracing/DataSourceObservationListenerTests.java b/datasource-micrometer/src/test/java/net/ttddyy/observation/tracing/DataSourceObservationListenerTests.java index 4cff4b3..3cc8295 100644 --- a/datasource-micrometer/src/test/java/net/ttddyy/observation/tracing/DataSourceObservationListenerTests.java +++ b/datasource-micrometer/src/test/java/net/ttddyy/observation/tracing/DataSourceObservationListenerTests.java @@ -1,5 +1,5 @@ /* - * Copyright 2022-2023 the original author or authors. + * Copyright 2022-2024 the original author or authors. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -328,6 +328,38 @@ void connection() throws Exception { .hasIpEqualTo("localhost").hasPortEqualTo(5555).hasEventWithNameEqualTo("acquired"); } + @Test + void getConnectionFailure() throws Exception { + this.registry.observationConfig().observationHandler(new ConnectionTracingObservationHandler(this.tracer)); + Method getConnection = DataSource.class.getMethod("getConnection"); + RuntimeException exception = new RuntimeException(); + + // when getConnection failed, result is null and thrown has exception + MethodExecutionContext executionContext = new MethodExecutionContext(); + executionContext.setMethod(getConnection); + executionContext.setTarget(mock(DataSource.class)); + executionContext.setThrown(exception); + + DataSourceObservationListener listener = new DataSourceObservationListener(this.registry); + listener.beforeMethod(executionContext); + assertThat(tracer.currentSpan()).isNotNull(); + listener.afterMethod(executionContext); + // when getConnection failed, it closes the connection span. + assertThat(tracer.currentSpan()).isNull(); + assertThat(tracer.getSpans()).hasSize(1); + + // @formatter:off + assertThat(this.tracer).onlySpan() + .hasNameEqualTo("connection") + .doesNotHaveRemoteServiceNameEqualTo("myDS") + .hasIpThatIsBlank() + .doesNotHavePortEqualTo(5555) + .doesNotHaveEventWithNameEqualTo("acquired") + .assertThatThrowable() + .isSameAs(exception); + // @formatter:on + } + @Test void commitAndRollback() throws Exception { this.registry.observationConfig().observationHandler(new ConnectionTracingObservationHandler(this.tracer));