fix(storage): Always create a reference when getting data from storage (fixes #132).

src/spider/client/Driver.hpp

@@ -228,6 +228,8 @@ class Driver {
 
         return Job<ReturnType>{
                 job_id,
+                Job<ReturnType>::JobSource::Driver,
+                m_id,
                 m_metadata_storage,
                 m_data_storage,
                 m_storage_factory,
@@ -291,6 +293,8 @@ class Driver {
 
         return Job<ReturnType>{
                 job_id,
+                Job<ReturnType>::JobSource::Driver,
+                m_id,
                 m_metadata_storage,
                 m_data_storage,
                 m_storage_factory,

src/spider/client/Job.hpp

@@ -157,21 +157,34 @@ class Job {
     }
 
 private:
-    Job(boost::uuids::uuid id,
+    enum class JobSource : uint8_t {
+        Driver,
+        Task,
+    };
+
+    Job(boost::uuids::uuid const id,
+        JobSource const source,
+        boost::uuids::uuid const source_id,
         std::shared_ptr<core::MetadataStorage> metadata_storage,
         std::shared_ptr<core::DataStorage> data_storage,
         std::shared_ptr<core::StorageFactory> storage_factory)
             : m_id{id},
+              m_source{source},
+              m_source_id{source_id},
               m_metadata_storage{std::move(metadata_storage)},
               m_data_storage{std::move(data_storage)},
               m_storage_factory{std::move(storage_factory)} {}
 
-    Job(boost::uuids::uuid id,
+    Job(boost::uuids::uuid const id,
+        JobSource const source,
+        boost::uuids::uuid const source_id,
         std::shared_ptr<core::MetadataStorage> metadata_storage,
         std::shared_ptr<core::DataStorage> data_storage,
         std::shared_ptr<core::StorageFactory> storage_factory,
         std::shared_ptr<core::StorageConnection> conn)
             : m_id{id},
+              m_source{source},
+              m_source_id{source_id},
               m_metadata_storage{std::move(metadata_storage)},
               m_data_storage{std::move(data_storage)},
               m_storage_factory{std::move(storage_factory)},
@@ -242,7 +255,21 @@ class Job {
                     if (!optional_data_id.has_value()) {
                         throw ConnectionException{fmt::format("Output data ID is missing")};
                     }
-                    err = m_data_storage->get_data(conn, optional_data_id.value(), &data);
+                    if (m_source == JobSource::Driver) {
+                        err = m_data_storage->get_driver_data(
+                                conn,
+                                m_source_id,
+                                optional_data_id.value(),
+                                &data
+                        );
+                    } else {
+                        err = m_data_storage->get_task_data(
+                                conn,
+                                m_source_id,
+                                optional_data_id.value(),
+                                &data
+                        );
+                    }
                     if (!err.success()) {
                         throw ConnectionException{
                                 fmt::format("Failed to get data: {}", err.description)
@@ -302,7 +329,21 @@ class Job {
                 if (!optional_data_id.has_value()) {
                     throw ConnectionException{fmt::format("Output data ID is missing")};
                 }
-                err = m_data_storage->get_data(conn, optional_data_id.value(), &data);
+                if (m_source == JobSource::Driver) {
+                    err = m_data_storage->get_driver_data(
+                            conn,
+                            m_source_id,
+                            optional_data_id.value(),
+                            &data
+                    );
+                } else {
+                    err = m_data_storage->get_task_data(
+                            conn,
+                            m_source_id,
+                            optional_data_id.value(),
+                            &data
+                    );
+                }
                 if (!err.success()) {
                     throw ConnectionException{
                             fmt::format("Failed to get data: {}", err.description)
@@ -337,6 +378,8 @@ class Job {
     // NOLINTEND(readability-function-cognitive-complexity)
 
     boost::uuids::uuid m_id;
+    JobSource m_source;
+    boost::uuids::uuid m_source_id;
     std::shared_ptr<core::MetadataStorage> m_metadata_storage;
     std::shared_ptr<core::DataStorage> m_data_storage;
     std::shared_ptr<core::StorageFactory> m_storage_factory;

src/spider/client/TaskContext.hpp

@@ -172,7 +172,14 @@ class TaskContext {
             throw ConnectionException(fmt::format("Failed to start job: {}", err.description));
         }
 
-        return Job<ReturnType>{job_id, m_metadata_store, m_data_store, m_storage_factory};
+        return Job<ReturnType>{
+                job_id,
+                Job<ReturnType>::JobSource::Task,
+                m_task_id,
+                m_metadata_store,
+                m_data_store,
+                m_storage_factory
+        };
     }
 
     /**
@@ -224,7 +231,14 @@ class TaskContext {
             throw ConnectionException(fmt::format("Failed to start job: {}", err.description));
         }
 
-        return Job<ReturnType>{job_id, m_metadata_store, m_data_store, m_storage_factory};
+        return Job<ReturnType>{
+                job_id,
+                Job<ReturnType>::JobSource::Task,
+                m_task_id,
+                m_metadata_store,
+                m_data_store,
+                m_storage_factory
+        };
     }
 
     /**

src/spider/storage/DataStorage.hpp

@@ -32,6 +32,38 @@ class DataStorage {
             = 0;
     virtual auto get_data(StorageConnection& conn, boost::uuids::uuid id, Data* data) -> StorageErr
             = 0;
+    /**
+     * Get a data object and register a reference for it from the given driver in a single
+     * transaction.
+     * @param conn
+     * @param driver_id
+     * @param data_id
+     * @param data output data
+     * @return StorageErr::Success if the transaction succeed. Error types otherwise.
+     */
+    virtual auto get_driver_data(
+            StorageConnection& conn,
+            boost::uuids::uuid driver_id,
+            boost::uuids::uuid data_id,
+            Data* data
+    ) -> StorageErr
+            = 0;
+    /**
+     * Get a data object and register a reference for it from the given task in a single
+     * transaction.
+     * @param conn
+     * @param task_id
+     * @param data_id
+     * @param data output data
+     * @return StorageErr::Success if the transaction succeed. Error types otherwise.
+     */
+    virtual auto get_task_data(
+            StorageConnection& conn,
+            boost::uuids::uuid task_id,
+            boost::uuids::uuid data_id,
+            Data* data
+    ) -> StorageErr
+            = 0;
     virtual auto set_data_locality(StorageConnection& conn, Data const& data) -> StorageErr = 0;
     virtual auto remove_data(StorageConnection& conn, boost::uuids::uuid id) -> StorageErr = 0;
     virtual auto

src/spider/storage/mysql/MySqlStorage.cpp

@@ -2074,42 +2074,114 @@ auto MySqlDataStorage::add_task_data(
     return StorageErr{};
 }
 
-auto MySqlDataStorage::get_data(StorageConnection& conn, boost::uuids::uuid id, Data* data)
+auto MySqlDataStorage::get_data_with_locality(
+        StorageConnection& conn,
+        boost::uuids::uuid const id,
+        Data* data
+) -> StorageErr {
+    std::unique_ptr<sql::PreparedStatement> statement(
+            static_cast<MySqlConnection&>(conn)->prepareStatement(
+                    "SELECT `id`, `value`, `hard_locality` FROM `data` WHERE `id` = ?"
+            )
+    );
+    sql::bytes id_bytes = uuid_get_bytes(id);
+    statement->setBytes(1, &id_bytes);
+    std::unique_ptr<sql::ResultSet> res(statement->executeQuery());
+    if (res->rowsCount() == 0) {
+        static_cast<MySqlConnection&>(conn)->rollback();
+        return StorageErr{
+                StorageErrType::KeyNotFoundErr,
+                fmt::format("no data with id {}", boost::uuids::to_string(id))
+        };
+    }
+    res->next();
+    *data = Data{id, get_sql_string(res->getString(2))};
+    data->set_hard_locality(res->getBoolean(3));
+
+    std::unique_ptr<sql::PreparedStatement> locality_statement(
+            static_cast<MySqlConnection&>(conn)->prepareStatement(
+                    "SELECT `address` FROM `data_locality` WHERE `id` = ?"
+            )
+    );
+    locality_statement->setBytes(1, &id_bytes);
+    std::unique_ptr<sql::ResultSet> const locality_res(locality_statement->executeQuery());
+    std::vector<std::string> locality;
+    while (locality_res->next()) {
+        locality.emplace_back(get_sql_string(locality_res->getString(1)));
+    }
+    if (!locality.empty()) {
+        data->set_locality(locality);
+    }
+    return StorageErr{};
+}
+
+auto MySqlDataStorage::get_data(StorageConnection& conn, boost::uuids::uuid const id, Data* data)
         -> StorageErr {
     try {
-        std::unique_ptr<sql::PreparedStatement> statement(
+        StorageErr const err = get_data_with_locality(conn, id, data);
+        if (false == err.success()) {
+            return err;
+        }
+    } catch (sql::SQLException& e) {
+        static_cast<MySqlConnection&>(conn)->rollback();
+        return StorageErr{StorageErrType::OtherErr, e.what()};
+    }
+    static_cast<MySqlConnection&>(conn)->commit();
+    return StorageErr{};
+}
+
+auto MySqlDataStorage::get_driver_data(
+        StorageConnection& conn,
+        boost::uuids::uuid const driver_id,
+        boost::uuids::uuid const data_id,
+        Data* data
+) -> StorageErr {
+    try {
+        StorageErr const err = get_data_with_locality(conn, data_id, data);
+        if (false == err.success()) {
+            return err;
+        }
+        // Add data reference from driver
+        std::unique_ptr<sql::PreparedStatement> statement{
                 static_cast<MySqlConnection&>(conn)->prepareStatement(
-                        "SELECT `id`, `value`, `hard_locality` FROM `data` WHERE `id` = ?"
+                        "INSERT INTO `data_ref_driver` (`id`, `driver_id`) VALUES (?, ?)"
                 )
-        );
-        sql::bytes id_bytes = uuid_get_bytes(id);
+        };
+        sql::bytes id_bytes = uuid_get_bytes(data_id);
+        sql::bytes driver_id_bytes = uuid_get_bytes(driver_id);
         statement->setBytes(1, &id_bytes);
-        std::unique_ptr<sql::ResultSet> res(statement->executeQuery());
-        if (res->rowsCount() == 0) {
-            static_cast<MySqlConnection&>(conn)->rollback();
-            return StorageErr{
-                    StorageErrType::KeyNotFoundErr,
-                    fmt::format("no data with id {}", boost::uuids::to_string(id))
-            };
-        }
-        res->next();
-        *data = Data{id, get_sql_string(res->getString(2))};
-        data->set_hard_locality(res->getBoolean(3));
+        statement->setBytes(2, &driver_id_bytes);
+        statement->executeUpdate();
+    } catch (sql::SQLException& e) {
+        static_cast<MySqlConnection&>(conn)->rollback();
+        return StorageErr{StorageErrType::OtherErr, e.what()};
+    }
+    static_cast<MySqlConnection&>(conn)->commit();
+    return StorageErr{};
+}
 
-        std::unique_ptr<sql::PreparedStatement> locality_statement(
+auto MySqlDataStorage::get_task_data(
+        StorageConnection& conn,
+        boost::uuids::uuid const task_id,
+        boost::uuids::uuid const data_id,
+        Data* data
+) -> StorageErr {
+    try {
+        StorageErr const err = get_data_with_locality(conn, data_id, data);
+        if (false == err.success()) {
+            return err;
+        }
+        // Add data reference from task
+        std::unique_ptr<sql::PreparedStatement> statement{
                 static_cast<MySqlConnection&>(conn)->prepareStatement(
-                        "SELECT `address` FROM `data_locality` WHERE `id` = ?"
+                        "INSERT INTO `data_ref_task` (`id`, `task_id`) VALUES (?, ?)"
                 )
-        );
-        locality_statement->setBytes(1, &id_bytes);
-        std::unique_ptr<sql::ResultSet> const locality_res(locality_statement->executeQuery());
-        std::vector<std::string> locality;
-        while (locality_res->next()) {
-            locality.emplace_back(get_sql_string(locality_res->getString(1)));
-        }
-        if (!locality.empty()) {
-            data->set_locality(locality);
-        }
+        };
+        sql::bytes id_bytes = uuid_get_bytes(data_id);
+        sql::bytes task_id_bytes = uuid_get_bytes(task_id);
+        statement->setBytes(1, &id_bytes);
+        statement->setBytes(2, &task_id_bytes);
+        statement->executeUpdate();
     } catch (sql::SQLException& e) {
         static_cast<MySqlConnection&>(conn)->rollback();
         return StorageErr{StorageErrType::OtherErr, e.what()};

src/spider/storage/mysql/MySqlStorage.hpp

@@ -148,6 +148,18 @@ class MySqlDataStorage : public DataStorage {
             -> StorageErr override;
     auto get_data(StorageConnection& conn, boost::uuids::uuid id, Data* data)
             -> StorageErr override;
+    auto get_driver_data(
+            StorageConnection& conn,
+            boost::uuids::uuid driver_id,
+            boost::uuids::uuid data_id,
+            Data* data
+    ) -> StorageErr override;
+    auto get_task_data(
+            StorageConnection& conn,
+            boost::uuids::uuid task_id,
+            boost::uuids::uuid data_id,
+            Data* data
+    ) -> StorageErr override;
     auto set_data_locality(StorageConnection& conn, Data const& data) -> StorageErr override;
     auto remove_data(StorageConnection& conn, boost::uuids::uuid id) -> StorageErr override;
     auto
@@ -187,6 +199,9 @@ class MySqlDataStorage : public DataStorage {
     ) -> StorageErr override;
 
 private:
+    static auto get_data_with_locality(StorageConnection& conn, boost::uuids::uuid id, Data* data)
+            -> StorageErr;
+
     MySqlDataStorage() = default;
 
     friend class MySqlStorageFactory;

src/spider/worker/FunctionManager.hpp

@@ -44,7 +44,8 @@ using ArgsBuffer = msgpack::sbuffer;
 
 using ResultBuffer = msgpack::sbuffer;
 
-using Function = std::function<ResultBuffer(TaskContext& context, ArgsBuffer const&)>;
+using Function = std::function<
+        ResultBuffer(TaskContext& context, boost::uuids::uuid task_id, ArgsBuffer const&)>;
 
 using FunctionMap = absl::flat_hash_map<std::string, Function>;
 
@@ -235,8 +236,12 @@ inline auto create_args_request(std::vector<msgpack::sbuffer> const& args_buffer
 template <class F>
 class FunctionInvoker {
 public:
-    static auto apply(F const& function, TaskContext& context, ArgsBuffer const& args_buffer)
-            -> ResultBuffer {
+    static auto apply(
+            F const& function,
+            TaskContext& context,
+            boost::uuids::uuid const task_id,
+            ArgsBuffer const& args_buffer
+    ) -> ResultBuffer {
         // NOLINTBEGIN(cppcoreguidelines-pro-type-union-access,cppcoreguidelines-pro-bounds-pointer-arithmetic)
         using ArgsTuple = signature<F>::args_t;
         using ReturnType = signature<F>::ret_t;
@@ -301,7 +306,7 @@ class FunctionInvoker {
                 if constexpr (cIsSpecializationV<T, spider::Data>) {
                     boost::uuids::uuid const data_id = arg.as<boost::uuids::uuid>();
                     std::unique_ptr<Data> data = std::make_unique<Data>();
-                    err = data_store->get_data(*conn, data_id, data.get());
+                    err = data_store->get_task_data(*conn, task_id, data_id, data.get());
                     if (!err.success()) {
                         return;
                     }
@@ -372,7 +377,8 @@ class FunctionManager {
                                 &FunctionInvoker<F>::apply,
                                 std::move(f),
                                 std::placeholders::_1,
-                                std::placeholders::_2
+                                std::placeholders::_2,
+                                std::placeholders::_3
                         )
                 )
                 .second;