Fix double-counting of Automaton in AutomatonQuery#ramBytesUsed (#16390)

Author: Sasilekha

lucene/CHANGES.txt

@@ -191,6 +191,8 @@ Bug Fixes
   been aborted, so IndexWriter#rollback and abortMerges no longer block until the entire graph is
   built. (Jeho Jeong)
 
+* GITHUB#16389: Fix double-counting of the underlying Automaton in AutomatonQuery#ramBytesUsed. (Sasilekha R)
+
 * GITHUB#14049: Randomize KNN codec params in RandomCodec. Fixes scalar quantization div-by-zero
   when all values are identical. (Mike Sokolov)
 

lucene/core/src/java/org/apache/lucene/search/AutomatonQuery.java

@@ -98,8 +98,10 @@ public AutomatonQuery(
     this.automatonIsBinary = isBinary;
     this.compiled = new CompiledAutomaton(automaton, false, true, isBinary);
 
+    // compiled may already reference the same Automaton instance; only count its bytes once.
+    long automatonBytes = compiled.sharesAutomaton(automaton) ? 0L : automaton.ramBytesUsed();
     this.ramBytesUsed =
-        BASE_RAM_BYTES + term.ramBytesUsed() + automaton.ramBytesUsed() + compiled.ramBytesUsed();
+        BASE_RAM_BYTES + term.ramBytesUsed() + automatonBytes + compiled.ramBytesUsed();
   }
 
   @Override

lucene/core/src/java/org/apache/lucene/util/automaton/CompiledAutomaton.java

@@ -492,6 +492,23 @@ public ByteRunnable getByteRunnable() {
     return nfaRunAutomaton;
   }
 
+  /**
+   * Returns {@code true} if this {@link CompiledAutomaton} internally holds the given {@link
+   * Automaton} instance. This is only the case on the {@link AUTOMATON_TYPE#NORMAL} paths built
+   * without a UTF-8 conversion (typically {@code isBinary=true} construction): the DFA path retains
+   * it as {@link #automaton}, and the NFA path retains it inside the internal NFA runner. Returns
+   * {@code false} for {@link AUTOMATON_TYPE#NONE}, {@link AUTOMATON_TYPE#ALL} and {@link
+   * AUTOMATON_TYPE#SINGLE} -- those short-circuit and keep no reference -- and for {@code
+   * isBinary=false} construction, where the internal automaton is a distinct UTF-8 conversion of
+   * the input.
+   *
+   * <p>Useful for callers that keep their own reference to the same {@link Automaton} and want to
+   * avoid double-counting it in their own {@link Accountable#ramBytesUsed()} calculation.
+   */
+  public boolean sharesAutomaton(Automaton a) {
+    return a == automaton || (nfaRunAutomaton != null && nfaRunAutomaton.getAutomaton() == a);
+  }
+
   /**
    * Get a {@link TransitionAccessor} instance, it will be different depending on whether a NFA or
    * DFA is passed in, and does not guarantee returning non-null object

lucene/core/src/test/org/apache/lucene/search/TestAutomatonQuery.java

@@ -34,13 +34,16 @@
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.tests.index.RandomIndexWriter;
 import org.apache.lucene.tests.util.LuceneTestCase;
+import org.apache.lucene.tests.util.RamUsageTester;
 import org.apache.lucene.tests.util.Rethrow;
 import org.apache.lucene.tests.util.TestUtil;
 import org.apache.lucene.tests.util.automaton.AutomatonTestUtil;
 import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.RamUsageEstimator;
 import org.apache.lucene.util.automaton.Automata;
 import org.apache.lucene.util.automaton.Automaton;
 import org.apache.lucene.util.automaton.Operations;
+import org.apache.lucene.util.automaton.RegExp;
 
 public class TestAutomatonQuery extends LuceneTestCase {
   private Directory directory;
@@ -252,4 +255,68 @@ public void testBiggishAutomaton() {
     Collections.sort(terms);
     new AutomatonQuery(new Term("foo", "bar"), Automata.makeStringUnion(terms));
   }
+
+  public void testRamBytesUsedDoesNotDoubleCountSharedDeterministicAutomaton() {
+    // Same shape as PrefixQuery: an already-deterministic binary automaton passed with
+    // isBinary=true. CompiledAutomaton keeps a reference to the same Automaton instance
+    // via runAutomaton.automaton; naively summing automaton + compiled would double-count.
+    Automaton prefix = PrefixQuery.toAutomaton(new BytesRef("prefix"));
+    AutomatonQuery q = new AutomatonQuery(new Term(FN, "prefix"), prefix, true);
+    assertSame(prefix, q.getCompiled().automaton);
+    assertRamBytesExcludesSharedAutomaton(q, prefix);
+  }
+
+  public void testRamBytesUsedDoesNotDoubleCountSharedNfaAutomaton() {
+    // isBinary=true with a non-deterministic input drives CompiledAutomaton down the NFA
+    // path, where nfaRunAutomaton wraps the same Automaton instance.
+    Automaton nfa = new Automaton();
+    int start = nfa.createState();
+    int a1 = nfa.createState();
+    int a2 = nfa.createState();
+    nfa.setAccept(a1, true);
+    nfa.setAccept(a2, true);
+    nfa.addTransition(start, a1, 'a', 'a');
+    nfa.addTransition(start, a2, 'a', 'a');
+    nfa.finishState();
+    assertFalse(nfa.isDeterministic());
+
+    AutomatonQuery q = new AutomatonQuery(new Term(FN, "nfa"), nfa, true);
+    assertNull(q.getCompiled().automaton);
+    assertTrue(q.getCompiled().sharesAutomaton(nfa));
+    assertRamBytesExcludesSharedAutomaton(q, nfa);
+  }
+
+  public void testRamBytesUsedIsBinaryFalseCountsOuterAutomaton() {
+    // isBinary=false path: CompiledAutomaton converts to UTF-8 internally, so the outer
+    // automaton and compiled hold distinct Automaton instances. Both must be counted --
+    // this test guards against a future refactor accidentally dropping the outer bytes.
+    Automaton a =
+        Operations.determinize(new RegExp("abc.*").toAutomaton(), DEFAULT_DETERMINIZE_WORK_LIMIT);
+    AutomatonQuery q = new AutomatonQuery(new Term(FN, "regex"), a);
+    assertFalse(q.getCompiled().sharesAutomaton(a));
+    long reported = q.ramBytesUsed();
+    assertTrue(
+        "outer automaton must still contribute on the isBinary=false path",
+        reported >= a.ramBytesUsed());
+    long actual = RamUsageTester.ramUsed(q);
+    assertEquals((double) actual, (double) reported, (double) actual * 0.10);
+  }
+
+  // Asserts that ramBytesUsed() does not include the shared Automaton twice.
+  // Reported bytes must equal (shallow AutomatonQuery + term + compiled), i.e., the
+  // shared automaton is accounted for only once (via compiled). If the bug were
+  // present, reported would be inflated by exactly sharedAutomaton.ramBytesUsed().
+  private static void assertRamBytesExcludesSharedAutomaton(
+      AutomatonQuery q, Automaton sharedAutomaton) {
+    long expected =
+        RamUsageEstimator.shallowSizeOfInstance(AutomatonQuery.class)
+            + q.term.ramBytesUsed()
+            + q.getCompiled().ramBytesUsed();
+    assertEquals(
+        "shared Automaton must not be counted twice (would over-report by "
+            + sharedAutomaton.ramBytesUsed()
+            + " bytes)",
+        expected,
+        q.ramBytesUsed());
+  }
 }