Add default terminate after max recall count

Yelp · Nov 18, 2024 · c340234 · c340234
1 parent 81311d8
commit c340234
Show file tree

Hide file tree

Showing 7 changed files with 59 additions and 0 deletions.
diff --git a/clientlib/src/main/proto/yelp/nrtsearch/luceneserver.proto b/clientlib/src/main/proto/yelp/nrtsearch/luceneserver.proto
@@ -458,6 +458,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 */
@@ -1157,6 +1159,8 @@ message IndexLiveSettings {
     google.protobuf.BoolValue parallelFetchByField = 16;
     // The number of documents/fields per parallel fetch task, default: 50
     google.protobuf.Int32Value parallelFetchChunkSize = 17;
+    // 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/handler/LiveSettingsHandler.java b/src/main/java/com/yelp/nrtsearch/server/handler/LiveSettingsHandler.java
@@ -133,6 +133,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/index/ImmutableIndexState.java b/src/main/java/com/yelp/nrtsearch/server/index/ImmutableIndexState.java
@@ -146,6 +146,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())
           .setParallelFetchByField(BoolValue.newBuilder().setValue(false).build())
@@ -167,6 +168,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;
   private final ParallelFetchConfig parallelFetchConfig;
@@ -262,6 +264,8 @@ public ImmutableIndexState(
     defaultSearchTimeoutCheckEvery =
         mergedLiveSettingsWithLocal.getDefaultSearchTimeoutCheckEvery().getValue();
     defaultTerminateAfter = mergedLiveSettingsWithLocal.getDefaultTerminateAfter().getValue();
+    defaultTerminateAfterMaxRecallCount =
+        mergedLiveSettingsWithLocal.getDefaultTerminateAfterMaxRecallCount().getValue();
     maxMergePreCopyDurationSec =
         mergedLiveSettingsWithLocal.getMaxMergePreCopyDurationSec().getValue();
     verboseMetrics = mergedLiveSettingsWithLocal.getVerboseMetrics().getValue();
@@ -718,6 +722,11 @@ public int getDefaultTerminateAfter() {
     return defaultTerminateAfter;
   }
 
+  @Override
+  public int getDefaultTerminateAfterMaxRecallCount() {
+    return defaultTerminateAfterMaxRecallCount;
+  }
+
   @Override
   public int getDefaultSearchTimeoutCheckEvery() {
     return defaultSearchTimeoutCheckEvery;
@@ -809,6 +818,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/index/IndexState.java b/src/main/java/com/yelp/nrtsearch/server/index/IndexState.java
@@ -487,6 +487,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/tools/cli/LiveSettingsV2Command.java b/src/main/java/com/yelp/nrtsearch/tools/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")
@@ -192,6 +197,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/test/java/com/yelp/nrtsearch/server/grpc/StateBackendServerTest.java b/src/test/java/com/yelp/nrtsearch/server/grpc/StateBackendServerTest.java
@@ -775,6 +775,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(
@@ -784,6 +786,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())
@@ -1647,6 +1650,7 @@ public void testLiveSettingsV1All() throws IOException {
             .setDefaultSearchTimeoutSec(13.0)
             .setDefaultSearchTimeoutCheckEvery(500)
             .setDefaultTerminateAfter(5000)
+            .setDefaultTerminateAfterMaxRecallCount(6000)
             .build();
 
     LiveSettingsResponse response = primaryClient.getBlockingStub().liveSettings(request);
@@ -1665,6 +1669,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())
             .setParallelFetchByField(BoolValue.newBuilder().setValue(false).build())

diff --git a/src/test/java/com/yelp/nrtsearch/server/index/ImmutableIndexStateTest.java b/src/test/java/com/yelp/nrtsearch/server/index/ImmutableIndexStateTest.java
@@ -781,6 +781,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());