HBASE-16194 Should count in MSLAB chunk allocation into heap size change when adding duplicate cells

hbase-server/src/main/java/org/apache/hadoop/hbase/regionserver/AbstractMemStore.java

@@ -109,7 +109,8 @@ public abstract class AbstractMemStore implements MemStore {
   @Override
   public long add(Cell cell) {
     Cell toAdd = maybeCloneWithAllocator(cell);
-    return internalAdd(toAdd);
+    boolean useMSLAB = (toAdd != cell);
+    return internalAdd(toAdd, useMSLAB);
   }
 
   /**
@@ -156,7 +157,8 @@ public abstract class AbstractMemStore implements MemStore {
   @Override
   public long delete(Cell deleteCell) {
     Cell toAdd = maybeCloneWithAllocator(deleteCell);
-    long s = internalAdd(toAdd);
+    boolean useMSLAB = (toAdd != deleteCell);
+    long s = internalAdd(toAdd, useMSLAB);
     return s;
   }
 
@@ -243,7 +245,7 @@ public abstract class AbstractMemStore implements MemStore {
     // hitting OOME - see TestMemStore.testUpsertMSLAB for a
     // test that triggers the pathological case if we don't avoid MSLAB
     // here.
-    long addedSize = internalAdd(cell);
+    long addedSize = internalAdd(cell, false);
 
     // Get the Cells for the row/family/qualifier regardless of timestamp.
     // For this case we want to clean up any other puts
@@ -387,9 +389,12 @@ public abstract class AbstractMemStore implements MemStore {
    * allocator, and doesn't take the lock.
    *
    * Callers should ensure they already have the read lock taken
+   * @param toAdd the cell to add
+   * @param useMSLAB whether using MSLAB
+   * @return the heap size change in bytes
    */
-  private long internalAdd(final Cell toAdd) {
+  private long internalAdd(final Cell toAdd, final boolean useMSLAB) {
-    long s = active.add(toAdd);
+    long s = active.add(toAdd, useMSLAB);
     setOldestEditTimeToNow();
     checkActiveSize();
     return s;

hbase-server/src/main/java/org/apache/hadoop/hbase/regionserver/HeapMemStoreLAB.java

@@ -29,6 +29,7 @@ import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hbase.util.ByteRange;
 import org.apache.hadoop.hbase.util.SimpleMutableByteRange;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 
 /**
@@ -206,6 +207,11 @@ public class HeapMemStoreLAB implements MemStoreLAB {
     }
   }
 
+  @VisibleForTesting
+  Chunk getCurrentChunk() {
+    return this.curChunk.get();
+  }
+
   /**
    * A chunk of memory out of which allocations are sliced.
    */
@@ -311,5 +317,10 @@ public class HeapMemStoreLAB implements MemStoreLAB {
         " allocs=" + allocCount.get() + "waste=" +
         (data.length - nextFreeOffset.get());
     }
+
+    @VisibleForTesting
+    int getNextFreeOffset() {
+      return this.nextFreeOffset.get();
+    }
   }
 }

hbase-server/src/main/java/org/apache/hadoop/hbase/regionserver/MemStoreCompactor.java

@@ -197,7 +197,8 @@ class MemStoreCompactor {
           // The scanner is doing all the elimination logic
           // now we just copy it to the new segment
           Cell newKV = result.maybeCloneWithAllocator(c);
-          result.internalAdd(newKV);
+          boolean useMSLAB = (newKV != c);
+          result.internalAdd(newKV, useMSLAB);
 
         }
         kvs.clear();

hbase-server/src/main/java/org/apache/hadoop/hbase/regionserver/MutableSegment.java

@@ -35,10 +35,12 @@ public class MutableSegment extends Segment {
 
   /**
    * Adds the given cell into the segment
+   * @param cell the cell to add
+   * @param useMSLAB whether using MSLAB
    * @return the change in the heap size
    */
-  public long add(Cell cell) {
+  public long add(Cell cell, boolean useMSLAB) {
-    return internalAdd(cell);
+    return internalAdd(cell, useMSLAB);
   }
 
   /**

hbase-server/src/main/java/org/apache/hadoop/hbase/regionserver/Segment.java

@@ -31,6 +31,8 @@ import org.apache.hadoop.hbase.classification.InterfaceAudience;
 import org.apache.hadoop.hbase.client.Scan;
 import org.apache.hadoop.hbase.util.ByteRange;
 
+import com.google.common.annotations.VisibleForTesting;
+
 /**
  * This is an abstraction of a segment maintained in a memstore, e.g., the active
  * cell set or its snapshot.
@@ -136,7 +138,7 @@ public abstract class Segment {
       return cell;
     }
 
-    int len = KeyValueUtil.length(cell);
+    int len = getCellLength(cell);
     ByteRange alloc = getMemStoreLAB().allocateBytes(len);
     if (alloc == null) {
       // The allocation was too large, allocator decided
@@ -150,6 +152,14 @@ public abstract class Segment {
     return newKv;
   }
 
+  /**
+   * Get cell length after serialized in {@link KeyValue}
+   */
+  @VisibleForTesting
+  int getCellLength(Cell cell) {
+    return KeyValueUtil.length(cell);
+  }
+
   public abstract boolean shouldSeek(Scan scan, long oldestUnexpiredTS);
 
   public abstract long getMinTimestamp();
@@ -240,9 +250,15 @@ public abstract class Segment {
     return comparator;
   }
 
-  protected long internalAdd(Cell cell) {
+  protected long internalAdd(Cell cell, boolean useMSLAB) {
     boolean succ = getCellSet().add(cell);
     long s = AbstractMemStore.heapSizeChange(cell, succ);
+    // If there's already a same cell in the CellSet and we are using MSLAB, we must count in the
+    // MSLAB allocation size as well, or else there will be memory leak (occupied heap size larger
+    // than the counted number)
+    if (!succ && useMSLAB) {
+      s += getCellLength(cell);
+    }
     updateMetaInfo(cell, s);
     return s;
   }
@@ -269,7 +285,8 @@ public abstract class Segment {
     return getCellSet().tailSet(firstCell);
   }
 
-  private MemStoreLAB getMemStoreLAB() {
+  @VisibleForTesting
+  public MemStoreLAB getMemStoreLAB() {
     return memStoreLAB;
   }
 

hbase-server/src/test/java/org/apache/hadoop/hbase/regionserver/TestDefaultMemStore.java

@@ -21,6 +21,7 @@ package org.apache.hadoop.hbase.regionserver;
 import com.google.common.base.Joiner;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
+
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.hadoop.conf.Configuration;
@@ -56,6 +57,7 @@ import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertNotNull;
+
 import org.junit.experimental.categories.Category;
 import org.junit.rules.TestName;
 import org.junit.rules.TestRule;
@@ -110,6 +112,30 @@ public class TestDefaultMemStore {
     assertTrue(Bytes.toString(found.getValueArray()), CellUtil.matchingValue(samekey, found));
   }
 
+  @Test
+  public void testPutSameCell() {
+    byte[] bytes = Bytes.toBytes(getName());
+    KeyValue kv = new KeyValue(bytes, bytes, bytes, bytes);
+    long sizeChangeForFirstCell = this.memstore.add(kv);
+    long sizeChangeForSecondCell = this.memstore.add(kv);
+    // make sure memstore size increase won't double-count MSLAB chunk size
+    assertEquals(AbstractMemStore.heapSizeChange(kv, true), sizeChangeForFirstCell);
+    Segment segment = this.memstore.getActive();
+    MemStoreLAB msLab = segment.getMemStoreLAB();
+    if (msLab != null) {
+      // make sure memstore size increased even when writing the same cell, if using MSLAB
+      assertEquals(segment.getCellLength(kv), sizeChangeForSecondCell);
+      // make sure chunk size increased even when writing the same cell, if using MSLAB
+      if (msLab instanceof HeapMemStoreLAB) {
+        assertEquals(2 * segment.getCellLength(kv),
+          ((HeapMemStoreLAB) msLab).getCurrentChunk().getNextFreeOffset());
+      }
+    } else {
+      // make sure no memstore size change w/o MSLAB
+      assertEquals(0, sizeChangeForSecondCell);
+    }
+  }
+
   /**
    * Test memstore snapshot happening while scanning.
    * @throws IOException