package net.rubyeye.xmemcached.test.unittest;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

import com.google.code.yanf4j.core.Session;
import com.google.code.yanf4j.core.impl.HandlerAdapter;
import com.google.code.yanf4j.nio.TCPController;

import net.rubyeye.xmemcached.HashAlgorithm;
import net.rubyeye.xmemcached.MemcachedClientBuilder;
import net.rubyeye.xmemcached.XMemcachedClient;
import net.rubyeye.xmemcached.XMemcachedClientBuilder;
import net.rubyeye.xmemcached.command.BinaryCommandFactory;
import net.rubyeye.xmemcached.impl.KetamaMemcachedSessionLocator;
import net.rubyeye.xmemcached.utils.AddrUtil;

public class ConnectionPoolMemcachedClientTest extends XMemcachedClientTest {
	private static final int CONNECTION_POOL_SIZE = 3;

	@Override
	public MemcachedClientBuilder createBuilder() throws Exception {
		MemcachedClientBuilder builder = new XMemcachedClientBuilder(AddrUtil
				.getAddresses(this.properties
						.getProperty("test.memcached.servers")));
		builder.setConnectionPoolSize(CONNECTION_POOL_SIZE);
		return builder;
	}

	/**
	 * Mock memcached server
	 * 
	 * @author boyan
	 * 
	 */
	static class MockServer {
		private static final int PORT = 9999;
		TCPController controller;
		AtomicInteger sessionCounter = new AtomicInteger(0);

		public void start() {
			controller = new TCPController();
			controller.setHandler(new HandlerAdapter() {

				@Override
				public void onSessionClosed(Session session) {
				}

				@Override
				public void onSessionCreated(Session session) {
					System.out.println("connection created,"
							+ sessionCounter.incrementAndGet());
				}

			});
			try {
				controller.bind(PORT);
			} catch (IOException e) {
				e.printStackTrace();
			}
		}

		public InetSocketAddress getServerAddress() {
			return new InetSocketAddress("localhost", PORT);
		}

		public void stop() {
			try {
				this.controller.stop();
			} catch (IOException e) {
				e.printStackTrace();
			}
		}
	}

	public void testHealSession() throws Exception {
		MockServer server = new MockServer();
		server.start();
		InetSocketAddress serverAddress = server.getServerAddress();
		XMemcachedClient client = new XMemcachedClient();
		client.setConnectionPoolSize(5);
		client.addServer(serverAddress);
		synchronized (this) {
			while (server.sessionCounter.get() < 5) {
				this.wait(1000);
			}
		}
		assertEquals(1, client.getAvaliableServers().size());
		assertEquals(5, client.getConnectionSizeBySocketAddress(serverAddress));
		assertEquals(5, server.sessionCounter.get());

		// stop mock server,try to heal sessions
		server.stop();

		Thread.sleep(10000);
		assertEquals(0, client.getAvaliableServers().size());
		assertEquals(0, client.getConnectionSizeBySocketAddress(serverAddress));
		// new server start
		server = new MockServer();
		server.start();
		Thread.sleep(20000);
		assertEquals(1, client.getAvaliableServers().size());
		assertEquals(5, client.getConnectionSizeBySocketAddress(serverAddress));
		assertEquals(5, server.sessionCounter.get());

		server.stop();
		client.shutdown();

	}

	@Override
	public MemcachedClientBuilder createWeightedBuilder() throws Exception {
		List<InetSocketAddress> addressList = AddrUtil
				.getAddresses(this.properties
						.getProperty("test.memcached.servers"));
		int[] weights = new int[addressList.size()];
		for (int i = 0; i < weights.length; i++) {
			weights[i] = i + 1;
		}

		MemcachedClientBuilder builder = new XMemcachedClientBuilder(
				addressList, weights);
		builder.setConnectionPoolSize(CONNECTION_POOL_SIZE);
		builder.setCommandFactory(new BinaryCommandFactory());
		builder.setSessionLocator(new KetamaMemcachedSessionLocator());
		return builder;
	}
}

@@ -128,364 +128,66 @@
 til.
-HashMap;%0D%0Aimport java.util.List;%0D%0Aimport java.util.Map;%0D%0Aimport java.util.concurrent.ConcurrentHashMap;%0D%0Aimport java.util.concurrent.atomic.AtomicInteger;%0D%0A%0D%0Aimport com.google.code.yanf4j.core.Session;%0D%0Aimport com.google.code.yanf4j.core.impl.HandlerAdapter;%0D%0Aimport com.google.code.yanf4j.nio.TCPController;%0D%0A%0D%0Aimport net.rubyeye.xmemcached.HashAlgorithm;
+List;%0D%0Aimport java.util.concurrent.atomic.AtomicInteger;%0D%0A
 %0D%0Aim
@@ -521,16 +521,170 @@
 til;%0D%0A%0D%0A
+import com.google.code.yanf4j.core.Session;%0D%0Aimport com.google.code.yanf4j.core.impl.HandlerAdapter;%0D%0Aimport com.google.code.yanf4j.nio.TCPController;%0D%0A%0D%0A
 public c
@@ -2331,24 +2331,61 @@
 olSize(5);%0D%0A
+%09%09client.setEnableHeartBeat(false);%0D%0A
 %09%09client.add
@@ -3008,9 +3008,9 @@
 eep(
-2
+1
 0000

repair ConnectionPoolMemcachedClientTest