From 012a03ff771e6736802607151dd0ff9740b2c06f Mon Sep 17 00:00:00 2001
From: Artem Bilan <abilan@vmware.com>
Date: Tue, 17 Oct 2023 13:25:05 -0400
Subject: [PATCH] Fix PostgresSubscribableChannel race condition

The `PostgresSubscribableChannel` uses a task executor
for dispatching messages.
Even if we stop `PostgresChannelMessageTableSubscriber` and
unsubscribe from the channel, the task might be ongoing.

* Use explicit `ThreadPoolTaskExecutor` in the test
to shout it down and wait for tasks to be completed before
verifying DB status
* Optimize `PostgresSubscribableChannel` to mark TX for rollback
when we got a message from DB, but no handlers subscribed

# Conflicts:
#	spring-integration-jdbc/src/test/java/org/springframework/integration/jdbc/channel/PostgresChannelMessageTableSubscriberTests.java
---
 .../channel/PostgresSubscribableChannel.java  | 11 ++++++-
 ...resChannelMessageTableSubscriberTests.java | 29 ++++++++++++++-----
 2 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/spring-integration-jdbc/src/main/java/org/springframework/integration/jdbc/channel/PostgresSubscribableChannel.java b/spring-integration-jdbc/src/main/java/org/springframework/integration/jdbc/channel/PostgresSubscribableChannel.java
index bcb9004b20..4d8d079b67 100644
--- a/spring-integration-jdbc/src/main/java/org/springframework/integration/jdbc/channel/PostgresSubscribableChannel.java
+++ b/spring-integration-jdbc/src/main/java/org/springframework/integration/jdbc/channel/PostgresSubscribableChannel.java
@@ -168,7 +168,16 @@ public class PostgresSubscribableChannel extends AbstractSubscribableChannel
 		if (this.hasHandlers) {
 			if (this.transactionTemplate != null) {
 				return this.retryTemplate.execute(context ->
-						this.transactionTemplate.execute(status -> pollMessage().map(this::dispatch)));
+						this.transactionTemplate.execute(status ->
+								pollMessage()
+										.filter(message -> {
+											if (!this.hasHandlers) {
+												status.setRollbackOnly();
+												return false;
+											}
+											return true;
+										})
+										.map(this::dispatch)));
 			}
 			else {
 				return pollMessage()
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 20f2b21c04..005d4330e8 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
@@ -48,6 +48,7 @@ import org.springframework.jdbc.datasource.init.ScriptUtils;
 import org.springframework.messaging.MessageHandler;
 import org.springframework.messaging.support.GenericMessage;
 import org.springframework.retry.support.RetryTemplate;
+import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
 import org.springframework.test.annotation.DirtiesContext;
 import org.springframework.test.context.junit.jupiter.SpringJUnitConfig;
 import org.springframework.transaction.PlatformTransactionManager;
@@ -109,6 +110,8 @@ public class PostgresChannelMessageTableSubscriberTests implements PostgresConta
 
 	private PostgresSubscribableChannel postgresSubscribableChannel;
 
+	private ThreadPoolTaskExecutor taskExecutor;
+
 	private String groupId;
 
 	@BeforeEach
@@ -121,15 +124,26 @@ public class PostgresChannelMessageTableSubscriberTests implements PostgresConta
 										POSTGRES_CONTAINER.getPassword())
 								.unwrap(PgConnection.class));
 
+
+		this.taskExecutor = new ThreadPoolTaskExecutor();
+		this.taskExecutor.setCorePoolSize(10);
+		this.taskExecutor.setWaitForTasksToCompleteOnShutdown(true);
+		this.taskExecutor.setAwaitTerminationSeconds(10);
+		this.taskExecutor.afterPropertiesSet();
+
 		this.groupId = testInfo.getDisplayName();
 
 		this.postgresSubscribableChannel =
 				new PostgresSubscribableChannel(messageStore, groupId, postgresChannelMessageTableSubscriber);
+		this.postgresSubscribableChannel.setBeanName("testPostgresChannel");
+		this.postgresSubscribableChannel.setDispatcherExecutor(this.taskExecutor);
+		this.postgresSubscribableChannel.afterPropertiesSet();
 	}
 
 	@AfterEach
 	void tearDown() {
 		this.postgresChannelMessageTableSubscriber.stop();
+		this.taskExecutor.shutdown();
 	}
 
 
@@ -172,13 +186,13 @@ public class PostgresChannelMessageTableSubscriberTests implements PostgresConta
 		postgresChannelMessageTableSubscriber.start();
 		MessageHandler messageHandler =
 				message -> {
-			try {
-				throw new RuntimeException("An error has occurred");
-			}
-			finally {
-				latch.countDown();
-			}
-		};
+					try {
+						throw new RuntimeException("An error has occurred");
+					}
+					finally {
+						latch.countDown();
+					}
+				};
 		postgresSubscribableChannel.subscribe(messageHandler);
 
 		messageStore.addMessageToGroup(groupId, new GenericMessage<>("1"));
@@ -189,6 +203,7 @@ public class PostgresChannelMessageTableSubscriberTests implements PostgresConta
 		// Stop subscriber to unlock records from TX for the next verification
 		postgresChannelMessageTableSubscriber.stop();
 		postgresSubscribableChannel.unsubscribe(messageHandler);
+		this.taskExecutor.shutdown();
 
 		assertThat(messageStore.messageGroupSize(groupId)).isEqualTo(2);
 		assertThat(messageStore.pollMessageFromGroup(groupId).getPayload()).isEqualTo("1");