Fix OutofMemory bug when the client can't connect to the server.

Author: mengli

src/main/com/schooner/MemCached/AscIIClient.java

@@ -442,7 +442,7 @@ private boolean set(String cmdname, String key, Object value, Date expiry, Integ
 				 */
 				return true;
 			}
-		} catch (IOException e) {
+		} catch (Exception e) {
 			// if we have an errorHandler, use its hook
 			if (errorHandler != null)
 				errorHandler.handleErrorOnSet(this, e, key);
@@ -599,7 +599,7 @@ private long incrdecr(String cmdname, String key, long inc, Integer hashCode) {
 				log.error(new StringBuffer().append("++++ error incr/decr key: ").append(key).toString());
 				log.error(new StringBuffer().append("++++ server response: ").append(line).toString());
 			}
-		} catch (IOException e) {
+		} catch (Exception e) {
 
 			// if we have an errorHandler, use its hook
 			if (errorHandler != null)
@@ -789,7 +789,7 @@ private Object get(String cmd, String key, Integer hashCode, boolean asString) {
 			}
 			sock.readBuf.clear();
 			return o;
-		} catch (IOException ce) {
+		} catch (Exception ce) {
 			// if we have an errorHandler, use its hook
 			if (errorHandler != null)
 				errorHandler.handleErrorOnGet(this, ce, key);
@@ -930,7 +930,7 @@ public MemcachedItem gets(String cmd, String key, Integer hashCode, boolean asSt
 			item.value = o;
 			return item;
 
-		} catch (IOException ce) {
+		} catch (Exception ce) {
 			// if we have an errorHandler, use its hook
 			if (errorHandler != null)
 				errorHandler.handleErrorOnGet(this, ce, key);
@@ -1417,7 +1417,7 @@ private Map<String, Map<String, String>> stats(String[] servers, String command,
 
 					statsMaps.put(servers[i], stats);
 				}
-			} catch (IOException e) {
+			} catch (Exception e) {
 
 				// if we have an errorHandler, use its hook
 				if (errorHandler != null)

src/main/com/schooner/MemCached/SchoonerSockIO.java

@@ -32,6 +32,8 @@
 import java.net.UnknownHostException;
 import java.nio.ByteBuffer;
 import java.nio.channels.ByteChannel;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import com.danga.MemCached.SockIOPool;
 
@@ -48,7 +50,6 @@ public abstract class SchoonerSockIO extends SockIOPool.SockIO {
 	public SchoonerSockIO(int bufferSize) throws UnknownHostException, IOException {
 		super(null, null, 0, 0, false);
 		this.bufferSize = bufferSize;
-		writeBuf = ByteBuffer.allocateDirect(bufferSize);
 	}
 
 	private int bufferSize = 1024 * 1025;
@@ -58,6 +59,10 @@ public SchoonerSockIO(int bufferSize) throws UnknownHostException, IOException {
 	public ByteBuffer writeBuf;
 
 	protected boolean isPooled = true;
+	
+	protected ConcurrentLinkedQueue<SchoonerSockIO> sockets;
+
+	protected AtomicInteger sockNum;
 
 	public abstract short preWrite();
 
@@ -86,4 +91,13 @@ public void setPooled(boolean isPooled) {
 	public boolean isPooled() {
 		return isPooled;
 	}
+	
+	public AtomicInteger getSockNum() {
+		return sockNum;
+	}
+
+	public void setSockNum(AtomicInteger sockNum) {
+		this.sockNum = sockNum;
+	}
+	
 }

src/main/com/schooner/MemCached/SchoonerSockIOPool.java

@@ -304,10 +304,10 @@ private void populateBuckets() {
 
 			// Create a socket pool for each host
 			socketPool.put(servers[i], new ConcurrentLinkedQueue<SchoonerSockIO>());
-			poolCurrentConn.put(servers[i], new AtomicInteger(initConn));
 
 			// Create the initial connections
-			for (int j = 0; j < initConn; j++) {
+			int j;
+			for (j = 0; j < initConn; j++) {
 				SchoonerSockIO socket = createSocket(servers[i]);
 				if (socket == null) {
 					log.error("++++ failed to create connection to: " + servers[i] + " -- only " + j + " created.");
@@ -317,6 +317,13 @@ private void populateBuckets() {
 				// Add this new connection to socket pool
 				addSocketToPool(servers[i], socket);
 			}
+			
+			ConcurrentLinkedQueue<SchoonerSockIO> sockets = socketPool.get(servers[i]);
+			AtomicInteger num = new AtomicInteger(j);
+			for (SchoonerSockIO schoonerSockIO : sockets) {
+				schoonerSockIO.setSockNum(num);
+			}
+			poolCurrentConn.put(servers[i], num);
 		}
 	}
 
@@ -351,9 +358,9 @@ private void populateConsistentBuckets() {
 
 			// Create a socket pool for each host
 			socketPool.put(servers[i], new ConcurrentLinkedQueue<SchoonerSockIO>());
-			poolCurrentConn.put(servers[i], new AtomicInteger(initConn));
 
-			for (int j = 0; j < initConn; j++) {
+			int j;
+			for (j = 0; j < initConn; j++) {
 				SchoonerSockIO socket = createSocket(servers[i]);
 				if (socket == null) {
 					log.error("++++ failed to create connection to: " + servers[i] + " -- only " + j + " created.");
@@ -363,6 +370,13 @@ private void populateConsistentBuckets() {
 				// Add this new connection to socket pool
 				addSocketToPool(servers[i], socket);
 			}
+			
+			ConcurrentLinkedQueue<SchoonerSockIO> sockets = socketPool.get(servers[i]);
+			AtomicInteger num = new AtomicInteger(j);
+			for (SchoonerSockIO schoonerSockIO : sockets) {
+				schoonerSockIO.setSockNum(num);
+			}
+			poolCurrentConn.put(servers[i], num);
 		}
 	}
 
@@ -409,6 +423,7 @@ protected final SchoonerSockIO createSocketWithAdd(String host) {
 			log.error("++++ failed to get SockIO obj for: " + host);
 			// log.error(ex.getMessage(), ex);
 			socket = null;
+			poolCurrentConn.get(host).decrementAndGet();
 		}
 
 		return socket;
@@ -539,6 +554,7 @@ public final SchoonerSockIO getConnection(String host) {
 		// if we have items in the pool then we can return it
 		ConcurrentLinkedQueue<SchoonerSockIO> sockets = socketPool.get(host);
 		SchoonerSockIO socket = sockets.poll();
+		System.out.println(poolCurrentConn.get(host).get());
 		if (socket == null) {
 			if (poolCurrentConn.get(host).get() < maxConn) {
 				socket = createSocketWithAdd(host);
@@ -1049,11 +1065,6 @@ public int getBufferSize() {
 
 	public static class UDPSockIO extends SchoonerSockIO {
 
-		// logger
-		// private static Logger log =
-		// Logger.getLogger(UDPSockIO.class.getName());
-		private ConcurrentLinkedQueue<SchoonerSockIO> sockets;
-
 		/**
 		 * 
 		 * <p>
@@ -1130,8 +1141,6 @@ public synchronized void addLength(int alength) {
 
 		private Selector selector;
 
-		private AtomicInteger sockNum;
-
 		@Override
 		public void trueClose() throws IOException {
 			if (selector != null) {
@@ -1154,8 +1163,9 @@ public UDPSockIO(SchoonerSockIOPool pool, String host, int bufferSize, int timeo
 			channel.socket().setSoTimeout(timeout);
 			selector = Selector.open();
 			((DatagramChannel) channel).register(selector, SelectionKey.OP_READ);
-			this.sockets = pool.socketPool.get(host);
-			this.sockNum = pool.poolCurrentConn.get(host);
+			writeBuf = ByteBuffer.allocateDirect(bufferSize);
+			sockets = pool.socketPool.get(host);
+			sockNum = pool.poolCurrentConn.get(host);
 		}
 
 		@Override
@@ -1341,10 +1351,6 @@ public static class TCPSockIO extends SchoonerSockIO {
 
 		private int hash = 0;
 
-		private ConcurrentLinkedQueue<SchoonerSockIO> sockets;
-
-		private AtomicInteger sockNum;
-
 		/**
 		 * creates a new SockIO object wrapping a socket connection to
 		 * host:port, and its input and output streams
@@ -1372,6 +1378,8 @@ public TCPSockIO(SchoonerSockIOPool pool, String host, int bufferSize, int timeo
 
 			// get socket: default is to use non-blocking connect
 			sock = getSocket(ip[0], Integer.parseInt(ip[1]), connectTimeout);
+			
+			writeBuf = ByteBuffer.allocateDirect(bufferSize);
 
 			if (timeout >= 0)
 				this.sock.setSoTimeout(timeout);
@@ -1383,8 +1391,8 @@ public TCPSockIO(SchoonerSockIOPool pool, String host, int bufferSize, int timeo
 			sockChannel = sock.getChannel();
 			hash = sock.hashCode();
 			this.host = host;
-			this.sockets = pool.socketPool.get(host);
-			this.sockNum = pool.poolCurrentConn.get(host);
+			sockets = pool.socketPool.get(host);
+			sockNum = pool.poolCurrentConn.get(host);
 		}
 
 		/**