GH-10362: Fix PostgresSubscribableChannel for race condition

Fixes: #10362

Unsubscribing message handlers while processing an in-flight message may lead to a `MessageDeliveryException`.
There is a possibility in the `PostgresSubscribableChannel` to have a polled message dispatched when no subscribers to the channel anymore.
The `doPollAndDispatchMessage()` can enter to the `pollMessage()` when it sees `hasHandlers`.
However, the last subscriber to the channel is unsubscribed in parallel.

* Fix `PostgresSubscribableChannel` using a `ReadWriteLock` to block `unsubscribe()` until currently polled message is processed.
Or opposite: do not let to `doPollAndDispatchMessage()` to enter `pollMessage()` logic until `unsubscribe()` is done.
This way the `hasHandlers` state is guarded by lock synchronization.

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

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

(cherry picked from commit 145e2a5)

Author: bertschneider

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

@@ -18,6 +18,7 @@
 
 import java.util.Optional;
 import java.util.concurrent.Executor;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
 
 import org.springframework.core.log.LogAccessor;
 import org.springframework.core.task.SimpleAsyncTaskExecutor;
@@ -47,6 +48,7 @@
  * @author Rafael Winterhalter
  * @author Artem Bilan
  * @author Igor Lovich
+ * @author Norbert Schneider
  *
  * @since 6.0
  */
@@ -65,6 +67,8 @@ public class PostgresSubscribableChannel extends AbstractSubscribableChannel
 
 	private final UnicastingDispatcher dispatcher = new UnicastingDispatcher();
 
+	private final ReentrantReadWriteLock hasHandlersLock = new ReentrantReadWriteLock();
+
 	private TransactionTemplate transactionTemplate;
 
 	private RetryTemplate retryTemplate = RetryTemplate.builder().maxAttempts(1).build();
@@ -156,12 +160,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.hasHandlersLock.writeLock().lock();
+		try {
+			boolean unsubscribed = super.unsubscribe(handle);
+			if (this.dispatcher.getHandlerCount() == 0) {
+				this.messageTableSubscriber.unsubscribe(this);
+				this.hasHandlers = false;
+			}
+			return unsubscribed;
+		}
+		finally {
+			this.hasHandlersLock.writeLock().unlock();
 		}
-		return unsubscribed;
 	}
 
 	@Override
@@ -201,25 +211,24 @@ private Optional<?> pollAndDispatchMessage() {
 	}
 
 	private Optional<?> doPollAndDispatchMessage() {
-		if (this.hasHandlers) {
-			if (this.transactionTemplate != null) {
-				return this.retryTemplate.execute(context ->
-						this.transactionTemplate.execute(status ->
-								pollMessage()
-										.filter(message -> {
-											if (!this.hasHandlers) {
-												status.setRollbackOnly();
-												return false;
-											}
-											return true;
-										})
-										.map(this::dispatch)));
-			}
-			else {
-				return pollMessage()
-						.map(message -> this.retryTemplate.execute(context -> dispatch(message)));
+		this.hasHandlersLock.readLock().lock();
+		try {
+			if (this.hasHandlers) {
+				TransactionTemplate transactionTemplateToUse = this.transactionTemplate;
+				if (transactionTemplateToUse != null) {
+					return this.retryTemplate.execute(context ->
+							transactionTemplateToUse.execute(status ->
+									pollMessage().map(this::dispatch)));
+				}
+				else {
+					return pollMessage()
+							.map(message -> this.retryTemplate.execute(context -> dispatch(message)));
+				}
 			}
 		}
+		finally {
+			this.hasHandlersLock.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;
 
@@ -70,6 +71,7 @@
  * @author Igor Lovich
  * @author Adama Sorho
  * @author Johannes Edmeier
+ * @author Norbert Schneider
  *
  * @since 6.0
  */
@@ -299,6 +301,49 @@ public void testRenewConnection() throws Exception {
 		assertThat(payloads).containsExactlyInAnyOrder("1", "2");
 	}
 
+	@Test
+	public void testUnsubscribeHandlerDuringDispatch() throws InterruptedException {
+		int numberOfMessages = 200;
+		CountDownLatch messagesUntilUnsubscribe = new CountDownLatch(10);
+		AtomicInteger receivedMessages = new AtomicInteger();
+		AtomicReference<Throwable> receivedException = new AtomicReference<>();
+
+		postgresChannelMessageTableSubscriber.start();
+		postgresSubscribableChannel.setErrorHandler(receivedException::set);
+
+		MessageHandler messageHandlerOne = message -> {
+			receivedMessages.getAndIncrement();
+			messagesUntilUnsubscribe.countDown();
+		};
+		postgresSubscribableChannel.subscribe(messageHandlerOne);
+
+		MessageHandler messageHandlerTwo = message -> {
+			receivedMessages.getAndIncrement();
+			messagesUntilUnsubscribe.countDown();
+		};
+		postgresSubscribableChannel.subscribe(messageHandlerTwo);
+
+		IntStream.range(0, numberOfMessages).forEach(i -> {
+			taskExecutor.execute(() ->
+				messageStore.addMessageToGroup(groupId, new GenericMessage<>(i)));
+		});
+
+		taskExecutor.execute(postgresSubscribableChannel::notifyUpdate);
+		taskExecutor.execute(postgresSubscribableChannel::notifyUpdate);
+		taskExecutor.execute(postgresSubscribableChannel::notifyUpdate);
+		taskExecutor.execute(postgresSubscribableChannel::notifyUpdate);
+		taskExecutor.execute(postgresSubscribableChannel::notifyUpdate);
+
+		assertThat(messagesUntilUnsubscribe.await(10, TimeUnit.SECONDS)).isTrue();
+		postgresSubscribableChannel.unsubscribe(messageHandlerOne);
+		postgresSubscribableChannel.unsubscribe(messageHandlerTwo);
+
+		int undeliveredMessages = messageStore.messageGroupSize(groupId);
+		int processedMessages = undeliveredMessages + receivedMessages.get();
+		assertThat(processedMessages).isEqualTo(numberOfMessages);
+		assertThat(receivedException).hasNullValue();
+	}
+
 	@Configuration
 	@EnableIntegration
 	public static class Config {