INT-3722: Fix Timing Issue TCP OG and Cached CF

JIRA: https://jira.spring.io/browse/INT-3722

The TCP outbound gateway correlates replies based on the connection id.

When a `CachingClientConnectionFactory` is being used, the connection is
returned to the pool too early, and can be reused. It is possible that
the current thread then removes the "next" pending reply from the correlation map.

Add code to the `CachingClientConnectionFactory` so that the "self" close from the
connection (called after `onMessage`) is deferred and the actual close (return to cache)
is controlled by the gateway itself.

This mechanism will no longer be needed when INT-3654 is resolved (removal of the "self"
closing by connections). At that time, connection users (such as the gateway) will be
in complete control.

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

@@ -34,6 +34,7 @@ import org.springframework.integration.handler.AbstractReplyProducingMessageHand
 import org.springframework.integration.ip.IpHeaders;
 import org.springframework.integration.ip.tcp.connection.AbstractClientConnectionFactory;
 import org.springframework.integration.ip.tcp.connection.AbstractConnectionFactory;
+import org.springframework.integration.ip.tcp.connection.CloseDeferrable;
 import org.springframework.integration.ip.tcp.connection.TcpConnection;
 import org.springframework.integration.ip.tcp.connection.TcpListener;
 import org.springframework.integration.ip.tcp.connection.TcpSender;
@@ -155,6 +156,9 @@ public class TcpOutboundGateway extends AbstractReplyProducingMessageHandler
 					logger.debug("released semaphore");
 				}
 			}
+			if (this.connectionFactory instanceof CloseDeferrable) {
+				((CloseDeferrable) this.connectionFactory).closeDeferred(connectionId);
+			}
 		}
 	}
 
@@ -221,6 +225,9 @@ public class TcpOutboundGateway extends AbstractReplyProducingMessageHandler
 
 	@Override
 	public void start() {
+		if (this.connectionFactory instanceof CloseDeferrable) {
+			((CloseDeferrable) this.connectionFactory).enableCloseDeferral(true);
+		}
 		this.connectionFactory.start();
 	}
 

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

@@ -16,6 +16,8 @@
 
 package org.springframework.integration.ip.tcp.connection;
 
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executor;
 
 import org.springframework.core.serializer.Deserializer;
@@ -37,12 +39,17 @@ import org.springframework.messaging.support.ErrorMessage;
  * @since 2.2
  *
  */
-public class CachingClientConnectionFactory extends AbstractClientConnectionFactory {
+public class CachingClientConnectionFactory extends AbstractClientConnectionFactory implements CloseDeferrable {
 
 	private final AbstractClientConnectionFactory targetConnectionFactory;
 
 	private final SimplePool<TcpConnectionSupport> pool;
 
+	private final Map<String, CachedConnection> deferredClosures =
+			new ConcurrentHashMap<String, CachedConnection>();
+
+	private volatile boolean deferClose;
+
 	/**
 	 * Construct a caching connection factory that delegates to the provided factory, with
 	 * the provided pool size.
@@ -135,6 +142,19 @@ public class CachingClientConnectionFactory extends AbstractClientConnectionFact
 		return cachedConnection;
 	}
 
+	@Override
+	public void enableCloseDeferral(boolean defer) {
+		this.deferClose = defer;
+	}
+
+	@Override
+	public void closeDeferred(String connectionId) {
+		CachedConnection deferred = this.deferredClosures.remove(connectionId);
+		if (deferred != null) {
+			deferred.doClose();
+		}
+	}
+
 	private class CachedConnection extends TcpConnectionInterceptorSupport {
 
 		private volatile boolean released;
@@ -146,6 +166,15 @@ public class CachingClientConnectionFactory extends AbstractClientConnectionFact
 
 		@Override
 		public synchronized void close() {
+			if (deferClose && !this.released) {
+				deferredClosures.put(getConnectionId(), this);
+			}
+			else {
+				doClose();
+			}
+		}
+
+		private synchronized void doClose() {
 			if (this.released) {
 				if (logger.isDebugEnabled()) {
 					logger.debug("Connection " + getConnectionId() + " has already been released");

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

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2015 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.integration.ip.tcp.connection;
+
+/**
+ * Temporary interface on the {@code CachingClientConnectionFactory} enabling the gateway
+ * to defer the implicit close after onMessage so the connection is not reused until after the
+ * gateway has completely finished with it. Will be removed when INT-3654 is resolved, whereby
+ * the gateway will be completely responsible for the close.
+ *
+ * @author Gary Russell
+ * @since 4.1.5
+ *
+ */
+public interface CloseDeferrable {
+
+	/**
+	 * Enable deferred closure.
+	 * @param defer true to defer.
+	 */
+	void enableCloseDeferral(boolean defer);
+
+	/**
+	 * Close (release) the connection if deferred.
+	 * @param connectionId the connection id.
+	 */
+	void closeDeferred(String connectionId);
+
+}

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

@@ -27,9 +27,11 @@ import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyInt;
+import static org.mockito.Matchers.anyString;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -46,6 +48,7 @@ import java.util.concurrent.Executors;
 import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import org.apache.commons.logging.Log;
 import org.junit.Test;
@@ -56,10 +59,14 @@ import org.mockito.invocation.InvocationOnMock;
 import org.mockito.stubbing.Answer;
 
 import org.springframework.beans.DirectFieldAccessor;
+import org.springframework.beans.factory.BeanFactory;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.context.ApplicationEvent;
 import org.springframework.context.ApplicationEventPublisher;
+import org.springframework.integration.channel.QueueChannel;
 import org.springframework.integration.ip.IpHeaders;
+import org.springframework.integration.ip.tcp.TcpOutboundGateway;
+import org.springframework.integration.ip.tcp.TcpSendingMessageHandler;
 import org.springframework.integration.ip.tcp.serializer.ByteArrayCrLfSerializer;
 import org.springframework.integration.ip.util.TestingUtilities;
 import org.springframework.integration.support.MessageBuilder;
@@ -590,6 +597,91 @@ public class CachingClientConnectionFactoryTests {
 		cache.stop();
 	}
 
+	@SuppressWarnings("unchecked")
+	@Test //INT-3722
+	public void testGatewayRelease() throws Exception {
+		int port = SocketUtils.findAvailableTcpPort();
+		TcpNetServerConnectionFactory in = new TcpNetServerConnectionFactory(port);
+		in.setApplicationEventPublisher(mock(ApplicationEventPublisher.class));
+		final TcpSendingMessageHandler handler = new TcpSendingMessageHandler();
+		handler.setConnectionFactory(in);
+		final AtomicInteger count = new AtomicInteger(2);
+		in.registerListener(new TcpListener() {
+
+			@Override
+			public boolean onMessage(Message<?> message) {
+				if (!(message instanceof ErrorMessage)) {
+					if (count.decrementAndGet() < 1) {
+						try {
+							Thread.sleep(1000);
+						}
+						catch (InterruptedException e) {
+							Thread.currentThread().interrupt();
+						}
+					}
+					handler.handleMessage(message);
+				}
+				return false;
+			}
+
+		});
+		handler.setBeanFactory(mock(BeanFactory.class));
+		handler.afterPropertiesSet();
+		handler.start();
+		int n = 0;
+		while (n++ < 100 && !in.isListening()) {
+			Thread.sleep(100);
+		}
+		assertTrue(in.isListening());
+		TcpNetClientConnectionFactory out = new TcpNetClientConnectionFactory("localhost", port);
+		out.setApplicationEventPublisher(mock(ApplicationEventPublisher.class));
+		CachingClientConnectionFactory cache = new CachingClientConnectionFactory(out, 2);
+		final TcpOutboundGateway gate = new TcpOutboundGateway();
+		gate.setConnectionFactory(cache);
+		QueueChannel outputChannel = new QueueChannel();
+		gate.setOutputChannel(outputChannel);
+		gate.setBeanFactory(mock(BeanFactory.class));
+		gate.afterPropertiesSet();
+		Log logger = spy(TestUtils.getPropertyValue(gate, "logger", Log.class));
+		new DirectFieldAccessor(gate).setPropertyValue("logger", logger);
+		when(logger.isDebugEnabled()).thenReturn(true);
+		doAnswer(new Answer<Void>() {
+
+			private final CountDownLatch latch = new CountDownLatch(2);
+
+			@Override
+			public Void answer(InvocationOnMock invocation) throws Throwable {
+				String log = (String) invocation.getArguments()[0];
+				if (log.startsWith("Response")) {
+					Executors.newSingleThreadScheduledExecutor().execute(new Runnable() {
+
+						@Override
+						public void run() {
+							gate.handleMessage(new GenericMessage<String>("bar"));
+						}
+					});
+					// hold up the first thread until the second has added its pending reply
+					latch.await(10, TimeUnit.SECONDS);
+				}
+				else if (log.startsWith("Added")) {
+					latch.countDown();
+				}
+				return null;
+			}
+		}).when(logger).debug(anyString());
+		gate.start();
+		gate.handleMessage(new GenericMessage<String>("foo"));
+		Message<byte[]> result = (Message<byte[]>) outputChannel.receive(10000);
+		assertNotNull(result);
+		assertEquals("foo", new String(result.getPayload()));
+		result = (Message<byte[]>) outputChannel.receive(10000);
+		assertNotNull(result);
+		assertEquals("bar", new String(result.getPayload()));
+		handler.stop();
+		gate.stop();
+		verify(logger, never()).error(anyString());
+	}
+
 	public TcpConnectionSupport makeMockConnection() {
 		TcpConnectionSupport connection = mock(TcpConnectionSupport.class);
 		when(connection.isOpen()).thenReturn(true);