From a145789d39c90533340ac20a30dcc5d8ad26be2b Mon Sep 17 00:00:00 2001 From: Luana Fragoso Date: Mon, 6 Jan 2025 18:18:40 -0800 Subject: [PATCH] Added output size for logging in diagnostics --- .../main/proto/yelp/nrtsearch/search.proto | 2 + .../server/handler/SearchHandler.java | 2 + .../nrtsearch/server/logging/HitsLogger.java | 3 +- .../server/logging/HitsLoggerFetchTask.java | 10 ++++- .../server/logging/HitsLoggerCreatorTest.java | 8 +++- .../server/logging/HitsLoggerTest.java | 40 ++++++++++++++++++- 6 files changed, 58 insertions(+), 7 deletions(-) diff --git a/clientlib/src/main/proto/yelp/nrtsearch/search.proto b/clientlib/src/main/proto/yelp/nrtsearch/search.proto index d3ffe0622..0cd018408 100644 --- a/clientlib/src/main/proto/yelp/nrtsearch/search.proto +++ b/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 diff --git a/src/main/java/com/yelp/nrtsearch/server/handler/SearchHandler.java b/src/main/java/com/yelp/nrtsearch/server/handler/SearchHandler.java index 6f9681b53..76d1e6154 100644 --- a/src/main/java/com/yelp/nrtsearch/server/handler/SearchHandler.java +++ b/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( diff --git a/src/main/java/com/yelp/nrtsearch/server/logging/HitsLogger.java b/src/main/java/com/yelp/nrtsearch/server/logging/HitsLogger.java index 509aa7766..59f1bb7bc 100644 --- a/src/main/java/com/yelp/nrtsearch/server/logging/HitsLogger.java +++ b/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 hits); + int log(SearchContext context, List hits); } diff --git a/src/main/java/com/yelp/nrtsearch/server/logging/HitsLoggerFetchTask.java b/src/main/java/com/yelp/nrtsearch/server/logging/HitsLoggerFetchTask.java index 9bd8715bd..e1a63970d 100644 --- a/src/main/java/com/yelp/nrtsearch/server/logging/HitsLoggerFetchTask.java +++ b/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 params) { } @Override - public void log(SearchContext context, List hits) {} + public int log(SearchContext context, List hits) { + return 0; + } } static class CustomHitsLogger2 implements HitsLogger { public CustomHitsLogger2(Map params) {} @Override - public void log(SearchContext context, List hits) {} + public int log(SearchContext context, List hits) { + return 0; + } } @Override diff --git a/src/test/java/com/yelp/nrtsearch/server/logging/HitsLoggerTest.java b/src/test/java/com/yelp/nrtsearch/server/logging/HitsLoggerTest.java index 0475f4922..fbec464d5 100644 --- a/src/test/java/com/yelp/nrtsearch/server/logging/HitsLoggerTest.java +++ b/src/test/java/com/yelp/nrtsearch/server/logging/HitsLoggerTest.java @@ -100,7 +100,7 @@ public CustomHitsLogger(Map params) { } @Override - public void log(SearchContext context, List hits) { + public int log(SearchContext context, List hits) { HitsLoggerTest.logMessage = "LOGGED "; for (SearchResponse.Hit.Builder hit : hits) { @@ -113,6 +113,8 @@ public void log(SearchContext context, List hits) { if (!params.isEmpty()) { HitsLoggerTest.logMessage += " " + params; } + + return HitsLoggerTest.logMessage.length(); } } @@ -124,7 +126,7 @@ public CustomHitsLogger2(Map params) { } @Override - public void log(SearchContext context, List hits) { + public int log(SearchContext context, List hits) { HitsLoggerTest.logMessage = "LOGGED_2 "; for (SearchResponse.Hit.Builder hit : hits) { @@ -137,6 +139,8 @@ public void log(SearchContext context, List 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()); + } }