OAK-11089: Allow SegmentReferences/RecordNumbers to provide their memory estimates

oak-segment-tar/src/main/java/org/apache/jackrabbit/oak/segment/ImmutableRecordNumbers.java

@@ -59,6 +59,11 @@ public int getOffset(int recordNumber) {
         }
     }
 
+    @Override
+    public int estimateMemoryUsage() {
+        return offsets.length * RecordNumbers.MEMORY_USAGE_PER_NUMBER;
+    }
+
     @NotNull
     @Override
     public Iterator<Entry> iterator() {

oak-segment-tar/src/main/java/org/apache/jackrabbit/oak/segment/MutableRecordNumbers.java

@@ -56,6 +56,11 @@ private static int getRecordEntry(int[] entries, int index) {
             : entries[index * 2];
     }
 
+    @Override
+    public int estimateMemoryUsage() {
+        return size / 2 * RecordNumbers.MEMORY_USAGE_PER_NUMBER;
+    }
+
     @NotNull
     @Override
     public synchronized Iterator<Entry> iterator() {

oak-segment-tar/src/main/java/org/apache/jackrabbit/oak/segment/MutableSegmentReferences.java

@@ -92,6 +92,11 @@ int size() {
         }
     }
 
+    @Override
+    public int estimateMemoryUsage() {
+        return SegmentReferences.MEMORY_USAGE_PER_REFERENCE * size();
+    }
+
     /**
      * Check if a reference exists for a provided segment ID.
      *

oak-segment-tar/src/main/java/org/apache/jackrabbit/oak/segment/RecordNumbers.java

@@ -30,6 +30,8 @@
  * A table to translate record numbers to offsets.
  */
 public interface RecordNumbers extends Iterable<Entry> {
+    /** The memory usage, in bytes, per record number stored. */
+    int MEMORY_USAGE_PER_NUMBER = 5;
 
     /**
      * An always empty {@code RecordNumber} table.
@@ -83,6 +85,22 @@ public Iterator<Entry> iterator() {
      */
     int getOffset(int recordNumber);
 
+
+    /**
+     * Estimates the memory usage, in bytes, of this table.
+     *
+     * <p>
+     * The default implementation iterates over all record numbers for backwards compatibility with all implementations,
+     * but implementations are encouraged to provide a more efficient implementation.
+     */
+    default int estimateMemoryUsage() {
+        int size = 0;
+        for (var number : this) {
+            size += MEMORY_USAGE_PER_NUMBER;
+        }
+        return size;
+    }
+
     /**
      * Represents an entry in the record table.
      */

oak-segment-tar/src/main/java/org/apache/jackrabbit/oak/segment/Segment.java

@@ -461,12 +461,8 @@ int estimateMemoryUsage() {
         size += 56; // 7 refs x 8 bytes
 
         if (id.isDataSegmentId()) {
-            int recordNumberCount = getRecordNumberCount();
-            size += 5 * recordNumberCount;
-
-            int referencedSegmentIdCount = getReferencedSegmentIdCount();
-            size += 8 * referencedSegmentIdCount;
-
+            size += recordNumbers.estimateMemoryUsage();
+            size += segmentReferences.estimateMemoryUsage();
             size += StringUtils.estimateMemoryUsage(info);
         }
         size += data.estimateMemoryUsage();

oak-segment-tar/src/main/java/org/apache/jackrabbit/oak/segment/SegmentReferences.java

@@ -30,6 +30,8 @@
  * Represents a list of segment IDs referenced from a segment.
  */
 public interface SegmentReferences extends Iterable<SegmentId> {
+    /** The memory usage, in bytes, per segment reference stored. */
+    int MEMORY_USAGE_PER_REFERENCE = 8;
 
     /** Builds a new instance of {@link SegmentReferences} from the provided {@link SegmentData}. */
     static @NotNull SegmentReferences fromSegmentData(@NotNull SegmentData data, @NotNull SegmentIdProvider idProvider) {
@@ -68,6 +70,11 @@ public SegmentId getSegmentId(int reference) {
                 return id;
             }
 
+            @Override
+            public int estimateMemoryUsage() {
+                return MEMORY_USAGE_PER_REFERENCE * referencedSegmentIdCount;
+            }
+
             @NotNull
             @Override
             public Iterator<SegmentId> iterator() {
@@ -96,4 +103,18 @@ protected SegmentId computeNext() {
      */
     SegmentId getSegmentId(int reference);
 
+    /**
+     * Estimates the memory usage, in bytes, of this list.
+     *
+     * <p>
+     * The default implementation iterates over all references for backwards compatibility with all implementations,
+     * but implementations are encouraged to provide a more efficient implementation.
+     */
+    default int estimateMemoryUsage() {
+        int size = 0;
+        for (var number : this) {
+            size += MEMORY_USAGE_PER_REFERENCE;
+        }
+        return size;
+    }
 }

oak-segment-tar/src/test/java/org/apache/jackrabbit/oak/segment/ImmutableRecordNumbersTest.java

@@ -24,6 +24,7 @@
 import java.util.Set;
 
 import org.apache.jackrabbit.oak.segment.RecordNumbers.Entry;
+import org.apache.jackrabbit.oak.segment.memory.MemoryStore;
 import org.junit.Test;
 
 public class ImmutableRecordNumbersTest {
@@ -85,6 +86,12 @@ public void iteratingShouldBeCorrect() {
         assertEquals(entries, iterated);
     }
 
+    @Test
+    public void shouldEstimateMemoryUsage() throws Exception {
+        RecordNumbers recordNumbers = new ImmutableRecordNumbers(new int[3], new byte[3]);
+        assertEquals(5 * 3, recordNumbers.estimateMemoryUsage());
+    }
+
     private Map<Integer, RecordEntry> recordEntries(Map<Integer, Integer> offsets) {
         Map<Integer, RecordEntry> entries = new HashMap<>();
 

oak-segment-tar/src/test/java/org/apache/jackrabbit/oak/segment/MutableSegmentReferencesTest.java

@@ -84,6 +84,19 @@ public void shouldMaintainSize() throws Exception {
         assertEquals(1, table.size());
     }
 
+    @Test
+    public void shouldEstimateMemoryUsage() throws Exception {
+        MemoryStore store = new MemoryStore();
+        MutableSegmentReferences table = new MutableSegmentReferences();
+
+        for (int i = 0; i < 3; i++) {
+            SegmentId id = store.getSegmentIdProvider().newDataSegmentId();
+            table.addOrReference(id);
+        }
+
+        assertEquals(8 * 3, table.estimateMemoryUsage());
+    }
+
     @Test
     public void shouldContainAddedSegment() throws Exception {
         MemoryStore store = new MemoryStore();

oak-segment-tar/src/test/java/org/apache/jackrabbit/oak/segment/RecordNumbersTest.java

@@ -19,9 +19,15 @@
 
 package org.apache.jackrabbit.oak.segment;
 
+import org.apache.jackrabbit.oak.segment.memory.MemoryStore;
+import org.jetbrains.annotations.NotNull;
 import org.apache.jackrabbit.oak.commons.Buffer;
 import org.apache.jackrabbit.oak.segment.data.SegmentData;
 import org.junit.Test;
+import org.mockito.Mockito;
+
+import java.util.Iterator;
+import java.util.List;
 
 import static org.junit.Assert.assertEquals;
 
@@ -45,4 +51,27 @@ public void shouldReadRecordNumbersFromSegmentData() {
         assertEquals(262080, recordNumbers.getOffset(3));
         assertEquals(-1, recordNumbers.getOffset(4));
     }
+
+    @Test
+    public void shouldEstimateMemoryUsageWithDefaultImplementation() throws Exception {
+        MemoryStore store = new MemoryStore();
+
+        RecordNumbers.Entry entry = Mockito.mock(RecordNumbers.Entry.class);
+        List<RecordNumbers.Entry> list = List.of(entry, entry, entry);
+
+        var customRecordNumbersImplementation = new RecordNumbers() {
+            @NotNull
+            @Override
+            public Iterator<Entry> iterator() {
+                return list.iterator();
+            }
+
+            @Override
+            public int getOffset(int recordNumber) {
+                throw new UnsupportedOperationException();
+            }
+        };
+
+        assertEquals(list.size() * 5, customRecordNumbersImplementation.estimateMemoryUsage());
+    }
 }

oak-segment-tar/src/test/java/org/apache/jackrabbit/oak/segment/SegmentReferencesTest.java

@@ -23,6 +23,8 @@
 
 import java.io.File;
 import java.io.IOException;
+import java.util.Iterator;
+import java.util.List;
 
 import org.apache.jackrabbit.oak.commons.Buffer;
 import org.apache.jackrabbit.oak.plugins.memory.EmptyNodeState;
@@ -31,6 +33,7 @@
 import org.apache.jackrabbit.oak.segment.file.FileStoreBuilder;
 import org.apache.jackrabbit.oak.segment.memory.MemoryStore;
 import org.apache.jackrabbit.oak.spi.state.NodeBuilder;
+import org.jetbrains.annotations.NotNull;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TemporaryFolder;
@@ -113,4 +116,29 @@ public void shouldReadSegmentReferencesFromSegmentData() throws IOException {
         assertEquals(segmentId.getMostSignificantBits(), -8306364159399214979L);
         assertEquals(segmentId.getLeastSignificantBits(), -6852035036232007869L);
     }
+
+    @Test
+    public void shouldEstimateMemoryUsageWithDefaultImplementation() throws Exception {
+        MemoryStore store = new MemoryStore();
+        List<SegmentId> segmentIds = List.of(
+            store.getSegmentIdProvider().newDataSegmentId(),
+            store.getSegmentIdProvider().newDataSegmentId(),
+            store.getSegmentIdProvider().newDataSegmentId()
+        );
+
+        var customSegmentReferencesImplementation = new SegmentReferences() {
+            @NotNull
+            @Override
+            public Iterator<SegmentId> iterator() {
+                return segmentIds.iterator();
+            }
+
+            @Override
+            public SegmentId getSegmentId(int reference) {
+                throw new UnsupportedOperationException();
+            }
+        };
+
+        assertEquals(segmentIds.size() * 8, customSegmentReferencesImplementation.estimateMemoryUsage());
+    }
 }