package com.kiven.test;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import org.apache.log4j.PropertyConfigurator;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.Watcher.Event.KeeperState;
import org.apache.zookeeper.ZooDefs.Ids;
import org.apache.zookeeper.ZooKeeper;
 
public class Test{
    /**
     * server列表, 以逗號分割
     */
    protected static String hosts = "172.16.217.148:2181";
    /**
     * 連接的超時時間, 毫秒
     */
    private final int SESSION_TIMEOUT = 5000;
    private CountDownLatch connectedSignal = new CountDownLatch(1);
    protected static ZooKeeper zk;
    private static String nodePath = "/Test/test1";
 
    //static String data = "a very long string about data to set to zookeeper";
    static int threads = 10;    //線程數
    static int runs = 1000;     //迭代次數
    static int start = 0;       //none
 
    static int size = 1024*4;   //寫入數據的大小,單位：字節
    static byte[] testdata;     //測試數據
 
    public static void main(String[] args) throws Exception {
    PropertyConfigurator.configure("log4j.properties");
    //生成寫入的數據，大小size字節
    testdata = new byte[size];
    for(int i=0;i<size;i++){
        testdata[i] = 'A';
    }
 
    Test test = new Test();
    //連接
    test.connect();
 
    WorkerStat[] statArray = new WorkerStat[threads];
    Thread[] threadArray = new Thread[threads];
 
    WorkerStat mainStat = new WorkerStat();
    mainStat.runs = runs * threads;
 
    long begin = System.currentTimeMillis();
    for (int i = 0; i < threads; i++) {
        statArray[i] = new WorkerStat();
        statArray[i].start = start + i * runs;
        statArray[i].runs = runs;
        threadArray[i] = new SetterThread(statArray[i]);
        threadArray[i].start();
    }
    for (int i = 0; i < threads; i++) {
        threadArray[i].join();
    }
    mainStat.setterTime = System.currentTimeMillis() - begin;
 
    begin = System.currentTimeMillis();
    for (int i = 0; i < threads; i++) {
        threadArray[i] = new GetterThread(statArray[i]);
        threadArray[i].start();
    }
    for (int i = 0; i < threads; i++) {
        threadArray[i].join();
    }
    mainStat.getterTime = System.currentTimeMillis() - begin;
 
    WorkerStat totalStat = new WorkerStat();
 
    System.out.println("Test over!!");
    System.out.println("Thread("+threads+")\truns\tset time(ms)\tget time(ms)");
 
    for (int i = 0; i < threads; i++) {
        totalStat.runs = totalStat.runs + statArray[i].runs;
        totalStat.setterTime = totalStat.setterTime + statArray[i].setterTime;
        totalStat.getterTime = totalStat.getterTime + statArray[i].getterTime;
    }
    System.out.println("Total\t\t" + totalStat.runs + "\t"+ totalStat.setterTime + "\t\t" + totalStat.getterTime);
    System.out.println("Avg\t\t" + runs + "\t" + totalStat.setterTime/ threads + "\t\t" + totalStat.getterTime / threads);      
    System.out.println("TPS\t\t\t" + 1000 * totalStat.runs/ totalStat.setterTime + "\t\t" + 1000 * totalStat.runs/ totalStat.getterTime);
    System.out.println("\nMain\t\t" + mainStat.runs + "\t"+ mainStat.setterTime + "\t\t" + mainStat.getterTime);
    System.out.println("TPS\t\t" + 1000 * mainStat.runs/ mainStat.setterTime + "\t\t" + 1000 * mainStat.runs/ mainStat.getterTime);
    }
 
    private static class WorkerStat {
    public int start, runs;
    public long setterTime, getterTime;
 
    public WorkerStat() {
        start = runs = 0;
        setterTime = getterTime = 0;
    }
    }
 
    private static class SetterThread extends Thread {
    private WorkerStat stat;
 
    SetterThread(WorkerStat stat) {
        this.stat = stat;
    }
 
    public void run() {
        long begin = System.currentTimeMillis();
        for (int i = stat.start; i < stat.start + stat.runs; i++) {
            //寫入節點數據
            try {
            zk.setData(nodePath, testdata, -1);
            } catch (Exception e) {
            e.printStackTrace();
            }
        }
        long end = System.currentTimeMillis();
        stat.setterTime = end - begin;
    }
    }
 
    private static class GetterThread extends Thread {
    private WorkerStat stat;
    GetterThread(WorkerStat stat) {
        this.stat = stat;
    }
    public void run() {
        long begin = System.currentTimeMillis();
        for (int i = stat.start; i < stat.start + stat.runs; i++) {
            //讀取節點數據
            try {
            zk.getData(nodePath, false, null);
        } catch (Exception e) {
            e.printStackTrace();
        } 
        }
        long end = System.currentTimeMillis();
        stat.getterTime = end - begin;
    }
    }
 
    //===============================================================================
    /**
     * 連接zookeeper server
     */
    public void connect() throws Exception {
        zk = new ZooKeeper(hosts, SESSION_TIMEOUT, new ConnWatcher());
        // 等待連接完成
        connectedSignal.await();
    }
 
    /**
     * 
     * @author Kiven
     *
     */
    public class ConnWatcher implements Watcher{
        public void process(WatchedEvent event) {
            // 連接建立, 回調process接口時, 其event.getState()為KeeperState.SyncConnected
            if (event.getState() == KeeperState.SyncConnected) {
                // 放開閘門, wait在connect方法上的線程將被喚醒
                connectedSignal.countDown();
            }
        }
    }
 
    /**
        以下為各個參數的詳細說明:
        path. znode的路徑.
        data. 與znode關聯的數據.
        acl. 指定權限信息, 如果不想指定權限, 可以傳入Ids.OPEN_ACL_UNSAFE.
        指定znode類型. CreateMode是一個枚舉類, 從中選擇一個成員傳入即可. 
    */
 
    /**
     * 創建持久化節點
     */
    public void create(String Path, byte[] data) throws Exception {
        zk.create(Path, data, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
        System.out.println("創建節點:"+Path);
        System.out.println("=================");
    }
 
    /** 
     *  
     *獲取節點信息 
     *@author kiven
     *@createDate 2013-01-16 15:17:22 
     *@param path 
     *@throws KeeperException 
     *@throws InterruptedException 
     */  
    public void getChild(String path) throws KeeperException, InterruptedException{     
        try{  
            List list=zk.getChildren(path, false);  
            if(list.isEmpty()){
                System.out.println(path+"中沒有節點");
            }else{
                System.out.println(path+"中存在節點");  
                for(String child:list){
                    System.out.println("節點為："+child);
                }
            }
        }catch (KeeperException.NoNodeException e) {  
            e.printStackTrace();
        }
    }
 
    /**
     * 設置節點數據
     * @throws Exception
     */
    public void setData(String path,String data) throws Exception{
        zk.setData(path, data.getBytes(), -1);
        System.out.println("set Data:"+"testSetData");
    }
 
    /**
     * 讀取節點數據
     * @throws Exception
     */
    public void getData() throws Exception{
        System.out.println("get Data:");
        zk.getData(nodePath, false, null);
    }
 
    /**
     * 刪除節點
     * @param path
     * @throws Exception
     */
    public void delete(String path) throws Exception{
        System.out.println("刪除節點:"+path);
        //如果版本號與znode的版本號不一致，將無法刪除，是一種樂觀加鎖機制；如果將版本號設置為-1，不會去檢測版本，直接刪除；
        zk.delete(path, -1);
    }
 
    /**
     * 關閉連接
     */
    public void close() {
        try {
            zk.close();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }
}