package ch.ethz.globis.distindex.client;

import ch.ethz.globis.distindex.client.pht.PHFactory;
import ch.ethz.globis.distindex.client.pht.PHTreeIndexProxy;
import ch.ethz.globis.distindex.client.pht.ZKPHFactory;
import ch.ethz.globis.distindex.middleware.api.Middleware;
import ch.ethz.globis.distindex.middleware.net.IndexMiddlewareFactory;

import ch.ethz.globis.distindex.middleware.util.MiddlewareUtil;
import org.apache.curator.test.TestingServer;
import org.openjdk.jmh.annotations.*;
import org.openjdk.jmh.runner.Runner;
import org.openjdk.jmh.runner.RunnerException;
import org.openjdk.jmh.runner.options.Options;
import org.openjdk.jmh.runner.options.OptionsBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.math.BigInteger;
import java.util.ArrayList;
import java.util.List;
import java.util.Random;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;


public class PhTreeBenchmark {

    private static Logger LOG = LoggerFactory.getLogger(PhTreeBenchmark.class);

    @State(Scope.Benchmark)
    public static class BenchmarkState {

        private static PHTreeIndexProxy<Object> indexProxy;

        private static final String ZK_HOST = "localhost";
        private static final int ZK_PORT = 2181;
        private static final int S_BASE_PORT = 7070;
        static final int NUMBER_OF_SERVERS = 16;

        private static ExecutorService threadPool;
        private static TestingServer zkServer;
        private static List<Middleware> middlewares = new ArrayList<Middleware>();

        static {
            init();
            indexProxy = setupIndex();
            indexProxy.setRangeKNN(true);
            insert();
        }

        private static PHTreeIndexProxy<Object> setupIndex() {
            PHFactory factory = new ZKPHFactory(ZK_HOST, ZK_PORT);
            PHTreeIndexProxy<Object> proxy = factory.createProxy(2, 64);
            return proxy;
        }

        private static void insert() {
            Random random = new Random();
            for (int i = 0; i < 100; i++) {
                long[] key = { random.nextLong(), random.nextLong() };
                indexProxy.put(key, new BigInteger(64, random));
            }
        }

        public static void init() {
            try {
                threadPool = Executors.newFixedThreadPool(NUMBER_OF_SERVERS * 2);
                zkServer = new TestingServer(ZK_PORT);
                zkServer.start();

                for (int i = 0; i < NUMBER_OF_SERVERS; i++) {
                    Middleware current = IndexMiddlewareFactory.newPhTree(ZK_HOST, S_BASE_PORT + i * 10, ZK_HOST, ZK_PORT);
                    MiddlewareUtil.startMiddleware(threadPool, current);
                    middlewares.add(current);
                }
            } catch (IOException e) {
                LOG.error("Failed to create the new testing utility.");
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    @Benchmark
    @BenchmarkMode({Mode.Throughput})
    @Warmup(iterations = 5, time = 100, timeUnit = TimeUnit.NANOSECONDS)
    @Measurement(iterations = 10, time = 100, timeUnit = TimeUnit.NANOSECONDS)
    public Object defaultBenchmark(BenchmarkState state) {
        return state.indexProxy.getNearestNeighbors(randomKey(), 10);
    }

    private static long[] randomKey() {
        Random random = new Random();
        return new long[] { random.nextLong(), random.nextLong() };
    }

    public static void main(String[] args) throws RunnerException {
        Options opt = new OptionsBuilder()
                .include(".*" + PhTreeBenchmark.class.getSimpleName() + ".*")
                .forks(1)
                .build();
        new Runner(opt).run();
    }
}
@@ -1440,10 +1440,9 @@
 S = 
-16
+4
 ;%0A%0A 
@@ -1630,22 +1630,49 @@
         
-static
+@Setup%0A        public void init()
  %7B%0A     
@@ -1678,20 +1678,28 @@
         
-init
+setupCluster
 ();%0A    
@@ -1751,33 +1751,235 @@
   in
-dexProxy.setRangeKNN(true
+sert();%0A        %7D%0A%0A        @TearDown%0A        public void close() throws IOException %7B%0A            for (Middleware middleware : middlewares) %7B%0A                middleware.close();%0A            %7D%0A            threadPool.shutdownNow(
 );%0A 
@@ -1981,38 +1981,46 @@
 ();%0A            
-insert
+zkServer.close
 ();%0A        %7D%0A%0A 
@@ -2563,36 +2563,44 @@
 lic static void 
-init
+setupCluster
 () %7B%0A           
@@ -3575,24 +3575,69 @@
 te state) %7B%0A
+        state.indexProxy.setRangeKNN(false);%0A
         retu

Index distribution Set up a teardown method for the benchmark suite.