GH-3993: Fix async race condition in TcpOutGateway (#3995)

* GH-3993: Fix async race condition in TcpOutGateway

Fixes #3993

When `TcpOutboundGateway` is in an `async` mode and `CCF` is configured
not for `singleUse` an `Semaphore` around an obtained `TcpConnection` is involved.
If we fail on `TcpConnection.send()`, resources are not clean up, including the
mentioned `Semaphore`: in async mode this happens only when we receive a reply.

* Catch an exception on the `TcpConnection.send()` and perform `cleanUp()` in async mode.
* Add `cleanUp()` into a scheduled task from the `TcpOutboundGateway.AsyncReply`
when no reply arrives in time.
* Optimize `TcpOutboundGateway.AsyncReply` behavior to cancel no-reply scheduled task
when reply arrives into a `CompletableFuture`

**Cherry-pick to `5.5.x`**

* * Call `cleanUp()` from no response scheduled task
only if `future.completeExceptionally()` is `true`

Author: artembilan

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2001-2021 the original author or authors.
+ * Copyright 2001-2023 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.
@@ -17,10 +17,12 @@
 package org.springframework.integration.ip.tcp;
 
 import java.io.IOException;
-import java.util.Date;
+import java.time.Instant;
 import java.util.Map;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
 
@@ -59,7 +61,6 @@
  * <p>
  * {@link org.springframework.context.Lifecycle} methods delegate to the underlying {@link AbstractConnectionFactory}.
  *
- *
  * @author Gary Russell
  * @author Artem Bilan
  *
@@ -223,7 +224,17 @@ protected Object handleRequestMessage(Message<?> requestMessage) {
 			this.pendingReplies.put(connectionId, reply);
 			String connectionIdToLog = connectionId;
 			logger.debug(() -> "Added pending reply " + connectionIdToLog);
-			connection.send(requestMessage);
+			try {
+				connection.send(requestMessage);
+			}
+			catch (Exception ex) {
+				// If it cannot send, then no reply for this connection.
+				// Therefor release resources for subsequent requests.
+				if (async) {
+					cleanUp(haveSemaphore, connection, connectionId);
+				}
+				throw ex;
+			}
 			if (this.closeStreamAfterSend) {
 				connection.shutdownOutput();
 			}
@@ -326,7 +337,7 @@ public boolean onMessage(Message<?> message) {
 		if (reply == null) {
 			if (message instanceof ErrorMessage) {
 				/*
-				 * Socket errors are sent here so they can be conveyed to any waiting thread.
+				 * Socket errors are sent here, so they can be conveyed to any waiting thread.
 				 * If there's not one, simply ignore.
 				 */
 				return false;
@@ -427,7 +438,11 @@ private final class AsyncReply {
 
 		private final boolean haveSemaphore;
 
-		private final SettableListenableFuture<Message<?>> future = new SettableListenableFuture<>();
+		private final ScheduledFuture<?> noResponseFuture;
+
+		private final CompletableFuture<Message<?>> future =
+				new CompletableFuture<Message<?>>()
+						.thenApply(this::cancelNoResponseFutureIfAny);
 
 		private volatile Message<?> reply;
 
@@ -440,12 +455,27 @@ private final class AsyncReply {
 			this.connection = connection;
 			this.haveSemaphore = haveSemaphore;
 			if (async && remoteTimeout > 0) {
-				getTaskScheduler().schedule(() -> {
-					TcpOutboundGateway.this.pendingReplies.remove(connection.getConnectionId());
-					this.future.setException(
-							new MessageTimeoutException(requestMessage, "Timed out waiting for response"));
-				}, new Date(System.currentTimeMillis() + remoteTimeout));
+				this.noResponseFuture =
+						getTaskScheduler()
+								.schedule(() -> {
+									if (this.future.completeExceptionally(
+											new MessageTimeoutException(requestMessage,
+													"Timed out waiting for response"))) {
+
+										cleanUp(this.haveSemaphore, this.connection, this.connection.getConnectionId());
+									}
+								}, Instant.now().plusMillis(remoteTimeout));
+			}
+			else {
+				this.noResponseFuture = null;
+			}
+		}
+
+		private Message<?> cancelNoResponseFutureIfAny(Message<?> message) {
+			if (this.noResponseFuture != null) {
+				this.noResponseFuture.cancel(true);
 			}
+			return message;
 		}
 
 		TcpConnection getConnection() {

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

@@ -61,6 +61,7 @@
 import org.springframework.integration.ip.tcp.connection.AbstractClientConnectionFactory;
 import org.springframework.integration.ip.tcp.connection.CachingClientConnectionFactory;
 import org.springframework.integration.ip.tcp.connection.FailoverClientConnectionFactory;
+import org.springframework.integration.ip.tcp.connection.TcpConnection;
 import org.springframework.integration.ip.tcp.connection.TcpConnectionSupport;
 import org.springframework.integration.ip.tcp.connection.TcpNetClientConnectionFactory;
 import org.springframework.integration.ip.tcp.connection.TcpNioClientConnectionFactory;
@@ -80,9 +81,14 @@
 import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
 import static org.assertj.core.api.Assertions.catchThrowable;
 import static org.assertj.core.api.Assertions.fail;
 import static org.awaitility.Awaitility.await;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.BDDMockito.given;
+import static org.mockito.BDDMockito.willReturn;
+import static org.mockito.BDDMockito.willThrow;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
@@ -397,7 +403,7 @@ private void testGoodNetGWTimeoutGuts(AbstractClientConnectionFactory ccf,
 
 		Expression remoteTimeoutExpression = Mockito.mock(Expression.class);
 
-		when(remoteTimeoutExpression.getValue(Mockito.any(EvaluationContext.class), Mockito.any(Message.class),
+		when(remoteTimeoutExpression.getValue(any(EvaluationContext.class), any(Message.class),
 				Mockito.eq(Long.class))).thenReturn(50L, 60000L);
 
 		gateway.setRemoteTimeoutExpression(remoteTimeoutExpression);
@@ -489,7 +495,7 @@ void testCachingFailover() throws Exception {
 		TcpConnectionSupport mockConn1 = makeMockConnection();
 		when(factory1.getConnection()).thenReturn(mockConn1);
 		doThrow(new UncheckedIOException(new IOException("fail")))
-				.when(mockConn1).send(Mockito.any(Message.class));
+				.when(mockConn1).send(any(Message.class));
 
 		AbstractClientConnectionFactory factory2 = new TcpNetClientConnectionFactory("localhost",
 				serverSocket.get().getLocalPort());
@@ -522,7 +528,7 @@ void testCachingFailover() throws Exception {
 		assertThat(reply.getPayload()).isEqualTo("bar");
 		done.set(true);
 		gateway.stop();
-		verify(mockConn1).send(Mockito.any(Message.class));
+		verify(mockConn1).send(any(Message.class));
 		factory2.stop();
 		serverSocket.get().close();
 	}
@@ -572,7 +578,7 @@ void testFailoverCached() throws Exception {
 		when(factory1.getConnection()).thenReturn(mockConn1);
 		when(factory1.isSingleUse()).thenReturn(true);
 		doThrow(new UncheckedIOException(new IOException("fail")))
-				.when(mockConn1).send(Mockito.any(Message.class));
+				.when(mockConn1).send(any(Message.class));
 		CachingClientConnectionFactory cachingFactory1 = new CachingClientConnectionFactory(factory1, 1);
 
 		AbstractClientConnectionFactory factory2 = new TcpNetClientConnectionFactory("localhost",
@@ -607,7 +613,7 @@ void testFailoverCached() throws Exception {
 		assertThat(reply.getPayload()).isEqualTo("bar");
 		done.set(true);
 		gateway.stop();
-		verify(mockConn1).send(Mockito.any(Message.class));
+		verify(mockConn1).send(any(Message.class));
 		factory2.stop();
 		serverSocket.get().close();
 	}
@@ -1080,4 +1086,37 @@ void testAsyncTimeout() throws Exception {
 		}
 	}
 
+	@Test
+	void semaphoreIsReleasedOnAsyncSendFailure() throws InterruptedException {
+		AbstractClientConnectionFactory ccf = mock(AbstractClientConnectionFactory.class);
+
+		TcpConnection connection = mock(TcpConnectionSupport.class);
+
+		given(connection.getConnectionId()).willReturn("testId");
+		willThrow(new RuntimeException("intentional"))
+				.given(connection)
+						.send(any(Message.class));
+
+		willReturn(connection)
+				.given(ccf)
+				.getConnection();
+
+		TcpOutboundGateway gateway = new TcpOutboundGateway();
+		gateway.setConnectionFactory(ccf);
+		gateway.setAsync(true);
+		gateway.setBeanFactory(mock(BeanFactory.class));
+		gateway.setRemoteTimeout(-1);
+		gateway.afterPropertiesSet();
+
+		assertThatExceptionOfType(MessageHandlingException.class)
+				.isThrownBy(() -> gateway.handleMessage(new GenericMessage<>("Test1")))
+				.withCauseExactlyInstanceOf(RuntimeException.class)
+				.withStackTraceContaining("intentional");
+
+		assertThatExceptionOfType(MessageHandlingException.class)
+				.isThrownBy(() -> gateway.handleMessage(new GenericMessage<>("Test2")))
+				.withCauseExactlyInstanceOf(RuntimeException.class)
+				.withStackTraceContaining("intentional");
+	}
+
 }