sss

Author: kewang1024

presto-native-execution/presto_cpp/main/CMakeLists.txt

@@ -63,6 +63,7 @@ target_link_libraries(
   presto_function_metadata
   presto_connectors
   presto_http
+  presto_thrift_server
   presto_operators
   presto_session_properties
   presto_velox_plan_conversion

presto-native-execution/presto_cpp/main/PrestoServer.cpp

@@ -43,6 +43,7 @@
 #include "presto_cpp/main/operators/ShuffleRead.h"
 #include "presto_cpp/main/types/PrestoToVeloxQueryPlan.h"
 #include "presto_cpp/main/types/VeloxPlanConversion.h"
+#include "thrift/server/ThriftServer.h"
 #include "velox/common/base/Counters.h"
 #include "velox/common/base/StatsReporter.h"
 #include "velox/common/caching/CacheTTLController.h"
@@ -619,6 +620,8 @@ void PrestoServer::run() {
     }
   };
 
+  startThriftServer(bindToNodeInternalAddressOnly, certPath, keyPath, ciphers);
+
   // Start everything. After the return from the following call we are shutting
   // down.
   httpServer_->start(getHttpServerFilters(), [&](proxygen::HTTPServer* server) {
@@ -679,6 +682,7 @@ void PrestoServer::run() {
   taskManager_.reset();
   PRESTO_SHUTDOWN_LOG(INFO) << "Destroying HTTP Server";
   httpServer_.reset();
+  thriftServer_.reset();
 
   unregisterFileReadersAndWriters();
   unregisterFileSystems();
@@ -1070,6 +1074,7 @@ void PrestoServer::stop() {
     httpServer_->stop();
     PRESTO_SHUTDOWN_LOG(INFO) << "HTTP Server stopped.";
   }
+  shutdownThriftServer();
 }
 
 size_t PrestoServer::numDriverThreads() const {
@@ -1492,7 +1497,7 @@ void PrestoServer::enableWorkerStatsReporting() {
 
 void PrestoServer::initVeloxPlanValidator() {
   VELOX_CHECK_NULL(planValidator_);
-  planValidator_ = std::make_unique<VeloxPlanValidator>();
+  planValidator_ = std::make_shared<VeloxPlanValidator>();
 }
 
 VeloxPlanValidator* PrestoServer::getVeloxPlanValidator() {
@@ -1793,6 +1798,75 @@ void PrestoServer::createTaskManager() {
       driverExecutor_.get(), httpSrvCpuExecutor_.get(), spillerExecutor_.get());
 }
 
+void PrestoServer::startThriftServer(
+    bool bindToNodeInternalAddressOnly,
+    const std::string& certPath,
+    const std::string& keyPath,
+    const std::string& ciphers) {
+  auto* systemConfig = SystemConfig::instance();
+  bool thriftServerEnabled = systemConfig
+                                 ->optionalProperty<bool>(std::string_view(
+                                     "presto.thrift-server.enabled"))
+                                 .value_or(true);
+
+  if (thriftServerEnabled) {
+    std::unique_ptr<thrift::ThriftConfig> thriftConfig;
+    folly::SocketAddress thriftAddress;
+    int thriftPort = systemConfig
+                         ->optionalProperty<int>(
+                             std::string_view("presto.thrift-server.port"))
+                         .value_or(9090);
+    if (bindToNodeInternalAddressOnly) {
+      thriftAddress.setFromHostPort(address_, thriftPort);
+    } else {
+      thriftAddress.setFromLocalPort(thriftPort);
+    }
+    thriftConfig = std::make_unique<thrift::ThriftConfig>(
+        thriftAddress, certPath, keyPath, ciphers);
+    thriftServer_ = std::make_unique<thrift::ThriftServer>(
+        std::move(thriftConfig),
+        pool_,
+        planValidator_,
+        taskManager_);
+
+    thriftServerFuture_ =
+        folly::via(folly::getGlobalCPUExecutor().get())
+            .thenTry([this](folly::Try<folly::Unit>) {
+              try {
+                PRESTO_STARTUP_LOG(INFO)
+                    << "Starting Thrift server asynchronously...";
+                thriftServer_->start();
+                PRESTO_STARTUP_LOG(INFO)
+                    << "Thrift server started successfully";
+              } catch (const std::exception& e) {
+                PRESTO_STARTUP_LOG(ERROR)
+                    << "Thrift server failed to start: " << e.what();
+                throw;
+              }
+            });
+  }
+}
+
+void PrestoServer::shutdownThriftServer() {
+  if (thriftServer_) {
+    PRESTO_SHUTDOWN_LOG(INFO) << "Stopping Thrift server";
+    thriftServer_->stop();
+
+    // Wait for Thrift server thread to complete with timeout
+    try {
+      std::move(thriftServerFuture_)
+          .within(std::chrono::seconds(5)) // 5-second timeout
+          .get();
+      PRESTO_SHUTDOWN_LOG(INFO) << "Thrift server stopped gracefully";
+    } catch (const std::exception& e) {
+      PRESTO_SHUTDOWN_LOG(WARNING)
+          << "Thrift server shutdown timeout or error: " << e.what();
+    }
+
+    thriftServer_.reset();
+  }
+}
+
 void PrestoServer::reportNodeStats(proxygen::ResponseHandler* downstream) {
   protocol::NodeStats nodeStats;
 

presto-native-execution/presto_cpp/main/PrestoServer.h

@@ -48,6 +48,10 @@ namespace facebook::presto::http {
 class HttpServer;
 }
 
+namespace facebook::presto::thrift {
+class ThriftServer;
+}
+
 namespace proxygen {
 class ResponseHandler;
 } // namespace proxygen
@@ -231,6 +235,16 @@ class PrestoServer {
 
   virtual void createTaskManager();
 
+  /// Utility method to start the Thrift server if enabled
+  void startThriftServer(
+      bool bindToNodeInternalAddressOnly,
+      const std::string& certPath,
+      const std::string& keyPath,
+      const std::string& ciphers);
+
+  /// Utility method to safely shutdown the Thrift server if running
+  void shutdownThriftServer();
+
   const std::string configDirectoryPath_;
 
   std::shared_ptr<CoordinatorDiscoverer> coordinatorDiscoverer_;
@@ -269,20 +283,22 @@ class PrestoServer {
   // Executor for spilling.
   std::unique_ptr<folly::CPUThreadPoolExecutor> spillerExecutor_;
 
-  std::unique_ptr<VeloxPlanValidator> planValidator_;
+  std::shared_ptr<VeloxPlanValidator> planValidator_;
 
   std::unique_ptr<http::HttpClientConnectionPool> exchangeSourceConnectionPool_;
 
   // If not null,  the instance of AsyncDataCache used for in-memory file cache.
   std::shared_ptr<velox::cache::AsyncDataCache> cache_;
 
   std::unique_ptr<http::HttpServer> httpServer_;
+  std::unique_ptr<thrift::ThriftServer> thriftServer_;
+  folly::Future<folly::Unit> thriftServerFuture_{folly::makeFuture()};
   std::unique_ptr<SignalHandler> signalHandler_;
   std::unique_ptr<Announcer> announcer_;
   std::unique_ptr<PeriodicHeartbeatManager> heartbeatManager_;
   std::shared_ptr<velox::memory::MemoryPool> pool_;
   std::shared_ptr<velox::memory::MemoryPool> nativeWorkerPool_;
-  std::unique_ptr<TaskManager> taskManager_;
+  std::shared_ptr<TaskManager> taskManager_;
   std::unique_ptr<TaskResource> taskResource_;
   std::atomic<NodeState> nodeState_{NodeState::kActive};
   folly::Synchronized<bool> shuttingDown_{false};

presto-native-execution/presto_cpp/main/common/tests/ConfigTest.cpp

@@ -191,6 +191,88 @@ TEST_F(ConfigTest, optionalSystemConfigs) {
   ASSERT_EQ(config.discoveryUri(), "my uri");
 }
 
+TEST_F(ConfigTest, thriftServerConfigs) {
+  SystemConfig config;
+
+  // Test default values (when no thrift server configs are provided)
+  init(config, {});
+
+  // Test with thrift server enabled
+  init(config, {{"presto.thrift-server.enabled", "true"}});
+  ASSERT_EQ(
+      config.optionalProperty<bool>(std::string_view("presto.thrift-server.enabled")).value_or(false),
+      true);
+
+  // Test thrift server port configuration
+  init(config, {{"presto.thrift-server.port", "9090"}});
+  ASSERT_EQ(
+      config.optionalProperty<int>(std::string_view("presto.thrift-server.port")).value_or(9090),
+      9090);
+
+  // Test thrift server max connections
+  init(config, {{"presto.thrift-server.max-connections", "5000"}});
+  ASSERT_EQ(
+      config.optionalProperty<int>(std::string_view("presto.thrift-server.max-connections")).value_or(10000),
+      5000);
+
+  // Test thrift server max requests
+  init(config, {{"presto.thrift-server.max-requests", "5000"}});
+  ASSERT_EQ(
+      config.optionalProperty<int>(std::string_view("presto.thrift-server.max-requests")).value_or(10000),
+      5000);
+
+  // Test thrift server idle timeout
+  init(config, {{"presto.thrift-server.idle-timeout", "600000"}});
+  ASSERT_EQ(
+      config.optionalProperty<int>(std::string_view("presto.thrift-server.idle-timeout")).value_or(300000),
+      600000);
+
+  // Test thrift server task expire time
+  init(config, {{"presto.thrift-server.task-expire-time-ms", "600000"}});
+  ASSERT_EQ(
+      config.optionalProperty<int>(std::string_view("presto.thrift-server.task-expire-time-ms")).value_or(300000),
+      600000);
+
+  // Test thrift server stream expire time
+  init(config, {{"presto.thrift-server.stream-expire-time", "600000"}});
+  ASSERT_EQ(
+      config.optionalProperty<int>(std::string_view("presto.thrift-server.stream-expire-time")).value_or(300000),
+      600000);
+
+  // Test multiple thrift server configs together
+  init(config, {
+      {"presto.thrift-server.enabled", "true"},
+      {"presto.thrift-server.port", "9091"},
+      {"presto.thrift-server.max-connections", "8000"},
+      {"presto.thrift-server.max-requests", "8000"},
+      {"presto.thrift-server.idle-timeout", "900000"},
+      {"presto.thrift-server.task-expire-time-ms", "900000"},
+      {"presto.thrift-server.stream-expire-time", "900000"}
+  });
+
+  ASSERT_EQ(
+      config.optionalProperty<bool>(std::string_view("presto.thrift-server.enabled")).value_or(false),
+      true);
+  ASSERT_EQ(
+      config.optionalProperty<int>(std::string_view("presto.thrift-server.port")).value_or(9090),
+      9091);
+  ASSERT_EQ(
+      config.optionalProperty<int>(std::string_view("presto.thrift-server.max-connections")).value_or(10000),
+      8000);
+  ASSERT_EQ(
+      config.optionalProperty<int>(std::string_view("presto.thrift-server.max-requests")).value_or(10000),
+      8000);
+  ASSERT_EQ(
+      config.optionalProperty<int>(std::string_view("presto.thrift-server.idle-timeout")).value_or(300000),
+      900000);
+  ASSERT_EQ(
+      config.optionalProperty<int>(std::string_view("presto.thrift-server.task-expire-time-ms")).value_or(300000),
+      900000);
+  ASSERT_EQ(
+      config.optionalProperty<int>(std::string_view("presto.thrift-server.stream-expire-time")).value_or(300000),
+      900000);
+}
+
 TEST_F(ConfigTest, optionalNodeConfigs) {
   NodeConfig config;
   init(config, {});

presto-native-execution/presto_cpp/main/thrift/CMakeLists.txt

@@ -70,3 +70,32 @@ add_dependencies(presto_thrift_extra presto_thrift-cpp2)
 if(PRESTO_ENABLE_TESTING)
   add_subdirectory(tests)
 endif()
+
+add_library(
+  presto_thrift_server
+        server/ThriftServer.cpp
+        server/PrestoThriftServiceHandler.cpp
+)
+
+target_include_directories(
+  presto_thrift_server PUBLIC 
+    ${presto_thrift_INCLUDES} 
+    ${THRIFT_INCLUDES}
+    ${GLOG_INCLUDE_DIR}
+    # Include main directory for PrestoThriftServiceHandler.h
+    ${CMAKE_SOURCE_DIR}/presto_cpp/main
+)
+
+target_link_libraries(
+  presto_thrift_server
+  presto_thrift
+  ${presto_thrift_LIBRARIES}
+  ${THRIFT_LIBRARIES}
+  ${GLOG_LIBRARY}
+  presto_cpp_main_common
+  presto_cpp_main_types
+  presto_task_lib
+  ${THRIFT_TRANSPORT}
+  xsimd
+)
+

presto-native-execution/presto_cpp/main/thrift/presto_thrift.thrift

@@ -705,6 +705,45 @@ struct TaskUpdateRequest {
   6: optional TableWriteInfo tableWriteInfo;
 }
 
+struct TaskResult {
+  1: i64 sequence;
+  2: i64 nextSequence;
+  3: optional IOBufPtr data;
+  4: bool complete;
+  5: optional list<i64> remainingBytes;
+}
+
 service PrestoThrift {
-  void fake();
+    /**
+     * Get task results - corresponds to /v1/task/{taskId}/results/{bufferId}/{token}
+     * @param taskId The ID of the task to get results for
+     * @param bufferId The buffer ID to get results from
+     * @param token Continuation token for paging
+     * @param maxSizeBytes Maximum number of bytes to return
+     * @param maxWaitMicros Maximum time to wait in microseconds
+     * @param getDataSize Two phase protocol: if true, return the size of the data in the first phrase
+     * @return TaskResult containing the data and metadata
+     */
+    TaskResult getTaskResults(
+      1: string taskId,
+      2: i64 bufferId,
+      3: i64 token,
+      4: i64 maxSizeBytes,
+      5: i64 maxWaitMicros,
+      6: bool getDataSize,
+    );
+
+    /**
+     * Acknowledge task results - corresponds to /v1/task/{taskId}/results/{bufferId}/{token}/acknowledge
+     * @param taskId The ID of the task to acknowledge results for
+     * @param bufferId The buffer ID to acknowledge results for
+     * @param token The token to acknowledge up to
+     */
+    void acknowledgeTaskResults(1: string taskId, 2: i64 bufferId, 3: i64 token);
+
+    /**
+     * Abort task results - corresponds to /v1/task/{taskId}/results
+     * @param taskId The ID of the task to abort results for
+     */
+    void abortTaskResults(1: string taskId, 2: i64 destination);
 }