Question

I’m trying to write a thread-safe Map[K, Set[V]] implementation in java.

1. If a unique key is added to the map, a new Set should be created (and added to)
2. If a non unique key is added to the map, the existing Set should be added to.
3. If a value is removed from a Set causing the Set to be empty, the entry should be removed from the map to avoid memory leaks.
4. I’d like to solve this without needing to synchronize the whole thing

I have included a failing test case below, please let me know if you have a solution.

package org.deleteme;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;

import junit.framework.Assert;

import org.junit.Test;

public class ConcurrentSetMapTest {
    public static class ConcurrentSetMap<K, V> {
        private final ConcurrentMap<K, Set<V>> map = new ConcurrentHashMap<K, Set<V>>();

        public void add(K key, V value) {
            Set<V> set = map.get(key);
            if (set != null) {
                set.add(value);
            } else {
                Set<V> candidateSet = createConcurrentSet(value);
                set = map.putIfAbsent(key, candidateSet);
                if (set != null) {
                    // candidate set not accepted, use existing
                    set.add(value);
                }
            }
        }

        public void remove(K key, V value) {
            Set<V> set = map.get(key);
            if (set != null) {
                boolean removed = set.remove(value);
                if (removed && set.isEmpty()) {
                    // this is not thread-safe and causes the test to fail
                    map.remove(key, set);
                }
            }
        }

        public boolean contains(K key, V value) {
            Set<V> set = map.get(key);
            if (set == null) {
                return false;
            }
            return set.contains(value);
        }

        protected Set<V> createConcurrentSet(V element) {
            Set<V> set = Collections.newSetFromMap(new ConcurrentHashMap<V, Boolean>());
            set.add(element);
            return set;
        }
    }

    @Test
    public void testThreadSafe() throws InterruptedException, ExecutionException {
        ConcurrentSetMap<String, String> setMap = new ConcurrentSetMap<String, String>();
        ExecutorService executors = Executors.newFixedThreadPool(4);
        List<Future<?>> futures = new ArrayList<Future<?>>();

        futures.add(executors.submit(new TestWorker(setMap, "key1")));
        futures.add(executors.submit(new TestWorker(setMap, "key1")));
        futures.add(executors.submit(new TestWorker(setMap, "key2")));
        futures.add(executors.submit(new TestWorker(setMap, "key2")));

        for (Future<?> future : futures) {
            future.get();
        }
    }

    public static class TestWorker implements Runnable {
        ConcurrentSetMap<String, String> setMap;
        String key;

        public TestWorker(ConcurrentSetMap<String, String> setMap, String key) {
            super();
            this.setMap = setMap;
            this.key = key;
        }

        public void run() {
            int sampleSize = 100000;
            for (int i = 0; i < sampleSize; ++ i) {
                // avoid value clashes with other threads
                String value = Thread.currentThread().getName() + i;

                Assert.assertFalse("Should not exist before add", setMap.contains(key, value));
                setMap.add(key, value);
                Assert.assertTrue("Should exist after add", setMap.contains(key, value));
                setMap.remove(key, value);
                Assert.assertFalse("Should not exist after remove", setMap.contains(key, value));
            }
        }
    }
}

Answer 1

I’ve managed to solve my problem 🙂

I failed to mention in my initial post that I require fast reads from the collection and I’m not too concerned about write speed. For this reason, I have come up with a solution that synchronizes write access but does not require synchronized read access. The code below now passes my test case.

Thanks to all for your suggestions.

public static class ConcurrentSetMap<K, V> {
    private final ConcurrentMap<K, Set<V>> map = new ConcurrentHashMap<K, Set<V>>();

    public synchronized void add(K key, V value) {
        Set<V> set = map.get(key);
        if (set != null) {
            set.add(value);
        } else {
            map.put(key, createConcurrentSet(value));
        }
    }

    public synchronized void remove(K key, V value) {
        Set<V> set = map.get(key);
        if (set != null) {
            set.remove(value);
            if (set.isEmpty()) {
                map.remove(key);
            }
        }
    }

    public boolean contains(K key, V value) {
        return get(key).contains(value);
    }

    public Set<V> get(K key) {
        Set<V> set = map.get(key);
        return set == null ? Collections.<V> emptySet() : set;
    }

    protected Set<V> createConcurrentSet(V value) {
        Set<V> set = Collections.newSetFromMap(new ConcurrentHashMap<V, Boolean>());
        set.add(value);
        return set;
    }
}