feat: Add batch job submission

src/spider/CMakeLists.txt

@@ -1,7 +1,7 @@
 # set variable as CACHE INTERNAL to access it from other scope
 set(SPIDER_CORE_SOURCES
-    storage/MySqlConnection.cpp
-    storage/MySqlStorage.cpp
+    storage/mysql/MySqlConnection.cpp
+    storage/mysql/MySqlStorage.cpp
     worker/FunctionManager.cpp
     worker/FunctionNameManager.cpp
     io/msgpack_message.cpp
@@ -25,8 +25,11 @@ set(SPIDER_CORE_HEADERS
     storage/MetadataStorage.hpp
     storage/DataStorage.hpp
     storage/StorageConnection.hpp
-    storage/MySqlConnection.hpp
-    storage/MySqlStorage.hpp
+    storage/mysql/mysql_stmt.hpp
+    storage/mysql/MySqlConnection.hpp
+    storage/mysql/MySqlStorage.hpp
+    storage/mysql/MySqlJobSubmissionBatch.hpp
+    storage/JobSubmissionBatch.hpp
     worker/FunctionManager.hpp
     worker/FunctionNameManager.hpp
     CACHE INTERNAL

src/spider/client/Data.hpp

@@ -15,7 +15,7 @@
 #include "../io/MsgPack.hpp"  // IWYU pragma: keep
 #include "../io/Serializer.hpp"
 #include "../storage/DataStorage.hpp"
-#include "../storage/MySqlConnection.hpp"
+#include "../storage/mysql/MySqlConnection.hpp"
 #include "Exception.hpp"
 
 namespace spider {

src/spider/client/Driver.cpp

@@ -6,6 +6,7 @@
 #include <stop_token>
 #include <string>
 #include <thread>
+#include <utility>
 #include <variant>
 
 #include <boost/uuid/random_generator.hpp>
@@ -15,8 +16,8 @@
 #include "../core/Error.hpp"
 #include "../core/KeyValueData.hpp"
 #include "../io/BoostAsio.hpp"  // IWYU pragma: keep
-#include "../storage/MySqlConnection.hpp"
-#include "../storage/MySqlStorage.hpp"
+#include "../storage/mysql/MySqlConnection.hpp"
+#include "../storage/mysql/MySqlStorage.hpp"
 #include "Exception.hpp"
 
 namespace spider {
@@ -33,9 +34,11 @@ Driver::Driver(std::string const& storage_url) {
     if (std::holds_alternative<core::StorageErr>(conn_result)) {
         throw ConnectionException(std::get<core::StorageErr>(conn_result).description);
     }
-    auto& conn = std::get<core::MySqlConnection>(conn_result);
+    m_conn = std::make_shared<core::MySqlConnection>(
+            std::get<core::MySqlConnection>(std::move(conn_result))
+    );
 
-    core::StorageErr const err = m_metadata_storage->add_driver(conn, core::Driver{m_id});
+    core::StorageErr const err = m_metadata_storage->add_driver(*m_conn, core::Driver{m_id});
     if (!err.success()) {
         if (core::StorageErrType::DuplicateKeyErr == err.type) {
             throw DriverIdInUseException(m_id);
@@ -71,9 +74,11 @@ Driver::Driver(std::string const& storage_url, boost::uuids::uuid const id) : m_
     if (std::holds_alternative<core::StorageErr>(conn_result)) {
         throw ConnectionException(std::get<core::StorageErr>(conn_result).description);
     }
-    auto& conn = std::get<core::MySqlConnection>(conn_result);
+    m_conn = std::make_shared<core::MySqlConnection>(
+            std::get<core::MySqlConnection>(std::move(conn_result))
+    );
 
-    core::StorageErr const err = m_metadata_storage->add_driver(conn, core::Driver{m_id});
+    core::StorageErr const err = m_metadata_storage->add_driver(*m_conn, core::Driver{m_id});
     if (!err.success()) {
         if (core::StorageErrType::DuplicateKeyErr == err.type) {
             throw DriverIdInUseException(m_id);
@@ -104,29 +109,15 @@ Driver::Driver(std::string const& storage_url, boost::uuids::uuid const id) : m_
 auto Driver::kv_store_insert(std::string const& key, std::string const& value) -> void {
     core::KeyValueData const kv_data{key, value, m_id};
 
-    std::variant<core::MySqlConnection, core::StorageErr> conn_result
-            = core::MySqlConnection::create(m_data_storage->get_url());
-    if (std::holds_alternative<core::StorageErr>(conn_result)) {
-        throw ConnectionException(std::get<core::StorageErr>(conn_result).description);
-    }
-    auto& conn = std::get<core::MySqlConnection>(conn_result);
-
-    core::StorageErr const err = m_data_storage->add_client_kv_data(conn, kv_data);
+    core::StorageErr const err = m_data_storage->add_client_kv_data(*m_conn, kv_data);
     if (!err.success()) {
         throw ConnectionException(err.description);
     }
 }
 
 auto Driver::kv_store_get(std::string const& key) -> std::optional<std::string> {
-    std::variant<core::MySqlConnection, core::StorageErr> conn_result
-            = core::MySqlConnection::create(m_data_storage->get_url());
-    if (std::holds_alternative<core::StorageErr>(conn_result)) {
-        throw ConnectionException(std::get<core::StorageErr>(conn_result).description);
-    }
-    auto& conn = std::get<core::MySqlConnection>(conn_result);
-
     std::string value;
-    core::StorageErr const err = m_data_storage->get_client_kv_data(conn, m_id, key, &value);
+    core::StorageErr const err = m_data_storage->get_client_kv_data(*m_conn, m_id, key, &value);
     if (!err.success()) {
         if (core::StorageErrType::KeyNotFoundErr == err.type) {
             return std::nullopt;

src/spider/client/Driver.hpp

@@ -8,7 +8,6 @@
 #include <thread>
 #include <tuple>
 #include <type_traits>
-#include <variant>
 #include <vector>
 
 #include <boost/uuid/random_generator.hpp>
@@ -18,7 +17,10 @@
 #include "../core/Error.hpp"
 #include "../core/TaskGraphImpl.hpp"
 #include "../io/Serializer.hpp"
-#include "../storage/MySqlConnection.hpp"
+#include "../storage/JobSubmissionBatch.hpp"
+#include "../storage/mysql/MySqlConnection.hpp"
+#include "../storage/mysql/MySqlJobSubmissionBatch.hpp"
+#include "../storage/StorageConnection.hpp"
 #include "../worker/FunctionManager.hpp"
 #include "../worker/FunctionNameManager.hpp"
 #include "Data.hpp"
@@ -133,6 +135,40 @@ class Driver {
         return TaskGraphType<ReturnType, Inputs...>{std::move(graph)};
     }
 
+    /**
+     * Begins a batch of `start` calls. This allows the driver to submit multiple jobs in a single
+     * batch, which can be more efficient than submitting jobs individually.
+     *
+     * Needs to be paired with `end_batch_start`.
+     *
+     * If a batch has already been started, this method is a no-op.
+     */
+    auto begin_batch_start() -> void {
+        if (nullptr != m_batch) {
+            return;
+        }
+        m_batch = std::make_shared<core::MySqlJobSubmissionBatch>(
+                // NOLINTNEXTLINE(cppcoreguidelines-pro-type-static-cast-downcast)
+                static_cast<core::MySqlConnection&>(*m_conn)
+        );
+    }
+
+    /**
+     * Ends a batch of `start` calls. This submits all jobs in the batch to Spider.
+     *
+     * @throw spider::ConnectionException
+     */
+    auto end_batch_start() -> void {
+        if (nullptr == m_batch) {
+            return;
+        }
+        core::StorageErr const err = m_batch->submit_batch(*m_conn);
+        m_batch = nullptr;
+        if (!err.success()) {
+            throw ConnectionException(fmt::format("Failed to start job: {}", err.description));
+        }
+    }
+
     /**
      * Starts running a task with the given inputs on Spider.
      *
@@ -175,18 +211,20 @@ class Driver {
         graph.add_task(new_task);
         graph.add_input_task(new_task.get_id());
         graph.add_output_task(new_task.get_id());
-        std::variant<core::MySqlConnection, core::StorageErr> conn_result
-                = core::MySqlConnection::create(m_metadata_storage->get_url());
-        if (std::holds_alternative<core::StorageErr>(conn_result)) {
-            throw ConnectionException(std::get<core::StorageErr>(conn_result).description);
-        }
-        auto& conn = std::get<core::MySqlConnection>(conn_result);
-        core::StorageErr err = m_metadata_storage->add_job(conn, job_id, m_id, graph);
-        if (!err.success()) {
-            throw ConnectionException(fmt::format("Failed to start job: {}", err.description));
+        if (nullptr != m_batch) {
+            core::StorageErr const err
+                    = m_metadata_storage->add_job_batch(*m_conn, *m_batch, job_id, m_id, graph);
+            if (!err.success()) {
+                throw ConnectionException(fmt::format("Failed to start job: {}", err.description));
+            }
+        } else {
+            core::StorageErr const err = m_metadata_storage->add_job(*m_conn, job_id, m_id, graph);
+            if (!err.success()) {
+                throw ConnectionException(fmt::format("Failed to start job: {}", err.description));
+            }
         }
 
-        return Job<ReturnType>{job_id, m_metadata_storage, m_data_storage};
+        return Job<ReturnType>{job_id, m_metadata_storage, m_data_storage, m_conn};
     }
 
     /**
@@ -224,19 +262,13 @@ class Driver {
         graph.m_impl->reset_ids();
         boost::uuids::random_generator gen;
         boost::uuids::uuid const job_id = gen();
-        std::variant<core::MySqlConnection, core::StorageErr> conn_result
-                = core::MySqlConnection::create(m_metadata_storage->get_url());
-        if (std::holds_alternative<core::StorageErr>(conn_result)) {
-            throw ConnectionException(std::get<core::StorageErr>(conn_result).description);
-        }
-        auto& conn = std::get<core::MySqlConnection>(conn_result);
         core::StorageErr const err
-                = m_metadata_storage->add_job(conn, job_id, m_id, graph.m_impl->get_graph());
+                = m_metadata_storage->add_job(*m_conn, job_id, m_id, graph.m_impl->get_graph());
         if (!err.success()) {
             throw ConnectionException(fmt::format("Failed to start job: {}", err.description));
         }
 
-        return Job<ReturnType>{job_id, m_metadata_storage, m_data_storage};
+        return Job<ReturnType>{job_id, m_metadata_storage, m_data_storage, m_conn};
     }
 
     /**
@@ -249,14 +281,8 @@ class Driver {
      */
     auto get_jobs() -> std::vector<boost::uuids::uuid> {
         std::vector<boost::uuids::uuid> job_ids;
-        std::variant<core::MySqlConnection, core::StorageErr> conn_result
-                = core::MySqlConnection::create(m_metadata_storage->get_url());
-        if (std::holds_alternative<spider::core::StorageErr>(conn_result)) {
-            throw ConnectionException(std::get<spider::core::StorageErr>(conn_result).description);
-        }
-        auto& conn = std::get<core::MySqlConnection>(conn_result);
         core::StorageErr const err
-                = m_metadata_storage->get_jobs_by_client_id(conn, m_id, &job_ids);
+                = m_metadata_storage->get_jobs_by_client_id(*m_conn, m_id, &job_ids);
         if (!err.success()) {
             throw ConnectionException("Failed to get jobs.");
         }
@@ -267,6 +293,8 @@ class Driver {
     boost::uuids::uuid m_id;
     std::shared_ptr<core::MetadataStorage> m_metadata_storage;
     std::shared_ptr<core::DataStorage> m_data_storage;
+    std::shared_ptr<core::StorageConnection> m_conn;
+    std::shared_ptr<core::JobSubmissionBatch> m_batch{nullptr};
     std::jthread m_heartbeat_thread;
 };
 }  // namespace spider