diff --git a/presto-native-execution/presto_cpp/main/PrestoServer.cpp b/presto-native-execution/presto_cpp/main/PrestoServer.cpp index 35e38749dbf04..85eaa5cbebef8 100644 --- a/presto-native-execution/presto_cpp/main/PrestoServer.cpp +++ b/presto-native-execution/presto_cpp/main/PrestoServer.cpp @@ -61,9 +61,6 @@ #include "velox/functions/prestosql/registration/RegistrationFunctions.h" #include "velox/functions/prestosql/window/WindowFunctionsRegistration.h" #include "velox/serializers/PrestoSerializer.h" -#include -#include -#include #ifdef PRESTO_ENABLE_REMOTE_FUNCTIONS @@ -154,15 +151,6 @@ std::string nodeState2String(NodeState nodeState) { return fmt::format(":{}>", static_cast(nodeState)); } -std::string getRequestBody(const std::vector>& body) { - std::string requestBody; - for (const auto& buf : body) { - // Convert IOBuf to fbstring and then to std::string - requestBody += buf->moveToFbString().toStdString(); - } - return requestBody; -} - PrestoServer::PrestoServer(const std::string& configDirectoryPath) : configDirectoryPath_(configDirectoryPath), signalHandler_(std::make_unique(this)), @@ -385,34 +373,8 @@ void PrestoServer::run() { proxygen::HTTPMessage* /*message*/, const std::vector>& body , proxygen::ResponseHandler* downstream) { - std::string requestBody=getRequestBody(body); - try { - folly::dynamic jsonBody = folly::parseJson(requestBody); - if (jsonBody == "SHUTTING_DOWN") { - // Print message and initiate shutdown - LOG(INFO) << "Shutdown requested"; - if (server->nodeState() == NodeState::kActive) { - // Call stop() to handle the shutdown process - std::thread([server]() { - server->stop(); - }).detach(); - } else { - // Print message and indicate that node is inactive or shutdown is already requested - LOG(INFO) << "Node is inactive or shutdown is already requested"; - } - http::sendOkResponse(downstream); - } else { - // Throw a 400 Bad Request error if the body has some other value - http::sendErrorResponse( - downstream, - "Bad Request: Body must contain 'SHUTTING_DOWN'.",http::kHttpBadRequest); - } - } - catch (const folly::json::parse_error& e) { - // Handle JSON parse errors and send a 400 Bad Request response - http::sendErrorResponse(downstream, e.what(), http::kHttpBadRequest); // 400 Bad Request - } - }); + server->handleGracefulShutdown(body,downstream); + }); httpServer_->registerGet( "/v1/status", [server = this]( @@ -1452,6 +1414,34 @@ void PrestoServer::reportNodeStatus(proxygen::ResponseHandler* downstream) { http::sendOkResponse(downstream, json(fetchNodeStatus())); } +void PrestoServer::handleGracefulShutdown(const std::vector>& body , proxygen::ResponseHandler* downstream){ + std::string requestBody; + for (const auto& buf : body) { + // Convert IOBuf to fbstring and then to std::string + requestBody += buf->moveToFbString().toStdString(); + } + requestBody.erase(requestBody.length()-2); + if (requestBody == "\"SHUTTING_DOWN\"") { + // Print message and initiate shutdown + LOG(INFO) << "Shutdown requested"; + if (nodeState() == NodeState::kActive) { + // Call stop() to handle the shutdown process + std::thread([this]() { + this->stop(); + }).detach(); + } else { + // Print message and indicate that node is inactive or shutdown is already requested + LOG(INFO) << "Node is inactive or shutdown is already requested"; + } + http::sendOkResponse(downstream); + } else { + // Throw a 400 Bad Request error if the body has some other value + http::sendErrorResponse( + downstream, + "Bad Request: Body must contain 'SHUTTING_DOWN'.",http::kHttpBadRequest); + } +} + protocol::NodeStatus PrestoServer::fetchNodeStatus() { auto systemConfig = SystemConfig::instance(); const int64_t nodeMemoryGb = systemConfig->systemMemoryGb(); diff --git a/presto-native-execution/presto_cpp/main/PrestoServer.h b/presto-native-execution/presto_cpp/main/PrestoServer.h index bee2d8d43391a..84b902368dc47 100644 --- a/presto-native-execution/presto_cpp/main/PrestoServer.h +++ b/presto-native-execution/presto_cpp/main/PrestoServer.h @@ -207,6 +207,8 @@ class PrestoServer { void reportNodeStatus(proxygen::ResponseHandler* downstream); + void handleGracefulShutdown(const std::vector>& body , proxygen::ResponseHandler* downstream); + protocol::NodeStatus fetchNodeStatus(); void populateMemAndCPUInfo();