Add search diagnostics to deadline exceeded exceptions (#621)

* Add search diagnostics to deadline exceeded exceptions --------- Co-authored-by: swekannan <shwethu.kannan93@gmail.com>
Yelp · Feb 7, 2024 · 5637580 · 5637580
1 parent b0a4763
commit 5637580
Show file tree

Hide file tree

Showing 3 changed files with 67 additions and 3 deletions.
diff --git a/src/main/java/com/yelp/nrtsearch/server/grpc/DeadlineUtils.java b/src/main/java/com/yelp/nrtsearch/server/grpc/DeadlineUtils.java
@@ -15,6 +15,7 @@
  */
 package com.yelp.nrtsearch.server.grpc;
 
+import com.yelp.nrtsearch.server.grpc.SearchResponse.Diagnostics;
 import com.yelp.nrtsearch.server.monitoring.DeadlineMetrics;
 import io.grpc.Context;
 import io.grpc.Deadline;
@@ -54,4 +55,21 @@ public static void checkDeadline(String message, String operation) {
       }
     }
   }
+
+  public static void checkDeadline(
+      String message, Diagnostics.Builder diagnostics, String operation) {
+    if (cancellationEnabled) {
+      Deadline deadline = Context.current().getDeadline();
+      if (deadline != null && deadline.isExpired()) {
+        DeadlineMetrics.nrtDeadlineCancelCount.labels(operation).inc();
+        throw Status.CANCELLED
+            .withDescription(
+                "Request deadline exceeded: "
+                    + message
+                    + ", Search Diagnostics: "
+                    + diagnostics.toString())
+            .asRuntimeException();
+      }
+    }
+  }
 }
diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/SearchHandler.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/SearchHandler.java
@@ -191,7 +191,7 @@ public SearchResponse handle(IndexState indexState, SearchRequest searchRequest)
 
       diagnostics.setFirstPassSearchTimeMs(((System.nanoTime() - searchStartTime) / 1000000.0));
 
-      DeadlineUtils.checkDeadline("SearchHandler: post recall", "SEARCH");
+      DeadlineUtils.checkDeadline("SearchHandler: post recall", diagnostics, "SEARCH");
 
       // add detailed timing metrics for query execution
       if (profileResultBuilder != null) {
@@ -206,7 +206,8 @@ public SearchResponse handle(IndexState indexState, SearchRequest searchRequest)
           hits = rescorer.rescore(hits, searchContext);
           long endNS = System.nanoTime();
           diagnostics.putRescorersTimeMs(rescorer.getName(), (endNS - startNS) / 1000000.0);
-          DeadlineUtils.checkDeadline("SearchHandler: post " + rescorer.getName(), "SEARCH");
+          DeadlineUtils.checkDeadline(
+              "SearchHandler: post " + rescorer.getName(), diagnostics, "SEARCH");
         }
         diagnostics.setRescoreTimeMs(((System.nanoTime() - rescoreStartTime) / 1000000.0));
       }
@@ -287,7 +288,7 @@ public SearchResponse handle(IndexState indexState, SearchRequest searchRequest)
     }
 
     // if we are out of time, don't bother with serialization
-    DeadlineUtils.checkDeadline("SearchHandler: end", "SEARCH");
+    DeadlineUtils.checkDeadline("SearchHandler: end", diagnostics, "SEARCH");
     SearchResponse searchResponse = searchContext.getResponseBuilder().build();
     if (!warming && searchContext.getIndexState().getVerboseMetrics()) {
       VerboseIndexCollector.updateSearchResponseMetrics(

diff --git a/src/test/java/com/yelp/nrtsearch/server/grpc/DeadlineUtilsTest.java b/src/test/java/com/yelp/nrtsearch/server/grpc/DeadlineUtilsTest.java
@@ -104,6 +104,51 @@ public void testWithDeadlineReached() {
     }
   }
 
+  @Test
+  public void testWithDeadlineReachedWithDiagnostics() {
+    var diagnostics = SearchResponse.Diagnostics.newBuilder();
+    diagnostics
+        .setFirstPassSearchTimeMs(100)
+        .setGetFieldsTimeMs(50)
+        .setRescoreTimeMs(20)
+        .setHighlightTimeMs(10);
+    DeadlineUtils.setCancellationEnabled(false);
+    CancellableContext context =
+        Context.current().withDeadlineAfter(1, TimeUnit.MILLISECONDS, executorService);
+    try {
+      context.run(
+          () -> {
+            try {
+              Thread.sleep(5);
+            } catch (InterruptedException ignored) {
+            }
+            DeadlineUtils.checkDeadline("test", diagnostics, "TEST");
+          });
+    } finally {
+      context.cancel(null);
+    }
+    DeadlineUtils.setCancellationEnabled(true);
+    context = Context.current().withDeadlineAfter(1, TimeUnit.MILLISECONDS, executorService);
+    try {
+      context.run(
+          () -> {
+            try {
+              Thread.sleep(5);
+            } catch (InterruptedException ignored) {
+            }
+            DeadlineUtils.checkDeadline("test", diagnostics, "TEST");
+          });
+      fail();
+    } catch (StatusRuntimeException e) {
+      assertEquals(Status.CANCELLED.getCode(), e.getStatus().getCode());
+      assertEquals(
+          "Request deadline exceeded: test, Search Diagnostics: firstPassSearchTimeMs: 100.0\nhighlightTimeMs: 10.0\ngetFieldsTimeMs: 50.0\nrescoreTimeMs: 20.0\n",
+          e.getStatus().getDescription());
+    } finally {
+      context.cancel(null);
+    }
+  }
+
   @Test
   public void testMetricsCounter() {
     int initialCount = getCancelMetricCount();