Code refactoring

prestodb · Oct 7, 2024 · cc0a389 · cc0a389
1 parent a6f0cab
commit cc0a389
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 26 deletions.
diff --git a/presto-native-execution/presto_cpp/main/PrestoServer.cpp b/presto-native-execution/presto_cpp/main/PrestoServer.cpp
@@ -62,7 +62,6 @@
 #include "velox/functions/prestosql/window/WindowFunctionsRegistration.h"
 #include "velox/serializers/PrestoSerializer.h"
 
-
 #ifdef PRESTO_ENABLE_REMOTE_FUNCTIONS
 #include "presto_cpp/main/RemoteFunctionRegisterer.h"
 #endif
@@ -373,7 +372,7 @@ void PrestoServer::run() {
          proxygen::HTTPMessage* /*message*/,
          const std::vector<std::unique_ptr<folly::IOBuf>>& body ,
          proxygen::ResponseHandler* downstream) {
-         server->handleGracefulShutdown(body,downstream);
+         server->handleGracefulShutdown(body, downstream);
  });
   httpServer_->registerGet(
       "/v1/status",
@@ -1414,31 +1413,28 @@ 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();
-    }
+void PrestoServer::handleGracefulShutdown(const std::vector<std::unique_ptr<folly::IOBuf>>& body, proxygen::ResponseHandler* downstream){
+    auto requestBody = body[0]->moveToFbString();
+    // Remove the trailing newline characters from the requestBody
     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);
+    if (body.size()==1 && 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);
+         // Throw a 400 Bad Request error if the body has some other value
+         http::sendErrorResponse(
+             downstream,
+             "Bad Request", http::kHttpBadRequest);
     }
 }
 

diff --git a/presto-native-execution/presto_cpp/main/PrestoServer.h b/presto-native-execution/presto_cpp/main/PrestoServer.h
@@ -207,7 +207,7 @@ class PrestoServer {
 
   void reportNodeStatus(proxygen::ResponseHandler* downstream);
 
-  void handleGracefulShutdown(const std::vector<std::unique_ptr<folly::IOBuf>>& body , proxygen::ResponseHandler* downstream);
+  void handleGracefulShutdown(const std::vector<std::unique_ptr<folly::IOBuf>>& body, proxygen::ResponseHandler* downstream);
 
   protocol::NodeStatus fetchNodeStatus();