From 161c0189274de85c6b480f911d433cbf0917c265 Mon Sep 17 00:00:00 2001
From: Allan Yang <allan163@apache.org>
Date: Mon, 13 Aug 2018 20:30:23 +0800
Subject: [PATCH] HBASE-21029 Miscount of memstore's heap/offheap size if same
 cell was put

---
 .../apache/hadoop/hbase/regionserver/Segment.java | 10 +++++++---
 .../hbase/regionserver/TestDefaultMemStore.java   | 15 +++++++++++++++
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/hbase-server/src/main/java/org/apache/hadoop/hbase/regionserver/Segment.java b/hbase-server/src/main/java/org/apache/hadoop/hbase/regionserver/Segment.java
index 7069bf831e5..1f414a7eed5 100644
--- a/hbase-server/src/main/java/org/apache/hadoop/hbase/regionserver/Segment.java
+++ b/hbase-server/src/main/java/org/apache/hadoop/hbase/regionserver/Segment.java
@@ -299,11 +299,15 @@ public abstract class Segment implements MemStoreSizing {
     // 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 || mslabUsed) {
+    boolean sizeChanged = succ || mslabUsed;
+    if (sizeChanged) {
       cellSize = getCellLength(cellToAdd);
     }
-    long heapSize = heapSizeChange(cellToAdd, succ);
-    long offHeapSize = offHeapSizeChange(cellToAdd, succ);
+    // same as above, if MSLAB is used, we need to inc the heap/offheap size, otherwise there will
+    // be a memory miscount. Since we are now use heapSize + offHeapSize to decide whether a flush
+    // is needed.
+    long heapSize = heapSizeChange(cellToAdd, sizeChanged);
+    long offHeapSize = offHeapSizeChange(cellToAdd, sizeChanged);
     incMemStoreSize(cellSize, heapSize, offHeapSize);
     if (memstoreSizing != null) {
       memstoreSizing.incMemStoreSize(cellSize, heapSize, offHeapSize);
diff --git a/hbase-server/src/test/java/org/apache/hadoop/hbase/regionserver/TestDefaultMemStore.java b/hbase-server/src/test/java/org/apache/hadoop/hbase/regionserver/TestDefaultMemStore.java
index 77f796f8c3b..ad0d2ffe984 100644
--- a/hbase-server/src/test/java/org/apache/hadoop/hbase/regionserver/TestDefaultMemStore.java
+++ b/hbase-server/src/test/java/org/apache/hadoop/hbase/regionserver/TestDefaultMemStore.java
@@ -138,6 +138,21 @@ public class TestDefaultMemStore {
     Segment segment = this.memstore.getActive();
     MemStoreLAB msLab = segment.getMemStoreLAB();
     if (msLab != null) {
+      if (msLab.isOnHeap()) {
+        assertTrue("HeapSize should always bigger or equal than data size",
+            sizeChangeForFirstCell.getHeapSize() >= sizeChangeForFirstCell
+                .getDataSize());
+        assertTrue("HeapSize should always bigger or equal than data size",
+            sizeChangeForSecondCell.getHeapSize() >= sizeChangeForSecondCell
+                .getDataSize());
+      } else {
+        assertTrue("OffHeapSize should always bigger or equal than data size",
+            sizeChangeForFirstCell.getOffHeapSize() >= sizeChangeForFirstCell
+                .getDataSize());
+        assertTrue("OffHeapSize should always bigger or equal than data size",
+            sizeChangeForSecondCell.getOffHeapSize() >= sizeChangeForSecondCell
+                .getDataSize());
+      }
       // make sure memstore size increased even when writing the same cell, if using MSLAB
       assertEquals(Segment.getCellLength(kv),
           sizeChangeForSecondCell.getMemStoreSize().getDataSize());