From 9465aa7d89834e45987cae5e1fde0124c4b1609e Mon Sep 17 00:00:00 2001 From: Akshat Jain Date: Mon, 5 Aug 2024 12:16:10 +0530 Subject: [PATCH 1/2] Memoize the redundant calls to overlord in sql statements endpoint --- .../sql/resources/SqlStatementResource.java | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java index 322727aea92d..4f504ffe0010 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java +++ b/extensions-core/multi-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,19 +589,19 @@ private Optional getStatementStatus( MSQControllerTask msqControllerTask = getMSQControllerTaskAndCheckPermission(queryId, authenticationResult, forAction); SqlStatementState sqlStatementState = SqlStatementResourceHelper.getSqlStatementState(statusPlus); - Supplier> msqTaskReportPayloadSupplier = () -> { + Supplier 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( @@ -608,7 +609,7 @@ private Optional getStatementStatus( taskResponse, statusPlus, sqlStatementState, - msqTaskReportPayloadSupplier.get().orElse(null), + msqTaskReportPayloadSupplier.get(), jsonMapper, detail ); @@ -627,9 +628,9 @@ private Optional 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 )); } } From 8ba104b4489b2b9ed6a34bfb4500216219050dbf Mon Sep 17 00:00:00 2001 From: Akshat Jain Date: Mon, 5 Aug 2024 13:03:12 +0530 Subject: [PATCH 2/2] Address review comment --- .../sql/resources/SqlStatementResource.java | 22 +++++++++---------- 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java index 4f504ffe0010..f56622804786 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java @@ -22,7 +22,6 @@ 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; @@ -114,7 +113,6 @@ import java.util.Map; import java.util.Objects; import java.util.Optional; -import java.util.function.Supplier; import java.util.stream.Collectors; @@ -589,19 +587,19 @@ private Optional getStatementStatus( MSQControllerTask msqControllerTask = getMSQControllerTaskAndCheckPermission(queryId, authenticationResult, forAction); SqlStatementState sqlStatementState = SqlStatementResourceHelper.getSqlStatementState(statusPlus); - Supplier msqTaskReportPayloadSupplier = Suppliers.memoize(() -> { + MSQTaskReportPayload taskReportPayload = null; + if (detail || SqlStatementState.FAILED == sqlStatementState) { try { - return SqlStatementResourceHelper.getPayload( + taskReportPayload = SqlStatementResourceHelper.getPayload( contactOverlord(overlordClient.taskReportAsMap(queryId), queryId) ); } catch (DruidException e) { - if (e.getErrorCode().equals("notFound") || e.getMessage().contains("Unable to contact overlord")) { - return null; + if (!e.getErrorCode().equals("notFound") && !e.getMessage().contains("Unable to contact overlord")) { + throw e; } - throw e; } - }); + } if (SqlStatementState.FAILED == sqlStatementState) { return SqlStatementResourceHelper.getExceptionPayload( @@ -609,7 +607,7 @@ private Optional getStatementStatus( taskResponse, statusPlus, sqlStatementState, - msqTaskReportPayloadSupplier.get(), + taskReportPayload, jsonMapper, detail ); @@ -628,9 +626,9 @@ private Optional getStatementStatus( msqControllerTask.getQuerySpec().getDestination() ).orElse(null) : null, null, - detail ? SqlStatementResourceHelper.getQueryStagesReport(msqTaskReportPayloadSupplier.get()) : null, - detail ? SqlStatementResourceHelper.getQueryCounters(msqTaskReportPayloadSupplier.get()) : null, - detail ? SqlStatementResourceHelper.getQueryWarningDetails(msqTaskReportPayloadSupplier.get()) : null + SqlStatementResourceHelper.getQueryStagesReport(taskReportPayload), + SqlStatementResourceHelper.getQueryCounters(taskReportPayload), + SqlStatementResourceHelper.getQueryWarningDetails(taskReportPayload) )); } }