Add output size for logging hits in diagnostics

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

@@ -988,6 +988,8 @@ message SearchResponse {
         double initialDeadlineMs = 14;
         // Time for logging hits
         double loggingHitsTimeMs = 15;
+        // The size of the data logged in bytes
+        int64 loggingHitsOutputSizeBytes = 16;
     }
 
     // Message for query document hit

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

@@ -315,6 +315,8 @@ public SearchResponse handle(IndexState indexState, SearchRequest searchRequest)
       if (searchContext.getFetchTasks().getHitsLoggerFetchTask() != null) {
         diagnostics.setLoggingHitsTimeMs(
             searchContext.getFetchTasks().getHitsLoggerFetchTask().getTimeTakenMs());
+        diagnostics.setLoggingHitsOutputSizeBytes(
+            searchContext.getFetchTasks().getHitsLoggerFetchTask().getOutputSize());
       }
       if (searchContext.getFetchTasks().getInnerHitFetchTaskList() != null) {
         diagnostics.putAllInnerHitsDiagnostics(

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

@@ -30,6 +30,7 @@ public interface HitsLogger {
    *
    * @param context the {@link SearchContext} to keep the contexts for this search request
    * @param hits query hits
+   * @return the size in bytes of the logged data
    */
-  void log(SearchContext context, List<SearchResponse.Hit.Builder> hits);
+  int log(SearchContext context, List<SearchResponse.Hit.Builder> hits);
 }

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

@@ -33,6 +33,8 @@ public class HitsLoggerFetchTask implements FetchTask {
   private final int hitsToLog;
   private final DoubleAdder timeTakenMs = new DoubleAdder();
 
+  private int outputSize;
+
   public HitsLoggerFetchTask(LoggingHits loggingHits) {
     this.hitsLogger = HitsLoggerCreator.getInstance().createHitsLogger(loggingHits);
     this.hitsToLog = loggingHits.getHitsToLog();
@@ -52,14 +54,18 @@ public void processAllHits(SearchContext searchContext, List<SearchResponse.Hit.
     // 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()));
+      outputSize = hitsLogger.log(searchContext, hits.subList(0, searchContext.getHitsToLog()));
     } else {
-      hitsLogger.log(searchContext, hits);
+      outputSize = hitsLogger.log(searchContext, hits);
     }
 
     timeTakenMs.add(((System.nanoTime() - startTime) / TEN_TO_THE_POWER_SIX));
   }
 
+  public int getOutputSize() {
+    return outputSize;
+  }
+
   /**
    * Get the total time taken so far to logging hits.
    *

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

@@ -58,14 +58,18 @@ public CustomHitsLogger(Map<String, Object> params) {
       }
 
       @Override
-      public void log(SearchContext context, List<SearchResponse.Hit.Builder> hits) {}
+      public int log(SearchContext context, List<SearchResponse.Hit.Builder> hits) {
+        return 0;
+      }
     }
 
     static class CustomHitsLogger2 implements HitsLogger {
       public CustomHitsLogger2(Map<String, Object> params) {}
 
       @Override
-      public void log(SearchContext context, List<SearchResponse.Hit.Builder> hits) {}
+      public int log(SearchContext context, List<SearchResponse.Hit.Builder> hits) {
+        return 0;
+      }
     }
 
     @Override

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

@@ -100,7 +100,7 @@ public CustomHitsLogger(Map<String, Object> params) {
       }
 
       @Override
-      public void log(SearchContext context, List<SearchResponse.Hit.Builder> hits) {
+      public int log(SearchContext context, List<SearchResponse.Hit.Builder> hits) {
         HitsLoggerTest.logMessage = "LOGGED ";
 
         for (SearchResponse.Hit.Builder hit : hits) {
@@ -113,6 +113,8 @@ public void log(SearchContext context, List<SearchResponse.Hit.Builder> hits) {
         if (!params.isEmpty()) {
           HitsLoggerTest.logMessage += " " + params;
         }
+
+        return HitsLoggerTest.logMessage.length();
       }
     }
 
@@ -124,7 +126,7 @@ public CustomHitsLogger2(Map<String, Object> params) {
       }
 
       @Override
-      public void log(SearchContext context, List<SearchResponse.Hit.Builder> hits) {
+      public int log(SearchContext context, List<SearchResponse.Hit.Builder> hits) {
         HitsLoggerTest.logMessage = "LOGGED_2 ";
 
         for (SearchResponse.Hit.Builder hit : hits) {
@@ -137,6 +139,8 @@ public void log(SearchContext context, List<SearchResponse.Hit.Builder> hits) {
         if (!params.isEmpty()) {
           HitsLoggerTest.logMessage += " " + params;
         }
+
+        return HitsLoggerTest.logMessage.length();
       }
     }
 
@@ -463,4 +467,36 @@ public void testLoggingTimeTaken() {
 
     assertTrue(response.getDiagnostics().getLoggingHitsTimeMs() > 0);
   }
+
+  @Test
+  public void testLoggingHitsOutputByteSize() {
+    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(1)
+                    .setParams(
+                        Struct.newBuilder()
+                            .putFields(
+                                "external_value", Value.newBuilder().setStringValue("abc").build()))
+                    .build())
+            .build();
+    SearchResponse response = getGrpcServer().getBlockingStub().search(request);
+
+    assertEquals(10, response.getTotalHits().getValue());
+    assertEquals(39, response.getDiagnostics().getLoggingHitsOutputSizeBytes());
+  }
 }