Add Reactive mode for AbstractPollingEndpoint (#2429)

* Add Reactive mode for AbstractPollingEndpoint

* When `SourcePollingChannelAdapter.outputChannel` is a
`ReactiveStreamsSubscribableChannel`, use `Flux.generate()` for polling
* Refactor `AbstractPollingEndpoint` to remove redundant `Poller` class
in favor of lambda
* Extract `pollForMessage()` method to handle TX states instead of
`Poller` class previously

* * Rebase and fix conflicts

* Polishing for GatewayProxyFactoryBeanTests

Author: artembilan

spring-integration-core/src/main/java/org/springframework/integration/endpoint/AbstractPollingEndpoint.java

@@ -16,14 +16,17 @@
 
 package org.springframework.integration.endpoint;
 
+import java.time.Duration;
 import java.util.Collection;
+import java.util.Date;
 import java.util.List;
 import java.util.concurrent.Callable;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ScheduledFuture;
 import java.util.stream.Collectors;
 
 import org.aopalliance.aop.Advice;
+import org.reactivestreams.Subscription;
 
 import org.springframework.aop.framework.ProxyFactory;
 import org.springframework.beans.factory.BeanClassLoaderAware;
@@ -43,6 +46,7 @@
 import org.springframework.messaging.MessagingException;
 import org.springframework.scheduling.Trigger;
 import org.springframework.scheduling.support.PeriodicTrigger;
+import org.springframework.scheduling.support.SimpleTriggerContext;
 import org.springframework.transaction.interceptor.TransactionInterceptor;
 import org.springframework.transaction.support.TransactionSynchronization;
 import org.springframework.transaction.support.TransactionSynchronizationManager;
@@ -51,6 +55,10 @@
 import org.springframework.util.CollectionUtils;
 import org.springframework.util.ErrorHandler;
 
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+import reactor.core.scheduler.Schedulers;
+
 /**
  * @author Mark Fisher
  * @author Oleg Zhurakousky
@@ -66,23 +74,27 @@ public abstract class AbstractPollingEndpoint extends AbstractEndpoint implement
 
 	private boolean syncExecutor = true;
 
+	private ClassLoader beanClassLoader = ClassUtils.getDefaultClassLoader();
+
+	private Trigger trigger = new PeriodicTrigger(10);
+
+	private long maxMessagesPerPoll = -1;
+
 	private ErrorHandler errorHandler;
 
 	private boolean errorHandlerIsDefault;
 
-	private Trigger trigger = new PeriodicTrigger(10);
-
 	private List<Advice> adviceChain;
 
-	private ClassLoader beanClassLoader = ClassUtils.getDefaultClassLoader();
+	private TransactionSynchronizationFactory transactionSynchronizationFactory;
 
-	private long maxMessagesPerPoll = -1;
+	private volatile Callable<Message<?>> pollingTask;
 
-	private TransactionSynchronizationFactory transactionSynchronizationFactory;
+	private volatile Flux<Message<?>> pollingFlux;
 
-	private volatile ScheduledFuture<?> runningTask;
+	private volatile Subscription subscription;
 
-	private volatile Runnable poller;
+	private volatile ScheduledFuture<?> runningTask;
 
 	private volatile boolean initialized;
 
@@ -167,6 +179,14 @@ protected boolean isReceiveOnlyAdvice(Advice advice) {
 	protected void applyReceiveOnlyAdviceChain(Collection<Advice> chain) {
 	}
 
+	protected boolean isReactive() {
+		return false;
+	}
+
+	protected Flux<Message<?>> getPollingFlux() {
+		return this.pollingFlux;
+	}
+
 	@Override
 	protected void onInit() {
 		synchronized (this.initializationMonitor) {
@@ -200,16 +220,38 @@ protected void onInit() {
 		}
 	}
 
+	// LifecycleSupport implementation
+
+	@Override // guarded by super#lifecycleLock
+	protected void doStart() {
+		if (!this.initialized) {
+			onInit();
+		}
+
+		this.pollingTask = createPollingTask();
+
+		if (isReactive()) {
+			this.pollingFlux = createFluxGenerator();
+		}
+		else {
+			Assert.state(getTaskScheduler() != null, "unable to start polling, no taskScheduler available");
+
+			this.runningTask =
+					getTaskScheduler()
+							.schedule(createPoller(), this.trigger);
+		}
+	}
+
 	@SuppressWarnings("unchecked")
-	private Runnable createPoller() throws Exception {
+	private Callable<Message<?>> createPollingTask() {
 		List<Advice> receiveOnlyAdviceChain = null;
 		if (!CollectionUtils.isEmpty(this.adviceChain)) {
 			receiveOnlyAdviceChain = this.adviceChain.stream()
 					.filter(this::isReceiveOnlyAdvice)
 					.collect(Collectors.toList());
 		}
 
-		Callable<Boolean> pollingTask = this::doPoll;
+		Callable<Message<?>> pollingTask = this::doPoll;
 
 		List<Advice> adviceChain = this.adviceChain;
 		if (!CollectionUtils.isEmpty(adviceChain)) {
@@ -219,65 +261,122 @@ private Runnable createPoller() throws Exception {
 						.filter(advice -> !isReceiveOnlyAdvice(advice))
 						.forEach(proxyFactory::addAdvice);
 			}
-			pollingTask = (Callable<Boolean>) proxyFactory.getProxy(this.beanClassLoader);
+			pollingTask = (Callable<Message<?>>) proxyFactory.getProxy(this.beanClassLoader);
 		}
 		if (!CollectionUtils.isEmpty(receiveOnlyAdviceChain)) {
 			applyReceiveOnlyAdviceChain(receiveOnlyAdviceChain);
 		}
-		return new Poller(pollingTask);
+
+		return pollingTask;
 	}
 
-	// LifecycleSupport implementation
+	private Runnable createPoller() {
+		return () ->
+				this.taskExecutor.execute(() -> {
+					int count = 0;
+					while (this.initialized && (this.maxMessagesPerPoll <= 0 || count < this.maxMessagesPerPoll)) {
+						if (pollForMessage() == null) {
+							break;
+						}
+						count++;
+					}
+				});
+	}
 
-	@Override // guarded by super#lifecycleLock
-	protected void doStart() {
-		if (!this.initialized) {
-			this.onInit();
-		}
-		Assert.state(this.getTaskScheduler() != null,
-				"unable to start polling, no taskScheduler available");
+	private Flux<Message<?>> createFluxGenerator() {
+		SimpleTriggerContext triggerContext = new SimpleTriggerContext();
+
+		return Flux
+				.<Duration>generate(sink -> {
+					Date date = this.trigger.nextExecutionTime(triggerContext);
+					if (date != null) {
+						triggerContext.update(date, null, null);
+						long millis = date.getTime() - System.currentTimeMillis();
+						sink.next(Duration.ofMillis(millis));
+					}
+					else {
+						sink.complete();
+					}
+				})
+				.concatMap(duration ->
+						Mono.delay(duration)
+								.doOnNext(l ->
+										triggerContext.update(triggerContext.lastScheduledExecutionTime(),
+												new Date(), null))
+								.flatMapMany(l ->
+										Flux
+												.<Message<?>>generate(fluxSink -> {
+													Message<?> message = pollForMessage();
+													if (message != null) {
+														fluxSink.next(message);
+													}
+													else {
+														fluxSink.complete();
+													}
+												})
+												.take(this.maxMessagesPerPoll)
+												.subscribeOn(Schedulers.fromExecutor(this.taskExecutor))
+												.doOnComplete(() ->
+														triggerContext.update(triggerContext.lastScheduledExecutionTime(),
+																triggerContext.lastActualExecutionTime(),
+																new Date())
+												)), 1)
+				.repeat(this::isRunning)
+				.doOnSubscribe(subscription -> this.subscription = subscription);
+	}
+
+	private Message<?> pollForMessage() {
 		try {
-			this.poller = createPoller();
+			return this.pollingTask.call();
 		}
 		catch (Exception e) {
-			this.initialized = false;
-			throw new MessagingException("Failed to create Poller", e);
+			if (e instanceof MessagingException) {
+				throw (MessagingException) e;
+			}
+			else {
+				Message<?> failedMessage = null;
+				if (this.transactionSynchronizationFactory != null) {
+					Object resource = TransactionSynchronizationManager.getResource(getResourceToBind());
+					if (resource instanceof IntegrationResourceHolder) {
+						failedMessage = ((IntegrationResourceHolder) resource).getMessage();
+					}
+				}
+				throw new MessagingException(failedMessage, e);
+			}
 		}
-		this.runningTask = this.getTaskScheduler().schedule(this.poller, this.trigger);
-	}
-
-	@Override // guarded by super#lifecycleLock
-	protected void doStop() {
-		if (this.runningTask != null) {
-			this.runningTask.cancel(true);
+		finally {
+			if (this.transactionSynchronizationFactory != null) {
+				Object resource = getResourceToBind();
+				if (TransactionSynchronizationManager.hasResource(resource)) {
+					TransactionSynchronizationManager.unbindResource(resource);
+				}
+			}
 		}
-		this.runningTask = null;
 	}
 
-	private boolean doPoll() {
-		IntegrationResourceHolder holder = this.bindResourceHolderIfNecessary(
-				this.getResourceKey(), this.getResourceToBind());
-		Message<?> message = null;
+	private Message<?> doPoll() {
+		IntegrationResourceHolder holder = bindResourceHolderIfNecessary(getResourceKey(), getResourceToBind());
+		Message<?> message;
 		try {
-			message = this.receiveMessage();
+			message = receiveMessage();
 		}
 		catch (Exception e) {
 			if (Thread.interrupted()) {
 				if (logger.isDebugEnabled()) {
 					logger.debug("Poll interrupted - during stop()? : " + e.getMessage());
 				}
-				return false;
+				return null;
 			}
 			else {
 				throw (RuntimeException) e;
 			}
 		}
-		boolean result;
+
 		if (message == null) {
 			if (this.logger.isDebugEnabled()) {
 				this.logger.debug("Received no Message during the poll, returning 'false'");
 			}
-			result = false;
+			return null;
 		}
 		else {
 			if (this.logger.isDebugEnabled()) {
@@ -286,20 +385,35 @@ private boolean doPoll() {
 			if (holder != null) {
 				holder.setMessage(message);
 			}
-			try {
-				this.handleMessage(message);
-			}
-			catch (Exception e) {
-				if (e instanceof MessagingException) {
-					throw new MessagingExceptionWrapper(message, (MessagingException) e);
+
+			if (!isReactive()) {
+				try {
+					handleMessage(message);
 				}
-				else {
-					throw new MessagingException(message, e);
+				catch (Exception e) {
+					if (e instanceof MessagingException) {
+						throw new MessagingExceptionWrapper(message, (MessagingException) e);
+					}
+					else {
+						throw new MessagingException(message, e);
+					}
 				}
 			}
-			result = true;
 		}
-		return result;
+
+		return message;
+	}
+
+	@Override // guarded by super#lifecycleLock
+	protected void doStop() {
+		if (this.runningTask != null) {
+			this.runningTask.cancel(true);
+		}
+		this.runningTask = null;
+
+		if (this.subscription != null) {
+			this.subscription.cancel();
+		}
 	}
 
 	/**
@@ -369,57 +483,4 @@ private IntegrationResourceHolder bindResourceHolderIfNecessary(String key, Obje
 		return null;
 	}
 
-	/**
-	 * Default Poller implementation
-	 */
-	private final class Poller implements Runnable {
-
-		private final Callable<Boolean> pollingTask;
-
-		Poller(Callable<Boolean> pollingTask) {
-			this.pollingTask = pollingTask;
-		}
-
-		@Override
-		public void run() {
-			AbstractPollingEndpoint.this.taskExecutor.execute(() -> {
-				int count = 0;
-				while (AbstractPollingEndpoint.this.initialized
-						&& (AbstractPollingEndpoint.this.maxMessagesPerPoll <= 0
-						|| count < AbstractPollingEndpoint.this.maxMessagesPerPoll)) {
-					try {
-						if (!Poller.this.pollingTask.call()) {
-							break;
-						}
-						count++;
-					}
-					catch (Exception e) {
-						if (e instanceof MessagingException) {
-							throw (MessagingException) e;
-						}
-						else {
-							Message<?> failedMessage = null;
-							if (AbstractPollingEndpoint.this.transactionSynchronizationFactory != null) {
-								Object resource = TransactionSynchronizationManager.getResource(getResourceToBind());
-								if (resource instanceof IntegrationResourceHolder) {
-									failedMessage = ((IntegrationResourceHolder) resource).getMessage();
-								}
-							}
-							throw new MessagingException(failedMessage, e);
-						}
-					}
-					finally {
-						if (AbstractPollingEndpoint.this.transactionSynchronizationFactory != null) {
-							Object resource = getResourceToBind();
-							if (TransactionSynchronizationManager.hasResource(resource)) {
-								TransactionSynchronizationManager.unbindResource(resource);
-							}
-						}
-					}
-				}
-			});
-		}
-
-	}
-
 }