From 3917d22591bc382e457212ecf388b99cdfd4c3c2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jean-Fran=C3=A7ois=20BOEUF?= <jeanfrancois.boeuf@gmail.com>
Date: Thu, 28 Mar 2024 10:34:46 +0100
Subject: [PATCH] Subtract deleted file size from the cache size of
 NRTCachingDirectory. (#13206)

---
 lucene/CHANGES.txt                            |  2 ++
 .../lucene/store/NRTCachingDirectory.java     | 13 ++++++++++
 .../lucene/store/TestNRTCachingDirectory.java | 24 +++++++++++++++++++
 3 files changed, 39 insertions(+)

diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index 8082384472f..776d1267147 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -151,6 +151,8 @@ Bug Fixes
 
 * GITHUB#12878: Fix the declared Exceptions of Expression#evaluate() to match those
   of DoubleValues#doubleValue(). (Uwe Schindler)
+  
+* GITHUB#13206: Subtract deleted file size from the cache size of NRTCachingDirectory. (Jean-François Boeuf)
 
 Changes in Backwards Compatibility Policy
 -----------------------------------------
diff --git a/lucene/core/src/java/org/apache/lucene/store/NRTCachingDirectory.java b/lucene/core/src/java/org/apache/lucene/store/NRTCachingDirectory.java
index 26006e9b2d0..16d6aa22ce5 100644
--- a/lucene/core/src/java/org/apache/lucene/store/NRTCachingDirectory.java
+++ b/lucene/core/src/java/org/apache/lucene/store/NRTCachingDirectory.java
@@ -73,6 +73,12 @@ public class NRTCachingDirectory extends FilterDirectory implements Accountable
           new SingleInstanceLockFactory(),
           ByteBuffersDataOutput::new,
           (fileName, content) -> {
+            // Defensive check to handle the case the file has been deleted before this lambda
+            // is called when the IndexOutput is closed. Unsafe in the unlikely case the deletion
+            // happens concurrently on another thread.
+            if (isCachedFile(fileName) == false) {
+              return null;
+            }
             cacheSize.addAndGet(content.size());
             return ByteBuffersDirectory.OUTPUT_AS_MANY_BUFFERS.apply(fileName, content);
           });
@@ -118,7 +124,10 @@ public class NRTCachingDirectory extends FilterDirectory implements Accountable
       System.out.println("nrtdir.deleteFile name=" + name);
     }
     if (cacheDirectory.fileExists(name)) {
+      long size = cacheDirectory.fileLength(name);
       cacheDirectory.deleteFile(name);
+      long newSize = cacheSize.addAndGet(-size);
+      assert newSize >= 0;
     } else {
       in.deleteFile(name);
     }
@@ -295,6 +304,10 @@ public class NRTCachingDirectory extends FilterDirectory implements Accountable
     }
   }
 
+  private synchronized boolean isCachedFile(String fileName) {
+    return cacheDirectory.fileExists(fileName);
+  }
+
   private void unCache(String fileName) throws IOException {
     // Must sync here because other sync methods have
     // if (cache.fileNameExists(name)) { ... } else { ... }:
diff --git a/lucene/core/src/test/org/apache/lucene/store/TestNRTCachingDirectory.java b/lucene/core/src/test/org/apache/lucene/store/TestNRTCachingDirectory.java
index b81c3aeed6b..1588baf9be9 100644
--- a/lucene/core/src/test/org/apache/lucene/store/TestNRTCachingDirectory.java
+++ b/lucene/core/src/test/org/apache/lucene/store/TestNRTCachingDirectory.java
@@ -34,6 +34,7 @@ import org.apache.lucene.tests.store.BaseDirectoryTestCase;
 import org.apache.lucene.tests.util.LineFileDocs;
 import org.apache.lucene.tests.util.TestUtil;
 import org.apache.lucene.util.BytesRef;
+import org.junit.Assert;
 
 public class TestNRTCachingDirectory extends BaseDirectoryTestCase {
 
@@ -167,4 +168,27 @@ public class TestNRTCachingDirectory extends BaseDirectoryTestCase {
 
     dir.close();
   }
+
+  public void testCacheSizeAfterDelete() throws IOException {
+    IOContext ioContext = new IOContext(new FlushInfo(3, 40));
+    String fn = "f1";
+    try (Directory dir = newDirectory();
+        NRTCachingDirectory nrt = new NRTCachingDirectory(dir, 1, 1); ) {
+      // deletes a closed file
+      try (IndexOutput out = nrt.createOutput(fn, ioContext)) {
+        for (int i = 0; i < 10; i++) out.writeInt(i);
+      }
+      Assert.assertEquals(40, nrt.ramBytesUsed());
+      nrt.deleteFile(fn);
+      Assert.assertEquals(0, nrt.ramBytesUsed());
+
+      // Deletes an unclosed file (write before and after deletion
+      try (IndexOutput out = nrt.createOutput(fn, ioContext)) {
+        for (int i = 0; i < 10; i++) out.writeInt(i);
+        nrt.deleteFile(fn);
+        for (int i = 0; i < 10; i++) out.writeInt(i);
+      }
+      Assert.assertEquals(0, nrt.ramBytesUsed());
+    }
+  }
 }