Add logging hits time taken in diagnostics

clientlib/src/main/proto/yelp/nrtsearch/search.proto

@@ -975,6 +975,8 @@ message SearchResponse {
         repeated VectorDiagnostics vectorDiagnostics = 13;
         // Amount of gRPC request deadline remaining at the start of search processing
         double initialDeadlineMs = 14;
+        // Time for logging hits
+        double loggingHitsTimeMs = 15;
     }
 
     // Message for query document hit

src/main/java/com/yelp/nrtsearch/server/handler/SearchHandler.java

@@ -300,6 +300,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()

src/main/java/com/yelp/nrtsearch/server/logging/HitsLoggerFetchTask.java

@@ -21,13 +21,16 @@
 import com.yelp.nrtsearch.server.search.FetchTasks.FetchTask;
 import com.yelp.nrtsearch.server.search.SearchContext;
 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();
   }
 }

src/test/java/com/yelp/nrtsearch/server/logging/HitsLoggerTest.java

@@ -16,6 +16,7 @@
 package com.yelp.nrtsearch.server.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);
+  }
 }