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(true);
+
 	private @Nullable TransactionTemplate transactionTemplate;
 
 	private RetryTemplate retryTemplate =
@@ -153,23 +157,35 @@ protected void onInit() {
 
 	@Override
 	public boolean subscribe(MessageHandler handler) {
-		boolean subscribed = super.subscribe(handler);
-		if (this.dispatcher.getHandlerCount() == 1) {
-			this.messageTableSubscriber.subscribe(this);
-			this.hasHandlers = true;
-			notifyUpdate();
+		this.lock.writeLock().lock();
+		try {
+			boolean subscribed = super.subscribe(handler);
+			if (this.dispatcher.getHandlerCount() == 1) {
+				this.messageTableSubscriber.subscribe(this);
+				this.hasHandlers = true;
+				notifyUpdate();
+			}
+			return subscribed;
+		}
+		finally {
+			this.lock.writeLock().unlock();
 		}
-		return subscribed;
 	}
 
 	@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
@@ -188,8 +204,14 @@ public void notifyUpdate() {
 		this.executor.execute(() -> {
 			Optional<?> dispatchedMessage;
 			do {
-				dispatchedMessage = pollAndDispatchMessage();
-			} while (dispatchedMessage.isPresent());
+				this.lock.readLock().lock();
+				try {
+					dispatchedMessage = pollAndDispatchMessage();
+				}
+				finally {
+					this.lock.readLock().unlock();
+				}
+			} while (this.hasHandlers && dispatchedMessage.isPresent());
 		});
 	}
 

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,36 @@ public void testRenewConnection() throws Exception {
 		assertThat(payloads).containsExactlyInAnyOrder("1", "2");
 	}
 
+	@Test
+	public void testUnsubscribeHandlerDuringDispatch() throws InterruptedException {
+		int numberOfMessages = 300;
+		CountDownLatch receivedMessageCount = new CountDownLatch(numberOfMessages);
+		AtomicReference<Throwable> receivedException = new AtomicReference<>();
+
+		postgresChannelMessageTableSubscriber.start();
+		postgresSubscribableChannel.setErrorHandler(receivedException::set);
+		MessageHandler messageHandler = message -> {
+			receivedMessageCount.countDown();
+		};
+		postgresSubscribableChannel.subscribe(messageHandler);
+
+		IntStream.range(0, numberOfMessages).forEach(i -> {
+			taskExecutor.execute(() -> {
+				messageStore.addMessageToGroup(groupId, new GenericMessage<>(String.valueOf(i)));
+			});
+		});
+
+		Thread.sleep(50);
+		taskExecutor.execute(postgresSubscribableChannel::notifyUpdate);
+		taskExecutor.execute(postgresSubscribableChannel::notifyUpdate);
+		taskExecutor.execute(postgresSubscribableChannel::notifyUpdate);
+		Thread.sleep(50);
+		postgresSubscribableChannel.unsubscribe(messageHandler);
+
+		assertThat(receivedMessageCount.await(2, TimeUnit.SECONDS)).isFalse();
+		assertThat(receivedException).hasNullValue();
+	}
+
 	@Configuration
 	@EnableIntegration
 	public static class Config {