INT-2815 Support Tx Synch In Pollable Consumers

Pull transaction synchronization code up from SPCA
to AbstractPollingEndpoint, providing support for
transaction synchronization in PollingConsumer.

Also, ensure headers are copied to any message generated by the
ExpressionEvaluatingTransactionSynchronizationProcessor.

spring-integration-core/src/main/java/org/springframework/integration/config/ConsumerEndpointFactoryBean.java

@@ -21,7 +21,6 @@ import java.util.List;
 import org.aopalliance.aop.Advice;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
-
 import org.springframework.aop.framework.Advised;
 import org.springframework.aop.framework.ProxyFactory;
 import org.springframework.aop.support.AopUtils;
@@ -229,6 +228,7 @@ public class ConsumerEndpointFactoryBean
 				pollingConsumer.setErrorHandler(this.pollerMetadata.getErrorHandler());
 
 				pollingConsumer.setReceiveTimeout(this.pollerMetadata.getReceiveTimeout());
+				pollingConsumer.setTransactionSynchronizationFactory(this.pollerMetadata.getTransactionSynchronizationFactory());
 				pollingConsumer.setBeanClassLoader(beanClassLoader);
 				pollingConsumer.setBeanFactory(beanFactory);
 				this.endpoint = pollingConsumer;

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2011 the original author or authors.
+ * Copyright 2002-2012 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.
@@ -22,19 +22,22 @@ import java.util.concurrent.Executor;
 import java.util.concurrent.ScheduledFuture;
 
 import org.aopalliance.aop.Advice;
-
 import org.springframework.aop.framework.ProxyFactory;
 import org.springframework.beans.factory.BeanClassLoaderAware;
 import org.springframework.core.task.SyncTaskExecutor;
+import org.springframework.integration.Message;
 import org.springframework.integration.MessageHandlingException;
 import org.springframework.integration.MessagingException;
 import org.springframework.integration.channel.MessagePublishingErrorHandler;
 import org.springframework.integration.message.ErrorMessage;
 import org.springframework.integration.scheduling.PollerMetadata;
 import org.springframework.integration.support.channel.BeanFactoryChannelResolver;
+import org.springframework.integration.transaction.IntegrationResourceHolder;
+import org.springframework.integration.transaction.TransactionSynchronizationFactory;
 import org.springframework.integration.util.ErrorHandlingTaskExecutor;
 import org.springframework.scheduling.Trigger;
 import org.springframework.scheduling.support.PeriodicTrigger;
+import org.springframework.transaction.support.TransactionSynchronizationManager;
 import org.springframework.util.Assert;
 import org.springframework.util.ClassUtils;
 import org.springframework.util.CollectionUtils;
@@ -43,11 +46,12 @@ import org.springframework.util.ErrorHandler;
 /**
  * @author Mark Fisher
  * @author Oleg Zhurakousky
+ * @author Gary Russell
  */
 public abstract class AbstractPollingEndpoint extends AbstractEndpoint implements BeanClassLoaderAware {
 
 	private volatile Executor taskExecutor = new SyncTaskExecutor();
-	
+
 	private volatile ErrorHandler errorHandler;
 
 	private volatile Trigger trigger = new PeriodicTrigger(10);
@@ -61,16 +65,17 @@ public abstract class AbstractPollingEndpoint extends AbstractEndpoint implement
 	private volatile Runnable poller;
 
 	private volatile boolean initialized;
-	
+
 	private volatile long maxMessagesPerPoll = -1;
-	
+
 	private final Object initializationMonitor = new Object();
 
+	private volatile TransactionSynchronizationFactory transactionSynchronizationFactory;
 
 	public AbstractPollingEndpoint() {
 		this.setPhase(Integer.MAX_VALUE);
 	}
-	
+
 	/**
 	 * @deprecated  As of release 2.0.2, use individual setters
 	 */
@@ -82,7 +87,7 @@ public abstract class AbstractPollingEndpoint extends AbstractEndpoint implement
 		this.setTaskExecutor(pollerMetadata.getTaskExecutor());
 		this.setTrigger(pollerMetadata.getTrigger());
 	}
-	
+
 	public void setTaskExecutor(Executor taskExecutor) {
 		this.taskExecutor = (taskExecutor != null ? taskExecutor : new SyncTaskExecutor());
 	}
@@ -107,6 +112,11 @@ public abstract class AbstractPollingEndpoint extends AbstractEndpoint implement
 		this.beanClassLoader = classLoader;
 	}
 
+	public void setTransactionSynchronizationFactory(
+			TransactionSynchronizationFactory transactionSynchronizationFactory) {
+		this.transactionSynchronizationFactory = transactionSynchronizationFactory;
+	}
+
 	@Override
 	protected void onInit() {
 		synchronized (this.initializationMonitor) {
@@ -119,7 +129,7 @@ public abstract class AbstractPollingEndpoint extends AbstractEndpoint implement
 				this.taskExecutor = providedExecutor;
 			}
 			if (this.taskExecutor != null) {
-				if (!(this.taskExecutor instanceof ErrorHandlingTaskExecutor)) {				
+				if (!(this.taskExecutor instanceof ErrorHandlingTaskExecutor)) {
 					if (this.errorHandler == null) {
 						Assert.notNull(this.getBeanFactory(), "BeanFactory is required");
 						this.errorHandler = new MessagePublishingErrorHandler(
@@ -131,7 +141,7 @@ public abstract class AbstractPollingEndpoint extends AbstractEndpoint implement
 			try {
 				this.poller = this.createPoller();
 				this.initialized = true;
-			} 
+			}
 			catch (Exception e) {
 				throw new MessagingException("Failed to create Poller", e);
 			}
@@ -140,13 +150,13 @@ public abstract class AbstractPollingEndpoint extends AbstractEndpoint implement
 
 	@SuppressWarnings("unchecked")
 	private Runnable createPoller() throws Exception {
-		
+
 		Callable<Boolean> pollingTask = new Callable<Boolean>() {
 			public Boolean call() throws Exception {
 				return doPoll();
 			}
 		};
-		
+
 		List<Advice> adviceChain = this.adviceChain;
 		if (!CollectionUtils.isEmpty(adviceChain)) {
 			ProxyFactory proxyFactory = new ProxyFactory(pollingTask);
@@ -160,7 +170,20 @@ public abstract class AbstractPollingEndpoint extends AbstractEndpoint implement
 		return new Poller(pollingTask);
 	}
 
-
+	/**
+	 * Synchronize with an existing transaction (if any) and receive
+	 * a message using {@link #doReceive()}.
+	 * @return The message (or null).
+	 */
+	protected final Message<?> syncIfTxAndReceive() {
+		IntegrationResourceHolder holder = bindResourceHolderIfNecessary(
+				this.getResourceKey(), this.getResourceToBind());
+		Message<?> message = this.doReceive();
+		if (holder != null && message != null) {
+			holder.setMessage(message);
+		}
+		return message;
+	}
 	// LifecycleSupport implementation
 
 	@Override // guarded by super#lifecycleLock
@@ -182,9 +205,29 @@ public abstract class AbstractPollingEndpoint extends AbstractEndpoint implement
 		this.initialized = false;
 	}
 
-	
+	private IntegrationResourceHolder bindResourceHolderIfNecessary(String key, Object resource) {
+		IntegrationResourceHolder holder = null;
+
+		if (this.transactionSynchronizationFactory != null) {
+			if (TransactionSynchronizationManager.isActualTransactionActive()) {
+				holder = new IntegrationResourceHolder();
+				if (key != null) {
+					holder.addAttribute(key, resource);
+				}
+				TransactionSynchronizationManager.bindResource(resource, holder);
+				TransactionSynchronizationManager.registerSynchronization(this.transactionSynchronizationFactory.create(resource));
+			}
+		}
+		return holder;
+	}
+
 	protected abstract boolean doPoll();
 
+	protected abstract Message<?> doReceive();
+
+	protected abstract Object getResourceToBind();
+
+	protected abstract String getResourceKey();
 
 	/**
 	 * Default Poller implementation
@@ -208,11 +251,11 @@ public abstract class AbstractPollingEndpoint extends AbstractEndpoint implement
 								break;
 							}
 							count++;
-						} 
+						}
 						catch (Exception e) {
 							if (e instanceof RuntimeException) {
 								throw (RuntimeException) e;
-							} 
+							}
 							else {
 								throw new MessageHandlingException(new ErrorMessage(e));
 							}

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

@@ -22,6 +22,7 @@ import org.springframework.context.Lifecycle;
 import org.springframework.integration.Message;
 import org.springframework.integration.core.MessageHandler;
 import org.springframework.integration.core.PollableChannel;
+import org.springframework.integration.transaction.IntegrationResourceHolder;
 import org.springframework.util.Assert;
 
 /**
@@ -74,9 +75,7 @@ public class PollingConsumer extends AbstractPollingEndpoint {
 
 	@Override
 	protected boolean doPoll() {
-		Message<?> message = (this.receiveTimeout >= 0)
-				? this.inputChannel.receive(this.receiveTimeout)
-				: this.inputChannel.receive();
+		Message<?> message = this.syncIfTxAndReceive();
 		if (this.logger.isDebugEnabled()){
 			this.logger.debug("Poll resulted in Message: " + message);
 		}
@@ -89,4 +88,22 @@ public class PollingConsumer extends AbstractPollingEndpoint {
 		this.handler.handleMessage(message);
 		return true;
 	}
+
+	@Override
+	protected Message<?> doReceive() {
+		Message<?> message = (this.receiveTimeout >= 0)
+				? this.inputChannel.receive(this.receiveTimeout)
+				: this.inputChannel.receive();
+		return message;
+	}
+
+	@Override
+	protected Object getResourceToBind() {
+		return this.inputChannel;
+	}
+
+	@Override
+	protected String getResourceKey() {
+		return IntegrationResourceHolder.INPUT_CHANNEL;
+	}
 }

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

@@ -25,8 +25,6 @@ import org.springframework.integration.core.MessagingTemplate;
 import org.springframework.integration.history.MessageHistory;
 import org.springframework.integration.history.TrackableComponent;
 import org.springframework.integration.transaction.IntegrationResourceHolder;
-import org.springframework.integration.transaction.TransactionSynchronizationFactory;
-import org.springframework.transaction.support.TransactionSynchronizationManager;
 import org.springframework.util.Assert;
 
 /**
@@ -47,13 +45,6 @@ public class SourcePollingChannelAdapter extends AbstractPollingEndpoint impleme
 
 	private final MessagingTemplate messagingTemplate = new MessagingTemplate();
 
-	private volatile TransactionSynchronizationFactory transactionSynchronizationFactory;
-
-	public void setTransactionSynchronizationFactory(
-			TransactionSynchronizationFactory transactionSynchronizationFactory) {
-		this.transactionSynchronizationFactory = transactionSynchronizationFactory;
-	}
-
 	/**
 	 * Specify the source to be polled for Messages.
 	 */
@@ -99,26 +90,12 @@ public class SourcePollingChannelAdapter extends AbstractPollingEndpoint impleme
 	@Override
 	protected boolean doPoll() {
 
-		Message<?> message;
-		IntegrationResourceHolder holder = null;
-
-		if (TransactionSynchronizationManager.isActualTransactionActive()) {
-			if (transactionSynchronizationFactory != null){
-				holder = new IntegrationResourceHolder();
-				holder.addAttribute(IntegrationResourceHolder.MESSAGE_SOURCE, source);
-				TransactionSynchronizationManager.bindResource(source, holder);
-				TransactionSynchronizationManager.registerSynchronization(transactionSynchronizationFactory.create(source));
-			}
-		}
-		message = this.source.receive();
+		Message<?> message = this.syncIfTxAndReceive();
 
 		if (this.logger.isDebugEnabled()){
 			this.logger.debug("Poll resulted in Message: " + message);
 		}
 		if (message != null) {
-			if (holder != null) {
-				holder.setMessage(message);
-			}
 			if (this.shouldTrack) {
 				message = MessageHistory.write(message, this);
 			}
@@ -140,4 +117,20 @@ public class SourcePollingChannelAdapter extends AbstractPollingEndpoint impleme
 		}
 		return false;
 	}
+
+	@Override
+	protected Message<?> doReceive() {
+		return this.source.receive();
+	}
+
+	@Override
+	protected Object getResourceToBind() {
+		return this.source;
+	}
+
+	@Override
+	protected String getResourceKey() {
+		return IntegrationResourceHolder.MESSAGE_SOURCE;
+	}
+
 }

spring-integration-core/src/main/java/org/springframework/integration/transaction/ExpressionEvaluatingTransactionSynchronizationProcessor.java

@@ -115,7 +115,9 @@ public class ExpressionEvaluatingTransactionSynchronizationProcessor extends Int
 								"as part of '" + expressionType + "' transaction synchronization");
 					}
 					try {
-						spelResultMessage = MessageBuilder.withPayload(value).build();
+						spelResultMessage = MessageBuilder.withPayload(value)
+								.copyHeaders(message.getHeaders())
+								.build();
 						this.sendMessage(messageChannel, spelResultMessage);
 					}
 					catch (Exception e) {

spring-integration-core/src/main/java/org/springframework/integration/transaction/IntegrationResourceHolder.java

@@ -32,6 +32,8 @@ public class IntegrationResourceHolder implements ResourceHolder {
 
 	public static final String MESSAGE_SOURCE = "messageSource";
 
+	public static final String INPUT_CHANNEL = "inputChannel";
+
 	private volatile Message<?> message;
 
 	private final Map<String, Object> attributes = new HashMap<String, Object>();