diff --git a/presto-native-execution/presto_cpp/main/PrestoServer.cpp b/presto-native-execution/presto_cpp/main/PrestoServer.cpp index 91e3a39ae13b4..68cf4e36ef201 100644 --- a/presto-native-execution/presto_cpp/main/PrestoServer.cpp +++ b/presto-native-execution/presto_cpp/main/PrestoServer.cpp @@ -365,6 +365,44 @@ void PrestoServer::run() { json infoStateJson = convertNodeState(server->nodeState()); http::sendOkResponse(downstream, infoStateJson); }); + httpServer_->registerPut( + "/v1/info/state", + [server = this]( + proxygen::HTTPMessage* /*message*/, + 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(); + } + 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 shutdown is already requested + LOG(INFO) << "Shutdown 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'.",400); + } + } + catch (const folly::json::parse_error& e) { + // Handle JSON parse errors and send a 400 Bad Request response + http::sendErrorResponse(downstream, e.what(), 400); // 400 Bad Request + } + }); httpServer_->registerGet( "/v1/status", [server = this](