INT-2519 Fix Memory Leak

When using NIO, a Map of connections is maintained (keyed
by the SocketChannel) and used to manage registrations
for NIO events, timeouts etc.

When a connection is closed, the corresponding entry
should be removed from the map. The code to do this
is in AbstractConnectionFactory.processNioSelections().
However, if no socket timeout (soTimeout) has been set,
or it was explicitly set to 0, connections are not
removed from the map. This was because the timeout
logic and map cleanup is done in the same iteration
loop.

Now, if soTimeout is not set, the clean up loop
operates every nioHarvestInterval milliseconds
(default 2000), we don't want to run it on
every selector event.

In addition, it runs if the selectionCount is
zero - this might occur when a socket
is closed, when the selector.wakeup() is called.

INT-2519 Polishing

Rename connections field in NIO connection factories. Name
collision caused TestUtils problem on some platforms.

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2011 the original author or authors.
+ * Copyright 2002-2012 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.
@@ -28,6 +28,7 @@ import java.util.Iterator;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Set;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
@@ -100,6 +101,12 @@ public abstract class AbstractConnectionFactory extends IntegrationObjectSupport
 
 	protected final Object lifecycleMonitor = new Object();
 
+	private volatile long nextCheckForClosedNioConnections;
+
+	private volatile int nioHarvestInterval = DEFAULT_NIO_HARVEST_INTERVAL;
+
+	private static final int DEFAULT_NIO_HARVEST_INTERVAL = 2000;
+
 	public AbstractConnectionFactory(int port) {
 		this.port = port;
 	}
@@ -383,6 +390,18 @@ public abstract class AbstractConnectionFactory extends IntegrationObjectSupport
 		return lookupHost;
 	}
 
+	/**
+	 * How often we clean up closed NIO connections if soTimeout is 0.
+	 * Ignored when soTimeout > 0 because the clean up
+	 * process is run as part of the timeout handling.
+	 * Default 2000 milliseconds.
+	 * @param nioHarvestInterval The interval in milliseconds.
+	 */
+	public void setNioHarvestInterval(int nioHarvestInterval) {
+		Assert.isTrue(nioHarvestInterval > 0, "NIO Harvest interval must be > 0");
+		this.nioHarvestInterval = nioHarvestInterval;
+	}
+
 	/**
 	 * Closes the server.
 	 */
@@ -490,24 +509,28 @@ public abstract class AbstractConnectionFactory extends IntegrationObjectSupport
 	/**
 	 *
 	 * Times out any expired connections then, if selectionCount > 0, processes the selected keys.
+	 * Removes closed connections from the connections field, and from the connections parameter.
 	 *
-	 * @param selectionCount
-	 * @param selector
-	 * @param connections
+	 * @param selectionCount Number of IO Events, if 0 we were probably woken up by a close.
+	 * @param selector The selector
+	 * @param connections Map of connections
 	 * @throws IOException
 	 */
 	protected void processNioSelections(int selectionCount, final Selector selector, ServerSocketChannel server,
 			Map<SocketChannel, TcpNioConnection> connections) throws IOException {
-		long now = 0;
-		if (this.soTimeout > 0) {
-			Iterator<SocketChannel> it = connections.keySet().iterator();
-			now = System.currentTimeMillis();
+		long now = System.currentTimeMillis();
+		if (this.soTimeout > 0 ||
+				now >= this.nextCheckForClosedNioConnections ||
+				selectionCount == 0) {
+			this.nextCheckForClosedNioConnections = now + this.nioHarvestInterval;
+			Iterator<Entry<SocketChannel, TcpNioConnection>> it = connections.entrySet().iterator();
 			while (it.hasNext()) {
-				SocketChannel channel = it.next();
+				SocketChannel channel = it.next().getKey();
 				if (!channel.isOpen()) {
 					logger.debug("Removing closed channel");
 					it.remove();
-				} else {
+				}
+				else if (soTimeout > 0) {
 					TcpNioConnection connection = connections.get(channel);
 					if (now - connection.getLastRead() > this.soTimeout) {
 						logger.warn("Timing out TcpNioConnection " +

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

@@ -44,7 +44,7 @@ public class TcpNioClientConnectionFactory extends
 	
 	private Selector selector;
 	
-	private Map<SocketChannel, TcpNioConnection> connections = new ConcurrentHashMap<SocketChannel, TcpNioConnection>();
+	private Map<SocketChannel, TcpNioConnection> channelMap = new ConcurrentHashMap<SocketChannel, TcpNioConnection>();
 	
 	private BlockingQueue<SocketChannel> newChannels = new LinkedBlockingQueue<SocketChannel>();
 
@@ -93,7 +93,7 @@ public class TcpNioClientConnectionFactory extends
 		if (this.getSoTimeout() > 0) {
 			connection.setLastRead(System.currentTimeMillis());
 		}
-		this.connections.put(socketChannel, connection);
+		this.channelMap.put(socketChannel, connection);
 		newChannels.add(socketChannel);
 		selector.wakeup();
 		return wrappedConnection;
@@ -134,14 +134,14 @@ public class TcpNioClientConnectionFactory extends
 				}
 				while ((newChannel = newChannels.poll()) != null) {
 					try {
-						newChannel.register(this.selector, SelectionKey.OP_READ, connections.get(newChannel));
+						newChannel.register(this.selector, SelectionKey.OP_READ, channelMap.get(newChannel));
 					} catch (ClosedChannelException cce) {
 						if (logger.isDebugEnabled()) {
 							logger.debug("Channel closed before registering with selector for reading");
 						}
 					}
 				}
-				this.processNioSelections(selectionCount, selector, null, this.connections);
+				this.processNioSelections(selectionCount, selector, null, this.channelMap);
 			}
 		} catch (Exception e) {
 			logger.error("Exception in read selector thread", e);
@@ -163,7 +163,7 @@ public class TcpNioClientConnectionFactory extends
 	 * @return the connections
 	 */
 	protected Map<SocketChannel, TcpNioConnection> getConnections() {
-		return connections;
+		return channelMap;
 	}
 
 	/**

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

@@ -44,7 +44,7 @@ public class TcpNioServerConnectionFactory extends AbstractServerConnectionFacto
 	
 	private boolean usingDirectBuffers;
 	
-	private Map<SocketChannel, TcpNioConnection> connections = new HashMap<SocketChannel, TcpNioConnection>();
+	private Map<SocketChannel, TcpNioConnection> channelMap = new HashMap<SocketChannel, TcpNioConnection>();
 
 	private Selector selector;
 	
@@ -123,7 +123,7 @@ public class TcpNioServerConnectionFactory extends AbstractServerConnectionFacto
 					logger.debug("CancelledKeyException during Selector.select()");
 				}
 			}
-			this.processNioSelections(selectionCount, selector, server, this.connections);			
+			this.processNioSelections(selectionCount, selector, server, this.channelMap);
 		}
 	}
 
@@ -148,7 +148,7 @@ public class TcpNioServerConnectionFactory extends AbstractServerConnectionFacto
 		}
 		connection.setTaskExecutor(this.getTaskExecutor());
 		connection.setLastRead(now);
-		connections.put(channel, connection);
+		channelMap.put(channel, connection);
 		channel.register(selector, SelectionKey.OP_READ, connection);
 	}
 
@@ -200,7 +200,7 @@ public class TcpNioServerConnectionFactory extends AbstractServerConnectionFacto
 	 * @return the connections
 	 */
 	protected Map<SocketChannel, TcpNioConnection> getConnections() {
-		return connections;
+		return channelMap;
 	}
 	
 

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2010 the original author or authors.
+ * Copyright 2002-2012 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.
@@ -16,14 +16,26 @@
 
 package org.springframework.integration.ip.tcp.connection;
 
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.lang.reflect.Field;
 import java.net.ServerSocket;
 import java.net.Socket;
 import java.net.SocketTimeoutException;
+import java.nio.channels.SelectionKey;
+import java.nio.channels.Selector;
+import java.nio.channels.SocketChannel;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
@@ -33,6 +45,10 @@ import javax.net.ServerSocketFactory;
 import org.junit.Test;
 import org.springframework.integration.ip.util.SocketTestUtils;
 import org.springframework.integration.support.MessageBuilder;
+import org.springframework.integration.test.util.TestUtils;
+import org.springframework.util.ReflectionUtils;
+import org.springframework.util.ReflectionUtils.FieldCallback;
+import org.springframework.util.ReflectionUtils.FieldFilter;
 
 
 /**
@@ -110,9 +126,112 @@ public class TcpNioConnectionTests {
 		} catch (Exception e) {
 			fail("Unexpected exception " + e);
 		}
-
 	}
 	
+	@Test
+	public void testMemoryLeak() throws Exception {
+		final int port = SocketTestUtils.findAvailableServerSocket();
+		TcpNioClientConnectionFactory factory = new TcpNioClientConnectionFactory("localhost", port);
+		factory.setNioHarvestInterval(100);
+		factory.start();
+		final CountDownLatch latch = new CountDownLatch(1);
+		Executors.newSingleThreadExecutor().execute(new Runnable() {
+			public void run() {
+				try {
+					ServerSocket server = ServerSocketFactory.getDefault().createServerSocket(port);
+					latch.countDown();
+					Socket socket = server.accept();
+					byte[] b = new byte[6];
+					readFully(socket.getInputStream(), b);
+				} catch (Exception e) {
+					e.printStackTrace();
+				}
+			}
+		});
+		assertTrue(latch.await(10000, TimeUnit.MILLISECONDS));
+		try {
+			TcpConnection connection = factory.getConnection();
+			Map<SocketChannel, TcpNioConnection> connections = factory.getConnections();
+			assertEquals(1, connections.size());
+			connection.close();
+			assertTrue(!connection.isOpen());
+			int n = 0;
+			while (connections.size() > 0) {
+				Thread.sleep(100);
+				if (n++ > 100) {
+					break;
+				}
+			}
+			assertEquals(0, connections.size());
+		} catch (Exception e) {
+			e.printStackTrace();
+			fail("Unexpected exception " + e);
+		}
+		factory.stop();
+	}
+
+	@Test
+	public void testCleanup() throws Exception {
+		TcpNioClientConnectionFactory factory = new TcpNioClientConnectionFactory("localhost", 0);
+		factory.setNioHarvestInterval(100);
+		Map<SocketChannel, TcpNioConnection> connections = new HashMap<SocketChannel, TcpNioConnection>();
+		SocketChannel chan1 = mock(SocketChannel.class);
+		SocketChannel chan2 = mock(SocketChannel.class);
+		SocketChannel chan3 = mock(SocketChannel.class);
+		TcpNioConnection conn1 = mock(TcpNioConnection.class);
+		TcpNioConnection conn2 = mock(TcpNioConnection.class);
+		TcpNioConnection conn3 = mock(TcpNioConnection.class);
+		connections.put(chan1, conn1);
+		connections.put(chan2, conn2);
+		connections.put(chan3, conn3);
+		final List<Field> fields = new ArrayList<Field>();
+		ReflectionUtils.doWithFields(SocketChannel.class, new FieldCallback() {
+
+			public void doWith(Field field) throws IllegalArgumentException,
+					IllegalAccessException {
+				field.setAccessible(true);
+				fields.add(field);
+			}
+		}, new FieldFilter() {
+
+			public boolean matches(Field field) {
+				return field.getName().equals("open");
+			}});
+		Field field = fields.get(0);
+		// Can't use Mockito because isOpen() is final
+		ReflectionUtils.setField(field, chan1, true);
+		ReflectionUtils.setField(field, chan2, true);
+		ReflectionUtils.setField(field, chan3, true);
+		Selector selector = mock(Selector.class);
+		HashSet<SelectionKey> keys = new HashSet<SelectionKey>();
+		when(selector.selectedKeys()).thenReturn(keys);
+		factory.processNioSelections(1, selector, null, connections);
+		assertEquals(3, connections.size()); // all open
+
+		ReflectionUtils.setField(field, chan1, false);
+		factory.processNioSelections(1, selector, null, connections);
+		assertEquals(3, connections.size()); // interval didn't pass
+		Thread.sleep(110);
+		factory.processNioSelections(1, selector, null, connections);
+		assertEquals(2, connections.size()); // first is closed
+
+		ReflectionUtils.setField(field, chan2, false);
+		factory.processNioSelections(1, selector, null, connections);
+		assertEquals(2, connections.size()); // interval didn't pass
+		Thread.sleep(110);
+		factory.processNioSelections(1, selector, null, connections);
+		assertEquals(1, connections.size()); // second is closed
+
+		ReflectionUtils.setField(field, chan3, false);
+		factory.processNioSelections(1, selector, null, connections);
+		assertEquals(1, connections.size()); // interval didn't pass
+		Thread.sleep(110);
+		factory.processNioSelections(1, selector, null, connections);
+		assertEquals(0, connections.size()); // third is closed
+
+		assertEquals(0, TestUtils.getPropertyValue(factory, "connections", List.class).size());
+	}
+
 	private void readFully(InputStream is, byte[] buff) throws IOException {
 		for (int i = 0; i < buff.length; i++) {
 			buff[i] = (byte) is.read();