Backward compatibility of #794 (#795)

* Backward compatibility of #794 * Addressed PR comment
Yelp · Dec 6, 2024 · 5b34fda · 5b34fda
1 parent 546f040
commit 5b34fda
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 0 deletions.
diff --git a/clientlib/src/main/proto/yelp/nrtsearch/search.proto b/clientlib/src/main/proto/yelp/nrtsearch/search.proto
@@ -766,6 +766,7 @@ message SearchResponse {
         double rescoreTimeMs = 10;
         map<string, double> rescorersTimeMs = 11;
         map<string, Diagnostics> innerHitsDiagnostics = 12;
+        double loggingHitsTimeMs = 15;
     }
 
     message Hit {

diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/SearchHandler.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/SearchHandler.java
@@ -249,6 +249,10 @@ public SearchResponse handle(IndexState indexState, SearchRequest searchRequest)
         diagnostics.setHighlightTimeMs(
             searchContext.getFetchTasks().getHighlightFetchTask().getTimeTakenMs());
       }
+      if (searchContext.getFetchTasks().getHitsLoggerFetchTask() != null) {
+        diagnostics.setLoggingHitsTimeMs(
+            searchContext.getFetchTasks().getHitsLoggerFetchTask().getTimeTakenMs());
+      }
       if (searchContext.getFetchTasks().getInnerHitFetchTaskList() != null) {
         diagnostics.putAllInnerHitsDiagnostics(
             searchContext.getFetchTasks().getInnerHitFetchTaskList().stream()

diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/logging/HitsLoggerFetchTask.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/logging/HitsLoggerFetchTask.java
@@ -21,13 +21,16 @@
 import com.yelp.nrtsearch.server.luceneserver.search.SearchContext;
 import com.yelp.nrtsearch.server.plugins.HitsLoggerPlugin;
 import java.util.List;
+import java.util.concurrent.atomic.DoubleAdder;
 
 /**
  * Implementation of {@link FetchTask} which holds the required context to be able to log hits for a
  * search request.
  */
 public class HitsLoggerFetchTask implements FetchTask {
+  private static final double TEN_TO_THE_POWER_SIX = Math.pow(10, 6);
   private final HitsLogger hitsLogger;
+  private final DoubleAdder timeTakenMs = new DoubleAdder();
 
   public HitsLoggerFetchTask(LoggingHits loggingHits) {
     this.hitsLogger = HitsLoggerCreator.getInstance().createHitsLogger(loggingHits);
@@ -42,6 +45,17 @@ public HitsLoggerFetchTask(LoggingHits loggingHits) {
    */
   @Override
   public void processAllHits(SearchContext searchContext, List<SearchResponse.Hit.Builder> hits) {
+    long startTime = System.nanoTime();
     hitsLogger.log(searchContext, hits);
+    timeTakenMs.add(((System.nanoTime() - startTime) / TEN_TO_THE_POWER_SIX));
+  }
+
+  /**
+   * Get the total time taken so far to logging hits.
+   *
+   * @return Total time taken to logging hits in ms.
+   */
+  public double getTimeTakenMs() {
+    return timeTakenMs.doubleValue();
   }
 }
diff --git a/src/test/java/com/yelp/nrtsearch/server/luceneserver/logging/HitsLoggerTest.java b/src/test/java/com/yelp/nrtsearch/server/luceneserver/logging/HitsLoggerTest.java
@@ -16,6 +16,7 @@
 package com.yelp.nrtsearch.server.luceneserver.logging;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 
 import com.google.protobuf.Struct;
 import com.google.protobuf.Value;
@@ -192,4 +193,34 @@ public void testCustomHitsLoggerWithoutParam() {
 
     assertEquals(expectedLogMessage, HitsLoggerTest.logMessage);
   }
+
+  @Test
+  public void testLoggingTimeTaken() {
+    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")
+                    .setParams(
+                        Struct.newBuilder()
+                            .putFields(
+                                "external_value", Value.newBuilder().setStringValue("abc").build()))
+                    .build())
+            .build();
+    SearchResponse response = getGrpcServer().getBlockingStub().search(request);
+
+    assertTrue(response.getDiagnostics().getLoggingHitsTimeMs() > 0);
+  }
 }