GH-3256: Support Testing TCP Connections

Resolves #3256

Author: garyrussell

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

@@ -16,10 +16,13 @@
 
 package org.springframework.integration.ip.tcp.connection;
 
+import java.io.IOException;
+import java.io.UncheckedIOException;
 import java.net.Socket;
 import java.time.Duration;
 import java.util.concurrent.locks.ReadWriteLock;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.function.Predicate;
 
 import org.springframework.context.ApplicationEventPublisher;
 import org.springframework.lang.Nullable;
@@ -44,6 +47,9 @@ public abstract class AbstractClientConnectionFactory extends AbstractConnection
 
 	private Duration connectTimeout = Duration.ofSeconds(DEFAULT_CONNECT_TIMEOUT);
 
+	@Nullable
+	private Predicate<TcpConnectionSupport> connectionTest;
+
 	private volatile TcpConnectionSupport theConnection;
 
 	/**
@@ -68,6 +74,7 @@ protected Duration getConnectTimeout() {
 		return this.connectTimeout;
 	}
 
+
 	/**
 	 * Set whether to automatically (default) or manually add a {@link TcpListener} to the
 	 * connections created by this factory. By default, the factory automatically configures
@@ -79,6 +86,27 @@ public void enableManualListenerRegistration() {
 		this.manualListenerRegistration = true;
 	}
 
+	/**
+	 * Get a {@link Predicate} that will be invoked to test a new connection; return true
+	 * to accept the connection, false the reject.
+	 * @return the predicate.
+	 * @since 5.3
+	 */
+	@Nullable
+	protected Predicate<TcpConnectionSupport> getConnectionTest() {
+		return this.connectionTest;
+	}
+
+	/**
+	 * Set a {@link Predicate} that will be invoked to test a new connection; return true
+	 * to accept the connection, false the reject.
+	 * @param connectionTest the predicate.
+	 * @since 5.3
+	 */
+	public void setConnectionTest(@Nullable Predicate<TcpConnectionSupport> connectionTest) {
+		this.connectionTest = connectionTest;
+	}
+
 	/**
 	 * Obtains a connection - if {@link #setSingleUse(boolean)} was called with
 	 * true, a new connection is returned; otherwise a single connection is
@@ -126,21 +154,12 @@ protected final TcpConnectionSupport obtainNewConnection() throws InterruptedExc
 			if (!singleUse) {
 				// Another write lock holder might have created a new one by now.
 				connection = obtainSharedConnection();
-				if (connection != null) {
+				if (connection != null && connection.isOpen()) {
 					return connection;
 				}
 			}
 
-			if (logger.isDebugEnabled()) {
-				logger.debug("Opening new socket connection to " + getHost() + ":" + getPort());
-			}
-
-			connection = buildNewConnection();
-			if (!singleUse) {
-				setTheConnection(connection);
-			}
-			connection.publishConnectionOpenEvent();
-			return connection;
+			return doObtain(singleUse);
 		}
 		catch (RuntimeException e) {
 			ApplicationEventPublisher applicationEventPublisher = getApplicationEventPublisher();
@@ -156,6 +175,24 @@ protected final TcpConnectionSupport obtainNewConnection() throws InterruptedExc
 		}
 	}
 
+	private TcpConnectionSupport doObtain(boolean singleUse) {
+		TcpConnectionSupport connection;
+		if (logger.isDebugEnabled()) {
+			logger.debug("Opening new socket connection to " + getHost() + ":" + getPort());
+		}
+
+		connection = buildNewConnection();
+		if (this.connectionTest != null && !this.connectionTest.test(connection)) {
+			connection.close();
+			throw new UncheckedIOException(new IOException("Connection test failed for " + connection));
+		}
+		if (!singleUse) {
+			setTheConnection(connection);
+		}
+		connection.publishConnectionOpenEvent();
+		return connection;
+	}
+
 	protected TcpConnectionSupport buildNewConnection() {
 		throw new UnsupportedOperationException(
 				"Factories that don't override this class' obtainConnection() must implement this method");
@@ -187,6 +224,9 @@ protected void initializeConnection(TcpConnectionSupport connection, Socket sock
 		connection.setMapper(getMapper());
 		connection.setDeserializer(getDeserializer());
 		connection.setSerializer(getSerializer());
+		if (this.connectionTest != null) {
+			connection.setNeedsTest(true);
+		}
 	}
 
 	/**

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

@@ -78,6 +78,8 @@ public abstract class TcpConnectionSupport implements TcpConnection {
 
 	private TcpListener listener;
 
+	private volatile TcpListener testListener;
+
 	private TcpSender sender;
 
 	private String connectionId;
@@ -92,6 +94,13 @@ public abstract class TcpConnectionSupport implements TcpConnection {
 
 	private boolean manualListenerRegistration;
 
+	/*
+	 * This boolean is to avoid looking for a temporary listener when not needed
+	 * to avoid a CPU cache flush. This does not have to be volatile because it
+	 * is reset by the thread that checks for the temporary listener.
+	 */
+	private boolean needsTest;
+
 	public TcpConnectionSupport() {
 		this(null);
 	}
@@ -238,6 +247,15 @@ public void setSerializer(Serializer<?> serializer) {
 		}
 	}
 
+	/**
+	 * Set to true to use a temporary listener for just the first incoming message.
+	 * @param needsTest true for a temporary listener.
+	 * @since 5.3
+	 */
+	public void setNeedsTest(boolean needsTest) {
+		this.needsTest = needsTest;
+	}
+
 	/**
 	 * Set the listener that will receive incoming Messages.
 	 * @param listener The listener.
@@ -247,6 +265,17 @@ public void registerListener(@Nullable TcpListener listener) {
 		this.listenerRegisteredLatch.countDown();
 	}
 
+	/**
+	 * Set a temporary listener to receive just the first incoming message.
+	 * Used in conjunction with a connectionTest in a client connection
+	 * factory.
+	 * @param tListener the test listener.
+	 * @since 5.3
+	 */
+	public void registerTestListener(TcpListener tListener) {
+		this.testListener = tListener;
+	}
+
 	/**
 	 * 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
@@ -282,6 +311,10 @@ public TcpListener getListener() {
 			}
 			waitForListenerRegistration();
 		}
+		if (this.needsTest && this.testListener != null) {
+			this.needsTest = false;
+			return this.testListener;
+		}
 		return this.listener;
 	}
 

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2019 the original author or authors.
+ * Copyright 2002-2020 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.
@@ -23,6 +23,7 @@
  * @since 2.0
  *
  */
+@FunctionalInterface
 public interface TcpSender {
 
 	/**
@@ -37,6 +38,7 @@ public interface TcpSender {
 	 * method is called each time a connection is closed.
 	 * @param connection The connection.
 	 */
-	void removeDeadConnection(TcpConnection connection);
+	default void removeDeadConnection(TcpConnection connection) {
+	}
 
 }