Add default terminate after max recall count (#790) (#793)

Add default terminate after max recall count
Yelp · Dec 6, 2024 · dfa5c79 · dfa5c79
1 parent 5b34fda
commit dfa5c79
Show file tree

Hide file tree

Showing 10 changed files with 106 additions and 2 deletions.
diff --git a/clientlib/src/main/proto/yelp/nrtsearch/luceneserver.proto b/clientlib/src/main/proto/yelp/nrtsearch/luceneserver.proto
@@ -473,6 +473,8 @@ message LiveSettingsRequest {
     int32 defaultSearchTimeoutCheckEvery = 13;
     //Terminate after value to use when not specified in the search request.
     int32 defaultTerminateAfter = 14;
+    //Terminate after max recall count value to use when not specified in the search request.
+    int32 defaultTerminateAfterMaxRecallCount = 15;
 }
 
 /* Response from Server to liveSettings */
@@ -1101,6 +1103,8 @@ message IndexLiveSettings {
     google.protobuf.UInt64Value maxMergePreCopyDurationSec = 14;
     // Collect and publish additional index metrics, which may be more expensive in terms of volume, memory and/or compute, default: false
     google.protobuf.BoolValue verboseMetrics = 15;
+    // Terminate after max recall count value to use when not specified in the search request, or 0 for none, default: 0
+    google.protobuf.Int32Value defaultTerminateAfterMaxRecallCount = 18;
 }
 
 message IndexStateInfo {

diff --git a/src/main/java/com/yelp/nrtsearch/server/cli/LiveSettingsV2Command.java b/src/main/java/com/yelp/nrtsearch/server/cli/LiveSettingsV2Command.java
@@ -106,6 +106,11 @@ public class LiveSettingsV2Command implements Callable<Integer> {
       description = "Terminate after to use when not provided by the request")
   private Integer defaultTerminateAfter;
 
+  @CommandLine.Option(
+      names = {"--defaultTerminateAfterMaxRecallCount"},
+      description = "Terminate after max recall count to use when not provided by the request")
+  private Integer defaultTerminateAfterMaxRecallCount;
+
   @CommandLine.Option(
       names = {"--maxMergePreCopyDurationSec"},
       description = "Maximum time allowed for merge precopy in seconds")
@@ -181,6 +186,10 @@ public Integer call() throws Exception {
         liveSettingsBuilder.setDefaultTerminateAfter(
             Int32Value.newBuilder().setValue(defaultTerminateAfter).build());
       }
+      if (defaultTerminateAfterMaxRecallCount != null) {
+        liveSettingsBuilder.setDefaultTerminateAfterMaxRecallCount(
+            Int32Value.newBuilder().setValue(defaultTerminateAfterMaxRecallCount).build());
+      }
       if (maxMergePreCopyDurationSec != null) {
         liveSettingsBuilder.setMaxMergePreCopyDurationSec(
             UInt64Value.newBuilder().setValue(maxMergePreCopyDurationSec));

diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/IndexState.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/IndexState.java
@@ -473,6 +473,9 @@ public abstract IndexWriterConfig getIndexWriterConfig(
   /** Get the default terminate after. */
   public abstract int getDefaultTerminateAfter();
 
+  /** Get the default terminate after max recall count. */
+  public abstract int getDefaultTerminateAfterMaxRecallCount();
+
   /** Get the default search timeout check every. */
   public abstract int getDefaultSearchTimeoutCheckEvery();
 

diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/LiveSettingsHandler.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/LiveSettingsHandler.java
@@ -119,6 +119,12 @@ private LiveSettingsResponse handleAsLiveSettingsV2(
                 .setValue(liveSettingsRequest.getDefaultTerminateAfter())
                 .build());
       }
+      if (liveSettingsRequest.getDefaultTerminateAfterMaxRecallCount() >= 0) {
+        settingsBuilder.setDefaultTerminateAfterMaxRecallCount(
+            Int32Value.newBuilder()
+                .setValue(liveSettingsRequest.getDefaultTerminateAfterMaxRecallCount())
+                .build());
+      }
       try {
         updatedSettings = indexStateManager.updateLiveSettings(settingsBuilder.build(), false);
       } catch (IOException e) {

diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/index/ImmutableIndexState.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/index/ImmutableIndexState.java
@@ -163,6 +163,7 @@ public class ImmutableIndexState extends IndexState {
           .setDefaultSearchTimeoutSec(DoubleValue.newBuilder().setValue(0).build())
           .setDefaultSearchTimeoutCheckEvery(Int32Value.newBuilder().setValue(0).build())
           .setDefaultTerminateAfter(Int32Value.newBuilder().setValue(0).build())
+          .setDefaultTerminateAfterMaxRecallCount(Int32Value.newBuilder().setValue(0).build())
           .setMaxMergePreCopyDurationSec(UInt64Value.newBuilder().setValue(0))
           .setVerboseMetrics(BoolValue.newBuilder().setValue(false).build())
           .build();
@@ -181,6 +182,7 @@ public class ImmutableIndexState extends IndexState {
   private final double defaultSearchTimeoutSec;
   private final int defaultSearchTimeoutCheckEvery;
   private final int defaultTerminateAfter;
+  private final int defaultTerminateAfterMaxRecallCount;
   private final long maxMergePreCopyDurationSec;
   private final boolean verboseMetrics;
 
@@ -273,6 +275,8 @@ public ImmutableIndexState(
     defaultSearchTimeoutCheckEvery =
         mergedLiveSettingsWithLocal.getDefaultSearchTimeoutCheckEvery().getValue();
     defaultTerminateAfter = mergedLiveSettingsWithLocal.getDefaultTerminateAfter().getValue();
+    defaultTerminateAfterMaxRecallCount =
+        mergedLiveSettingsWithLocal.getDefaultTerminateAfterMaxRecallCount().getValue();
     maxMergePreCopyDurationSec =
         mergedLiveSettingsWithLocal.getMaxMergePreCopyDurationSec().getValue();
     verboseMetrics = mergedLiveSettingsWithLocal.getVerboseMetrics().getValue();
@@ -823,6 +827,11 @@ public int getDefaultTerminateAfter() {
     return defaultTerminateAfter;
   }
 
+  @Override
+  public int getDefaultTerminateAfterMaxRecallCount() {
+    return defaultTerminateAfterMaxRecallCount;
+  }
+
   @Override
   public int getDefaultSearchTimeoutCheckEvery() {
     return defaultSearchTimeoutCheckEvery;
@@ -924,6 +933,9 @@ static void validateLiveSettings(IndexLiveSettings liveSettings) {
     if (liveSettings.getDefaultTerminateAfter().getValue() < 0) {
       throw new IllegalArgumentException("defaultTerminateAfter must be >= 0");
     }
+    if (liveSettings.getDefaultTerminateAfterMaxRecallCount().getValue() < 0) {
+      throw new IllegalArgumentException("defaultTerminateAfterMaxRecallCount must be >= 0");
+    }
     if (liveSettings.getMaxMergePreCopyDurationSec().getValue() < 0) {
       throw new IllegalArgumentException("maxMergePreCopyDurationSec must be >= 0");
     }

diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/TerminateAfterWrapper.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/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/luceneserver/search/collectors/DocCollector.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/collectors/DocCollector.java
@@ -175,7 +175,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/grpc/StateBackendServerTest.java b/src/test/java/com/yelp/nrtsearch/server/grpc/StateBackendServerTest.java
@@ -823,6 +823,8 @@ public void testSetIndexLiveSettings() throws IOException {
                         IndexLiveSettings.newBuilder()
                             .setDefaultTerminateAfter(
                                 Int32Value.newBuilder().setValue(1000).build())
+                            .setDefaultTerminateAfterMaxRecallCount(
+                                Int32Value.newBuilder().setValue(1000).build())
                             .setSegmentsPerTier(Int32Value.newBuilder().setValue(4).build())
                             .setSliceMaxSegments(Int32Value.newBuilder().setValue(50).build())
                             .setDefaultSearchTimeoutSec(
@@ -832,6 +834,7 @@ public void testSetIndexLiveSettings() throws IOException {
     IndexLiveSettings expectedSettings =
         ImmutableIndexState.DEFAULT_INDEX_LIVE_SETTINGS.toBuilder()
             .setDefaultTerminateAfter(Int32Value.newBuilder().setValue(1000).build())
+            .setDefaultTerminateAfterMaxRecallCount(Int32Value.newBuilder().setValue(1000).build())
             .setSegmentsPerTier(Int32Value.newBuilder().setValue(4).build())
             .setSliceMaxSegments(Int32Value.newBuilder().setValue(50).build())
             .setDefaultSearchTimeoutSec(DoubleValue.newBuilder().setValue(5.1).build())
@@ -1697,6 +1700,7 @@ public void testLiveSettingsV1All() throws IOException {
             .setDefaultSearchTimeoutSec(13.0)
             .setDefaultSearchTimeoutCheckEvery(500)
             .setDefaultTerminateAfter(5000)
+            .setDefaultTerminateAfterMaxRecallCount(6000)
             .build();
 
     LiveSettingsResponse response = primaryClient.getBlockingStub().liveSettings(request);
@@ -1715,6 +1719,7 @@ public void testLiveSettingsV1All() throws IOException {
             .setDefaultSearchTimeoutSec(DoubleValue.newBuilder().setValue(13.0).build())
             .setDefaultSearchTimeoutCheckEvery(Int32Value.newBuilder().setValue(500).build())
             .setDefaultTerminateAfter(Int32Value.newBuilder().setValue(5000).build())
+            .setDefaultTerminateAfterMaxRecallCount(Int32Value.newBuilder().setValue(6000).build())
             .setMaxMergePreCopyDurationSec(UInt64Value.newBuilder().setValue(0))
             .setVerboseMetrics(BoolValue.newBuilder().setValue(false).build())
             .build();

diff --git a/src/test/java/com/yelp/nrtsearch/server/luceneserver/index/ImmutableIndexStateTest.java b/src/test/java/com/yelp/nrtsearch/server/luceneserver/index/ImmutableIndexStateTest.java
@@ -777,6 +777,26 @@ public void testDefaultTerminateAfter_invalid() throws IOException {
     assertLiveSettingException(expectedMsg, b -> b.setDefaultTerminateAfter(wrap(-1)));
   }
 
+  @Test
+  public void testDefaultTerminateAfterMaxRecallCount_default() throws IOException {
+    assertEquals(0, getIndexState(getEmptyState()).getDefaultTerminateAfterMaxRecallCount());
+  }
+
+  @Test
+  public void testDefaultTerminateAfterMaxRecallCount_set() throws IOException {
+    verifyIntLiveSetting(
+        100,
+        ImmutableIndexState::getDefaultTerminateAfterMaxRecallCount,
+        b -> b.setDefaultTerminateAfterMaxRecallCount(wrap(100)));
+  }
+
+  @Test
+  public void testDefaultTerminateAfterMaxRecallCount_invalid() throws IOException {
+    String expectedMsg = "defaultTerminateAfterMaxRecallCount must be >= 0";
+    assertLiveSettingException(
+        expectedMsg, b -> b.setDefaultTerminateAfterMaxRecallCount(wrap(-1)));
+  }
+
   @Test
   public void testMaxMergePreCopyDurationSec_default() throws IOException {
     assertEquals(0, getIndexState(getEmptyState()).getMaxMergePreCopyDurationSec());

diff --git a/src/test/java/com/yelp/nrtsearch/server/luceneserver/search/collectors/DocCollectorTest.java b/src/test/java/com/yelp/nrtsearch/server/luceneserver/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 =