import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.CyclicBarrier;

public class ConcurrentHashMapTest {

    static final int K = 100;
    static final int SIZE = 1000;

    public static void main(String[] args) throws InterruptedException {
        for (int i = 0; i < K; i++) {
            System.out.println(i + " of " + K);
            test();
        }
    }

    private static void test() throws InterruptedException {
        ConcurrentHashMap<Integer, Integer> map = new ConcurrentHashMap<>();
        // put 0's
        for (int i = 0; i < SIZE; i++) {
            map.put(i, 0);
        }

        CyclicBarrier barrier = new CyclicBarrier(2); // to start working simultaneously
        CountDownLatch latch = new CountDownLatch(2); // to wait until both threads finished

        // this thread put 1's into map
        new Thread(new Changer(map, barrier, latch)).start();

        // this thread remove all 0's from map
        new Thread(new Remover(map, barrier, latch)).start();

        latch.await();

        // there should be SIZE 1's in map
        if (map.size() != SIZE) {
            System.out.println("Number of 1's: " + map.values().stream().filter(v -> v == 1).count());
            throw new IllegalStateException(String.format("Size should be %d, but it is %d", SIZE, map.size()));
        }
    }

    static class Changer implements Runnable {

        private final ConcurrentHashMap<Integer, Integer> map;
        private final CyclicBarrier barrier;
        private final CountDownLatch latch;

        public Changer(ConcurrentHashMap<Integer, Integer> map, CyclicBarrier barrier, CountDownLatch latch) {
            this.map = map;
            this.barrier = barrier;
            this.latch = latch;
        }

        @Override
        public void run() {
            try {
                barrier.await();
            } catch (Exception e) {}

            for (int i = 0; i < SIZE; i++) {
                map.put(i, 1);
            }

            latch.countDown();
        }
    }

    static class Remover implements Runnable {

        private final ConcurrentHashMap<Integer, Integer> map;
        private final CyclicBarrier barrier;
        private final CountDownLatch latch;

        public Remover(ConcurrentHashMap<Integer, Integer> map, CyclicBarrier barrier, CountDownLatch latch) {
            this.map = map;
            this.barrier = barrier;
            this.latch = latch;
        }

        @Override
        public void run() {
            try {
                barrier.await();
            } catch (Exception e) {}

            map.entrySet().removeIf(e -> e.getValue() == 0);

            latch.countDown();
        }
    }
}