[native] Migrate TaskStatus, TaskUpdateRequest, TaskInfo to Thrift with JSON fields in CPP

presto-native-execution/presto_cpp/main/TaskResource.cpp

@@ -208,27 +208,37 @@ proxygen::RequestHandler* TaskResource::createOrUpdateTaskImpl(
     const std::vector<std::string>& pathMatch,
     const std::function<std::unique_ptr<protocol::TaskInfo>(
         const protocol::TaskId& taskId,
-        const std::string& updateJson,
+        const std::string& requestBody,
         const bool summarize,
-        long startProcessCpuTime)>& createOrUpdateFunc) {
+        long startProcessCpuTime,
+        bool receiveThrift)>& createOrUpdateFunc) {
   protocol::TaskId taskId = pathMatch[1];
   bool summarize = message->hasQueryParam("summarize");
+
+  auto& headers = message->getHeaders();
+  const auto& acceptHeader = headers.getSingleOrEmpty(proxygen::HTTP_HEADER_ACCEPT);
+  const auto sendThrift =
+      acceptHeader.find(http::kMimeTypeApplicationThrift) != std::string::npos;
+  const auto& contentHeader = headers.getSingleOrEmpty(proxygen::HTTP_HEADER_CONTENT_TYPE);
+  const auto receiveThrift =
+      contentHeader.find(http::kMimeTypeApplicationThrift) != std::string::npos;
+
   return new http::CallbackRequestHandler(
-      [this, taskId, summarize, createOrUpdateFunc](
+      [this, taskId, summarize, createOrUpdateFunc, sendThrift, receiveThrift](
           proxygen::HTTPMessage* /*message*/,
           const std::vector<std::unique_ptr<folly::IOBuf>>& body,
           proxygen::ResponseHandler* downstream,
           std::shared_ptr<http::CallbackRequestHandlerState> handlerState) {
         folly::via(
             httpSrvCpuExecutor_,
-            [this, &body, taskId, summarize, createOrUpdateFunc]() {
+            [this, &body, taskId, summarize, createOrUpdateFunc, receiveThrift]() {
               const auto startProcessCpuTimeNs = util::getProcessCpuTimeNs();
-              std::string updateJson = util::extractMessageBody(body);
+              std::string requestBody = util::extractMessageBody(body);
 
               std::unique_ptr<protocol::TaskInfo> taskInfo;
               try {
                 taskInfo = createOrUpdateFunc(
-                    taskId, updateJson, summarize, startProcessCpuTimeNs);
+                    taskId, requestBody, summarize, startProcessCpuTimeNs, receiveThrift);
               } catch (const velox::VeloxException& e) {
                 // Creating an empty task, putting errors inside so that next
                 // status fetch from coordinator will catch the error and well
@@ -243,12 +253,19 @@ proxygen::RequestHandler* TaskResource::createOrUpdateTaskImpl(
                   throw;
                 }
               }
-              return json(*taskInfo);
+              return taskInfo;
             })
             .via(folly::EventBaseManager::get()->getEventBase())
-            .thenValue([downstream, handlerState](auto&& taskInfoJson) {
+            .thenValue([downstream, handlerState, sendThrift](auto taskInfo) {
               if (!handlerState->requestExpired()) {
-                http::sendOkResponse(downstream, taskInfoJson);
+                if (sendThrift) {
+                  thrift::TaskInfo thriftTaskInfo;
+                  toThrift(*taskInfo, thriftTaskInfo);
+                  http::sendOkThriftResponse(
+                      downstream, thriftWrite(thriftTaskInfo));
+                } else {
+                  http::sendOkResponse(downstream, json(*taskInfo));
+                }
               }
             })
             .thenError(
@@ -275,11 +292,12 @@ proxygen::RequestHandler* TaskResource::createOrUpdateBatchTask(
       message,
       pathMatch,
       [&](const protocol::TaskId& taskId,
-          const std::string& updateJson,
+          const std::string& requestBody,
           const bool summarize,
-          long startProcessCpuTime) {
+          long startProcessCpuTime,
+          bool /*receiveThrift*/) {
         protocol::BatchTaskUpdateRequest batchUpdateRequest =
-            json::parse(updateJson);
+            json::parse(requestBody);
         auto updateRequest = batchUpdateRequest.taskUpdateRequest;
         VELOX_USER_CHECK_NOT_NULL(updateRequest.fragment);
 
@@ -327,16 +345,22 @@ proxygen::RequestHandler* TaskResource::createOrUpdateTask(
       message,
       pathMatch,
       [&](const protocol::TaskId& taskId,
-          const std::string& updateJson,
+          const std::string& requestBody,
           const bool summarize,
-          long startProcessCpuTime) {
-        protocol::TaskUpdateRequest updateRequest = json::parse(updateJson);
+          long startProcessCpuTime,
+          bool receiveThrift) {
+        protocol::TaskUpdateRequest updateRequest;
+        if (receiveThrift) {
+          auto thriftTaskUpdateRequest = std::make_shared<thrift::TaskUpdateRequest>();
+          thriftRead(requestBody, thriftTaskUpdateRequest);
+          fromThrift(*thriftTaskUpdateRequest, updateRequest);
+        } else {
+          updateRequest = json::parse(requestBody);
+        }
         velox::core::PlanFragment planFragment;
         std::shared_ptr<velox::core::QueryCtx> queryCtx;
         if (updateRequest.fragment) {
-          auto fragment =
-              velox::encoding::Base64::decode(*updateRequest.fragment);
-          protocol::PlanFragment prestoPlan = json::parse(fragment);
+          protocol::PlanFragment prestoPlan = json::parse(receiveThrift ? *updateRequest.fragment : velox::encoding::Base64::decode(*updateRequest.fragment));
 
           queryCtx =
               taskManager_.getQueryContextManager()->findOrCreateQueryCtx(
@@ -510,12 +534,12 @@ proxygen::RequestHandler* TaskResource::getTaskStatus(
   auto maxWait = getMaxWait(message);
 
   auto& headers = message->getHeaders();
-  auto acceptHeader = headers.getSingleOrEmpty(proxygen::HTTP_HEADER_ACCEPT);
-  auto useThrift =
+  const auto& acceptHeader = headers.getSingleOrEmpty(proxygen::HTTP_HEADER_ACCEPT);
+  const auto sendThrift =
       acceptHeader.find(http::kMimeTypeApplicationThrift) != std::string::npos;
 
   return new http::CallbackRequestHandler(
-      [this, useThrift, taskId, currentState, maxWait](
+      [this, sendThrift, taskId, currentState, maxWait](
           proxygen::HTTPMessage* /*message*/,
           const std::vector<std::unique_ptr<folly::IOBuf>>& /*body*/,
           proxygen::ResponseHandler* downstream,
@@ -525,7 +549,7 @@ proxygen::RequestHandler* TaskResource::getTaskStatus(
             httpSrvCpuExecutor_,
             [this,
              evb,
-             useThrift,
+             sendThrift,
              taskId,
              currentState,
              maxWait,
@@ -535,10 +559,10 @@ proxygen::RequestHandler* TaskResource::getTaskStatus(
                   .getTaskStatus(taskId, currentState, maxWait, handlerState)
                   .via(evb)
                   .thenValue(
-                      [useThrift, downstream, taskId, handlerState](
+                      [sendThrift, downstream, taskId, handlerState](
                           std::unique_ptr<protocol::TaskStatus> taskStatus) {
                         if (!handlerState->requestExpired()) {
-                          if (useThrift) {
+                          if (sendThrift) {
                             thrift::TaskStatus thriftTaskStatus;
                             toThrift(*taskStatus, thriftTaskStatus);
                             http::sendOkThriftResponse(

presto-native-execution/presto_cpp/main/TaskResource.h

@@ -76,7 +76,8 @@ class TaskResource {
           const protocol::TaskId&,
           const std::string&,
           const bool,
-          long)>& createOrUpdateFunc);
+          long,
+          const bool)>& createOrUpdateFunc);
 
   proxygen::RequestHandler* deleteTask(
       proxygen::HTTPMessage* message,

presto-native-execution/presto_cpp/main/common/tests/test_json.h

@@ -14,11 +14,14 @@
 #pragma once
 
 #include <fstream>
-#include <ios>
 #include <iosfwd>
+#include <boost/filesystem.hpp>
+#include <boost/algorithm/string.hpp>
 
 #include "presto_cpp/presto_protocol/core/presto_protocol_core.h"
 
+namespace fs = boost::filesystem;
+
 namespace nlohmann {
 
 // This is required avoid stack overflow when a gtest error printer is invoked.
@@ -48,3 +51,19 @@ inline std::string slurp(const std::string& path) {
   buf << input.rdbuf();
   return buf.str();
 }
+
+inline std::string getDataPath(const std::string& dirUnderFbcode, const std::string& fileName) {
+  std::string currentPath = fs::current_path().c_str();
+  if (boost::algorithm::ends_with(currentPath, "fbcode")) {
+    return currentPath + dirUnderFbcode + fileName;
+  }
+
+  // CLion runs the tests from cmake-build-release/ or cmake-build-debug/
+  // directory. Hard-coded json files are not copied there and test fails with
+  // file not found. Fixing the path so that we can trigger these tests from
+  // CLion.
+  boost::algorithm::replace_all(currentPath, "cmake-build-release/", "");
+  boost::algorithm::replace_all(currentPath, "cmake-build-debug/", "");
+
+  return currentPath + "/data/" + fileName;
+}

presto-native-execution/presto_cpp/main/tests/CMakeLists.txt

@@ -22,7 +22,10 @@ add_executable(
   ServerOperationTest.cpp
   SessionPropertiesTest.cpp
   TaskManagerTest.cpp
-  QueryContextManagerTest.cpp)
+  QueryContextManagerTest.cpp
+  TaskInfoTest.cpp
+  TaskStatusTest.cpp
+  TaskUpdateRequestTest.cpp)
 
 if(DEFINED PRESTO_MEMORY_CHECKER_TYPE AND PRESTO_MEMORY_CHECKER_TYPE STREQUAL
                                           "LINUX_MEMORY_CHECKER")

presto-native-execution/presto_cpp/main/tests/TaskInfoTest.cpp

@@ -0,0 +1,103 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <gtest/gtest.h>
+#include "presto_cpp/main/thrift/ProtocolToThrift.h"
+#include "presto_cpp/presto_protocol/core/Duration.h"
+#include "presto_cpp/main/common/tests/test_json.h"
+#include "presto_cpp/main/connectors/PrestoToVeloxConnector.h"
+
+using namespace facebook;
+using namespace facebook::presto::protocol;
+
+class TaskInfoTest : public ::testing::Test {};
+
+const std::string BASE_DATA_PATH = "/github/presto-trunk/presto-native-execution/presto_cpp/main/tests/data/";
+
+TEST_F(TaskInfoTest, duration) {
+  double thrift = 0;
+  facebook::presto::thrift::toThrift(Duration(123, TimeUnit::MILLISECONDS), thrift);
+  ASSERT_EQ(thrift, 123);
+}
+
+TEST_F(TaskInfoTest, binaryMetadataUpdates) {
+  std::string str = slurp(getDataPath(BASE_DATA_PATH, "MetadataUpdates.json"));
+  json j = json::parse(str);
+  registerPrestoToVeloxConnector(std::make_unique<facebook::presto::HivePrestoToVeloxConnector>("hive"));
+  MetadataUpdates metadataUpdates = j;
+  std::unique_ptr<std::string> thriftMetadataUpdates = std::make_unique<std::string>();
+  facebook::presto::thrift::toThrift(metadataUpdates, *thriftMetadataUpdates);
+
+  json thriftJson = json::parse(*thriftMetadataUpdates);
+  ASSERT_EQ(j, thriftJson);
+
+  presto::unregisterPrestoToVeloxConnector("hive");
+}
+
+TEST_F(TaskInfoTest, taskInfo) {
+  std::string str = slurp(getDataPath(BASE_DATA_PATH, "TaskInfo.json"));
+  json j = json::parse(str);
+  registerPrestoToVeloxConnector(std::make_unique<facebook::presto::HivePrestoToVeloxConnector>("hive"));
+  TaskInfo taskInfo = j;
+  facebook::presto::thrift::TaskInfo thriftTaskInfo;
+  facebook::presto::thrift::toThrift(taskInfo, thriftTaskInfo);
+
+  json thriftJson = json::parse(*thriftTaskInfo.metadataUpdates()->metadataUpdates());
+  ASSERT_EQ(taskInfo.metadataUpdates, thriftJson);
+  ASSERT_EQ(thriftTaskInfo.needsPlan(), false);
+  ASSERT_EQ(thriftTaskInfo.outputBuffers()->buffers()->size(), 2);
+  ASSERT_EQ(thriftTaskInfo.outputBuffers()->buffers()[0].bufferId()->id(), 100);
+  ASSERT_EQ(thriftTaskInfo.outputBuffers()->buffers()[1].bufferId()->id(), 200);
+  ASSERT_EQ(thriftTaskInfo.stats()->blockedReasons()->count(facebook::presto::thrift::BlockedReason::WAITING_FOR_MEMORY), 1);
+  ASSERT_EQ(thriftTaskInfo.stats()->runtimeStats()->metrics()->size(), 2);
+  ASSERT_EQ(thriftTaskInfo.stats()->runtimeStats()->metrics()["test_metric1"].sum(), 123);
+  ASSERT_EQ(thriftTaskInfo.stats()->runtimeStats()->metrics()["test_metric2"].name(), "test_metric2");
+
+  presto::unregisterPrestoToVeloxConnector("hive");
+}
+
+TEST_F(TaskInfoTest, taskId) {
+  TaskId taskId = "queryId.1.2.3.4";
+  facebook::presto::thrift::TaskId thriftTaskId;
+  facebook::presto::thrift::toThrift(taskId, thriftTaskId);
+
+  ASSERT_EQ(thriftTaskId.stageExecutionId()->stageId()->queryId(), "queryId");
+  ASSERT_EQ(thriftTaskId.stageExecutionId()->stageId()->id(), 1);
+  ASSERT_EQ(thriftTaskId.stageExecutionId()->id(), 2);
+  ASSERT_EQ(thriftTaskId.id(), 3);
+  ASSERT_EQ(thriftTaskId.attemptNumber(), 4);
+}
+
+
+TEST_F(TaskInfoTest, operatorStatsEmptyBlockedReason) {
+  std::string str = slurp(getDataPath(BASE_DATA_PATH, "OperatorStatsEmptyBlockedReason.json"));
+  json j = json::parse(str);
+  OperatorStats operatorStats = j;
+  facebook::presto::thrift::OperatorStats thriftOperatorStats;
+  facebook::presto::thrift::toThrift(operatorStats, thriftOperatorStats);
+
+  ASSERT_EQ(thriftOperatorStats.blockedReason().has_value(), false);
+  ASSERT_EQ(thriftOperatorStats.blockedWall(), 80);
+  ASSERT_EQ(thriftOperatorStats.finishCpu(), 1000);
+}
+
+TEST_F(TaskInfoTest, operatorStats) {
+  std::string str = slurp(getDataPath(BASE_DATA_PATH, "OperatorStats.json"));
+  json j = json::parse(str);
+  OperatorStats operatorStats = j;
+  facebook::presto::thrift::OperatorStats thriftOperatorStats;
+  facebook::presto::thrift::toThrift(operatorStats, thriftOperatorStats);
+
+  ASSERT_EQ(thriftOperatorStats.blockedReason(), facebook::presto::thrift::BlockedReason::WAITING_FOR_MEMORY);
+}