GH-10362: Block on PostgresSubscribableChannel.unsubscribe

Unsubscribing message handlers while processing an in-flight message may lead to
a MessageDeliveryException.

Fixes: GH-10362

Signed-off-by: Norbert Schneider <[email protected]>

Author: bertschneider

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

@@ -19,6 +19,7 @@
 import java.time.Duration;
 import java.util.Optional;
 import java.util.concurrent.Executor;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
 
 import org.jspecify.annotations.Nullable;
 
@@ -53,6 +54,7 @@
  * @author Rafael Winterhalter
  * @author Artem Bilan
  * @author Igor Lovich
+ * @author Norbert Schneider
  *
  * @since 6.0
  */
@@ -71,6 +73,8 @@ public class PostgresSubscribableChannel extends AbstractSubscribableChannel
 
 	private final UnicastingDispatcher dispatcher = new UnicastingDispatcher();
 
+	private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
+
 	private @Nullable TransactionTemplate transactionTemplate;
 
 	private RetryTemplate retryTemplate =
@@ -164,12 +168,18 @@ public boolean subscribe(MessageHandler handler) {
 
 	@Override
 	public boolean unsubscribe(MessageHandler handle) {
-		boolean unsubscribed = super.unsubscribe(handle);
-		if (this.dispatcher.getHandlerCount() == 0) {
-			this.messageTableSubscriber.unsubscribe(this);
-			this.hasHandlers = false;
+		this.lock.writeLock().lock();
+		try {
+			boolean unsubscribed = super.unsubscribe(handle);
+			if (this.dispatcher.getHandlerCount() == 0) {
+				this.messageTableSubscriber.unsubscribe(this);
+				this.hasHandlers = false;
+			}
+			return unsubscribed;
+		}
+		finally {
+			this.lock.writeLock().unlock();
 		}
-		return unsubscribed;
 	}
 
 	@Override
@@ -209,26 +219,24 @@ private Optional<?> pollAndDispatchMessage() {
 	}
 
 	private Optional<?> doPollAndDispatchMessage() {
-		if (this.hasHandlers) {
-			TransactionTemplate transactionTemplateToUse = this.transactionTemplate;
-			if (transactionTemplateToUse != null) {
-				return executeWithRetry(() ->
-						transactionTemplateToUse.execute(status ->
-								pollMessage()
-										.filter(message -> {
-											if (!this.hasHandlers) {
-												status.setRollbackOnly();
-												return false;
-											}
-											return true;
-										})
-										.map(this::dispatch)));
-			}
-			else {
-				return pollMessage()
-						.map(message -> executeWithRetry(() -> dispatch(message)));
+		this.lock.readLock().lock();
+		try {
+			if (this.hasHandlers) {
+				TransactionTemplate transactionTemplateToUse = this.transactionTemplate;
+				if (transactionTemplateToUse != null) {
+					return executeWithRetry(() ->
+							transactionTemplateToUse.execute(status ->
+									pollMessage().map(this::dispatch)));
+				}
+				else {
+					return pollMessage()
+							.map(message -> executeWithRetry(() -> dispatch(message)));
+				}
 			}
 		}
+		finally {
+			this.lock.readLock().unlock();
+		}
 		return Optional.empty();
 	}
 

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

@@ -27,6 +27,7 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
+import java.util.stream.IntStream;
 
 import javax.sql.DataSource;
 
@@ -72,7 +73,7 @@
  * @author Igor Lovich
  * @author Adama Sorho
  * @author Johannes Edmeier
- *
+ * @author Norbert Schneider
  * @since 6.0
  */
 @SpringJUnitConfig
@@ -310,6 +311,41 @@ public void testRenewConnection() throws Exception {
 		assertThat(payloads).containsExactlyInAnyOrder("1", "2");
 	}
 
+	@Test
+	public void testUnsubscribeHandlerDuringDispatch() throws InterruptedException {
+		int numberOfMessages = 150;
+		CountDownLatch messagesUntilUnsubscribe = new CountDownLatch(10);
+		AtomicInteger receivedMessages = new AtomicInteger();
+		AtomicReference<Throwable> receivedException = new AtomicReference<>();
+
+		postgresChannelMessageTableSubscriber.start();
+		postgresSubscribableChannel.setErrorHandler(receivedException::set);
+		MessageHandler messageHandler = message -> {
+			receivedMessages.getAndIncrement();
+			messagesUntilUnsubscribe.countDown();
+		};
+		postgresSubscribableChannel.subscribe(messageHandler);
+
+		IntStream.range(0, numberOfMessages).forEach(i -> {
+			taskExecutor.execute(() -> {
+				messageStore.addMessageToGroup(groupId, new GenericMessage<>(String.valueOf(i)));
+			});
+		});
+
+		taskExecutor.execute(postgresSubscribableChannel::notifyUpdate);
+		taskExecutor.execute(postgresSubscribableChannel::notifyUpdate);
+		taskExecutor.execute(postgresSubscribableChannel::notifyUpdate);
+		taskExecutor.execute(postgresSubscribableChannel::notifyUpdate);
+		taskExecutor.execute(postgresSubscribableChannel::notifyUpdate);
+
+		assertThat(messagesUntilUnsubscribe.await(30, TimeUnit.SECONDS)).isTrue();
+		postgresSubscribableChannel.unsubscribe(messageHandler);
+
+		int undeliveredMessages = messageStore.messageGroupSize(groupId);
+		assertThat(numberOfMessages).isEqualTo(undeliveredMessages + receivedMessages.get());
+		assertThat(receivedException).hasNullValue();
+	}
+
 	@Configuration
 	@EnableIntegration
 	public static class Config {