From 7c394621df55eba74b85a71daec5a58e603bbdc6 Mon Sep 17 00:00:00 2001
From: swethakann <shwethu.kannan93@gmail.com>
Date: Tue, 26 Nov 2024 13:56:14 -0800
Subject: [PATCH] PR changes

---
 .../server/search/TerminateAfterWrapper.java  |  5 +++
 .../search/collectors/DocCollector.java       |  2 +-
 .../search/collectors/DocCollectorTest.java   | 42 ++++++++++++++++++-
 3 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/src/main/java/com/yelp/nrtsearch/server/search/TerminateAfterWrapper.java b/src/main/java/com/yelp/nrtsearch/server/search/TerminateAfterWrapper.java
index b5fbac459..6de20cabd 100644
--- a/src/main/java/com/yelp/nrtsearch/server/search/TerminateAfterWrapper.java
+++ b/src/main/java/com/yelp/nrtsearch/server/search/TerminateAfterWrapper.java
@@ -97,6 +97,11 @@ public int getTerminateAfter() {
     return terminateAfter;
   }
 
+  /** Max documents to count beyond terminateAfter. */
+  public int getTerminateAfterMaxRecallCount() {
+    return terminateAfterMaxRecallCount;
+  }
+
   /**
    * {@link Collector} implementation that wraps another collector and terminates collection after a
    * certain global count of documents is reached.
diff --git a/src/main/java/com/yelp/nrtsearch/server/search/collectors/DocCollector.java b/src/main/java/com/yelp/nrtsearch/server/search/collectors/DocCollector.java
index d170f3999..133cb1b1b 100644
--- a/src/main/java/com/yelp/nrtsearch/server/search/collectors/DocCollector.java
+++ b/src/main/java/com/yelp/nrtsearch/server/search/collectors/DocCollector.java
@@ -179,7 +179,7 @@ <C extends Collector> CollectorManager<? extends Collector, SearcherResult> wrap
     int terminateAfterMaxRecallCount =
         request.getTerminateAfterMaxRecallCount() > 0
             ? request.getTerminateAfterMaxRecallCount()
-            : 0;
+            : indexState.getDefaultTerminateAfterMaxRecallCount();
     if (terminateAfter > 0) {
       wrapped =
           new TerminateAfterWrapper<>(
diff --git a/src/test/java/com/yelp/nrtsearch/server/search/collectors/DocCollectorTest.java b/src/test/java/com/yelp/nrtsearch/server/search/collectors/DocCollectorTest.java
index 7de62e096..09e24ba0a 100644
--- a/src/test/java/com/yelp/nrtsearch/server/search/collectors/DocCollectorTest.java
+++ b/src/test/java/com/yelp/nrtsearch/server/search/collectors/DocCollectorTest.java
@@ -212,12 +212,21 @@ public void testNumHitsToCollect() {
 
   @Test
   public void testHasTerminateAfterWrapper() {
-    SearchRequest request = SearchRequest.newBuilder().setTopHits(10).setTerminateAfter(5).build();
+    SearchRequest request =
+        SearchRequest.newBuilder()
+            .setTopHits(10)
+            .setTerminateAfter(5)
+            .setTerminateAfterMaxRecallCount(10)
+            .build();
     TestDocCollector docCollector = new TestDocCollector(request);
     assertTrue(docCollector.getManager() instanceof TestDocCollector.TestCollectorManager);
     assertTrue(docCollector.getWrappedManager() instanceof TerminateAfterWrapper);
     assertEquals(
         5, ((TerminateAfterWrapper<?>) docCollector.getWrappedManager()).getTerminateAfter());
+    assertEquals(
+        10,
+        ((TerminateAfterWrapper<?>) docCollector.getWrappedManager())
+            .getTerminateAfterMaxRecallCount());
   }
 
   @Test
@@ -246,6 +255,37 @@ public void testOverrideDefaultTerminateAfter() {
         75, ((TerminateAfterWrapper<?>) docCollector.getWrappedManager()).getTerminateAfter());
   }
 
+  @Test
+  public void testUsesDefaultTerminateAfterMaxRecallCount() {
+    IndexState indexState = Mockito.mock(IndexState.class);
+    when(indexState.getDefaultTerminateAfter()).thenReturn(100);
+    when(indexState.getDefaultTerminateAfterMaxRecallCount()).thenReturn(1000);
+
+    SearchRequest request = SearchRequest.newBuilder().setTopHits(10).build();
+    TestDocCollector docCollector = new TestDocCollector(request, indexState);
+    assertEquals(
+        1000,
+        ((TerminateAfterWrapper<?>) docCollector.getWrappedManager())
+            .getTerminateAfterMaxRecallCount());
+  }
+
+  @Test
+  public void testOverrideDefaultTerminateAfterMaxRecallCount() {
+    IndexState indexState = Mockito.mock(IndexState.class);
+    when(indexState.getDefaultTerminateAfter()).thenReturn(100);
+    when(indexState.getDefaultTerminateAfterMaxRecallCount()).thenReturn(1000);
+
+    SearchRequest request =
+        SearchRequest.newBuilder().setTopHits(10).setTerminateAfterMaxRecallCount(75).build();
+    TestDocCollector docCollector = new TestDocCollector(request, indexState);
+    assertTrue(docCollector.getManager() instanceof TestDocCollector.TestCollectorManager);
+    assertTrue(docCollector.getWrappedManager() instanceof TerminateAfterWrapper);
+    assertEquals(
+        75,
+        ((TerminateAfterWrapper<?>) docCollector.getWrappedManager())
+            .getTerminateAfterMaxRecallCount());
+  }
+
   @Test
   public void testWithAllWrappers() {
     SearchRequest request =