INT-3728: TCP: Fix Early Receive with Caching CCF

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

The `CachingClientConnectionFactory` uses a temporary (rejecting) listener until
the connection is established and then replaces the listener in the connection.

There is a race condition in that if the server starts sending messages before the
listener is replaced, the message is rejected.

Add a mechanism to delay `onMessage` calls until the real listener has been
registered.

Also fix `onMessage` in the cached connection so an `ErrorMessage` is propagated correctly.

To reproduce: revert src/main; add `Thread.sleep(1000)` before `registerListener(tcpListener);`
in `CachedConnection` ctor and run the new test.

To introduce a similar timing hole with the new code, add the sleep before
`this.theConnection.registerListener(this);` in `TcpConnectionInterceptorSupport`.

Summary of changes:

`CachingClientConnectionFactory`
- register the underlying connection's listener in the ctor, utilizing the
    `TcpConnectionInterceptorSupport.registerListener()` method.
- fix `ErrorMessage` propagation.

`AbstractClientConnectionFactory`
- propagate the `enableManualListenerRegistration` to connections.

`TcpConnectionSuport`
- implement delay when manual listener registration is enabled.

Add test case.

Fix Failover Tests

`FailoverClientConnectionFactory`
- propagate enable manual listener registration to underlying factories

Polishing
(cherry picked from commit 9e1119b)

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2014 the original author or authors.
+ * Copyright 2002-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.
@@ -34,6 +34,8 @@ public abstract class AbstractClientConnectionFactory extends AbstractConnection
 
 	private volatile TcpConnectionSupport theConnection;
 
+	private volatile boolean manualListenerRegistration;
+
 	/**
 	 * Constructs a factory that will established connections to the host and port.
 	 * @param host The host.
@@ -43,6 +45,17 @@ public abstract class AbstractClientConnectionFactory extends AbstractConnection
 		super(host, port);
 	}
 
+	/**
+	 * Set whether to automatically (default) or manually add a {@link TcpListener} to the
+	 * connections created by this factory. By default, the factory automatically configures
+	 * the listener. When manual registration is in place, incoming messages will be delayed
+	 * until the listener is registered.
+	 * @since 1.4.5
+	 */
+	public void enableManualListenerRegistration() {
+		this.manualListenerRegistration = true;
+	}
+
 	/**
 	 * Obtains a connection - if {@link #setSingleUse(boolean)} was called with
 	 * true, a new connection is returned; otherwise a single connection is
@@ -126,9 +139,14 @@ public abstract class AbstractClientConnectionFactory extends AbstractConnection
 	 * @param socket The new socket.
 	 */
 	protected void initializeConnection(TcpConnectionSupport connection, Socket socket) {
-		TcpListener listener = this.getListener();
-		if (listener != null) {
-			connection.registerListener(listener);
+		if (this.manualListenerRegistration) {
+			connection.enableManualListenerRegistration();
+		}
+		else {
+			TcpListener listener = this.getListener();
+			if (listener != null) {
+				connection.registerListener(listener);
+			}
 		}
 		TcpSender sender = this.getSender();
 		if (sender != null) {

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

@@ -16,6 +16,7 @@
 
 package org.springframework.integration.ip.tcp.connection;
 
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
@@ -141,9 +142,7 @@ public class CachingClientConnectionFactory extends AbstractClientConnectionFact
 
 	@Override
 	public TcpConnectionSupport obtainConnection() throws Exception {
-		CachedConnection cachedConnection = new CachedConnection(this.pool.getItem());
-		cachedConnection.registerListener(getListener());
-		return cachedConnection;
+		return new CachedConnection(this.pool.getItem(), getListener());
 	}
 
 	@Override
@@ -168,9 +167,9 @@ public class CachingClientConnectionFactory extends AbstractClientConnectionFact
 
 		private volatile boolean released;
 
-		public CachedConnection(TcpConnectionSupport connection) {
+		public CachedConnection(TcpConnectionSupport connection, TcpListener tcpListener) {
 			super.setTheConnection(connection);
-			connection.registerListener(this);
+			registerListener(tcpListener);
 		}
 
 		@Override
@@ -225,17 +224,30 @@ public class CachingClientConnectionFactory extends AbstractClientConnectionFact
 		 */
 		@Override
 		public boolean onMessage(Message<?> message) {
-			AbstractIntegrationMessageBuilder<?> messageBuilder =
-					CachingClientConnectionFactory.this.getMessageBuilderFactory()
-							.fromMessage(message)
-							.setHeader(IpHeaders.CONNECTION_ID, getConnectionId());
-			if (message.getHeaders().get(IpHeaders.ACTUAL_CONNECTION_ID) == null) {
-				messageBuilder.setHeader(IpHeaders.ACTUAL_CONNECTION_ID,
-						message.getHeaders().get(IpHeaders.CONNECTION_ID));
+			Message<?> modifiedMessage;
+			if (message instanceof ErrorMessage) {
+				Map<String, Object> headers = new HashMap<String, Object>(message.getHeaders());
+				headers.put(IpHeaders.CONNECTION_ID, getConnectionId());
+				if (headers.get(IpHeaders.ACTUAL_CONNECTION_ID) == null) {
+					headers.put(IpHeaders.ACTUAL_CONNECTION_ID,
+							message.getHeaders().get(IpHeaders.CONNECTION_ID));
+				}
+				modifiedMessage = new ErrorMessage((Throwable) message.getPayload(), headers);
+			}
+			else {
+				AbstractIntegrationMessageBuilder<?> messageBuilder =
+						CachingClientConnectionFactory.this.getMessageBuilderFactory()
+								.fromMessage(message)
+								.setHeader(IpHeaders.CONNECTION_ID, getConnectionId());
+				if (message.getHeaders().get(IpHeaders.ACTUAL_CONNECTION_ID) == null) {
+					messageBuilder.setHeader(IpHeaders.ACTUAL_CONNECTION_ID,
+							message.getHeaders().get(IpHeaders.CONNECTION_ID));
+				}
+				modifiedMessage = messageBuilder.build();
 			}
 			TcpListener listener = getListener();
 			if (listener != null) {
-				listener.onMessage(messageBuilder.build());
+				listener.onMessage(modifiedMessage);
 			}
 			else {
 				if (logger.isDebugEnabled()) {
@@ -413,17 +425,7 @@ public class CachingClientConnectionFactory extends AbstractClientConnectionFact
 	@Override
 	public void registerListener(TcpListener listener) {
 		super.registerListener(listener);
-		this.targetConnectionFactory.registerListener(new TcpListener() {
-
-			@Override
-			public boolean onMessage(Message<?> message) {
-				if (!(message instanceof ErrorMessage)) {
-					throw new UnsupportedOperationException("This should never be called");
-				}
-				return false;
-			}
-
-		});
+		this.targetConnectionFactory.enableManualListenerRegistration();
 	}
 
 	@Override

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2014 the original author or authors.
+ * Copyright 2002-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.
@@ -93,6 +93,13 @@ public class FailoverClientConnectionFactory extends AbstractClientConnectionFac
 		}
 	}
 
+	@Override
+	public void enableManualListenerRegistration() {
+		for (AbstractClientConnectionFactory factory : this.factories) {
+			factory.enableManualListenerRegistration();
+		}
+	}
+
 	@Override
 	public void registerSender(TcpSender sender) {
 		for (AbstractClientConnectionFactory factory : this.factories) {

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2001-2014 the original author or authors.
+ * Copyright 2001-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.
@@ -22,6 +22,8 @@ import java.net.SocketException;
 import java.util.Collections;
 import java.util.Map;
 import java.util.UUID;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
 
@@ -34,6 +36,7 @@ import org.springframework.core.serializer.Serializer;
 import org.springframework.integration.ip.IpHeaders;
 import org.springframework.integration.ip.tcp.serializer.AbstractByteArraySerializer;
 import org.springframework.messaging.Message;
+import org.springframework.messaging.MessagingException;
 import org.springframework.messaging.support.ErrorMessage;
 import org.springframework.util.Assert;
 
@@ -50,6 +53,8 @@ public abstract class TcpConnectionSupport implements TcpConnection {
 
 	protected final Log logger = LogFactory.getLog(this.getClass());
 
+	private final CountDownLatch listenerRegisteredLatch = new CountDownLatch(1);
+
 	@SuppressWarnings("rawtypes")
 	private volatile Deserializer deserializer;
 
@@ -88,6 +93,8 @@ public abstract class TcpConnectionSupport implements TcpConnection {
 
 	private volatile boolean noReadErrorOnClose;
 
+	private volatile boolean manualListenerRegistration;
+
 	public TcpConnectionSupport() {
 		this(null);
 	}
@@ -171,22 +178,24 @@ public abstract class TcpConnectionSupport implements TcpConnection {
 	/**
 	 * If we have been intercepted, propagate the close from the outermost interceptor;
 	 * otherwise, just call close().
-	 * 
+	 *
 	 * @param isException true when this call is the result of an Exception.
 	 */
 	protected void closeConnection(boolean isException) {
-		if (!(this.listener instanceof TcpConnectionInterceptor)) {
+		TcpListener listener = getListener();
+		if (!(listener instanceof TcpConnectionInterceptor)) {
 			close();
-			return;
 		}
-		TcpConnectionInterceptor outerInterceptor = (TcpConnectionInterceptor) this.listener;
-		while (outerInterceptor.getListener() instanceof TcpConnectionInterceptor) {
-			outerInterceptor = (TcpConnectionInterceptor) outerInterceptor.getListener();
-		}
-		outerInterceptor.close();
-		if (isException) {
-			// ensure physical close in case the interceptor did not close
-			this.close();
+		else {
+			TcpConnectionInterceptor outerInterceptor = (TcpConnectionInterceptor) listener;
+			while (outerInterceptor.getListener() instanceof TcpConnectionInterceptor) {
+				outerInterceptor = (TcpConnectionInterceptor) outerInterceptor.getListener();
+			}
+			outerInterceptor.close();
+			if (isException) {
+				// ensure physical close in case the interceptor did not close
+				this.close();
+			}
 		}
 	}
 
@@ -245,7 +254,7 @@ public abstract class TcpConnectionSupport implements TcpConnection {
 	}
 
 	/**
-	 * Sets the listener that will receive incoming Messages.
+	 * Set the listener that will receive incoming Messages.
 	 * @param listener The listener.
 	 */
 	public void registerListener(TcpListener listener) {
@@ -253,13 +262,33 @@ public abstract class TcpConnectionSupport implements TcpConnection {
 		// Determine the actual listener for this connection
 		if (!(this.listener instanceof TcpConnectionInterceptor)) {
 			this.actualListener = this.listener;
-		} else {
+		}
+		else {
 			TcpConnectionInterceptor outerInterceptor = (TcpConnectionInterceptor) this.listener;
 			while (outerInterceptor.getListener() instanceof TcpConnectionInterceptor) {
 				outerInterceptor = (TcpConnectionInterceptor) outerInterceptor.getListener();
 			}
 			this.actualListener = outerInterceptor.getListener();
 		}
+		this.listenerRegisteredLatch.countDown();
+	}
+
+	/**
+	 * Set whether or not automatic or manual registration of the {@link TcpListener} is to be
+	 * used. (Default automatic). When manual registration is in place, incoming messages will
+	 * be delayed until the listener is registered.
+	 * @since 1.4.5
+	 */
+	public void enableManualListenerRegistration() {
+		this.manualListenerRegistration = true;
+		this.listener = new TcpListener() {
+
+			@Override
+			public boolean onMessage(Message<?> message) {
+				return getListener().onMessage(message);
+			}
+
+		};
 	}
 
 	/**
@@ -280,9 +309,23 @@ public abstract class TcpConnectionSupport implements TcpConnection {
 	 */
 	@Override
 	public TcpListener getListener() {
+		if (this.manualListenerRegistration) {
+			waitForListenerRegistration();
+		}
 		return this.listener;
 	}
 
+	private void waitForListenerRegistration() {
+		try {
+			Assert.state(listenerRegisteredLatch.await(1, TimeUnit.MINUTES), "TcpListener not registered");
+			manualListenerRegistration = false;
+		}
+		catch (InterruptedException e) {
+			Thread.currentThread().interrupt();
+			throw new MessagingException("Interrupted while waiting for listener registration", e);
+		}
+	}
+
 	/**
 	 * @return the sender
 	 */

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

@@ -35,6 +35,7 @@ import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.net.Socket;
@@ -49,6 +50,7 @@ import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 
 import org.apache.commons.logging.Log;
 import org.junit.Test;
@@ -681,7 +683,48 @@ public class CachingClientConnectionFactoryTests {
 		verify(logger, never()).error(anyString());
 	}
 
-	public TcpConnectionSupport makeMockConnection() {
+	@Test // INT-3728
+	public void testEarlyReceive() throws Exception {
+		final CountDownLatch latch = new CountDownLatch(1);
+		final AbstractClientConnectionFactory factory = new TcpNetClientConnectionFactory("", 0) {
+
+			@Override
+			protected Socket createSocket(String host, int port) throws IOException {
+				Socket mock = mock(Socket.class);
+				when(mock.getInputStream()).thenReturn(new ByteArrayInputStream("foo\r\n".getBytes()));
+				return mock;
+			}
+
+			@Override
+			public boolean isActive() {
+				return true;
+			}
+
+		};
+		factory.setApplicationEventPublisher(mock(ApplicationEventPublisher.class));
+		final CachingClientConnectionFactory cachingFactory = new CachingClientConnectionFactory(factory, 1);
+		final AtomicReference<Message<?>> received = new AtomicReference<Message<?>>();
+		cachingFactory.registerListener(new TcpListener() {
+
+			@Override
+			public boolean onMessage(Message<?> message) {
+				if (!(message instanceof ErrorMessage)) {
+					received.set(message);
+					latch.countDown();
+				}
+				return false;
+			}
+		});
+		cachingFactory.start();
+
+		cachingFactory.getConnection();
+		assertTrue(latch.await(10, TimeUnit.SECONDS));
+		assertNotNull(received.get());
+		assertNotNull(received.get().getHeaders().get(IpHeaders.ACTUAL_CONNECTION_ID));
+		cachingFactory.stop();
+	}
+
+	private TcpConnectionSupport makeMockConnection() {
 		TcpConnectionSupport connection = mock(TcpConnectionSupport.class);
 		when(connection.isOpen()).thenReturn(true);
 		return connection;