diff --git a/spring-integration-jdbc/src/main/java/org/springframework/integration/jdbc/channel/PostgresChannelMessageTableSubscriber.java b/spring-integration-jdbc/src/main/java/org/springframework/integration/jdbc/channel/PostgresChannelMessageTableSubscriber.java index d0a5cd8290..9075ebedc3 100644 --- a/spring-integration-jdbc/src/main/java/org/springframework/integration/jdbc/channel/PostgresChannelMessageTableSubscriber.java +++ b/spring-integration-jdbc/src/main/java/org/springframework/integration/jdbc/channel/PostgresChannelMessageTableSubscriber.java @@ -18,6 +18,7 @@ package org.springframework.integration.jdbc.channel; 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 @@ import org.springframework.util.Assert; * @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 final class PostgresChannelMessageTableSubscriber implements SmartLifecyc 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 @@ public final class PostgresChannelMessageTableSubscriber implements SmartLifecyc 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 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 subscriptions = this.subscriptionsMap.get(parameter); + if (subscriptions == null) { + continue; + } + for (Subscription subscription : subscriptions) { + subscription.notifyUpdate(); } } + } } finally { diff --git a/spring-integration-jdbc/src/test/java/org/springframework/integration/jdbc/channel/PostgresChannelMessageTableSubscriberTests.java b/spring-integration-jdbc/src/test/java/org/springframework/integration/jdbc/channel/PostgresChannelMessageTableSubscriberTests.java index bf88968390..7a74cd7202 100644 --- a/spring-integration-jdbc/src/test/java/org/springframework/integration/jdbc/channel/PostgresChannelMessageTableSubscriberTests.java +++ b/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 @@ import static org.assertj.core.api.Assertions.assertThat; * @author Artem Bilan * @author Igor Lovich * @author Adama Sorho + * @author Johannes Edmeier * * @since 6.0 */ @@ -102,15 +105,14 @@ public class PostgresChannelMessageTableSubscriberTests implements PostgresConta 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 @@ public class PostgresChannelMessageTableSubscriberTests implements PostgresConta assertThat(payloads).containsExactly("1"); } + @Test + public void testRenewConnection() throws Exception { + CountDownLatch latch = new CountDownLatch(2); + List 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 class PostgresChannelMessageTableSubscriberTests implements PostgresConta } + 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; + } + + } }