From 4737575bbdec2678ed349c994c62104b77923c9e Mon Sep 17 00:00:00 2001
From: Nick Ginther <nickginther@fullstory.com>
Date: Mon, 16 Sep 2024 11:04:56 -0500
Subject: [PATCH] add FileSizeMap

---
 .../solr/storage/SizeAwareDirectory.java      | 54 ++++++++++++++++---
 1 file changed, 46 insertions(+), 8 deletions(-)

diff --git a/solr/core/src/java/org/apache/solr/storage/SizeAwareDirectory.java b/solr/core/src/java/org/apache/solr/storage/SizeAwareDirectory.java
index ebc789144f7..9592c8707bc 100644
--- a/solr/core/src/java/org/apache/solr/storage/SizeAwareDirectory.java
+++ b/solr/core/src/java/org/apache/solr/storage/SizeAwareDirectory.java
@@ -56,7 +56,7 @@ public class SizeAwareDirectory extends FilterDirectory
         this.onDiskSize.add(onDiskSize);
       };
 
-  private final ConcurrentHashMap<String, Sizes> fileSizeMap = new ConcurrentHashMap<>();
+  private final FileSizeMap fileSizeMap = new FileSizeMap();
 
   private final ConcurrentHashMap<String, SizeAccountingIndexOutput> liveOutputs =
       new ConcurrentHashMap<>();
@@ -68,6 +68,44 @@ private interface SizeWriter {
     void apply(long size, long onDiskSize, String name);
   }
 
+  private static class FileSizeMap implements Accountable {
+    private final long RAM_BYTES_USED =
+            RamUsageEstimator.shallowSizeOfInstance(FileSizeMap.class);
+    private final ConcurrentHashMap<String, Sizes> fileSizeMap = new ConcurrentHashMap<>();
+    private volatile LongAdder onDiskSize = new LongAdder();
+
+    public Sizes add(String name, Sizes sizes) {
+      Sizes extant = fileSizeMap.put(name, sizes);
+      onDiskSize.add(sizes.onDiskSize);
+      return extant;
+    }
+
+    public Sizes remove(String name) {
+      Sizes sizes = fileSizeMap.remove(name);
+      if (sizes != null) {
+        onDiskSize.add(-sizes.onDiskSize);
+      }
+      return sizes;
+    }
+
+    public long onDiskSize() {
+      return onDiskSize.sum();
+    }
+
+    public Sizes get(String key) {
+      return fileSizeMap.get(key);
+    }
+
+    public long size() {
+      return fileSizeMap.size();
+    }
+
+    @Override
+    public long ramBytesUsed() {
+      return RAM_BYTES_USED;
+    }
+  }
+
   private static class Sizes implements Accountable {
     private static final long RAM_BYTES_USED =
         RamUsageEstimator.shallowSizeOfInstance(SizeAccountingIndexOutput.class);
@@ -99,7 +137,7 @@ public SizeAwareDirectory(Directory in, long reconcileTTLNanos) {
   @Override
   public long ramBytesUsed() {
     return BASE_RAM_BYTES_USED
-        + RamUsageEstimator.sizeOfMap(fileSizeMap)
+        + fileSizeMap.ramBytesUsed()
         + RamUsageEstimator.sizeOfMap(liveOutputs);
   }
 
@@ -154,7 +192,7 @@ public long onDiskSize() throws IOException {
     if (initialized
         && (reconcileThreshold == null
             || System.nanoTime() - reconciledTimeNanos < reconcileTTLNanos)) {
-      return onDiskSize.sum();
+      return fileSizeMap.onDiskSize();
     }
     return initSize().onDiskSize;
   }
@@ -211,7 +249,7 @@ private Sizes initSize() throws IOException {
           if (fileSize > 0) {
             // whether the file exists or not, we don't care about it if it has zero size.
             // more often though, 0 size means the file isn't there.
-            fileSizeMap.put(file, sizes);
+            fileSizeMap.add(file, sizes);
             if (DirectoryFactory.sizeOf(in, file) == 0) {
               // during reconciliation, we have to check for file presence _after_ adding
               // to the map, to prevent a race condition that could leak entries into `fileSizeMap`
@@ -342,7 +380,7 @@ private static final class SizeAccountingIndexOutput extends IndexOutput impleme
 
     private final IndexOutput backing;
 
-    private final ConcurrentHashMap<String, Sizes> fileSizeMap;
+    private final FileSizeMap fileSizeMap;
 
     private final ConcurrentHashMap<String, SizeAccountingIndexOutput> liveOutputs;
 
@@ -353,7 +391,7 @@ private static final class SizeAccountingIndexOutput extends IndexOutput impleme
     private SizeAccountingIndexOutput(
         String name,
         IndexOutput backing,
-        ConcurrentHashMap<String, Sizes> fileSizeMap,
+        FileSizeMap fileSizeMap,
         ConcurrentHashMap<String, SizeAccountingIndexOutput> liveOutputs,
         SizeWriter sizeWriter) {
       super("byteSize(" + name + ")", name);
@@ -394,7 +432,7 @@ public void close() throws IOException {
           // on-disk size here
           sizeWriter.apply(0, finalBytesWritten - lastBytesWritten, name);
         }
-        fileSizeMap.put(name, new Sizes(backing.getFilePointer(), onDiskSize));
+        fileSizeMap.add(name, new Sizes(backing.getFilePointer(), onDiskSize));
         liveOutputs.remove(name);
       }
     }
@@ -442,7 +480,7 @@ public long ramBytesUsed() {
   @Override
   public void rename(String source, String dest) throws IOException {
     in.rename(source, dest);
-    Sizes extant = fileSizeMap.put(dest, fileSizeMap.remove(source));
+    Sizes extant = fileSizeMap.add(dest, fileSizeMap.remove(source));
     assert extant == null; // it's illegal for dest to already exist
   }
 }