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 fb622ebf62..1ba20713c6 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 @@ -169,70 +169,84 @@ public final class PostgresChannelMessageTableSubscriber implements SmartLifecyc 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"); - } - catch (Exception ex) { - try { - conn.close(); - } - catch (Exception suppressed) { - ex.addSuppressed(suppressed); - } - throw ex; - } - this.subscriptionsMap.values() - .forEach(subscriptions -> subscriptions.forEach(Subscription::notifyUpdate)); - try { - this.connection = conn; - while (isActive()) { - startingLatch.countDown(); + this.future = executorToUse.submit(() -> doStart(startingLatch)); - 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 || 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 { + 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 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"); + catch (Exception suppressed) { + ex.addSuppressed(suppressed); } + 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; + } + 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; + } + 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 { + 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)) { 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 1dda46b363..551aa028a0 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 @@ -23,8 +23,10 @@ import java.util.ArrayList; 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 class PostgresChannelMessageTableSubscriberTests implements PostgresConta CountDownLatch latch = new CountDownLatch(2); List 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 class PostgresChannelMessageTableSubscriberTests implements PostgresConta private static class ConnectionSupplier implements PgConnectionSupplier { - Runnable onGetConnection; + Consumer onGetConnection; @Override public PgConnection get() throws SQLException { @@ -349,10 +362,11 @@ public class PostgresChannelMessageTableSubscriberTests implements PostgresConta POSTGRES_CONTAINER.getPassword()) .unwrap(PgConnection.class); if (this.onGetConnection != null) { - this.onGetConnection.run(); + this.onGetConnection.accept(conn); } return conn; } } + }