always use ByteCountingLRUMap.remove() to remove

server/src/main/java/io/druid/client/cache/ByteCountingLRUMap.java

@@ -22,6 +22,7 @@ package io.druid.client.cache;
 import com.metamx.common.logger.Logger;
 
 import java.nio.ByteBuffer;
+import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.Set;
@@ -110,13 +111,13 @@ class ByteCountingLRUMap extends LinkedHashMap<ByteBuffer, byte[]>
   }
 
   /**
-   * We want keySet().iterator().remove() to account for object removal
-   * The underlying Map calls this.remove(key) so we do not need to override this
+   * Don't allow key removal using the underlying keySet iterator
+   * All removal operations must use ByteCountingLRUMap.remove()
    */
   @Override
   public Set<ByteBuffer> keySet()
   {
-    return super.keySet();
+    return Collections.unmodifiableSet(super.keySet());
   }
 
   @Override

server/src/main/java/io/druid/client/cache/MapCache.java

@@ -19,12 +19,14 @@
 
 package io.druid.client.cache;
 
+import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.google.common.primitives.Ints;
 
 import java.nio.ByteBuffer;
 import java.util.Collections;
 import java.util.Iterator;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
@@ -121,6 +123,7 @@ public class MapCache implements Cache
     }
     synchronized (clearLock) {
       Iterator<ByteBuffer> iter = baseMap.keySet().iterator();
+      List<ByteBuffer> toRemove = Lists.newLinkedList();
       while (iter.hasNext()) {
         ByteBuffer next = iter.next();
 
@@ -128,9 +131,12 @@ public class MapCache implements Cache
             && next.get(1) == idBytes[1]
             && next.get(2) == idBytes[2]
             && next.get(3) == idBytes[3]) {
-          iter.remove();
+          toRemove.add(next);
         }
       }
+      for(ByteBuffer key : toRemove) {
+        baseMap.remove(key);
+      }
     }
   }
 

server/src/test/java/io/druid/client/cache/ByteCountingLRUMapTest.java

@@ -19,12 +19,14 @@
 
 package io.druid.client.cache;
 
+import com.google.common.collect.Lists;
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 
 import java.nio.ByteBuffer;
 import java.util.Iterator;
+import java.util.List;
 
 /**
  */
@@ -68,12 +70,16 @@ public class ByteCountingLRUMapTest
     Assert.assertEquals(oneByte, ByteBuffer.wrap(map.get(twoByte)));
 
     Iterator<ByteBuffer> it = map.keySet().iterator();
+    List<ByteBuffer> toRemove = Lists.newLinkedList();
     while(it.hasNext()) {
       ByteBuffer buf = it.next();
       if(buf.remaining() == 10) {
-        it.remove();
+        toRemove.add(buf);
       }
     }
+    for(ByteBuffer buf : toRemove) {
+      map.remove(buf);
+    }
     assertMapValues(1, 3, 2);
 
     map.remove(twoByte);