InMemorySetStore.java

package net.morimekta.providence.storage;

import net.morimekta.util.concurrent.ReadWriteMutex;
import net.morimekta.util.concurrent.ReentrantReadWriteMutex;

import javax.annotation.Nonnull;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.function.Function;

/**
 * Simple in-memory set storage of providence messages. Uses a local hash map for
 * storing the instances. The store is thread safe through using re-entrant
 * read-write mutex handling, so reading can happen in parallel.
 */
public class InMemorySetStore<K, M> implements ReadWriteSetStore<K,M> {
    private final Map<K, M>      map;
    private final ReadWriteMutex mutex;
    private final Function<M, K> messageToKey;

    public InMemorySetStore(Function<M,K> messageToKey) {
        this.messageToKey = messageToKey;
        this.map   = new HashMap<>();
        this.mutex = new ReentrantReadWriteMutex();
    }

    @Nonnull
    @Override
    public Map<K, M> getAll(@Nonnull Collection<K> keys) {
        return mutex.lockForReading(() -> {
            Map<K, M> out = new HashMap<>();
            for (K key : keys) {
                if (map.containsKey(key)) {
                    out.put(key, map.get(key));
                }
            }
            return out;
        });
    }

    @Override
    public boolean containsKey(@Nonnull K key) {
        return mutex.lockForReading(() -> map.containsKey(key));
    }

    @Override @Nonnull
    public Collection<K> keys() {
        return mutex.lockForReading(() -> new HashSet<>(map.keySet()));
    }

    @Override
    public int size() {
        return mutex.lockForReading(map::size);
    }

    @Override
    public void putAll(@Nonnull Collection<M> values) {
        mutex.lockForWriting(() -> {
            for (M entry : values) {
                map.put(messageToKey.apply(entry), entry);
            }
        });
    }

    @Override
    public void removeAll(Collection<K> keys) {
        mutex.lockForWriting(() -> map.keySet().removeAll(keys));
    }
}