PostgresChannelMessageTableSubscriber: Renew connection only if invalid

Fixes: #9111

An evolution of the #9061: renew the connection only when we need to.

(cherry picked from commit da29e2d)

# Conflicts:
#	spring-integration-jdbc/src/main/java/org/springframework/integration/jdbc/channel/PostgresChannelMessageTableSubscriber.java

Author: joshiste

spring-integration-jdbc/src/main/java/org/springframework/integration/jdbc/channel/PostgresChannelMessageTableSubscriber.java

@@ -169,70 +169,84 @@ public synchronized void start() {
 		this.latch = new CountDownLatch(1);
 
 		CountDownLatch startingLatch = new CountDownLatch(1);
-		this.future = executorToUse.submit(() -> {
-			try {
-				while (isActive()) {
-					try {
-						PgConnection conn = this.connectionSupplier.get();
-						try (Statement stmt = conn.createStatement()) {
-							stmt.execute("LISTEN " + this.tablePrefix.toLowerCase() + "channel_message_notify");
+		this.future = executorToUse.submit(() -> doStart(startingLatch));
+
+		try {
+			if (!startingLatch.await(5, TimeUnit.SECONDS)) {
+				throw new IllegalStateException("Failed to start " + this);
+			}
+		}
+		catch (InterruptedException ex) {
+			Thread.currentThread().interrupt();
+			throw new IllegalStateException("Failed to start " + this, ex);
+		}
+	}
+
+	private void doStart(CountDownLatch startingLatch) {
+		try {
+			while (isActive()) {
+				try {
+					PgConnection conn = this.connectionSupplier.get();
+					try (Statement stmt = conn.createStatement()) {
+						stmt.execute("LISTEN " + this.tablePrefix.toLowerCase() + "channel_message_notify");
+					}
+					catch (Exception ex) {
+						try {
+							conn.close();
+						}
+						catch (Exception suppressed) {
+							ex.addSuppressed(suppressed);
 						}
-						catch (Exception ex) {
-							try {
-								conn.close();
+						throw ex;
+					}
+					this.subscriptionsMap.values()
+							.forEach(subscriptions -> subscriptions.forEach(Subscription::notifyUpdate));
+					try {
+						this.connection = conn;
+						while (isActive()) {
+							startingLatch.countDown();
+
+							PGNotification[] notifications = conn.getNotifications((int) this.notificationTimeout.toMillis());
+							// Unfortunately, there is no good way of interrupting a notification
+							// poll but by closing its connection.
+							if (!isActive()) {
+								return;
 							}
-							catch (Exception suppressed) {
-								ex.addSuppressed(suppressed);
+							if ((notifications == null || notifications.length == 0) && !conn.isValid(1)) {
+								//We did not receive any notifications within the timeout period.
+								//If the connection is still valid, we will continue polling
+								//Otherwise, we will close the connection and re-establish it.
+								break;
 							}
-							throw ex;
-						}
-						this.subscriptionsMap.values()
-								.forEach(subscriptions -> subscriptions.forEach(Subscription::notifyUpdate));
-						try {
-							this.connection = conn;
-							while (isActive()) {
-								startingLatch.countDown();
-
-								PGNotification[] notifications = conn.getNotifications((int) this.notificationTimeout.toMillis());
-								// Unfortunately, there is no good way of interrupting a notification
-								// poll but by closing its connection.
-								if (!isActive()) {
-									return;
+							for (PGNotification notification : notifications) {
+								String parameter = notification.getParameter();
+								Set<Subscription> subscriptions = this.subscriptionsMap.get(parameter);
+								if (subscriptions == null) {
+									continue;
 								}
-								if (notifications == null || notifications.length == 0) {
-									//We did not receive any notifications within the timeout period.
-									//We will close the connection and re-establish it.
-									break;
-								}
-								for (PGNotification notification : notifications) {
-									String parameter = notification.getParameter();
-									Set<Subscription> subscriptions = this.subscriptionsMap.get(parameter);
-									if (subscriptions == null) {
-										continue;
-									}
-									for (Subscription subscription : subscriptions) {
-										subscription.notifyUpdate();
-									}
+								for (Subscription subscription : subscriptions) {
+									subscription.notifyUpdate();
 								}
 							}
-						}
-						finally {
-							conn.close();
+
 						}
 					}
-					catch (Exception e) {
-						// The getNotifications method does not throw a meaningful message on interruption.
-						// Therefore, we do not log an error, unless it occurred while active.
-						if (isActive()) {
-							LOGGER.error(e, "Failed to poll notifications from Postgres database");
-						}
+					finally {
+						conn.close();
+					}
+				}
+				catch (Exception e) {
+					// The getNotifications method does not throw a meaningful message on interruption.
+					// Therefore, we do not log an error, unless it occurred while active.
+					if (isActive()) {
+						LOGGER.error(e, "Failed to poll notifications from Postgres database");
 					}
 				}
 			}
-			finally {
-				this.latch.countDown();
-			}
-		});
+		}
+		finally {
+			this.latch.countDown();
+		}
 
 		try {
 			if (!startingLatch.await(5, TimeUnit.SECONDS)) {

spring-integration-jdbc/src/test/java/org/springframework/integration/jdbc/channel/PostgresChannelMessageTableSubscriberTests.java

@@ -23,8 +23,10 @@
 import java.util.List;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
 
 import javax.sql.DataSource;
 
@@ -284,7 +286,18 @@ public void testRenewConnection() throws Exception {
 		CountDownLatch latch = new CountDownLatch(2);
 		List<Object> payloads = new ArrayList<>();
 		CountDownLatch connectionLatch = new CountDownLatch(2);
-		connectionSupplier.onGetConnection = connectionLatch::countDown;
+		AtomicBoolean connectionCloseState = new AtomicBoolean();
+		connectionSupplier.onGetConnection = conn -> {
+			connectionLatch.countDown();
+			if (connectionCloseState.compareAndSet(false, true)) {
+				try {
+					conn.close();
+				}
+				catch (Exception e) {
+					//nop
+				}
+			}
+		};
 		postgresChannelMessageTableSubscriber.start();
 		postgresSubscribableChannel.subscribe(message -> {
 			payloads.add(message.getPayload());
@@ -340,7 +353,7 @@ public JdbcChannelMessageStore jdbcChannelMessageStore(DataSource dataSource) {
 
 	private static class ConnectionSupplier implements PgConnectionSupplier {
 
-		Runnable onGetConnection;
+		Consumer<PgConnection> onGetConnection;
 
 		@Override
 		public PgConnection get() throws SQLException {
@@ -349,10 +362,11 @@ public PgConnection get() throws SQLException {
 							POSTGRES_CONTAINER.getPassword())
 					.unwrap(PgConnection.class);
 			if (this.onGetConnection != null) {
-				this.onGetConnection.run();
+				this.onGetConnection.accept(conn);
 			}
 			return conn;
 		}
 
 	}
+
 }