[CASSANDRA-21472][4.0] Fix memtable on-heap accounting drift in BTree.update and BTreeRow.merge

src/java/org/apache/cassandra/db/Columns.java

@@ -436,7 +436,7 @@ public long unsharedHeapSize()
         if(this == NONE)
             return 0;
 
-        return EMPTY_SIZE + BTree.sizeOfStructureOnHeap(columns);
+        return EMPTY_SIZE + BTree.sizeOnHeapOf(columns);
     }
 
     @Override

src/java/org/apache/cassandra/db/rows/BTreeRow.java

@@ -520,7 +520,7 @@ public long unsharedHeapSizeExcludingData()
                         + clustering.unsharedHeapSizeExcludingData()
                         + primaryKeyLivenessInfo.unsharedHeapSize()
                         + deletion.unsharedHeapSize()
-                        + BTree.sizeOfStructureOnHeap(btree);
+                        + BTree.sizeOnHeapOf(btree);
 
         return accumulate((cd, v) -> v + cd.unsharedHeapSizeExcludingData(), heapSize);
     }
@@ -588,10 +588,20 @@ else if (rowDeletion.isShadowedBy(livenessInfo))
                 {
                     // The update's deletion shadows part of the existing row. Those cells ARE owned by
                     // the memtable, so record their removal via retain().
-                    existingBtree = BTree.transformAndFilter(existingBtree, reconciler::retain);
+                    Object[] retained = BTree.transformAndFilter(existingBtree, reconciler::retain);
+                    if (existingBtree != retained)
+                    {
+                        reconcileF.onAllocatedOnHeap(BTree.sizeOnHeapOf(retained) - BTree.sizeOnHeapOf(existingBtree));
+                        existingBtree = retained;
+                    }
                 }
             }
             Object[] tree = BTree.update(existingBtree, updateBtree, ColumnData.comparator, reconciler);
+            // BTree.update and the reconciler only account the column data (cells and column-tree nodes); the row's
+            // own LivenessInfo/Deletion are not. When they change (e.g. a row tombstone supersedes a live row) the
+            // new objects become memtable-owned and the old ones are released, so account that delta here.
+            reconcileF.onAllocatedOnHeap((livenessInfo.unsharedHeapSize() + rowDeletion.unsharedHeapSize())
+                                         - (existing.primaryKeyLivenessInfo().unsharedHeapSize() + existing.deletion().unsharedHeapSize()));
             return new BTreeRow(existing.clustering, livenessInfo, rowDeletion, tree, minDeletionTime(tree, livenessInfo, deletion));
         }
     }

src/java/org/apache/cassandra/db/rows/ColumnData.java

@@ -151,6 +151,7 @@ public ColumnData merge(ColumnData existing, ColumnData update)
                     }
                     cells = BTree.update(existingTree, updateTree, existingComplex.column.cellComparator(), (UpdateFunction) reconciler);
                 }
+                onAllocatedOnHeap(maxComplexDeletion.unsharedHeapSize() - existingDeletion.unsharedHeapSize());
                 return new ComplexColumnData(existingComplex.column, cells, maxComplexDeletion);
             }
         }
@@ -212,8 +213,28 @@ private ColumnData removeShadowed(ColumnData existing, PostReconciliationFunctio
                 ComplexColumnData existingComplex = (ComplexColumnData) existing;
                 if (activeDeletion.supersedes(existingComplex.complexDeletion()))
                 {
-                    Object[] cells = BTree.transformAndFilter(existingComplex.tree(), (ColumnData cd) -> removeShadowed(cd, recordDeletion));
-                    return BTree.isEmpty(cells) ? null : new ComplexColumnData(existingComplex.column, cells, DeletionTime.LIVE);
+                    Object[] existingTree = existingComplex.tree();
+                    Object[] cells = BTree.transformAndFilter(existingTree, (ColumnData cd) -> removeShadowed(cd, recordDeletion));
+                    ComplexColumnData result = BTree.isEmpty(cells) ? null
+                                                                    : new ComplexColumnData(existingComplex.column, cells, DeletionTime.LIVE);
+                    // The shadowed inner cells are released through recordDeletion.delete above, but that does not cover
+                    // the complex column's own structure: its cell tree (which can span multiple BTree nodes), its
+                    // complex deletion, and, when the column is dropped entirely, its wrapper. All were counted as
+                    // owned when the column was first written (ComplexColumnData.unsharedHeapSizeExcludingData), so release that
+                    // delta here. The rewritten column carries DeletionTime.LIVE, so the dropped complex
+                    // deletion's heap is released too, matching the swap Reconciler.merge accounts on its path.
+                    // On the update side (recordDeletion == noOp) this is a no-op, so skip it entirely.
+                    if (recordDeletion != ColumnData.noOp)
+                    {
+                        long structureBefore = ComplexColumnData.EMPTY_SIZE
+                                               + existingComplex.complexDeletion().unsharedHeapSize()
+                                               + BTree.sizeOnHeapOf(existingTree);
+                        long structureAfter = result == null ? 0 : ComplexColumnData.EMPTY_SIZE
+                                                                   + DeletionTime.LIVE.unsharedHeapSize()
+                                                                   + BTree.sizeOnHeapOf(cells);
+                        recordDeletion.onAllocatedOnHeap(structureAfter - structureBefore);
+                    }
+                    return result;
                 }
             }
 

src/java/org/apache/cassandra/db/rows/ComplexColumnData.java

@@ -51,7 +51,7 @@ public class ComplexColumnData extends ColumnData implements Iterable<Cell<?>>
 {
     static final Cell<?>[] NO_CELLS = new Cell<?>[0];
 
-    private static final long EMPTY_SIZE = ObjectSizes.measure(new ComplexColumnData(ColumnMetadata.regularColumn("", "", "", SetType.getInstance(ByteType.instance, true)), NO_CELLS, new DeletionTime(0, 0)));
+    static final long EMPTY_SIZE = ObjectSizes.measure(new ComplexColumnData(ColumnMetadata.regularColumn("", "", "", SetType.getInstance(ByteType.instance, true)), NO_CELLS, new DeletionTime(0, 0)));
 
     // The cells for 'column' sorted by cell path.
     private final Object[] cells;
@@ -138,7 +138,7 @@ public int dataSize()
 
     public long unsharedHeapSizeExcludingData()
     {
-        long heapSize = EMPTY_SIZE + BTree.sizeOnHeapOf(cells);
+        long heapSize = EMPTY_SIZE + BTree.sizeOnHeapOf(cells) + complexDeletion.unsharedHeapSize();
         // TODO: this can be turned into a simple multiplication, at least while we have only one Cell implementation
         for (Cell<?> cell : this)
             heapSize += cell.unsharedHeapSizeExcludingData();

src/java/org/apache/cassandra/utils/btree/BTree.java

@@ -947,19 +947,6 @@ public static int size(Object[] tree)
         return ((int[]) tree[length - 1])[(length / 2) - 1];
     }
 
-    public static long sizeOfStructureOnHeap(Object[] tree)
-    {
-        if (tree == EMPTY_LEAF)
-            return 0;
-
-        long size = ObjectSizes.sizeOfArray(tree);
-        if (isLeaf(tree))
-            return size;
-        for (int i = getChildStart(tree); i < getChildEnd(tree); i++)
-            size += sizeOfStructureOnHeap((Object[]) tree[i]);
-        return size;
-    }
-
     /**
      * Checks is the node is a leaf.
      *
@@ -2218,6 +2205,23 @@ private static long sizeOnHeapOfLeaf(Object[] tree)
         return ObjectSizes.sizeOfArray(tree);
     }
 
+    /**
+     * The heap occupied by a single node (its backing array, plus the sizeMap for a branch), <em>not</em> including
+     * its children -- the per-node contribution to {@link #sizeOnHeapOf(Object[])}. The update builders use it to keep
+     * their running {@code allocated} total net: every newly built node adds its shallow heap and every source node it
+     * replaces subtracts its shallow heap, so reused subtrees cancel and the total equals the change in tree heap.
+     */
+    private static long shallowHeapOf(Object[] node)
+    {
+        if (isEmpty(node))
+            return 0;
+
+        long size = ObjectSizes.sizeOfArray(node);
+        if (!isLeaf(node))
+            size += ObjectSizes.sizeOfArray(sizeMap(node));
+        return size;
+    }
+
     // Arbitrary boundaries
     private static Object POSITIVE_INFINITY = new Object();
     private static Object NEGATIVE_INFINITY = new Object();
@@ -2416,17 +2420,28 @@ final boolean isEmpty()
          * @return the root of the constructed tree.
          */
         public Object[] completeBuild()
+        {
+            return completeBuild(null);
+        }
+
+        /**
+         * As {@link #completeBuild()}, but {@code rootUnode} is the source node at this (the root) level, so that an
+         * unchanged root is reused and a rebuilt root releases its predecessor from the running heap total -- the
+         * root-level analogue of the per-child accounting {@link Updater#updateRecursive} performs via drainAndPropagate.
+         */
+        public Object[] completeBuild(Object[] rootUnode)
         {
             LeafOrBranchBuilder level = this;
             while (true)
             {
                 if (!level.hasOverflow())
-                    return level.drain();
+                    return level.drainAndPropagate(rootUnode, null);
 
                 BranchBuilder parent = level.ensureParent();
-                level.drainAndPropagate(null, parent);
+                level.drainAndPropagate(rootUnode, parent);
                 if (level.savedBuffer != null)
                     Arrays.fill(level.savedBuffer, null);
+                rootUnode = null;
                 level = parent;
             }
         }
@@ -2508,6 +2523,7 @@ static boolean areIdentical(int[] a, int aOffset, int[] b, int bOffset, int coun
      */
     private static abstract class LeafBuilder extends LeafOrBranchBuilder
     {
+        static final long DISABLED = Long.MIN_VALUE;
         long allocated;
 
         LeafBuilder()
@@ -2516,6 +2532,12 @@ private static abstract class LeafBuilder extends LeafOrBranchBuilder
             buffer = new Object[MAX_KEYS];
         }
 
+        final void addAllocated(long bytes)
+        {
+            if (allocated != DISABLED)
+                allocated += bytes;
+        }
+
         /**
          * Add {@code nextKey} to the buffer, overflowing if necessary
          */
@@ -2673,8 +2695,7 @@ Object[] redistributeAndDrain(Object[] pred, int predSize, Object predNextKey)
             int newPredecessorCount = predSize - steal;
             Object[] newPredecessor = new Object[newPredecessorCount | 1];
             System.arraycopy(pred, 0, newPredecessor, 0, newPredecessorCount);
-            if (allocated >= 0)
-                allocated += ObjectSizes.sizeOfReferenceArray(newPredecessorCount | 1);
+            addAllocated(ObjectSizes.sizeOfReferenceArray(newPredecessorCount | 1));
             ensureParent().addChildAndNextKey(newPredecessor, newPredecessorCount, pred[newPredecessorCount]);
             return newLeaf;
         }
@@ -2726,8 +2747,7 @@ void propagateOverflow()
         {
             // propagate the leaf we have saved in savedBuffer
             // precondition: savedLeafCount == MAX_KEYS
-            if (allocated >= 0)
-                allocated += ObjectSizes.sizeOfReferenceArray(MAX_KEYS);
+            addAllocated(ObjectSizes.sizeOfReferenceArray(MAX_KEYS));
             ensureParent().addChildAndNextKey(savedBuffer, MAX_KEYS, savedNextKey);
             savedBuffer = null;
             savedNextKey = null;
@@ -2749,6 +2769,9 @@ Object[] drainAndPropagate(Object[] unode, BranchBuilder propagateTo)
                 // we have too few items, so spread the two buffers across two new nodes
                 leaf = redistributeOverflowAndDrain();
                 sizeOfLeaf = MIN_KEYS;
+                // redistributeAndDrain accounted the propagated predecessor; account this returned leaf and release
+                // the source node it (together with the predecessor) replaces
+                addAllocated(ObjectSizes.sizeOfReferenceArray(MIN_KEYS) - (unode == null ? 0 : sizeOnHeapOfLeaf(unode)));
             }
             else if (!hasOverflow() && unode != null && count == sizeOfLeaf(unode) && areIdentical(buffer, 0, unode, 0, count))
             {
@@ -2764,8 +2787,9 @@ else if (!hasOverflow() && unode != null && count == sizeOfLeaf(unode) && areIde
 
                 sizeOfLeaf = count;
                 leaf = drain();
-                if (allocated >= 0 && sizeOfLeaf > 0)
-                    allocated += ObjectSizes.sizeOfReferenceArray(sizeOfLeaf | 1) - (unode == null ? 0 : sizeOnHeapOfLeaf(unode));
+                // account the new leaf (if any) and release the source it replaces, keeping the total a net delta
+                addAllocated((sizeOfLeaf > 0 ? ObjectSizes.sizeOfReferenceArray(sizeOfLeaf | 1) : 0)
+                             - (unode == null ? 0 : sizeOnHeapOfLeaf(unode)));
             }
 
             count = 0;
@@ -2870,10 +2894,10 @@ void addChildAndNextKey(Object[] newChild, int newChildSize, Object nextKey)
          */
         void propagateOverflow()
         {
-            // propagate the leaf we have saved in leaf().savedBuffer
-            if (leaf.allocated >= 0)
-                leaf.allocated += ObjectSizes.sizeOfReferenceArray(2 * (1 + MAX_KEYS));
+            // propagate the branch we have saved in savedBuffer; it is a brand-new node, so account it in full
+            // (array + sizeMap) with nothing to release
             int size = setOverflowSizeMap(savedBuffer, MAX_KEYS);
+            leaf.addAllocated(shallowHeapOf(savedBuffer));
             ensureParent().addChildAndNextKey(savedBuffer, size, savedNextKey);
             savedBuffer = null;
             savedNextKey = null;
@@ -2928,8 +2952,7 @@ Object[] redistributeOverflowAndDrain()
             System.arraycopy(savedBuffer, 0, savedBranch, 0, savedBranchCount);
             System.arraycopy(savedBuffer, MAX_KEYS, savedBranch, savedBranchCount, savedBranchCount + 1);
             int savedBranchSize = setOverflowSizeMap(savedBranch, savedBranchCount);
-            if (leaf.allocated >= 0)
-                leaf.allocated += ObjectSizes.sizeOfReferenceArray(2 * (1 + savedBranchCount));
+            leaf.addAllocated(shallowHeapOf(savedBranch));
             ensureParent().addChildAndNextKey(savedBranch, savedBranchSize, savedBuffer[savedBranchCount]);
             savedNextKey = null;
 
@@ -3012,6 +3035,9 @@ Object[] drainAndPropagate(Object[] unode, BranchBuilder propagateTo)
             {
                 branch = redistributeOverflowAndDrain();
                 sizeOfBranch = sizeOfBranch(branch);
+                // redistributeOverflowAndDrain accounted the propagated branch; account this returned branch and
+                // release the source it (together with that branch) replaces
+                leaf.addAllocated(shallowHeapOf(branch) - (unode == null ? 0 : shallowHeapOf(unode)));
             }
             else
             {
@@ -3035,6 +3061,8 @@ && areIdentical(buffer, MAX_KEYS, unode, usz, usz + 1))
                     System.arraycopy(buffer, 0, branch, 0, count);
                     System.arraycopy(buffer, MAX_KEYS, branch, count, count + 1);
                     sizeOfBranch = setDrainSizeMap(unode, usz, branch, count);
+                    // account the new branch (array + sizeMap) and release the source branch it replaces
+                    leaf.addAllocated(shallowHeapOf(branch) - (unode == null ? 0 : shallowHeapOf(unode)));
                 }
             }
 
@@ -3329,7 +3357,7 @@ public static class FastBuilder<V> extends AbstractFastBuilder implements AutoCl
 
         FastBuilder()
         {
-            allocated = -1;
+            allocated = DISABLED;
         } // disable allocation tracking
 
         public void add(V value)
@@ -3437,7 +3465,7 @@ Object[] update(Object[] update, Object[] insert, Comparator<? super Compare> co
             this.insert.init(insert);
             this.updateF = updateF;
             this.comparator = comparator;
-            this.allocated = isSimple(updateF) ? -1 : 0;
+            this.allocated = isSimple(updateF) ? DISABLED : 0;
             int leafDepth = BTree.depth(update) - 1;
             LeafOrBranchBuilder builder = leaf();
             for (int i = 0; i < leafDepth; ++i)
@@ -3446,9 +3474,9 @@ Object[] update(Object[] update, Object[] insert, Comparator<? super Compare> co
             Insert ik = this.insert.next();
             ik = updateRecursive(ik, update, null, builder);
             assert ik == null;
-            Object[] result = builder.completeBuild();
+            Object[] result = builder.completeBuild(update);
 
-            if (allocated > 0)
+            if (allocated != DISABLED)
                 updateF.onAllocatedOnHeap(allocated);
 
             return result;
@@ -3633,7 +3661,7 @@ private static abstract class AbstractTransformer<I, O> extends AbstractFastBuil
 
         AbstractTransformer()
         {
-            allocated = -1;
+            allocated = DISABLED;
             ensureParent();
             parent.inUse = false;
         }

src/java/org/apache/cassandra/utils/btree/UpdateFunction.java

@@ -44,7 +44,7 @@ public interface UpdateFunction<K, V>
     V merge(V replacing, K update);
 
     /**
-     * @param heapSize extra heap space allocated (over previous tree)
+     * @param heapSize heap space signed delta allocated (over previous tree), can be negative
      */
     void onAllocatedOnHeap(long heapSize);