Skip to content

Commit

Permalink
Code Refactoring
Browse files Browse the repository at this point in the history
  • Loading branch information
Shakyan Kushwaha authored and Shakyan Kushwaha committed Oct 4, 2024
1 parent d54c6a2 commit f787527
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 40 deletions.
70 changes: 30 additions & 40 deletions presto-native-execution/presto_cpp/main/PrestoServer.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -61,9 +61,6 @@
#include "velox/functions/prestosql/registration/RegistrationFunctions.h"
#include "velox/functions/prestosql/window/WindowFunctionsRegistration.h"
#include "velox/serializers/PrestoSerializer.h"
#include <folly/json_patch.h>
#include <folly/dynamic.h>
#include <folly/json.h>


#ifdef PRESTO_ENABLE_REMOTE_FUNCTIONS
Expand Down Expand Up @@ -154,15 +151,6 @@ std::string nodeState2String(NodeState nodeState) {
return fmt::format("<unknown>:{}>", static_cast<int>(nodeState));
}

std::string getRequestBody(const std::vector<std::unique_ptr<folly::IOBuf>>& 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<SignalHandler>(this)),
Expand Down Expand Up @@ -385,34 +373,8 @@ void PrestoServer::run() {
proxygen::HTTPMessage* /*message*/,
const std::vector<std::unique_ptr<folly::IOBuf>>& 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](
Expand Down Expand Up @@ -1452,6 +1414,34 @@ void PrestoServer::reportNodeStatus(proxygen::ResponseHandler* downstream) {
http::sendOkResponse(downstream, json(fetchNodeStatus()));
}

void PrestoServer::handleGracefulShutdown(const std::vector<std::unique_ptr<folly::IOBuf>>& 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();
Expand Down
2 changes: 2 additions & 0 deletions presto-native-execution/presto_cpp/main/PrestoServer.h
Original file line number Diff line number Diff line change
Expand Up @@ -207,6 +207,8 @@ class PrestoServer {

void reportNodeStatus(proxygen::ResponseHandler* downstream);

void handleGracefulShutdown(const std::vector<std::unique_ptr<folly::IOBuf>>& body , proxygen::ResponseHandler* downstream);

protocol::NodeStatus fetchNodeStatus();

void populateMemAndCPUInfo();
Expand Down

0 comments on commit f787527

Please sign in to comment.