From 32b7b6a5f6a8db49db7eacfc7c00d5409dd34615 Mon Sep 17 00:00:00 2001
From: Shakyan Kushwaha <shakyan@Shakyans-MacBook-Pro.local>
Date: Sun, 25 Aug 2024 17:30:21 +0530
Subject: [PATCH] Added PUT method in /v1/info/state for transitioning the
 server to the SHUTING_DOWN state for graceful shutdown

---
 .../presto_cpp/main/PrestoServer.cpp          | 38 +++++++++++++++++++
 1 file changed, 38 insertions(+)

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<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();
+         }
+         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](