From c3f97fbdc11cf29e17a4e715981108dda7ba3aea Mon Sep 17 00:00:00 2001
From: Adrien Grand <jpountz@gmail.com>
Date: Mon, 28 Sep 2020 15:08:08 +0200
Subject: [PATCH] Compute RAM usage ByteBuffersDataOutput on the fly. (#1919)

This helps remove the assumption that all blocks have the same size.
---
 .../lucene/store/ByteBuffersDataOutput.java   | 15 +++----
 .../store/TestByteBuffersDataOutput.java      | 42 +++++++++++++++++++
 2 files changed, 50 insertions(+), 7 deletions(-)

diff --git a/lucene/core/src/java/org/apache/lucene/store/ByteBuffersDataOutput.java b/lucene/core/src/java/org/apache/lucene/store/ByteBuffersDataOutput.java
index ac1dc7579e5..7711302c5e4 100644
--- a/lucene/core/src/java/org/apache/lucene/store/ByteBuffersDataOutput.java
+++ b/lucene/core/src/java/org/apache/lucene/store/ByteBuffersDataOutput.java
@@ -117,6 +117,9 @@ public final class ByteBuffersDataOutput extends DataOutput implements Accountab
    */
   private final ArrayDeque<ByteBuffer> blocks = new ArrayDeque<>();
 
+  /** Cumulative RAM usage across all blocks. */
+  private long ramBytesUsed;
+
   /**
    * The current-or-next write block.
    */
@@ -400,13 +403,8 @@ public final class ByteBuffersDataOutput extends DataOutput implements Accountab
   public long ramBytesUsed() {
     // Return a rough estimation for allocated blocks. Note that we do not make
     // any special distinction for direct memory buffers.
-    ByteBuffer first = blocks.peek();
-    if (first == null) {
-      return 0L;
-    } else {
-      // All blocks have the same capacity.
-      return (first.capacity() + RamUsageEstimator.NUM_BYTES_OBJECT_REF) * blocks.size();
-    }
+    assert ramBytesUsed == blocks.stream().mapToLong(ByteBuffer::capacity).sum() + blocks.size() * RamUsageEstimator.NUM_BYTES_OBJECT_REF;
+    return ramBytesUsed;
   }
 
   /**
@@ -422,6 +420,7 @@ public final class ByteBuffersDataOutput extends DataOutput implements Accountab
       blocks.forEach(blockReuse);
     }
     blocks.clear();
+    ramBytesUsed = 0;
     currentBlock = EMPTY;
   }
 
@@ -455,6 +454,7 @@ public final class ByteBuffersDataOutput extends DataOutput implements Accountab
     currentBlock = blockAllocate.apply(requiredBlockSize);
     assert currentBlock.capacity() == requiredBlockSize;
     blocks.add(currentBlock);
+    ramBytesUsed += RamUsageEstimator.NUM_BYTES_OBJECT_REF + currentBlock.capacity();
   }
 
   private void rewriteToBlockSize(int targetBlockBits) {
@@ -476,6 +476,7 @@ public final class ByteBuffersDataOutput extends DataOutput implements Accountab
     assert blocks.isEmpty();
     this.blockBits = targetBlockBits;
     blocks.addAll(cloned.blocks);
+    ramBytesUsed = cloned.ramBytesUsed;
   }
 
   private static int computeBlockSizeBitsFor(long bytes) {
diff --git a/lucene/core/src/test/org/apache/lucene/store/TestByteBuffersDataOutput.java b/lucene/core/src/test/org/apache/lucene/store/TestByteBuffersDataOutput.java
index 8a9012acc48..d67e9de0267 100644
--- a/lucene/core/src/test/org/apache/lucene/store/TestByteBuffersDataOutput.java
+++ b/lucene/core/src/test/org/apache/lucene/store/TestByteBuffersDataOutput.java
@@ -20,10 +20,12 @@ import static org.junit.Assert.*;
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.util.List;
 import java.util.Random;
 import java.util.concurrent.atomic.AtomicInteger;
 import org.apache.lucene.util.ArrayUtil;
 import org.apache.lucene.util.LuceneTestCase;
+import org.apache.lucene.util.RamUsageEstimator;
 import org.junit.Assert;
 import org.junit.Test;
 
@@ -159,4 +161,44 @@ public final class TestByteBuffersDataOutput extends BaseDataOutputTestCase<Byte
       assertTrue(bb.hasArray()); // heap-based by default, so array should be there.
     }
   }
+
+  @Test
+  public void testRamBytesUsed() {
+    ByteBuffersDataOutput out = new ByteBuffersDataOutput();
+    // Empty output requires no RAM
+    assertEquals(0, out.ramBytesUsed());
+
+    // Non-empty buffer requires RAM
+    out.writeInt(4);
+    assertEquals(out.ramBytesUsed(), computeRamBytesUsed(out));
+
+    // Make sure this keeps working with multiple backing buffers
+    while (out.toBufferList().size() < 2) {
+      out.writeLong(42);
+    }
+    assertEquals(out.ramBytesUsed(), computeRamBytesUsed(out));
+
+    // Make sure this keeps working when increasing the block size
+    int currentBlockCapacity = out.toBufferList().get(0).capacity();
+    do {
+      out.writeLong(42);
+    } while (out.toBufferList().get(0).capacity() == currentBlockCapacity);
+    assertEquals(out.ramBytesUsed(), computeRamBytesUsed(out));
+
+    // Back to zero after a clear
+    out.reset();
+    assertEquals(0, out.ramBytesUsed());
+
+    // And back to non-empty
+    out.writeInt(4);
+    assertEquals(out.ramBytesUsed(), computeRamBytesUsed(out));
+  }
+
+  private static long computeRamBytesUsed(ByteBuffersDataOutput out) {
+    if (out.size() == 0) {
+      return 0;
+    }
+    List<ByteBuffer> buffers = out.toBufferList();
+    return buffers.stream().mapToLong(ByteBuffer::capacity).sum() + buffers.size() * RamUsageEstimator.NUM_BYTES_OBJECT_REF;
+  }
 }