From cc0a389cd2338e831959eb76d64fafba590771af Mon Sep 17 00:00:00 2001
From: Shakyan Kushwaha <shakyan@Shakyans-MacBook-Pro.local>
Date: Mon, 7 Oct 2024 11:14:34 +0530
Subject: [PATCH] Code refactoring

---
 .../presto_cpp/main/PrestoServer.cpp          | 46 +++++++++----------
 .../presto_cpp/main/PrestoServer.h            |  2 +-
 2 files changed, 22 insertions(+), 26 deletions(-)

diff --git a/presto-native-execution/presto_cpp/main/PrestoServer.cpp b/presto-native-execution/presto_cpp/main/PrestoServer.cpp
index 85eaa5cbebef8..18acb9d2cbf0d 100644
--- 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
index 84b902368dc47..4777003bb2466 100644
--- 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();