spring-projectsGH-8582:Transactional support in PostgresSubscribableChannel

Author: igorlovich

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 @@
  *
  * @author Rafael Winterhalter
  * @author Artem Bilan
+ * @author Igor Lovich
  *
  * @since 6.0
  */
@@ -149,6 +150,8 @@ public synchronized void start() {
 			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 synchronized void start() {
 						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,18 @@ public synchronized void start() {
 				this.latch.countDown();
 			}
 		});
+
+		try {
+			if (!startingLatch.await(5, TimeUnit.SECONDS)) {
+				throw new IllegalStateException("Failed to start "
+						+ PostgresChannelMessageTableSubscriber.class.getName());
+			}
+		}
+		catch (InterruptedException e) {
+			Thread.currentThread().interrupt();
+			throw new IllegalStateException("Failed to start "
+					+ PostgresChannelMessageTableSubscriber.class.getName(), e);
+		}
 	}
 
 	private boolean isActive() {

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,15 +16,20 @@
 
 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;
 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 @@
  *
  * @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 TransactionTemplate transactionTemplate;
+
 	private final Object groupId;
 
 	private final PostgresChannelMessageTableSubscriber messageTableSubscriber;
 
-	private UnicastingDispatcher dispatcher = new UnicastingDispatcher(new SimpleAsyncTaskExecutor());
+	private final UnicastingDispatcher dispatcher = new UnicastingDispatcher();
+
+	private RetryTemplate retryTemplate = RetryTemplate.builder().maxAttempts(1).build();
+
+	private Executor executor = new SimpleAsyncTaskExecutor();
 
 	/**
 	 * Create a subscribable channel for a Postgres database.
@@ -75,7 +89,26 @@ public PostgresSubscribableChannel(JdbcChannelMessageStore jdbcChannelMessageSto
 	 */
 	public void setDispatcherExecutor(Executor executor) {
 		Assert.notNull(executor, "An executor must be provided.");
-		this.dispatcher = new UnicastingDispatcher(executor);
+		this.executor = executor;
+	}
+
+	/**
+	 * Sets the transaction manager to use for message processing. Each message will be processed in a
+	 * separate transaction
+	 * @param transactionManager The transaction manager to use
+	 */
+	public void setTransactionManager(PlatformTransactionManager transactionManager) {
+		Assert.notNull(transactionManager, "A platform transaction manager must be provided.");
+		this.transactionTemplate = new TransactionTemplate(transactionManager);
+	}
+
+	/**
+	 * Sets retry template to use for retries in case of exception in downstream processing
+	 * @param retryTemplate The retry template to use
+	 */
+	public void setRetryTemplate(RetryTemplate retryTemplate) {
+		Assert.notNull(retryTemplate, "A retry template must be provided.");
+		this.retryTemplate = retryTemplate;
 	}
 
 	@Override
@@ -110,10 +143,37 @@ protected boolean doSend(Message<?> message, long timeout) {
 
 	@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 ->
+										Optional.ofNullable(this.jdbcChannelMessageStore.pollMessageFromGroup(this.groupId))
+												.map(this::dispatch)
+								)
+						);
+					}
+					else {
+						dispatchedMessage = Optional.ofNullable(this.jdbcChannelMessageStore.pollMessageFromGroup(this.groupId))
+								.map(message ->
+										this.retryTemplate.execute(context -> dispatch(message))
+								);
+					}
+
+				} while (dispatchedMessage.isPresent());
+			}
+			catch (Exception e) {
+				LOGGER.error(e, "Exception during message dispatch");
+			}
+		});
+	}
+
+	private Message<?> dispatch(Message<?> message) {
+		this.dispatcher.dispatch(message);
+		return message;
 	}
 
 	@Override