From dfa5c79d2f0725ab9f25ac9a0056986ed8966189 Mon Sep 17 00:00:00 2001 From: swethakann Date: Fri, 6 Dec 2024 14:05:12 -0800 Subject: [PATCH] Add default terminate after max recall count (#790) (#793) Add default terminate after max recall count --- .../proto/yelp/nrtsearch/luceneserver.proto | 4 ++ .../server/cli/LiveSettingsV2Command.java | 9 ++++ .../server/luceneserver/IndexState.java | 3 ++ .../luceneserver/LiveSettingsHandler.java | 6 +++ .../index/ImmutableIndexState.java | 12 ++++++ .../search/TerminateAfterWrapper.java | 5 +++ .../search/collectors/DocCollector.java | 2 +- .../server/grpc/StateBackendServerTest.java | 5 +++ .../index/ImmutableIndexStateTest.java | 20 +++++++++ .../search/collectors/DocCollectorTest.java | 42 ++++++++++++++++++- 10 files changed, 106 insertions(+), 2 deletions(-) diff --git a/clientlib/src/main/proto/yelp/nrtsearch/luceneserver.proto b/clientlib/src/main/proto/yelp/nrtsearch/luceneserver.proto index 3529a28f6..34cfab994 100644 --- 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 index dbeef4338..5e26e1e7d 100644 --- 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 { 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 index 49b55f04d..8c8078538 100644 --- 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 index 014a9ebbb..396ef1e16 100644 --- 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 index 64beee0ae..2c24b90d5 100644 --- 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 index 6b576375e..f4ae60ddd 100644 --- 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 index 9d2db8059..ac6a3c11d 100644 --- 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 @@ CollectorManager 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 index 09a12670f..9e9a1a69f 100644 --- 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 index 0b0d4a4ea..6a5aa878f 100644 --- 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 index b1f2a398e..3aba08b6e 100644 --- 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 =