Memoize the redundant calls to overlord in sql statements endpoint

apache · Aug 5, 2024 · 9465aa7 · 9465aa7
1 parent c7eacd0
commit 9465aa7
Showing 1 changed file with 10 additions and 9 deletions.
diff --git a/...ti-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java b/...ti-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java
@@ -22,6 +22,7 @@
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Suppliers;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.io.CountingOutputStream;
 import com.google.common.util.concurrent.ListenableFuture;
@@ -588,27 +589,27 @@ private Optional<SqlStatementResult> getStatementStatus(
     MSQControllerTask msqControllerTask = getMSQControllerTaskAndCheckPermission(queryId, authenticationResult, forAction);
     SqlStatementState sqlStatementState = SqlStatementResourceHelper.getSqlStatementState(statusPlus);
 
-    Supplier<Optional<MSQTaskReportPayload>> msqTaskReportPayloadSupplier = () -> {
+    Supplier<MSQTaskReportPayload> msqTaskReportPayloadSupplier = Suppliers.memoize(() -> {
       try {
-        return Optional.ofNullable(SqlStatementResourceHelper.getPayload(
+        return SqlStatementResourceHelper.getPayload(
             contactOverlord(overlordClient.taskReportAsMap(queryId), queryId)
-        ));
+        );
       }
       catch (DruidException e) {
         if (e.getErrorCode().equals("notFound") || e.getMessage().contains("Unable to contact overlord")) {
-          return Optional.empty();
+          return null;
         }
         throw e;
       }
-    };
+    });
 
     if (SqlStatementState.FAILED == sqlStatementState) {
       return SqlStatementResourceHelper.getExceptionPayload(
           queryId,
           taskResponse,
           statusPlus,
           sqlStatementState,
-          msqTaskReportPayloadSupplier.get().orElse(null),
+          msqTaskReportPayloadSupplier.get(),
           jsonMapper,
           detail
       );
@@ -627,9 +628,9 @@ private Optional<SqlStatementResult> getStatementStatus(
               msqControllerTask.getQuerySpec().getDestination()
           ).orElse(null) : null,
           null,
-          detail ? SqlStatementResourceHelper.getQueryStagesReport(msqTaskReportPayloadSupplier.get().orElse(null)) : null,
-          detail ? SqlStatementResourceHelper.getQueryCounters(msqTaskReportPayloadSupplier.get().orElse(null)) : null,
-          detail ? SqlStatementResourceHelper.getQueryWarningDetails(msqTaskReportPayloadSupplier.get().orElse(null)) : null
+          detail ? SqlStatementResourceHelper.getQueryStagesReport(msqTaskReportPayloadSupplier.get()) : null,
+          detail ? SqlStatementResourceHelper.getQueryCounters(msqTaskReportPayloadSupplier.get()) : null,
+          detail ? SqlStatementResourceHelper.getQueryWarningDetails(msqTaskReportPayloadSupplier.get()) : null
       ));
     }
   }