Merge pull request #686 from garyrussell/INT-2839

* INT-2839: Fix TCP Caching Connections with Gateways

spring-integration-ip/src/main/java/org/springframework/integration/ip/tcp/connection/CachingClientConnectionFactory.java

@@ -21,7 +21,10 @@ import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.springframework.core.serializer.Deserializer;
 import org.springframework.core.serializer.Serializer;
+import org.springframework.integration.Message;
 import org.springframework.integration.MessagingException;
+import org.springframework.integration.ip.IpHeaders;
+import org.springframework.integration.support.MessageBuilder;
 import org.springframework.integration.util.SimplePool;
 
 /**
@@ -37,6 +40,8 @@ public class CachingClientConnectionFactory extends AbstractClientConnectionFact
 
 	private final SimplePool<TcpConnection> pool;
 
+	private volatile TcpListener listener;
+
 	public CachingClientConnectionFactory(AbstractClientConnectionFactory target, int poolSize) {
 		super("", 0);
 		// override single-use to true to force "close" after use
@@ -140,6 +145,27 @@ public class CachingClientConnectionFactory extends AbstractClientConnectionFact
 			return this.getConnectionId();
 		}
 
+		@Override
+		public TcpListener getListener() {
+			return CachingClientConnectionFactory.this.listener;
+		}
+
+		/**
+		 * We have to intercept the message to replace the connectionId header with
+		 * ours so the listener can correlate a response with a request. We supply
+		 * the actual connectionId in another header for convenience and tracing
+		 * purposes.
+		 */
+		@Override
+		public boolean onMessage(Message<?> message) {
+			CachingClientConnectionFactory.this.listener.onMessage(MessageBuilder.fromMessage(message)
+					.setHeader(IpHeaders.CONNECTION_ID, this.getConnectionId())
+					.setHeader(IpHeaders.ACTUAL_CONNECTION_ID, message.getHeaders().get(IpHeaders.CONNECTION_ID))
+					.build());
+			close(); // return to pool after response is received
+			return true; // true so the single-use connection doesn't close itself
+		}
+
 	}
 
 ///////////////// DELEGATE METHODS ///////////////////////
@@ -281,6 +307,7 @@ public class CachingClientConnectionFactory extends AbstractClientConnectionFact
 
 	@Override
 	public void registerListener(TcpListener listener) {
+		this.listener = listener;
 		targetConnectionFactory.registerListener(listener);
 	}
 

spring-integration-ip/src/test/java/org/springframework/integration/ip/tcp/connection/CachingClientConnectionFactoryTests-context.xml

@@ -22,6 +22,12 @@
 		<int:queue/>
 	</int:channel>
 
+	<int-ip:tcp-outbound-channel-adapter
+		connection-factory="scf"
+		channel="replies" />
+
+	<int:channel id="replies" />
+
 	<int-ip:tcp-connection-factory
 		id="ccf"
 		type="client"
@@ -43,4 +49,29 @@
 
 	<int:channel id="outbound" />
 
+	<int-ip:tcp-connection-factory
+		id="gateway.ccf"
+		type="client"
+		host="localhost"
+		port="9876"
+		so-timeout="60000"
+	/>
+
+	<bean id="gateway.caching.ccf" class="org.springframework.integration.ip.tcp.connection.CachingClientConnectionFactory">
+		<constructor-arg ref="gateway.ccf" />
+		<constructor-arg value="10" />
+		<property name="connectionWaitTimeout" value="10000"/>
+	</bean>
+
+	<int-ip:tcp-outbound-gateway
+		id="ob.gw"
+		connection-factory="gateway.caching.ccf"
+		request-channel="toGateway"
+		reply-channel="fromGateway" />
+
+	<int:channel id="toGateway" />
+
+	<int:channel id="fromGateway">
+		<int:queue/>
+	</int:channel>
 </beans>

spring-integration-ip/src/test/java/org/springframework/integration/ip/tcp/connection/CachingClientConnectionFactoryTests.java

@@ -25,7 +25,11 @@ import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.Executors;
 import java.util.concurrent.Semaphore;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -39,6 +43,7 @@ import org.springframework.integration.core.SubscribableChannel;
 import org.springframework.integration.ip.IpHeaders;
 import org.springframework.integration.ip.util.TestingUtilities;
 import org.springframework.integration.message.GenericMessage;
+import org.springframework.integration.support.MessageBuilder;
 import org.springframework.integration.test.util.TestUtils;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
@@ -61,6 +66,15 @@ public class CachingClientConnectionFactoryTests {
 	@Autowired
 	AbstractServerConnectionFactory serverCf;
 
+	@Autowired
+	SubscribableChannel toGateway;
+
+	@Autowired
+	SubscribableChannel replies;
+
+	@Autowired
+	PollableChannel fromGateway;
+
 	@Test
 	public void testReuse() throws Exception {
 		AbstractClientConnectionFactory factory = mock(AbstractClientConnectionFactory.class);
@@ -275,10 +289,46 @@ public class CachingClientConnectionFactoryTests {
 		Message<?> m = inbound.receive(1000);
 		assertNotNull(m);
 		String connectionId = m.getHeaders().get(IpHeaders.CONNECTION_ID, String.class);
+
+		// assert we use the same connection from the pool
 		outbound.send(new GenericMessage<String>("Hello, world!"));
 		m = inbound.receive(1000);
 		assertNotNull(m);
 		assertEquals(connectionId, m.getHeaders().get(IpHeaders.CONNECTION_ID, String.class));
+	}
 
+	@Test
+	public void gatewayIntegrationTest() throws Exception {
+		final List<String> connectionIds = new ArrayList<String>();
+		final AtomicBoolean okToRun = new AtomicBoolean(true);
+		Executors.newSingleThreadExecutor().execute(new Runnable() {
+			public void run() {
+				while (okToRun.get()) {
+					Message<?> m = inbound.receive(1000);
+					if (m != null) {
+						connectionIds.add((String) m.getHeaders().get(IpHeaders.CONNECTION_ID));
+						replies.send(MessageBuilder.withPayload("foo:" + new String((byte[]) m.getPayload()))
+								.copyHeaders(m.getHeaders())
+								.build());
+					}
+				}
+			}
+		});
+		TestingUtilities.waitListening(serverCf, null);
+		toGateway.send(new GenericMessage<String>("Hello, world!"));
+		Message<?> m = fromGateway.receive(1000);
+		assertNotNull(m);
+		assertEquals("foo:" + "Hello, world!", new String((byte[]) m.getPayload()));
+
+		// assert we use the same connection from the pool
+		toGateway.send(new GenericMessage<String>("Hello, world2!"));
+		m = fromGateway.receive(1000);
+		assertNotNull(m);
+		assertEquals("foo:" + "Hello, world2!", new String((byte[]) m.getPayload()));
+
+		assertEquals(2, connectionIds.size());
+		assertEquals(connectionIds.get(0), connectionIds.get(1));
+
+		okToRun.set(false);
 	}
 }