GH-9061: renew connection in PostgresChannelMessageTableSubscriber

Fixes: #9061

`PostgresChannelMessageTableSubscriber` never renews the connection.
This causes problems on DB failover.
With this change the connection is renewed when notifications are not received for a certain time.

(cherry picked from commit 642278d)

Author: joshiste

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

@@ -18,6 +18,7 @@
 
 import java.sql.SQLException;
 import java.sql.Statement;
+import java.time.Duration;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
@@ -64,6 +65,7 @@
  * @author Artem Bilan
  * @author Igor Lovich
  * @author Christian Tzolov
+ * @author Johannes Edmeier
  *
  * @since 6.0
  */
@@ -88,6 +90,8 @@ public final class PostgresChannelMessageTableSubscriber implements SmartLifecyc
 	@Nullable
 	private volatile PgConnection connection;
 
+	private Duration notificationTimeout = Duration.ofSeconds(60);
+
 	/**
 	 * Create a new subscriber using the {@link JdbcChannelMessageStore#DEFAULT_TABLE_PREFIX}.
 	 * @param connectionSupplier The connection supplier for the targeted Postgres database.
@@ -136,6 +140,19 @@ public void setTaskExecutor(AsyncTaskExecutor taskExecutor) {
 		this.taskExecutor = taskExecutor;
 	}
 
+	/**
+	 * Set the timeout for the notification polling.
+	 * If for the specified duration no notificiation are received the underlying connection is closed and re-established.
+	 * Setting a value of {@code Duration.ZERO} will disable the timeout and wait forever.
+	 * This might cause problems in DB failover scenarios.
+	 * @param notificationTimeout the timeout for the notification polling.
+	 * @since 6.1.8
+	 */
+	public void setNotificationTimeout(Duration notificationTimeout) {
+		Assert.notNull(notificationTimeout, "'notificationTimeout' must not be null.");
+		this.notificationTimeout = notificationTimeout;
+	}
+
 	/**
 	 * Add a new subscription to this subscriber.
 	 * @param subscription The subscription to register.
@@ -213,24 +230,28 @@ private void doStart(CountDownLatch startingLatch) {
 						while (isActive()) {
 							startingLatch.countDown();
 
-							PGNotification[] notifications = conn.getNotifications(0);
+							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;
 							}
-							if (notifications != null) {
-								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();
-									}
+							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();
 								}
 							}
+
 						}
 					}
 					finally {

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

@@ -17,6 +17,8 @@
 package org.springframework.integration.jdbc.channel;
 
 import java.sql.DriverManager;
+import java.sql.SQLException;
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.CountDownLatch;
@@ -62,6 +64,7 @@
  * @author Artem Bilan
  * @author Igor Lovich
  * @author Adama Sorho
+ * @author Johannes Edmeier
  *
  * @since 6.0
  */
@@ -102,15 +105,14 @@ CREATE FUNCTION INT_CHANNEL_MESSAGE_NOTIFY_FCT()
 
 	private String groupId;
 
+	private ConnectionSupplier connectionSupplier;
+
 	@BeforeEach
 	void setUp(TestInfo testInfo) {
 		// Not initiated as a bean to allow for registrations prior and post the life cycle
-		this.postgresChannelMessageTableSubscriber =
-				new PostgresChannelMessageTableSubscriber(() ->
-						DriverManager.getConnection(POSTGRES_CONTAINER.getJdbcUrl(),
-										POSTGRES_CONTAINER.getUsername(),
-										POSTGRES_CONTAINER.getPassword())
-								.unwrap(PgConnection.class));
+		this.connectionSupplier = new ConnectionSupplier();
+		this.postgresChannelMessageTableSubscriber = new PostgresChannelMessageTableSubscriber(connectionSupplier);
+		this.postgresChannelMessageTableSubscriber.setNotificationTimeout(Duration.ofSeconds(5));
 
 
 		this.taskExecutor = new ThreadPoolTaskExecutor();
@@ -263,6 +265,26 @@ void testRetryOnErrorDuringDispatch(boolean transactionsEnabled) throws Interrup
 		assertThat(payloads).containsExactly("1");
 	}
 
+	@Test
+	public void testRenewConnection() throws Exception {
+		CountDownLatch latch = new CountDownLatch(2);
+		List<Object> payloads = new ArrayList<>();
+		CountDownLatch connectionLatch = new CountDownLatch(2);
+		connectionSupplier.onGetConnection = connectionLatch::countDown;
+		postgresChannelMessageTableSubscriber.start();
+		postgresSubscribableChannel.subscribe(message -> {
+			payloads.add(message.getPayload());
+			latch.countDown();
+		});
+
+		assertThat(connectionLatch.await(10, TimeUnit.SECONDS)).isTrue();
+
+		messageStore.addMessageToGroup(groupId, new GenericMessage<>("1"));
+		messageStore.addMessageToGroup(groupId, new GenericMessage<>("2"));
+		assertThat(latch.await(10, TimeUnit.SECONDS)).isTrue();
+		assertThat(payloads).containsExactlyInAnyOrder("1", "2");
+	}
+
 	@Configuration
 	@EnableIntegration
 	public static class Config {
@@ -302,4 +324,21 @@ public JdbcChannelMessageStore jdbcChannelMessageStore(DataSource dataSource) {
 
 	}
 
+	private static class ConnectionSupplier implements PgConnectionSupplier {
+
+		Runnable onGetConnection;
+
+		@Override
+		public PgConnection get() throws SQLException {
+			var conn = DriverManager.getConnection(POSTGRES_CONTAINER.getJdbcUrl(),
+							POSTGRES_CONTAINER.getUsername(),
+							POSTGRES_CONTAINER.getPassword())
+					.unwrap(PgConnection.class);
+			if (this.onGetConnection != null) {
+				this.onGetConnection.run();
+			}
+			return conn;
+		}
+
+	}
 }