import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import org.apache.pulsar.common.util.collections.ConcurrentOpenHashMap;
import org.openjdk.jmh.annotations.Benchmark;
import org.openjdk.jmh.annotations.Level;
import org.openjdk.jmh.annotations.Measurement;
import org.openjdk.jmh.annotations.Scope;
import org.openjdk.jmh.annotations.Setup;
import org.openjdk.jmh.annotations.State;
import org.openjdk.jmh.annotations.TearDown;
import org.openjdk.jmh.annotations.Warmup;
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;

@State(Scope.Benchmark)
public class ConcurrentHashMapGetTest {

    static final int minKey = 0;
    static final int maxKey = 10000;

    @State(Scope.Benchmark)
    public static class BenchmarkState {
        static final int numThreads = 4;
        ExecutorService executor;
        final List<Future<Void>> futures = new ArrayList<>();
        final ConcurrentHashMap<Integer, Boolean> map = new ConcurrentHashMap<>();
        final ConcurrentOpenHashMap<Integer, Boolean> map2 = ConcurrentOpenHashMap.<Integer, Boolean>newBuilder()
                .build();
        final Map<Integer, Boolean> map3 = Collections.synchronizedMap(new HashMap<>());

        void setUp() {
            final var random = new Random();
            for (int i = 0; i < 5000; i++) {
                final var key = random.nextInt(maxKey + 1);
                map.put(key, true);
                map2.put(key, true);
                map3.put(key, true);
            }
            executor = Executors.newFixedThreadPool(numThreads);
        }

        void tearDown() {
            futures.clear();
            map.clear();
            map2.clear();
            map3.clear();
            executor.shutdown();
        }
    }

    @Setup(Level.Trial)
    public void setUp(BenchmarkState state) {
        state.setUp();
    }

    @TearDown(Level.Trial)
    public void tearDown(BenchmarkState state) {
        state.tearDown();
    }

    @Warmup(iterations = 10, time = 500, timeUnit = TimeUnit.MILLISECONDS)
    @Measurement(iterations = 5, time = 500, timeUnit = TimeUnit.MILLISECONDS)
    @Benchmark
    public void testOfficialConcurrentHashMap(BenchmarkState state) throws Exception {
        for (int i = 0; i < BenchmarkState.numThreads; i++) {
            state.futures.add(state.executor.submit(() -> {
                for (int key = minKey; key <= maxKey; key++) {
                    state.map.get(key);
                }
                return null;
            }));
        }
        for (var future : state.futures) {
            future.get();
        }
    }

    @Warmup(iterations = 10, time = 500, timeUnit = TimeUnit.MILLISECONDS)
    @Measurement(iterations = 5, time = 500, timeUnit = TimeUnit.MILLISECONDS)
    @Benchmark
    public void testConcurrentOpenHashMap(BenchmarkState state) throws Exception {
        for (int i = 0; i < BenchmarkState.numThreads; i++) {
            state.futures.add(state.executor.submit(() -> {
                for (int key = minKey; key <= maxKey; key++) {
                    state.map2.get(key);
                }
                return null;
            }));
        }
        for (var future : state.futures) {
            future.get();
        }
    }

    @Warmup(iterations = 10, time = 500, timeUnit = TimeUnit.MILLISECONDS)
    @Measurement(iterations = 5, time = 500, timeUnit = TimeUnit.MILLISECONDS)
    @Benchmark
    public void testSynchronizedHashMap(BenchmarkState state) throws Exception {
        for (int i = 0; i < BenchmarkState.numThreads; i++) {
            state.futures.add(state.executor.submit(() -> {
                for (int key = minKey; key <= maxKey; key++) {
                    state.map3.get(key);
                }
                return null;
            }));
        }
        for (var future : state.futures) {
            future.get();
        }
    }

    public static void main(String[] args) throws RunnerException {
        Options opt = new OptionsBuilder()
                .include(ConcurrentHashMapGetTest.class.getSimpleName())
                .forks(1)
                .build();
        new Runner(opt).run();
    }
}