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;