GH-8582: Add TX support for PostgresSubChannel

Fixes https://github.com/spring-projects/spring-integration/issues/8582

* Introduce a `PostgresSubscribableChannel.setTransactionManager()`
to wrap a message polling and dispatching operation into a transaction
* In addition add a `RetryTemplate` support around transaction attempts

**Cherry-pick to `6.0.x`**

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2022 the original author or authors.
+ * Copyright 2022-2023 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -59,6 +59,7 @@ import org.springframework.util.Assert;
  *
  * @author Rafael Winterhalter
  * @author Artem Bilan
+ * @author Igor Lovich
  *
  * @since 6.0
  */
@@ -149,6 +150,8 @@ public final class PostgresChannelMessageTableSubscriber implements SmartLifecyc
 			this.executor = executorToUse;
 		}
 		this.latch = new CountDownLatch(1);
+
+		CountDownLatch startingLatch = new CountDownLatch(1);
 		this.future = executorToUse.submit(() -> {
 			try {
 				while (isActive()) {
@@ -171,6 +174,8 @@ public final class PostgresChannelMessageTableSubscriber implements SmartLifecyc
 						try {
 							this.connection = conn;
 							while (isActive()) {
+								startingLatch.countDown();
+
 								PGNotification[] notifications = conn.getNotifications(0);
 								// Unfortunately, there is no good way of interrupting a notification
 								// poll but by closing its connection.
@@ -208,6 +213,16 @@ public final class PostgresChannelMessageTableSubscriber implements SmartLifecyc
 				this.latch.countDown();
 			}
 		});
+
+		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 boolean isActive() {
@@ -234,8 +249,7 @@ public final class PostgresChannelMessageTableSubscriber implements SmartLifecyc
 		}
 		try {
 			if (!this.latch.await(5, TimeUnit.SECONDS)) {
-				throw new IllegalStateException("Failed to stop "
-						+ PostgresChannelMessageTableSubscriber.class.getName());
+				throw new IllegalStateException("Failed to stop " + this);
 			}
 		}
 		catch (InterruptedException ignored) {

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2022 the original author or authors.
+ * Copyright 2022-2023 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,8 +16,10 @@
 
 package org.springframework.integration.jdbc.channel;
 
+import java.util.Optional;
 import java.util.concurrent.Executor;
 
+import org.springframework.core.log.LogAccessor;
 import org.springframework.core.task.SimpleAsyncTaskExecutor;
 import org.springframework.integration.channel.AbstractSubscribableChannel;
 import org.springframework.integration.dispatcher.MessageDispatcher;
@@ -25,6 +27,9 @@ import org.springframework.integration.dispatcher.UnicastingDispatcher;
 import org.springframework.integration.jdbc.store.JdbcChannelMessageStore;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageHandler;
+import org.springframework.retry.support.RetryTemplate;
+import org.springframework.transaction.PlatformTransactionManager;
+import org.springframework.transaction.support.TransactionTemplate;
 import org.springframework.util.Assert;
 
 /**
@@ -39,19 +44,28 @@ import org.springframework.util.Assert;
  *
  * @author Rafael Winterhalter
  * @author Artem Bilan
+ * @author Igor Lovich
  *
  * @since 6.0
  */
 public class PostgresSubscribableChannel extends AbstractSubscribableChannel
 		implements PostgresChannelMessageTableSubscriber.Subscription {
 
+	private static final LogAccessor LOGGER = new LogAccessor(PostgresSubscribableChannel.class);
+
 	private final JdbcChannelMessageStore jdbcChannelMessageStore;
 
 	private final Object groupId;
 
 	private final PostgresChannelMessageTableSubscriber messageTableSubscriber;
 
-	private UnicastingDispatcher dispatcher = new UnicastingDispatcher(new SimpleAsyncTaskExecutor());
+	private final UnicastingDispatcher dispatcher = new UnicastingDispatcher();
+
+	private TransactionTemplate transactionTemplate;
+
+	private RetryTemplate retryTemplate = RetryTemplate.builder().maxAttempts(1).build();
+
+	private Executor executor = new SimpleAsyncTaskExecutor();
 
 	/**
 	 * Create a subscribable channel for a Postgres database.
@@ -75,7 +89,30 @@ public class PostgresSubscribableChannel extends AbstractSubscribableChannel
 	 */
 	public void setDispatcherExecutor(Executor executor) {
 		Assert.notNull(executor, "An executor must be provided.");
-		this.dispatcher = new UnicastingDispatcher(executor);
+		this.executor = executor;
+	}
+
+	/**
+	 * Set the transaction manager to use for message processing. Each message will be processed in a
+	 * separate transaction
+	 * @param transactionManager The transaction manager to use
+	 * @since 6.0.5
+	 * @see PlatformTransactionManager
+	 */
+	public void setTransactionManager(PlatformTransactionManager transactionManager) {
+		Assert.notNull(transactionManager, "A platform transaction manager must be provided.");
+		this.transactionTemplate = new TransactionTemplate(transactionManager);
+	}
+
+	/**
+	 * Set the retry template to use for retries in case of exception in downstream processing
+	 * @param retryTemplate The retry template to use
+	 * @since 6.0.5
+	 * @see RetryTemplate
+	 */
+	public void setRetryTemplate(RetryTemplate retryTemplate) {
+		Assert.notNull(retryTemplate, "A retry template must be provided.");
+		this.retryTemplate = retryTemplate;
 	}
 
 	@Override
@@ -110,10 +147,37 @@ public class PostgresSubscribableChannel extends AbstractSubscribableChannel
 
 	@Override
 	public void notifyUpdate() {
-		Message<?> message;
-		while ((message = this.jdbcChannelMessageStore.pollMessageFromGroup(this.groupId)) != null) {
-			this.dispatcher.dispatch(message);
-		}
+		this.executor.execute(() -> {
+			try {
+				Optional<Message<?>> dispatchedMessage;
+				do {
+					if (this.transactionTemplate != null) {
+						dispatchedMessage =
+								this.retryTemplate.execute(context ->
+										this.transactionTemplate.execute(status ->
+												pollMessage()
+														.map(this::dispatch)));
+					}
+					else {
+						dispatchedMessage =
+								pollMessage()
+										.map(message -> this.retryTemplate.execute(context -> dispatch(message)));
+					}
+				} while (dispatchedMessage.isPresent());
+			}
+			catch (Exception ex) {
+				LOGGER.error(ex, "Exception during message dispatch");
+			}
+		});
+	}
+
+	private Optional<Message<?>> pollMessage() {
+		return Optional.ofNullable(this.jdbcChannelMessageStore.pollMessageFromGroup(this.groupId));
+	}
+
+	private Message<?> dispatch(Message<?> message) {
+		this.dispatcher.dispatch(message);
+		return message;
 	}
 
 	@Override

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2022 the original author or authors.
+ * Copyright 2022-2023 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -21,12 +21,17 @@ import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import javax.sql.DataSource;
 
 import org.apache.commons.dbcp2.BasicDataSource;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestInfo;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.postgresql.jdbc.PgConnection;
 
 import org.springframework.beans.factory.annotation.Autowired;
@@ -36,18 +41,22 @@ import org.springframework.core.io.ByteArrayResource;
 import org.springframework.integration.config.EnableIntegration;
 import org.springframework.integration.jdbc.store.JdbcChannelMessageStore;
 import org.springframework.integration.jdbc.store.channel.PostgresChannelMessageStoreQueryProvider;
+import org.springframework.jdbc.datasource.DataSourceTransactionManager;
 import org.springframework.jdbc.datasource.init.DataSourceInitializer;
 import org.springframework.jdbc.datasource.init.ResourceDatabasePopulator;
 import org.springframework.jdbc.datasource.init.ScriptUtils;
 import org.springframework.messaging.support.GenericMessage;
+import org.springframework.retry.support.RetryTemplate;
 import org.springframework.test.annotation.DirtiesContext;
 import org.springframework.test.context.junit.jupiter.SpringJUnitConfig;
+import org.springframework.transaction.PlatformTransactionManager;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
 /**
  * @author Rafael Winterhalter
  * @author Artem Bilan
+ * @author Igor Lovich
  *
  * @since 6.0
  */
@@ -92,41 +101,49 @@ public class PostgresChannelMessageTableSubscriberTests implements PostgresConta
 	@Autowired
 	private JdbcChannelMessageStore messageStore;
 
+	@Autowired
+	private PlatformTransactionManager transactionManager;
+
 	private PostgresChannelMessageTableSubscriber postgresChannelMessageTableSubscriber;
 
+	private PostgresSubscribableChannel postgresSubscribableChannel;
+
+	private String groupId;
+
 	@BeforeEach
-	void setUp() {
+	void setUp(TestInfo testInfo) {
 		// Not initiated as a bean to allow for registrations prior and post the life cycle
-		this.postgresChannelMessageTableSubscriber = new PostgresChannelMessageTableSubscriber(
-				() -> DriverManager.getConnection(POSTGRES_CONTAINER.getJdbcUrl(),
-								POSTGRES_CONTAINER.getUsername(),
-								POSTGRES_CONTAINER.getPassword())
-						.unwrap(PgConnection.class)
-		);
+		this.postgresChannelMessageTableSubscriber =
+				new PostgresChannelMessageTableSubscriber(() ->
+						DriverManager.getConnection(POSTGRES_CONTAINER.getJdbcUrl(),
+										POSTGRES_CONTAINER.getUsername(),
+										POSTGRES_CONTAINER.getPassword())
+								.unwrap(PgConnection.class));
+
+		this.groupId = testInfo.getDisplayName();
+
+		this.postgresSubscribableChannel =
+				new PostgresSubscribableChannel(messageStore, groupId, postgresChannelMessageTableSubscriber);
 	}
 
+	@AfterEach
+	void tearDown() {
+		this.postgresChannelMessageTableSubscriber.stop();
+	}
+
+
 	@Test
 	public void testMessagePollMessagesAddedAfterStart() throws Exception {
 		CountDownLatch latch = new CountDownLatch(2);
 		List<Object> payloads = new ArrayList<>();
 		postgresChannelMessageTableSubscriber.start();
-		try {
-			PostgresSubscribableChannel channel = new PostgresSubscribableChannel(messageStore,
-					"testMessagePollMessagesAddedAfterStart",
-					postgresChannelMessageTableSubscriber);
-			channel.subscribe(message -> {
-				payloads.add(message.getPayload());
-				latch.countDown();
-			});
-			messageStore.addMessageToGroup("testMessagePollMessagesAddedAfterStart", new GenericMessage<>("1"));
-			messageStore.addMessageToGroup("testMessagePollMessagesAddedAfterStart", new GenericMessage<>("2"));
-			assertThat(latch.await(3, TimeUnit.SECONDS))
-					.as("Expected Postgres notification within 3 seconds")
-					.isTrue();
-		}
-		finally {
-			postgresChannelMessageTableSubscriber.stop();
-		}
+		postgresSubscribableChannel.subscribe(message -> {
+			payloads.add(message.getPayload());
+			latch.countDown();
+		});
+		messageStore.addMessageToGroup(groupId, new GenericMessage<>("1"));
+		messageStore.addMessageToGroup(groupId, new GenericMessage<>("2"));
+		assertThat(latch.await(10, TimeUnit.SECONDS)).isTrue();
 		assertThat(payloads).containsExactly("1", "2");
 	}
 
@@ -134,28 +151,80 @@ public class PostgresChannelMessageTableSubscriberTests implements PostgresConta
 	public void testMessagePollMessagesAddedBeforeStart() throws InterruptedException {
 		CountDownLatch latch = new CountDownLatch(2);
 		List<Object> payloads = new ArrayList<>();
-		PostgresSubscribableChannel channel =
-				new PostgresSubscribableChannel(messageStore,
-						"testMessagePollMessagesAddedBeforeStart",
-						postgresChannelMessageTableSubscriber);
-		channel.subscribe(message -> {
+
+		postgresSubscribableChannel.subscribe(message -> {
 			payloads.add(message.getPayload());
 			latch.countDown();
 		});
-		messageStore.addMessageToGroup("testMessagePollMessagesAddedBeforeStart", new GenericMessage<>("1"));
-		messageStore.addMessageToGroup("testMessagePollMessagesAddedBeforeStart", new GenericMessage<>("2"));
+		messageStore.addMessageToGroup(groupId, new GenericMessage<>("1"));
+		messageStore.addMessageToGroup(groupId, new GenericMessage<>("2"));
 		postgresChannelMessageTableSubscriber.start();
-		try {
-			assertThat(latch.await(3, TimeUnit.SECONDS))
-					.as("Expected Postgres notification within 3 seconds")
-					.isTrue();
-		}
-		finally {
-			postgresChannelMessageTableSubscriber.stop();
-		}
+		assertThat(latch.await(10, TimeUnit.SECONDS)).isTrue();
 		assertThat(payloads).containsExactly("1", "2");
 	}
 
+	@Test
+	void testMessagesDispatchedInTransaction() throws InterruptedException {
+		CountDownLatch latch = new CountDownLatch(2);
+		postgresSubscribableChannel.setTransactionManager(transactionManager);
+
+		postgresChannelMessageTableSubscriber.start();
+		postgresSubscribableChannel.subscribe(message -> {
+			try {
+				throw new RuntimeException("An error has occurred");
+			}
+			finally {
+				latch.countDown();
+			}
+		});
+
+		messageStore.addMessageToGroup(groupId, new GenericMessage<>("1"));
+		messageStore.addMessageToGroup(groupId, new GenericMessage<>("2"));
+
+		assertThat(latch.await(10, TimeUnit.SECONDS)).isTrue();
+
+		assertThat(messageStore.messageGroupSize(groupId)).isEqualTo(2);
+		assertThat(messageStore.pollMessageFromGroup(groupId).getPayload()).isEqualTo("1");
+		assertThat(messageStore.pollMessageFromGroup(groupId).getPayload()).isEqualTo("2");
+	}
+
+	@ParameterizedTest
+	@ValueSource(booleans = {true, false})
+	void testRetryOnErrorDuringDispatch(boolean transactionsEnabled) throws InterruptedException {
+		CountDownLatch latch = new CountDownLatch(2);
+		List<Object> payloads = new ArrayList<>();
+		AtomicInteger actualTries = new AtomicInteger();
+
+		int maxAttempts = 2;
+		postgresSubscribableChannel.setRetryTemplate(RetryTemplate.builder().maxAttempts(maxAttempts).build());
+
+		if (transactionsEnabled) {
+			postgresSubscribableChannel.setTransactionManager(transactionManager);
+		}
+
+		postgresChannelMessageTableSubscriber.start();
+
+		postgresSubscribableChannel.subscribe(message -> {
+			try {
+				//fail once
+				if (actualTries.getAndIncrement() == 0) {
+					throw new RuntimeException("An error has occurred");
+				}
+				payloads.add(message.getPayload());
+			}
+			finally {
+				latch.countDown();
+			}
+		});
+
+		messageStore.addMessageToGroup(groupId, new GenericMessage<>("1"));
+
+		assertThat(latch.await(10, TimeUnit.SECONDS)).isTrue();
+
+		assertThat(actualTries.get()).isEqualTo(maxAttempts);
+		assertThat(payloads).containsExactly("1");
+	}
+
 	@Configuration
 	@EnableIntegration
 	public static class Config {
@@ -181,6 +250,11 @@ public class PostgresChannelMessageTableSubscriberTests implements PostgresConta
 			return dataSourceInitializer;
 		}
 
+		@Bean
+		PlatformTransactionManager transactionManager(DataSource dataSource) {
+			return new DataSourceTransactionManager(dataSource);
+		}
+
 		@Bean
 		public JdbcChannelMessageStore jdbcChannelMessageStore(DataSource dataSource) {
 			JdbcChannelMessageStore messageStore = new JdbcChannelMessageStore(dataSource);

src/reference/asciidoc/jdbc.adoc

@@ -621,6 +621,18 @@ public PostgresSubscribableChannel channel(
   return new PostgresSubscribableChannel(messageStore, "some group", subscriber);
 }
 ----
+
+*Transaction support*
+
+Starting with version 6.0.5, specifying a `PlatformTransactionManager` on a `PostgresSubscribableChannel` will notify subscribers in a transaction.
+An exception in a subscriber will cause the transaction to be rolled back and the message to be put back in the message store.
+Transactional support is not activated by default.
+
+*Retries*
+
+Starting with version 6.0.5, a retry policy can be specified by providing a `RetryTemplate` to the `PostgresSubscribableChannel`.
+By default, no retries are performed.
+
 ====
 
 [IMPORTANT]
@@ -632,6 +644,7 @@ Such connection pools do normally expect that issued connections are closed with
 For this need of an exclusive connection, it is also recommended that a JVM only runs a single `PostgresChannelMessageTableSubscriber` which can be used to register any number of subscriptions.
 ====
 
+
 [[stored-procedures]]
 === Stored Procedures