@BenchmarkMode(Mode.AverageTime)
@OutputTimeUnit(TimeUnit.MILLISECONDS)
@State(Scope.Benchmark)
public class MapBenchmark {

    private static final int CACHE_SIZE = 10;
    private static final int ITERATIONS = 100_000;

    private Map<Integer, String> lockedCache;
    private ConcurrentHashMap<Integer, String> concurrentCache;
    private ReentrantLock lock;
    private int[] keys;

    @Setup(Level.Trial)
    public void setup() {
        lockedCache = new HashMap<>(CACHE_SIZE);
        concurrentCache = new ConcurrentHashMap<>(CACHE_SIZE);
        lock = new ReentrantLock();
        keys = new int[CACHE_SIZE];

        for (int i = 0; i < CACHE_SIZE; i++) {
            keys[i] = i;
        }

        for (int key : keys) {
            String value = expensiveOperation(key);
            lockedCache.put(key, value);
            concurrentCache.put(key, value);
        }
    }

    public String expensiveOperation(int key) {
        return "Value-" + key;
    }

    @Benchmark
    @Threads(8)
    public void testReentrantLock() {
        for (int i = 0; i < ITERATIONS; i++) {
            int key = keys[ThreadLocalRandom.current().nextInt(CACHE_SIZE)];
            lock.lock();
            try {
                if (lockedCache.containsKey(key)) {
                    lockedCache.get(key);
                } else {
                    lockedCache.put(key, expensiveOperation(key));
                }
            } finally {
                lock.unlock();
            }
        }
    }

    @Benchmark
    @Threads(8)
    public void testConcurrentHashMap() {
        for (int i = 0; i < ITERATIONS; i++) {
            int key = keys[ThreadLocalRandom.current().nextInt(CACHE_SIZE)];
            concurrentCache.computeIfAbsent(key, this::expensiveOperation);
        }
    }

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