From 5c9198347b10973db56bf2d8db6813d097457826 Mon Sep 17 00:00:00 2001
From: Himanshu <g.himanshu@gmail.com>
Date: Thu, 13 Apr 2017 17:06:34 -0500
Subject: [PATCH] fix over counting of bytes in ByteCountingLRUMap (#4168)

---
 .../io/druid/client/cache/ByteCountingLRUMap.java   | 10 +++++++---
 .../druid/client/cache/ByteCountingLRUMapTest.java  | 13 +++++++++++++
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/server/src/main/java/io/druid/client/cache/ByteCountingLRUMap.java b/server/src/main/java/io/druid/client/cache/ByteCountingLRUMap.java
index 65be2d2630d..09cb0eedab2 100644
--- a/server/src/main/java/io/druid/client/cache/ByteCountingLRUMap.java
+++ b/server/src/main/java/io/druid/client/cache/ByteCountingLRUMap.java
@@ -19,6 +19,8 @@
 
 package io.druid.client.cache;
 
+import io.druid.java.util.common.logger.Logger;
+
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -29,8 +31,6 @@ import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicLong;
 
-import io.druid.java.util.common.logger.Logger;
-
 /**
 */
 class ByteCountingLRUMap extends LinkedHashMap<ByteBuffer, byte[]>
@@ -104,7 +104,11 @@ class ByteCountingLRUMap extends LinkedHashMap<ByteBuffer, byte[]>
       remove(keyToRemove);
     }
 
-    return super.put(key, value);
+    byte[] old = super.put(key, value);
+    if (old != null) {
+      numBytes.addAndGet(-key.remaining() - old.length);
+    }
+    return old;
   }
 
   @Override
diff --git a/server/src/test/java/io/druid/client/cache/ByteCountingLRUMapTest.java b/server/src/test/java/io/druid/client/cache/ByteCountingLRUMapTest.java
index a02c8c547de..09709809529 100644
--- a/server/src/test/java/io/druid/client/cache/ByteCountingLRUMapTest.java
+++ b/server/src/test/java/io/druid/client/cache/ByteCountingLRUMapTest.java
@@ -86,6 +86,19 @@ public class ByteCountingLRUMapTest
     assertMapValues(0, 0, 3);
   }
 
+  @Test
+  public void testSameKeyUpdate() throws Exception
+  {
+    final ByteBuffer k = ByteBuffer.allocate(1);
+
+    assertMapValues(0, 0, 0);
+    map.put(k, new byte[1]);
+    map.put(k, new byte[2]);
+    map.put(k, new byte[5]);
+    map.put(k, new byte[3]);
+    assertMapValues(1, 4, 0);
+  }
+
   private void assertMapValues(final int size, final int numBytes, final int evictionCount)
   {
     Assert.assertEquals(size, map.size());