title: Concurrent map of sets - Removing the set when empty
body: <p>I'm trying to write a thread-safe Map[K, Set[V]] implementation in java. </p>
  
  <ol>
  <li>If a unique key is added to the map, a new Set should be created (and added to)</li>
  <li>If a non unique key is added to the map, the existing Set should be added to.</li>
  <li>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.</li>
  <li>I'd like to solve this without needing to synchronize the whole thing </li>
  </ol>
  
  <p>I have included a failing test case below, please let me know if you have a solution.</p>
  
  <pre><code>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&lt;K, V&gt; {
          private final ConcurrentMap&lt;K, Set&lt;V&gt;&gt; map = new ConcurrentHashMap&lt;K, Set&lt;V&gt;&gt;();
  
          public void add(K key, V value) {
              Set&lt;V&gt; set = map.get(key);
              if (set != null) {
                  set.add(value);
              } else {
                  Set&lt;V&gt; 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&lt;V&gt; set = map.get(key);
              if (set != null) {
                  boolean removed = set.remove(value);
                  if (removed &amp;&amp; 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&lt;V&gt; set = map.get(key);
              if (set == null) {
                  return false;
              }
              return set.contains(value);
          }
  
          protected Set&lt;V&gt; createConcurrentSet(V element) {
              Set&lt;V&gt; set = Collections.newSetFromMap(new ConcurrentHashMap&lt;V, Boolean&gt;());
              set.add(element);
              return set;
          }
      }
  
      @Test
      public void testThreadSafe() throws InterruptedException, ExecutionException {
          ConcurrentSetMap&lt;String, String&gt; setMap = new ConcurrentSetMap&lt;String, String&gt;();
          ExecutorService executors = Executors.newFixedThreadPool(4);
          List&lt;Future&lt;?&gt;&gt; futures = new ArrayList&lt;Future&lt;?&gt;&gt;();
  
          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&lt;?&gt; future : futures) {
              future.get();
          }
      }
  
      public static class TestWorker implements Runnable {
          ConcurrentSetMap&lt;String, String&gt; setMap;
          String key;
  
          public TestWorker(ConcurrentSetMap&lt;String, String&gt; setMap, String key) {
              super();
              this.setMap = setMap;
              this.key = key;
          }
  
          public void run() {
              int sampleSize = 100000;
              for (int i = 0; i &lt; 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));
              }
          }
      }
  }
  </code></pre>