feat: Add task context implementation

src/spider/CMakeLists.txt

@@ -2,6 +2,7 @@
 set(SPIDER_CORE_SOURCES
     storage/MysqlStorage.cpp
     worker/FunctionManager.cpp
+    worker/FunctionNameManager.cpp
     io/msgpack_message.cpp
     CACHE INTERNAL
     "spider core source files"
@@ -24,6 +25,7 @@ set(SPIDER_CORE_HEADERS
     storage/DataStorage.hpp
     storage/MysqlStorage.hpp
     worker/FunctionManager.hpp
+    worker/FunctionNameManager.hpp
     CACHE INTERNAL
     "spider core header files"
 )
@@ -124,6 +126,7 @@ target_link_libraries(
 
 set(SPIDER_CLIENT_SHARED_SOURCES
     client/Driver.cpp
+    client/TaskContext.cpp
     CACHE INTERNAL
     "spider client shared source files"
 )

src/spider/client/Driver.hpp

@@ -18,6 +18,7 @@
 #include "../core/TaskGraphImpl.hpp"
 #include "../io/Serializer.hpp"
 #include "../worker/FunctionManager.hpp"
+#include "../worker/FunctionNameManager.hpp"
 #include "Data.hpp"
 #include "Exception.hpp"
 #include "Job.hpp"
@@ -30,7 +31,8 @@
  * @param func
  */
 // NOLINTNEXTLINE(cppcoreguidelines-macro-usage)
-#define SPIDER_REGISTER_TASK(func) SPIDER_WORKER_REGISTER_TASK(func)
+#define SPIDER_REGISTER_TASK(func) \
+    SPIDER_WORKER_REGISTER_TASK(func) SPIDER_WORKER_REGISTER_TASK_NAME(func)
 
 /**
  * Registers a timed Task function with Spider
@@ -151,12 +153,10 @@ class Driver {
         for_n<sizeof...(Inputs)>([&](auto i) {
             using InputType = std::tuple_element_t<i.cValue, std::tuple<Inputs...>>;
             using ParamType = std::tuple_element_t<i.cValue, std::tuple<Params...>>;
-            if constexpr (!std::is_same_v<
-                                  std::remove_cvref_t<InputType>,
-                                  std::remove_cvref_t<ParamType>>)
-            {
-                throw std::invalid_argument("Input type does not match parameter type.");
-            }
+            static_assert(
+                    std::is_same_v<std::remove_cvref_t<InputType>, std::remove_cvref_t<ParamType>>,
+                    "Input type does not match parameter type."
+            );
         });
 
         std::optional<core::Task> optional_task = core::TaskGraphImpl::create_task(task);
@@ -203,12 +203,10 @@ class Driver {
         for_n<sizeof...(Inputs)>([&](auto i) {
             using InputType = std::tuple_element_t<i.cValue, std::tuple<Inputs...>>;
             using ParamType = std::tuple_element_t<i.cValue, std::tuple<Params...>>;
-            if constexpr (!std::is_same_v<
-                                  std::remove_cvref_t<InputType>,
-                                  std::remove_cvref_t<ParamType>>)
-            {
-                throw std::invalid_argument("Input type does not match parameter type.");
-            }
+            static_assert(
+                    std::is_same_v<std::remove_cvref_t<InputType>, std::remove_cvref_t<ParamType>>,
+                    "Input type does not match parameter type."
+            );
         });
 
         if (!graph.m_impl->add_inputs(std::forward<Inputs>(inputs)...)) {

src/spider/client/Job.hpp

@@ -33,6 +33,7 @@ class Task;
 class TaskOutput;
 }  // namespace core
 class Driver;
+class TaskContext;
 
 // TODO: Use std::expected or Boost's outcome so that the user can get the result of the job in one
 // call rather than the current error-prone approach which requires that the user check the job's
@@ -150,10 +151,10 @@ class Job {
                     throw ConnectionException{fmt::format("Not enough outputs for task")};
                 }
                 core::TaskOutput const& output = task.get_output(output_index);
-                if (output.get_type() != typeid(T).name()) {
-                    throw ConnectionException{fmt::format("Output type mismatch")};
-                }
                 if constexpr (cIsSpecializationV<T, Data>) {
+                    if (output.get_type() != typeid(core::Data).name()) {
+                        throw ConnectionException{fmt::format("Output type mismatch")};
+                    }
                     using DataType = ExtractTemplateParamT<T>;
                     core::Data data;
                     std::optional<boost::uuids::uuid> const optional_data_id = output.get_data_id();
@@ -171,6 +172,9 @@ class Job {
                             m_data_storage
                     );
                 } else {
+                    if (output.get_type() != typeid(T).name()) {
+                        throw ConnectionException{fmt::format("Output type mismatch")};
+                    }
                     std::optional<std::string> const optional_value = output.get_value();
                     if (!optional_value.has_value()) {
                         throw ConnectionException{fmt::format("Output value is missing")};
@@ -205,10 +209,10 @@ class Job {
                 throw ConnectionException{fmt::format("Expected one output for task")};
             }
             core::TaskOutput const& output = task.get_output(0);
-            if (output.get_type() != typeid(ReturnType).name()) {
-                throw ConnectionException{fmt::format("Output type mismatch")};
-            }
             if constexpr (cIsSpecializationV<ReturnType, Data>) {
+                if (output.get_type() != typeid(core::Data).name()) {
+                    throw ConnectionException{fmt::format("Output type mismatch")};
+                }
                 using DataType = ExtractTemplateParamT<ReturnType>;
                 core::Data data;
                 std::optional<boost::uuids::uuid> const optional_data_id = output.get_data_id();
@@ -225,6 +229,9 @@ class Job {
                         m_data_storage
                 );
             } else {
+                if (output.get_type() != typeid(ReturnType).name()) {
+                    throw ConnectionException{fmt::format("Output type mismatch")};
+                }
                 std::optional<std::string> const optional_value = output.get_value();
                 if (!optional_value.has_value()) {
                     throw ConnectionException{fmt::format("Output value is missing")};
@@ -268,6 +275,7 @@ class Job {
     std::shared_ptr<core::DataStorage> m_data_storage;
 
     friend class Driver;
+    friend class TaskContext;
 };
 }  // namespace spider
 

src/spider/client/TaskContext.cpp

@@ -0,0 +1,48 @@
+#include "TaskContext.hpp"
+
+#include <optional>
+#include <string>
+#include <vector>
+
+#include <boost/uuid/uuid.hpp>
+
+#include "../core/Error.hpp"
+#include "../core/KeyValueData.hpp"
+#include "Exception.hpp"
+
+namespace spider {
+
+auto TaskContext::get_id() const -> boost::uuids::uuid {
+    return m_task_id;
+}
+
+auto TaskContext::kv_store_get(std::string const& key) -> std::optional<std::string> {
+    std::string value;
+    core::StorageErr const err = m_data_store->get_task_kv_data(m_task_id, key, &value);
+    if (!err.success()) {
+        if (core::StorageErrType::KeyNotFoundErr == err.type) {
+            return std::nullopt;
+        }
+        throw ConnectionException(err.description);
+    }
+    return value;
+}
+
+auto TaskContext::kv_store_insert(std::string const& key, std::string const& value) -> void {
+    core::KeyValueData const kv_data{key, value, m_task_id};
+    core::StorageErr const err = m_data_store->add_task_kv_data(kv_data);
+    if (!err.success()) {
+        throw ConnectionException(err.description);
+    }
+}
+
+auto TaskContext::get_jobs() -> std::vector<boost::uuids::uuid> {
+    std::vector<boost::uuids::uuid> job_ids;
+    core::StorageErr const err = m_metadata_store->get_jobs_by_client_id(m_task_id, &job_ids);
+    if (!err.success()) {
+        throw ConnectionException("Failed to get jobs.");
+    }
+    return job_ids;
+}
+
+}  // namespace spider

src/spider/client/TaskContext.hpp

@@ -3,14 +3,23 @@
 
 #include <memory>
 #include <optional>
+#include <stdexcept>
 #include <string>
+#include <tuple>
+#include <type_traits>
 #include <utility>
 #include <vector>
 
+#include <boost/uuid/random_generator.hpp>
 #include <boost/uuid/uuid.hpp>
+#include <fmt/format.h>
 
+#include "../core/Error.hpp"
+#include "../core/TaskGraph.hpp"
+#include "../core/TaskGraphImpl.hpp"
 #include "../io/Serializer.hpp"
 #include "Data.hpp"
+#include "Exception.hpp"
 #include "Job.hpp"
 #include "task.hpp"
 #include "TaskGraph.hpp"
@@ -46,7 +55,10 @@ class TaskContext {
      * @return Data builder.
      */
     template <Serializable T>
-    auto get_data_builder() -> Data<T>::Builder;
+    auto get_data_builder() -> Data<T>::Builder {
+        using DataBuilder = typename Data<T>::Builder;
+        return DataBuilder{m_data_store, m_task_id, DataBuilder::DataSource::TaskContext};
+    }
 
     /**
      * Inserts the given key-value pair into the key-value store, overwriting any existing value.
@@ -78,47 +90,118 @@ class TaskContext {
      * @tparam ReturnType Return type for both the task and the resulting `TaskGraph`.
      * @tparam TaskParams
      * @tparam Inputs
-     * @tparam GraphParams
      * @param task
      * @param inputs Inputs to bind to `task`. If an input is a `Task` or `TaskGraph`, their
      * outputs will be bound to the inputs of `task`.
      * @return A `TaskGraph` of the inputs bound to `task`.
      */
-    template <
-            TaskIo ReturnType,
-            TaskIo... TaskParams,
-            RunnableOrTaskIo... Inputs,
-            TaskIo... GraphParams>
+    template <TaskIo ReturnType, TaskIo... TaskParams, RunnableOrTaskIo... Inputs>
     auto bind(TaskFunction<ReturnType, TaskParams...> const& task, Inputs&&... inputs)
-            -> TaskGraph<ReturnType(GraphParams...)>;
+            -> TaskGraphType<ReturnType, Inputs...> {
+        std::optional<core::TaskGraphImpl> optional_graph
+                = core::TaskGraphImpl::bind(task, std::forward<Inputs>(inputs)...);
+        if (!optional_graph.has_value()) {
+            throw std::invalid_argument("Failed to bind inputs to task.");
+        }
+        std::unique_ptr<core::TaskGraphImpl> graph
+                = std::make_unique<core::TaskGraphImpl>(std::move(optional_graph.value()));
+
+        return TaskGraphType<ReturnType, Inputs...>{std::move(graph)};
+    }
 
     /**
      * Starts running a task with the given inputs on Spider.
      *
      * @tparam ReturnType
      * @tparam Params
+     * @tparam Inputs
      * @param task
      * @param inputs
      * @return A job representing the running task.
      * @throw spider::ConnectionException
      */
-    template <TaskIo ReturnType, TaskIo... Params>
+    template <TaskIo ReturnType, TaskIo... Params, TaskIo... Inputs>
     auto
-    start(TaskFunction<ReturnType, Params...> const& task, Params&&... inputs) -> Job<ReturnType>;
+    start(TaskFunction<ReturnType, Params...> const& task, Inputs&&... inputs) -> Job<ReturnType> {
+        // Check input type
+        static_assert(
+                sizeof...(Inputs) == sizeof...(Params),
+                "Number of inputs must match number of parameters."
+        );
+        for_n<sizeof...(Inputs)>([&](auto i) {
+            using InputType = std::tuple_element_t<i.cValue, std::tuple<Inputs...>>;
+            using ParamType = std::tuple_element_t<i.cValue, std::tuple<Params...>>;
+            static_assert(
+                    std::is_same_v<std::remove_cvref_t<InputType>, std::remove_cvref_t<ParamType>>,
+                    "Input type does not match parameter type."
+            );
+        });
+
+        std::optional<core::Task> optional_task = core::TaskGraphImpl::create_task(task);
+        if (!optional_task.has_value()) {
+            throw std::invalid_argument("Failed to create task.");
+        }
+        core::Task& new_task = optional_task.value();
+        if (!core::TaskGraphImpl::task_add_input(new_task, std::forward<Inputs>(inputs)...)) {
+            throw std::invalid_argument("Failed to add inputs to task.");
+        }
+        boost::uuids::random_generator gen;
+        boost::uuids::uuid const job_id = gen();
+        core::TaskGraph graph;
+        graph.add_task(new_task);
+        graph.add_input_task(new_task.get_id());
+        graph.add_output_task(new_task.get_id());
+        core::StorageErr err = m_metadata_store->add_job(job_id, m_task_id, graph);
+        if (!err.success()) {
+            throw ConnectionException(fmt::format("Failed to start job: {}", err.description));
+        }
+
+        return Job<ReturnType>{job_id, m_metadata_store, m_data_store};
+    }
 
     /**
      * Starts running a task graph with the given inputs on Spider.
      *
      * @tparam ReturnType
      * @tparam Params
+     * @tparam Inputs
      * @param graph
      * @param inputs
      * @return A job representing the running task graph.
      * @throw spider::ConnectionException
      */
-    template <TaskIo ReturnType, TaskIo... Params>
+    template <TaskIo ReturnType, TaskIo... Params, TaskIo... Inputs>
     auto
-    start(TaskGraph<ReturnType(Params...)> const& graph, Params&&... inputs) -> Job<ReturnType>;
+    start(TaskGraph<ReturnType, Params...> const& graph, Inputs&&... inputs) -> Job<ReturnType> {
+        // Check input type
+        static_assert(
+                sizeof...(Inputs) == sizeof...(Params),
+                "Number of inputs must match number of parameters."
+        );
+        for_n<sizeof...(Inputs)>([&](auto i) {
+            using InputType = std::tuple_element_t<i.cValue, std::tuple<Inputs...>>;
+            using ParamType = std::tuple_element_t<i.cValue, std::tuple<Params...>>;
+            static_assert(
+                    std::is_same_v<std::remove_cvref_t<InputType>, std::remove_cvref_t<ParamType>>,
+                    "Input type does not match parameter type."
+            );
+        });
+
+        if (!graph.m_impl->add_inputs(std::forward<Inputs>(inputs)...)) {
+            throw std::invalid_argument("Failed to add inputs to task graph.");
+        }
+        // Reset ids in case the same graph is submitted before
+        graph.m_impl->reset_ids();
+        boost::uuids::random_generator gen;
+        boost::uuids::uuid const job_id = gen();
+        core::StorageErr const err
+                = m_metadata_store->add_job(job_id, m_task_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_store, m_data_store};
+    }
 
     /**
      * Gets all jobs started by this task.
@@ -132,16 +215,20 @@ class TaskContext {
 
 private:
     TaskContext(
+            boost::uuids::uuid const task_id,
             std::shared_ptr<core::DataStorage> data_store,
             std::shared_ptr<core::MetadataStorage> metadata_store
     )
-            : m_data_store{std::move(data_store)},
+            : m_task_id{task_id},
+              m_data_store{std::move(data_store)},
               m_metadata_store{std::move(metadata_store)} {}
 
     auto get_data_store() -> std::shared_ptr<core::DataStorage> { return m_data_store; }
 
     auto get_metadata_store() -> std::shared_ptr<core::MetadataStorage> { return m_metadata_store; }
 
+    boost::uuids::uuid m_task_id;
+
     std::shared_ptr<core::DataStorage> m_data_store;
     std::shared_ptr<core::MetadataStorage> m_metadata_store;
 

src/spider/core/DataImpl.hpp

@@ -18,7 +18,7 @@ class DataImpl {
     }
 
     template <class T>
-    static auto get_impl(spider::Data<T> const& data) -> std::shared_ptr<DataStorage> {
+    static auto get_impl(spider::Data<T> const& data) -> std::unique_ptr<Data> const& {
         return data.get_impl();
     }
 };