perf: add connection pooling and parallel shard execution (opt-in, benefits large result sets)

Author: renecannao

include/sql_engine/connection_pool.h

@@ -0,0 +1,98 @@
+#ifndef SQL_ENGINE_CONNECTION_POOL_H
+#define SQL_ENGINE_CONNECTION_POOL_H
+
+// Thread-safe connection pool for MySQL backends.
+//
+// Each backend name maps to a stack of idle MYSQL* handles. Threads check out
+// a handle (creating one on demand), use it, and check it back in. The pool
+// grows dynamically -- there is no hard cap -- so N concurrent calls to the
+// same backend will use N connections.
+
+#include "sql_engine/backend_config.h"
+#include <mysql/mysql.h>
+#include <mutex>
+#include <unordered_map>
+#include <vector>
+#include <string>
+#include <stdexcept>
+
+namespace sql_engine {
+
+class ConnectionPool {
+public:
+    ConnectionPool() = default;
+
+    ~ConnectionPool() {
+        std::lock_guard<std::mutex> lk(mu_);
+        for (auto& kv : idle_) {
+            for (MYSQL* c : kv.second) {
+                if (c) mysql_close(c);
+            }
+        }
+        idle_.clear();
+    }
+
+    // Register a backend. Must be called before any checkout for that name.
+    void add_backend(const BackendConfig& config) {
+        std::lock_guard<std::mutex> lk(mu_);
+        configs_[config.name] = config;
+    }
+
+    // Obtain a connection for the named backend. Creates a new one if none is
+    // idle. The caller MUST call checkin() when done.
+    MYSQL* checkout(const std::string& backend) {
+        std::lock_guard<std::mutex> lk(mu_);
+        auto& stack = idle_[backend];
+        if (!stack.empty()) {
+            MYSQL* c = stack.back();
+            stack.pop_back();
+            // Quick liveness check (non-blocking).
+            if (mysql_ping(c) != 0) {
+                mysql_close(c);
+                return create_connection(backend);
+            }
+            return c;
+        }
+        return create_connection(backend);
+    }
+
+    // Return a connection to the pool for reuse.
+    void checkin(const std::string& backend, MYSQL* conn) {
+        if (!conn) return;
+        std::lock_guard<std::mutex> lk(mu_);
+        idle_[backend].push_back(conn);
+    }
+
+private:
+    std::mutex mu_;
+    std::unordered_map<std::string, BackendConfig> configs_;
+    std::unordered_map<std::string, std::vector<MYSQL*>> idle_;
+
+    // Must be called with mu_ held.
+    MYSQL* create_connection(const std::string& backend) {
+        auto it = configs_.find(backend);
+        if (it == configs_.end()) {
+            throw std::runtime_error("ConnectionPool: unknown backend: " + backend);
+        }
+        const BackendConfig& cfg = it->second;
+        MYSQL* c = mysql_init(nullptr);
+        if (!c) throw std::runtime_error("mysql_init failed for " + backend);
+
+        unsigned int timeout = 5;
+        mysql_options(c, MYSQL_OPT_CONNECT_TIMEOUT, &timeout);
+
+        if (!mysql_real_connect(c, cfg.host.c_str(), cfg.user.c_str(),
+                                cfg.password.c_str(), cfg.database.c_str(),
+                                cfg.port, nullptr, 0)) {
+            std::string err = mysql_error(c);
+            mysql_close(c);
+            throw std::runtime_error("ConnectionPool connect failed for " + backend + ": " + err);
+        }
+        mysql_set_character_set(c, "utf8mb4");
+        return c;
+    }
+};
+
+} // namespace sql_engine
+
+#endif // SQL_ENGINE_CONNECTION_POOL_H

include/sql_engine/operators/merge_aggregate_op.h

@@ -152,9 +152,9 @@ class MergeAggregateOperator : public Operator {
     std::vector<Operator*> children_;
     uint16_t group_key_count_;
     std::vector<uint8_t> merge_ops_;
-    bool parallel_open_;
     uint16_t merge_op_count_;
     sql_parser::Arena& arena_;
+    bool parallel_open_;
 
     struct GroupState {
         std::vector<Value> group_values;

include/sql_engine/operators/set_op_op.h

@@ -6,17 +6,30 @@
 #include <unordered_set>
 #include <string>
 #include <vector>
+#include <future>
 
 namespace sql_engine {
 
 class SetOpOperator : public Operator {
 public:
-    SetOpOperator(Operator* left, Operator* right, uint8_t op, bool all)
-        : left_(left), right_(right), op_(op), all_(all) {}
+    SetOpOperator(Operator* left, Operator* right, uint8_t op, bool all,
+                  bool parallel_open = false)
+        : left_(left), right_(right), op_(op), all_(all),
+          parallel_open_(parallel_open) {}
 
     void open() override {
-        left_->open();
-        right_->open();
+        if (parallel_open_) {
+            // Parallel open (#26): launch left and right opens concurrently.
+            // This is beneficial when both children are RemoteScan operators
+            // (each performing a network round-trip in open()).
+            auto fl = std::async(std::launch::async, [this]{ left_->open(); });
+            auto fr = std::async(std::launch::async, [this]{ right_->open(); });
+            fl.get();
+            fr.get();
+        } else {
+            left_->open();
+            right_->open();
+        }
         reading_left_ = true;
         seen_.clear();
 
@@ -99,6 +112,7 @@ class SetOpOperator : public Operator {
     Operator* right_;
     uint8_t op_;
     bool all_;
+    bool parallel_open_;
     bool reading_left_ = true;
     std::unordered_set<std::string> seen_;
     std::unordered_set<std::string> right_set_;

include/sql_engine/plan_executor.h

@@ -80,6 +80,13 @@ class PlanExecutor {
         remote_executor_ = exec;
     }
 
+    // Enable parallel opening of RemoteScan children in merge/set operators.
+    // Only safe when the RemoteExecutor is thread-safe (e.g. ThreadSafeMultiRemoteExecutor
+    // with connection pooling). Disabled by default for backward compatibility.
+    void set_parallel_open(bool enabled) {
+        parallel_open_enabled_ = enabled;
+    }
+
     // Access the subquery executor (for operators that need it)
     SubqueryExecutor<D>* subquery_executor() { return &subquery_exec_; }
 
@@ -294,6 +301,7 @@ class PlanExecutor {
     std::unordered_map<std::string, MutableDataSource*> mutable_sources_;
     std::vector<std::unique_ptr<Operator>> operators_;
     RemoteExecutor* remote_executor_ = nullptr;
+    bool parallel_open_enabled_ = false;
     DistributeFn distribute_fn_;
     SubqueryExecutor<D> subquery_exec_;
     sql_parser::Arena subquery_plan_arena_{65536, 1048576};
@@ -1011,8 +1019,12 @@ class PlanExecutor {
         Operator* right = build_operator(node->right);
         if (!left || !right) return nullptr;
 
+        // Enable parallel open when both children are remote scans and executor is thread-safe
+        bool parallel = parallel_open_enabled_ &&
+                         (node->left && node->left->type == PlanNodeType::REMOTE_SCAN &&
+                          node->right && node->right->type == PlanNodeType::REMOTE_SCAN);
         auto op = std::make_unique<SetOpOperator>(
-            left, right, node->set_op.op, node->set_op.all);
+            left, right, node->set_op.op, node->set_op.all, parallel);
         Operator* ptr = op.get();
         operators_.push_back(std::move(op));
         return ptr;
@@ -1037,12 +1049,16 @@ class PlanExecutor {
         }
         if (children.empty()) return nullptr;
 
+        // Enable parallel open when children are RemoteScans and executor is thread-safe
+        bool parallel = parallel_open_enabled_ &&
+                         (children.size() > 1) && has_remote_scan_children(node);
         auto op = std::make_unique<MergeAggregateOperator>(
             std::move(children),
             node->merge_aggregate.group_key_count,
             node->merge_aggregate.merge_ops,
             node->merge_aggregate.merge_op_count,
-            arena_);
+            arena_,
+            parallel);
         Operator* ptr = op.get();
         operators_.push_back(std::move(op));
         return ptr;
@@ -1074,16 +1090,40 @@ class PlanExecutor {
             sort_dirs.push_back(node->merge_sort.directions[i]);
         }
 
+        // Enable parallel open when children are RemoteScans and executor is thread-safe
+        bool parallel = parallel_open_enabled_ &&
+                         (children.size() > 1) && has_remote_scan_children_merge_sort(node);
         auto op = std::make_unique<MergeSortOperator>(
             std::move(children),
             sort_col_indices.data(),
             sort_dirs.data(),
-            node->merge_sort.key_count);
+            node->merge_sort.key_count,
+            parallel);
         Operator* ptr = op.get();
         operators_.push_back(std::move(op));
         return ptr;
     }
 
+    // Check whether all children of a MERGE_AGGREGATE node are REMOTE_SCAN.
+    static bool has_remote_scan_children(const PlanNode* node) {
+        if (!node || node->type != PlanNodeType::MERGE_AGGREGATE) return false;
+        for (uint16_t i = 0; i < node->merge_aggregate.child_count; ++i) {
+            if (node->merge_aggregate.children[i]->type != PlanNodeType::REMOTE_SCAN)
+                return false;
+        }
+        return node->merge_aggregate.child_count > 0;
+    }
+
+    // Check whether all children of a MERGE_SORT node are REMOTE_SCAN.
+    static bool has_remote_scan_children_merge_sort(const PlanNode* node) {
+        if (!node || node->type != PlanNodeType::MERGE_SORT) return false;
+        for (uint16_t i = 0; i < node->merge_sort.child_count; ++i) {
+            if (node->merge_sort.children[i]->type != PlanNodeType::REMOTE_SCAN)
+                return false;
+        }
+        return node->merge_sort.child_count > 0;
+    }
+
     uint16_t resolve_column_index(const sql_parser::AstNode* key, const TableInfo* table) {
         if (!key || !table) return 0;
         sql_parser::StringRef col_name;

include/sql_engine/session.h

@@ -54,6 +54,12 @@ class Session {
         remote_executor_ = exec;
     }
 
+    // Enable parallel opening of RemoteScan children. Only safe when the
+    // RemoteExecutor is thread-safe (e.g. ThreadSafeMultiRemoteExecutor).
+    void set_parallel_open(bool enabled) {
+        parallel_open_enabled_ = enabled;
+    }
+
     void set_shard_map(const ShardMap* sm) {
         shard_map_ = sm;
     }
@@ -193,6 +199,7 @@ class Session {
     Optimizer<D> optimizer_;
     RemoteExecutor* remote_executor_ = nullptr;
     const ShardMap* shard_map_ = nullptr;
+    bool parallel_open_enabled_ = false;
     std::unordered_map<std::string, DataSource*> sources_;
     std::unordered_map<std::string, MutableDataSource*> mutable_sources_;
 
@@ -203,6 +210,8 @@ class Session {
             executor.add_mutable_data_source(kv.first.c_str(), kv.second);
         if (remote_executor_)
             executor.set_remote_executor(remote_executor_);
+        if (parallel_open_enabled_)
+            executor.set_parallel_open(true);
         // If sharding is configured, provide a distribute callback so that
         // subqueries also go through the distributed planner.
         if (shard_map_ && remote_executor_) {