WaitForOne.java

package niffler.test;


import javax.annotation.Nullable;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

public class WaitForOne<K, V> {

    private final Map<K, WaitForOne<K, V>.SyncSubject> map = new ConcurrentHashMap<>();

    public void provide(K key, V value) {
        if (value != null && key != null) {
            WaitForOne<K, V>.SyncSubject subject = this.map.computeIfAbsent(key, SyncSubject::new);
            subject.provideIfNotProvided(value);
        } else {
            throw new IllegalArgumentException("Null is not allowed as key or value");
        }
    }

    @Nullable
    public V wait(K key, long timeoutMs) {
        SyncSubject subject = this.map.computeIfAbsent(key, SyncSubject::new);

        try {
            return subject.latch.await(timeoutMs, TimeUnit.MILLISECONDS)
                    ? subject.value
                    : null;
        } catch (InterruptedException var6) {
            return null;
        }
    }

    public boolean check(K key, WaitForOne.Checker<V> checker, long timeoutMs) {
        SyncSubject subject = this.map.computeIfAbsent(key, SyncSubject::new);

        try {
            while (!checker.check(subject.value)) {
                if (timeoutMs <= 0L) {
                    return false;
                }
                synchronized (subject) {
                    long startWait = System.currentTimeMillis();
                    subject.wait(timeoutMs);
                    timeoutMs -= System.currentTimeMillis() - startWait;
                }
            }
            return true;
        } catch (InterruptedException e) {
            return false;
        }
    }

    public Map<K, V> getAsMap() {
        return this.map.entrySet()
                .stream()
                .collect(Collectors.toMap(Map.Entry::getKey, (v) -> v.getValue().value));
    }

    private class SyncSubject {
        private CountDownLatch latch;
        private K key;
        private V value;

        private SyncSubject(K key) {
            this.key = key;
            this.latch = new CountDownLatch(1);
        }

        private synchronized void provideIfNotProvided(V value) {
            if (this.latch.getCount() != 0L) {
                this.value = value;
                this.latch.countDown();
            }
        }
    }

    @FunctionalInterface
    public interface Checker<V> {
        boolean check(V checkedValue);
    }
}