From c2a7160dd001de24f26521fa08f64a26f84fd53c Mon Sep 17 00:00:00 2001 From: Elena Felder <41136058+elefeint@users.noreply.github.com> Date: Tue, 25 Jan 2022 15:25:02 -0500 Subject: [PATCH] Close Spanner health check ResultSet to prevent session leak (#902) * close result set to prevent session leak * restrict visibility of test methods * use try-with-resources to close the resultset --- .../health/SpannerHealthIndicator.java | 8 ++-- .../health/SpannerHealthIndicatorTests.java | 43 ++++++++++++------- 2 files changed, 33 insertions(+), 18 deletions(-) diff --git a/spring-cloud-gcp-autoconfigure/src/main/java/com/google/cloud/spring/autoconfigure/spanner/health/SpannerHealthIndicator.java b/spring-cloud-gcp-autoconfigure/src/main/java/com/google/cloud/spring/autoconfigure/spanner/health/SpannerHealthIndicator.java index 0c572dfc1e..d70c56e0b3 100644 --- a/spring-cloud-gcp-autoconfigure/src/main/java/com/google/cloud/spring/autoconfigure/spanner/health/SpannerHealthIndicator.java +++ b/spring-cloud-gcp-autoconfigure/src/main/java/com/google/cloud/spring/autoconfigure/spanner/health/SpannerHealthIndicator.java @@ -52,9 +52,11 @@ public SpannerHealthIndicator(final SpannerTemplate spannerTemplate, String vali @Override protected void doHealthCheck(Builder builder) throws Exception { - ResultSet resultSet = spannerTemplate.executeQuery(validationStatement, null); - // Touch the record - resultSet.next(); + try (ResultSet resultSet = spannerTemplate.executeQuery(validationStatement, null)) { + // Touch the record + resultSet.next(); + } + builder.up(); } } diff --git a/spring-cloud-gcp-autoconfigure/src/test/java/com/google/cloud/spring/autoconfigure/spanner/health/SpannerHealthIndicatorTests.java b/spring-cloud-gcp-autoconfigure/src/test/java/com/google/cloud/spring/autoconfigure/spanner/health/SpannerHealthIndicatorTests.java index 82c6084120..99b5045e1b 100644 --- a/spring-cloud-gcp-autoconfigure/src/test/java/com/google/cloud/spring/autoconfigure/spanner/health/SpannerHealthIndicatorTests.java +++ b/spring-cloud-gcp-autoconfigure/src/test/java/com/google/cloud/spring/autoconfigure/spanner/health/SpannerHealthIndicatorTests.java @@ -17,6 +17,7 @@ package com.google.cloud.spring.autoconfigure.spanner.health; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; @@ -25,10 +26,10 @@ import com.google.cloud.spanner.ResultSet; import com.google.cloud.spanner.Statement; import com.google.cloud.spring.data.spanner.core.SpannerTemplate; -import org.junit.Test; -import org.junit.runner.RunWith; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; -import org.mockito.junit.MockitoJUnitRunner; +import org.mockito.junit.jupiter.MockitoExtension; import org.springframework.boot.actuate.health.Health; import org.springframework.boot.actuate.health.Status; @@ -37,8 +38,8 @@ * * @since 2.0.6 */ -@RunWith(MockitoJUnitRunner.class) -public class SpannerHealthIndicatorTests { +@ExtendWith(MockitoExtension.class) +class SpannerHealthIndicatorTests { @Mock private SpannerTemplate spannerTemplate; @@ -47,7 +48,7 @@ public class SpannerHealthIndicatorTests { private static final String QUERY = "SELECT 2"; @Test - public void testdoHealthCheckUp() throws Exception { + void testdoHealthCheckUp() throws Exception { SpannerHealthIndicator spannerHealthIndicator = new SpannerHealthIndicator(spannerTemplate, QUERY); @@ -61,10 +62,12 @@ public void testdoHealthCheckUp() throws Exception { assertThat(builder.build().getStatus()).isSameAs(Status.UP); verify(spannerTemplate).executeQuery(Statement.of(QUERY), null); verify(resultSet).next(); + // make sure Spanner ResultSet is closed to avoid session leak. + verify(resultSet).close(); } - @Test(expected = Exception.class) - public void testdoHealthCheckDownSpannerTemplate() throws Exception { + @Test + void testdoHealthCheckDownSpannerTemplate() throws Exception { SpannerHealthIndicator spannerHealthIndicator = new SpannerHealthIndicator(spannerTemplate, QUERY); @@ -73,11 +76,15 @@ public void testdoHealthCheckDownSpannerTemplate() throws Exception { Health.Builder builder = new Health.Builder(); - spannerHealthIndicator.doHealthCheck(builder); + assertThatThrownBy(() -> spannerHealthIndicator.doHealthCheck(builder)) + .isInstanceOf(RuntimeException.class) + .hasMessage("Cloud Spanner is down!!!"); + + verify(resultSet, never()).close(); } - @Test(expected = Exception.class) - public void testdoHealthCheckDownResultSet() throws Exception { + @Test + void testdoHealthCheckDownResultSet() throws Exception { SpannerHealthIndicator spannerHealthIndicator = new SpannerHealthIndicator(spannerTemplate, QUERY); @@ -86,11 +93,14 @@ public void testdoHealthCheckDownResultSet() throws Exception { Health.Builder builder = new Health.Builder(); - spannerHealthIndicator.doHealthCheck(builder); + assertThatThrownBy(() -> spannerHealthIndicator.doHealthCheck(builder)) + .isInstanceOf(RuntimeException.class) + .hasMessage("Cloud Spanner is down!!!"); + verify(resultSet).close(); } @Test - public void testHealthy() { + void testHealthy() { SpannerHealthIndicator spannerHealthIndicator = new SpannerHealthIndicator(spannerTemplate, QUERY); @@ -100,10 +110,11 @@ public void testHealthy() { assertThat(spannerHealthIndicator.health().getStatus()).isSameAs(Status.UP); verify(spannerTemplate).executeQuery(Statement.of(QUERY), null); verify(resultSet).next(); + verify(resultSet).close(); } @Test - public void testUnhealthySpannerTemplate() { + void testUnhealthySpannerTemplate() { SpannerHealthIndicator spannerHealthIndicator = new SpannerHealthIndicator(spannerTemplate, QUERY); @@ -113,10 +124,11 @@ public void testUnhealthySpannerTemplate() { assertThat(spannerHealthIndicator.health().getStatus()).isEqualTo(Status.DOWN); verify(spannerTemplate).executeQuery(Statement.of(QUERY), null); verify(resultSet, never()).next(); + verify(resultSet, never()).close(); } @Test - public void testUnhealthyResultSet() { + void testUnhealthyResultSet() { SpannerHealthIndicator spannerHealthIndicator = new SpannerHealthIndicator(spannerTemplate, QUERY); @@ -126,5 +138,6 @@ public void testUnhealthyResultSet() { assertThat(spannerHealthIndicator.health().getStatus()).isEqualTo(Status.DOWN); verify(spannerTemplate).executeQuery(Statement.of(QUERY), null); verify(resultSet).next(); + verify(resultSet).close(); } }