Reactor2TcpClient cleans up TcpClient instances

Issue: SPR-14231

Author: rstoyanchev

spring-messaging/src/main/java/org/springframework/messaging/tcp/reactor/Reactor2TcpClient.java

@@ -163,7 +163,8 @@ private static NioEventLoopGroup initEventLoopGroup() {
 	public ListenableFuture<Void> connect(final TcpConnectionHandler<P> connectionHandler) {
 		Assert.notNull(connectionHandler, "TcpConnectionHandler must not be null");
 
-		TcpClient<Message<P>, Message<P>> tcpClient;
+		final TcpClient<Message<P>, Message<P>> tcpClient;
+		Runnable cleanupTask;
 		synchronized (this.tcpClients) {
 			if (this.stopping) {
 				IllegalStateException ex = new IllegalStateException("Shutting down.");
@@ -172,9 +173,18 @@ public ListenableFuture<Void> connect(final TcpConnectionHandler<P> connectionHa
 			}
 			tcpClient = NetStreams.tcpClient(REACTOR_TCP_CLIENT_TYPE, this.tcpClientSpecFactory);
 			this.tcpClients.add(tcpClient);
+			cleanupTask = new Runnable() {
+				@Override
+				public void run() {
+					synchronized (tcpClients) {
+						tcpClients.remove(tcpClient);
+					}
+				}
+			};
 		}
 
-		Promise<Void> promise = tcpClient.start(new MessageChannelStreamHandler<P>(connectionHandler));
+		Promise<Void> promise = tcpClient.start(
+				new MessageChannelStreamHandler<P>(connectionHandler, cleanupTask));
 
 		return new PassThroughPromiseToListenableFutureAdapter<Void>(
 				promise.onError(new Consumer<Throwable>() {
@@ -191,7 +201,8 @@ public ListenableFuture<Void> connect(TcpConnectionHandler<P> connectionHandler,
 		Assert.notNull(connectionHandler, "TcpConnectionHandler must not be null");
 		Assert.notNull(strategy, "ReconnectStrategy must not be null");
 
-		TcpClient<Message<P>, Message<P>> tcpClient;
+		final TcpClient<Message<P>, Message<P>> tcpClient;
+		Runnable cleanupTask;
 		synchronized (this.tcpClients) {
 			if (this.stopping) {
 				IllegalStateException ex = new IllegalStateException("Shutting down.");
@@ -200,10 +211,18 @@ public ListenableFuture<Void> connect(TcpConnectionHandler<P> connectionHandler,
 			}
 			tcpClient = NetStreams.tcpClient(REACTOR_TCP_CLIENT_TYPE, this.tcpClientSpecFactory);
 			this.tcpClients.add(tcpClient);
+			cleanupTask = new Runnable() {
+				@Override
+				public void run() {
+					synchronized (tcpClients) {
+						tcpClients.remove(tcpClient);
+					}
+				}
+			};
 		}
 
 		Stream<Tuple2<InetSocketAddress, Integer>> stream = tcpClient.start(
-				new MessageChannelStreamHandler<P>(connectionHandler),
+				new MessageChannelStreamHandler<P>(connectionHandler, cleanupTask),
 				new ReactorReconnectAdapter(strategy));
 
 		return new PassThroughPromiseToListenableFutureAdapter<Void>(stream.next().after());
@@ -249,6 +268,7 @@ public void operationComplete(Future<Object> future) throws Exception {
 			});
 			promise = eventLoopPromise;
 		}
+
 		return new PassThroughPromiseToListenableFutureAdapter<Void>(promise);
 	}
 
@@ -278,8 +298,11 @@ private static class MessageChannelStreamHandler<P>
 
 		private final TcpConnectionHandler<P> connectionHandler;
 
-		public MessageChannelStreamHandler(TcpConnectionHandler<P> connectionHandler) {
+		private final Runnable cleanupTask;
+
+		public MessageChannelStreamHandler(TcpConnectionHandler<P> connectionHandler, Runnable cleanupTask) {
 			this.connectionHandler = connectionHandler;
+			this.cleanupTask = cleanupTask;
 		}
 
 		@Override
@@ -290,6 +313,7 @@ public Publisher<Void> apply(ChannelStream<Message<P>, Message<P>> channelStream
 					.finallyDo(new Consumer<Signal<Message<P>>>() {
 						@Override
 						public void accept(Signal<Message<P>> signal) {
+							cleanupTask.run();
 							if (signal.isOnError()) {
 								connectionHandler.handleFailure(signal.getThrowable());
 							}