From 89705c4fe0f2cc5dccc9bedf58b46a903dc5eadf Mon Sep 17 00:00:00 2001
From: Gary Russell <grussell@vmware.com>
Date: Mon, 24 May 2010 14:36:09 +0000
Subject: [PATCH] INT-1145 Retry Write After Socket Close and Clean up Reader
 in o/b Gateway

---
 .../tcp/AbstractTcpSendingMessageHandler.java | 12 +++++-
 .../ip/tcp/SimpleTcpNetOutboundGateway.java   |  3 +-
 .../tcp/SimpleTcpNetOutboundGatewayTests.java | 41 +++++++++++++++++++
 3 files changed, 54 insertions(+), 2 deletions(-)

diff --git a/org.springframework.integration.ip/src/main/java/org/springframework/integration/ip/tcp/AbstractTcpSendingMessageHandler.java b/org.springframework.integration.ip/src/main/java/org/springframework/integration/ip/tcp/AbstractTcpSendingMessageHandler.java
index 532b3e0d1d..d7a5dddbcb 100644
--- a/org.springframework.integration.ip/src/main/java/org/springframework/integration/ip/tcp/AbstractTcpSendingMessageHandler.java
+++ b/org.springframework.integration.ip/src/main/java/org/springframework/integration/ip/tcp/AbstractTcpSendingMessageHandler.java
@@ -15,6 +15,7 @@
  */
 package org.springframework.integration.ip.tcp;
 
+import java.io.IOException;
 import java.net.Socket;
 import java.net.SocketException;
 
@@ -95,7 +96,16 @@ public abstract class AbstractTcpSendingMessageHandler extends
 	 */
 	public void handleMessage(final Message<?> message) throws MessageRejectedException,
 			MessageHandlingException, MessageDeliveryException {
-		doWrite(message);
+		try {
+			doWrite(message);
+		} catch (MessageMappingException e) {
+			// retry - socket may have closed
+			if (e.getCause() instanceof IOException) {
+				doWrite(message);
+			} else {
+				throw e;
+			}
+		}
 	}
 
 	/**
diff --git a/org.springframework.integration.ip/src/main/java/org/springframework/integration/ip/tcp/SimpleTcpNetOutboundGateway.java b/org.springframework.integration.ip/src/main/java/org/springframework/integration/ip/tcp/SimpleTcpNetOutboundGateway.java
index a29b5aa8aa..c2618a71dc 100644
--- a/org.springframework.integration.ip/src/main/java/org/springframework/integration/ip/tcp/SimpleTcpNetOutboundGateway.java
+++ b/org.springframework.integration.ip/src/main/java/org/springframework/integration/ip/tcp/SimpleTcpNetOutboundGateway.java
@@ -67,7 +67,8 @@ public class SimpleTcpNetOutboundGateway extends
 	@Override
 	protected synchronized Object handleRequestMessage(Message<?> requestMessage) {
 		this.handler.handleMessage(requestMessage);
-		if (this.reader == null) {
+		if (this.reader == null ||
+			this.reader.getSocket() != handler.getSocket()) {
 			Socket socket = this.handler.getSocket();
 			this.reader = SocketIoUtils.createNetReader(this.messageFormat,
 				this.customSocketReaderClass, socket, this.maxMessageSize,
diff --git a/org.springframework.integration.ip/src/test/java/org/springframework/integration/ip/tcp/SimpleTcpNetOutboundGatewayTests.java b/org.springframework.integration.ip/src/test/java/org/springframework/integration/ip/tcp/SimpleTcpNetOutboundGatewayTests.java
index b58ea651bd..fea8728eff 100644
--- a/org.springframework.integration.ip/src/test/java/org/springframework/integration/ip/tcp/SimpleTcpNetOutboundGatewayTests.java
+++ b/org.springframework.integration.ip/src/test/java/org/springframework/integration/ip/tcp/SimpleTcpNetOutboundGatewayTests.java
@@ -15,6 +15,13 @@
  */
 package org.springframework.integration.ip.tcp;
 
+import java.net.ServerSocket;
+import java.net.Socket;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
+import javax.net.ServerSocketFactory;
+
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -22,6 +29,7 @@ import org.springframework.beans.factory.annotation.Qualifier;
 import org.springframework.integration.channel.PollableChannel;
 import org.springframework.integration.channel.SubscribableChannel;
 import org.springframework.integration.core.Message;
+import org.springframework.integration.ip.util.SocketUtils;
 import org.springframework.integration.message.MessageBuilder;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
@@ -109,4 +117,37 @@ public class SimpleTcpNetOutboundGatewayTests {
 		byte[] bytes = (byte[]) replyChannel.receive().getPayload();
 		assertEquals("echo:test", new String(bytes).trim());
 	}
+	
+	@Test
+	public void testOutboundClose() throws Exception {
+		final int port = SocketUtils.findAvailableServerSocket();
+		final CountDownLatch latch = new CountDownLatch(1);
+		Thread t = new Thread(new Runnable() {
+			public void run() {
+				try {
+					ServerSocket ss = ServerSocketFactory.getDefault().createServerSocket(port);
+					latch.countDown();
+					while (true) {
+						Socket s = ss.accept();
+						byte[] b = new byte[1024];
+						s.getInputStream().read(b);
+						s.getOutputStream().write("OK\r\n".getBytes());
+						s.close();
+					}
+				} catch (Exception e) {
+					e.printStackTrace();
+				}
+			}});
+		t.start();
+		latch.await(2000, TimeUnit.MILLISECONDS);
+		SimpleTcpNetOutboundGateway gateway = new SimpleTcpNetOutboundGateway
+			("localhost", port);
+		gateway.setMessageFormat(MessageFormats.FORMAT_CRLF);
+		Message<String> message = MessageBuilder.withPayload("test").build();
+		byte[] bytes = (byte[]) gateway.handleRequestMessage(message);
+		assertEquals("OK", new String(bytes));
+		bytes = (byte[]) gateway.handleRequestMessage(message);
+		assertEquals("OK", new String(bytes));
+	}
+	
 }