Add some conditions to skip hits logging (#817)

* Skip logging when warming * Added unit test
Yelp · Feb 3, 2025 · f12984b · f12984b
1 parent 5a0619f
commit f12984b
Show file tree

Hide file tree

Showing 7 changed files with 66 additions and 14 deletions.
diff --git a/src/main/java/com/yelp/nrtsearch/server/handler/SearchHandler.java b/src/main/java/com/yelp/nrtsearch/server/handler/SearchHandler.java
@@ -169,7 +169,7 @@ public SearchResponse handle(IndexState indexState, SearchRequest searchRequest)
 
       searchContext =
           SearchRequestProcessor.buildContextForRequest(
-              searchRequest, indexState, shardState, s, diagnostics, profileResultBuilder);
+              searchRequest, indexState, shardState, s, diagnostics, profileResultBuilder, warming);
 
       long searchStartTime = System.nanoTime();
 

diff --git a/src/main/java/com/yelp/nrtsearch/server/logging/HitsLoggerFetchTask.java b/src/main/java/com/yelp/nrtsearch/server/logging/HitsLoggerFetchTask.java
@@ -47,17 +47,20 @@ public HitsLoggerFetchTask(LoggingHits loggingHits) {
    */
   @Override
   public void processAllHits(SearchContext searchContext, List<SearchResponse.Hit.Builder> hits) {
-    long startTime = System.nanoTime();
+    if (!searchContext.isWarming() && searchContext.getHitsToLog() > 0) {
+      long startTime = System.nanoTime();
 
-    // hits list can contain extra hits that don't need to be logged, otherwise, pass all hits that
-    // can be logged
-    if (searchContext.getHitsToLog() < hits.size()) {
-      hitsLogger.log(searchContext, hits.subList(0, searchContext.getHitsToLog()));
-    } else {
-      hitsLogger.log(searchContext, hits);
-    }
+      // hits list can contain extra hits that don't need to be logged, otherwise, pass all hits
+      // that
+      // can be logged
+      if (searchContext.getHitsToLog() < hits.size()) {
+        hitsLogger.log(searchContext, hits.subList(0, searchContext.getHitsToLog()));
+      } else {
+        hitsLogger.log(searchContext, hits);
+      }
 
-    timeTakenMs.add(((System.nanoTime() - startTime) / TEN_TO_THE_POWER_SIX));
+      timeTakenMs.add(((System.nanoTime() - startTime) / TEN_TO_THE_POWER_SIX));
+    }
   }
 
   /**

diff --git a/src/main/java/com/yelp/nrtsearch/server/search/SearchContext.java b/src/main/java/com/yelp/nrtsearch/server/search/SearchContext.java
@@ -58,6 +58,7 @@ public enum VectorScoringMode {
   }
 
   private final boolean explain;
+  private final boolean warming;
 
   private SearchContext(Builder builder, boolean validate) {
     this.indexState = builder.indexState;
@@ -78,6 +79,7 @@ private SearchContext(Builder builder, boolean validate) {
     this.extraContext = builder.extraContext;
     this.queryNestedPath = builder.queryNestedPath;
     this.explain = builder.explain;
+    this.warming = builder.warming;
 
     if (validate) {
       validate();
@@ -192,6 +194,10 @@ public boolean isExplain() {
     return explain;
   }
 
+  public boolean isWarming() {
+    return warming;
+  }
+
   /** Get new context builder instance * */
   public static Builder newBuilder() {
     return new Builder();
@@ -249,6 +255,7 @@ public static class Builder {
     private Map<String, Object> extraContext;
     private String queryNestedPath;
     private boolean explain;
+    private boolean warming;
 
     private Builder() {}
 
@@ -360,6 +367,11 @@ public Builder setExplain(boolean explain) {
       return this;
     }
 
+    public Builder setWarming(boolean warming) {
+      this.warming = warming;
+      return this;
+    }
+
     /**
      * Use builder to create new search context. Skipping validation is possible, but mainly
      * intended for tests that do not require a complete context.

diff --git a/src/main/java/com/yelp/nrtsearch/server/search/SearchRequestProcessor.java b/src/main/java/com/yelp/nrtsearch/server/search/SearchRequestProcessor.java
@@ -116,7 +116,8 @@ public static SearchContext buildContextForRequest(
       ShardState shardState,
       SearcherTaxonomyManager.SearcherAndTaxonomy searcherAndTaxonomy,
       SearchResponse.Diagnostics.Builder diagnostics,
-      ProfileResult.Builder profileResult)
+      ProfileResult.Builder profileResult,
+      boolean warming)
       throws IOException {
 
     SearchContext.Builder contextBuilder = SearchContext.newBuilder();
@@ -130,7 +131,8 @@ public static SearchContext buildContextForRequest(
         .setTimestampSec(System.currentTimeMillis() / 1000)
         .setStartHit(searchRequest.getStartHit())
         .setTopHits(searchRequest.getTopHits())
-        .setExplain(searchRequest.getExplain());
+        .setExplain(searchRequest.getExplain())
+        .setWarming(warming);
 
     Map<String, FieldDef> queryVirtualFields = getVirtualFields(indexState, searchRequest);
     Map<String, FieldDef> queryRuntimeFields = getRuntimeFields(indexState, searchRequest);

diff --git a/src/test/java/com/yelp/nrtsearch/server/grpc/TimeoutTest.java b/src/test/java/com/yelp/nrtsearch/server/grpc/TimeoutTest.java
@@ -469,7 +469,8 @@ private TopDocs queryWithFunction(SearchRequest request, Function<SearchContext,
               shardState,
               s,
               Diagnostics.newBuilder(),
-              ProfileResult.newBuilder());
+              ProfileResult.newBuilder(),
+              false);
       return func.apply(context);
     } finally {
       if (s != null) {

diff --git a/src/test/java/com/yelp/nrtsearch/server/logging/HitsLoggerTest.java b/src/test/java/com/yelp/nrtsearch/server/logging/HitsLoggerTest.java
@@ -434,6 +434,39 @@ public void testResponseSizeReductionWithStartHitAndHitsToLogLessThanHitsCount()
     assertEquals(5, response.getHitsCount());
   }
 
+  @Test
+  public void testLoggingWithZeroHitsToLog() {
+    SearchRequest request =
+        SearchRequest.newBuilder()
+            .setTopHits(1)
+            .setStartHit(0)
+            .setIndexName(DEFAULT_TEST_INDEX)
+            .addRetrieveFields("doc_id")
+            .setQuery(
+                Query.newBuilder()
+                    .setTermQuery(
+                        TermQuery.newBuilder()
+                            .setField("vendor_name")
+                            .setTextValue("vendor")
+                            .build())
+                    .build())
+            .setLoggingHits(
+                LoggingHits.newBuilder()
+                    .setName("custom_logger")
+                    .setHitsToLog(0)
+                    .setParams(
+                        Struct.newBuilder()
+                            .putFields(
+                                "external_value", Value.newBuilder().setStringValue("abc").build()))
+                    .build())
+            .build();
+    SearchResponse response = getGrpcServer().getBlockingStub().search(request);
+
+    assertEquals("", HitsLoggerTest.logMessage);
+    assertEquals(10, response.getTotalHits().getValue());
+    assertEquals(1, response.getHitsCount());
+  }
+
   @Test
   public void testLoggingTimeTaken() {
     SearchRequest request =
@@ -453,6 +486,7 @@ public void testLoggingTimeTaken() {
             .setLoggingHits(
                 LoggingHits.newBuilder()
                     .setName("custom_logger")
+                    .setHitsToLog(1)
                     .setParams(
                         Struct.newBuilder()
                             .putFields(

diff --git a/src/test/java/com/yelp/nrtsearch/server/query/QueryNodeMapperTest.java b/src/test/java/com/yelp/nrtsearch/server/query/QueryNodeMapperTest.java
@@ -198,7 +198,7 @@ public void testFunctionScoreQueryField() throws IOException {
       s = shardState.acquire();
       SearchContext context =
           SearchRequestProcessor.buildContextForRequest(
-              request, indexState, shardState, s, Diagnostics.newBuilder(), null);
+              request, indexState, shardState, s, Diagnostics.newBuilder(), null, false);
       org.apache.lucene.queries.function.FunctionScoreQuery query =
           (org.apache.lucene.queries.function.FunctionScoreQuery) context.getQuery();
       assertNotNull(query);