Concurrent map of sets - Removing the set when empty

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

Don't write such a map, use somebody else's. I would use one of Guava's SetMultimap implementations such as HashMultimap and synchronize it using Multimaps.synchronizedSetMultimap.

Answer 2

Here is a fully concurrent and thread-safe implementation:

public 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> curSet = _map.get(key);
    while(true) {

      if((curSet != null) && curSet.contains(value)) {
        break;
      }

      Set<V> newSet = new HashSet<V>();
      newSet.add(value);

      if(curSet == null) {

        curSet = _map.putIfAbsent(key, newSet);
        if(curSet != null) {
          continue;
        }

      } else {

        newSet.addAll(curSet);
        if(!_map.replace(key, curSet, newSet)) {
          curSet = _map.get(key);
          continue;
        }
      }

      break;
    }
  }

  public void remove(K key, V value) {
    Set<V> curSet = _map.get(key);

    while(true) {
      if((curSet == null) || !curSet.contains(value))  {
        break;
      }

      if(curSet.size() == 1) {

        if(!_map.remove(key, curSet)) {
          curSet = _map.get(key);
          continue;
        }

      } else {

        Set<V> newSet = new HashSet<V>();
        newSet.addAll(curSet);
        newSet.remove(value);
        if(!_map.replace(key, curSet, newSet)) {
          curSet = _map.get(key);
          continue;
        }
      }

      break;
    }
  }

  public boolean contains(K key, V value) {
    Set<V> set = _map.get(key);
    return set != null && set.contains(value);
  }
}

Comparing timing to @PeterLawrey's answer (on my box), his takes 2.9 seconds, this takes 1.4 seconds.

Answer 3

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;
    }
} 

Answer 4

You need to use locking as you are performing multiple operations which need to be atomic collectively.

public class SynchronousMultiMap<K, V> {
    private final Map<K, Set<V>> map = new LinkedHashMap<K, Set<V>>();

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

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

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

    @Test
    public void testThreadSafe() throws ExecutionException, InterruptedException {
        ExecutorService executors = Executors.newFixedThreadPool(3);
        List<Future<?>> futures = new ArrayList<Future<?>>();
        SynchronousMultiMap<String, Integer> setMap = new SynchronousMultiMap<String, Integer>();
        int sampleSize = 1000000;

        String[] keys = "key1,key2,key3,key4".split(",");
        for (int i = 0; i < 3; i++)
            futures.add(executors.submit(new TestWorker(setMap, keys, sampleSize, i)));

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

    static class TestWorker implements Runnable {
        final SynchronousMultiMap<String, Integer> setMap;
        final String[] keys;
        final int sampleSize;
        final int value;

        public TestWorker(SynchronousMultiMap<String, Integer> setMap, String[] keys, int sampleSize, int value) {
            super();
            this.setMap = setMap;
            this.keys = keys;
            this.sampleSize = sampleSize;
            this.value = value;
        }

        public void run() {
            for (int i = 0; i < sampleSize; i += keys.length) {
                for (String key : keys) {
                    boolean contains = setMap.contains(key, value);
                    if (contains)
                        Assert.assertFalse("Should not exist before add", contains);
                    setMap.add(key, value);
                    boolean contains2 = setMap.contains(key, value);
                    if (!contains2)
                        Assert.assertTrue("Should exist after add", contains2);
                    setMap.remove(key, value);
                    boolean contains3 = setMap.contains(key, value);
                    if (contains3)
                        Assert.assertFalse("Should not exist after remove", contains3);
                }
            }
        }
    }
}

takes 0.35 seconds to run. With a sampleSize=1000000 it takes < 8 seconds.

Answer 5

Use ConcurrentMap of ConcurrentSkipListSet.

Call ConcurrentMap.putIfAbsent() to add a set for a key, if it doesn't exist.

Call ConcurrentMap.remove( key, emptySet ), where emptySet is an empty ConcurrentSkipListSet.

If the current value, corresponding to key, is empty, it'll be removed.