WIP Async writing

Author: franzpoeschel

include/openPMD/toolkit/Aws.hpp

@@ -4,25 +4,47 @@
 
 #include <aws/s3/S3Client.h>
 
+#include <condition_variable>
+
 namespace openPMD::internal
 {
+struct AwsAsyncHandler
+{
+    std::mutex mutex;
+    std::condition_variable event;
+    std::size_t request_counter = 0;
+    // Upon C++20, we can use a std::atomic for this and ditch the
+    // condition_variable + mutex approach
+    std::size_t completion_counter = 0;
+
+    void wait();
+    void add_task();
+    void add_and_notify_result();
+
+    ~AwsAsyncHandler();
+};
+
 struct ExternalBlockStorageAws : ExternalBlockStorageBackend
 {
 private:
     Aws::S3::S3Client m_client;
     std::string m_bucketName;
     std::optional<std::string> m_endpoint;
+    std::optional<AwsAsyncHandler> m_async;
 
 public:
     ExternalBlockStorageAws(
         Aws::S3::S3Client,
         std::string bucketName,
-        std::optional<std::string> endpoint);
+        std::optional<std::string> endpoint,
+        bool async);
     auto put(std::string const &identifier, void const *data, size_t len)
         -> std::string override;
     void get(std::string const &external_ref, void *data, size_t len) override;
     [[nodiscard]] auto externalStorageLocation() const
         -> nlohmann::json override;
+    void sync() override;
+
     ~ExternalBlockStorageAws() override;
 };
 } // namespace openPMD::internal

include/openPMD/toolkit/ExternalBlockStorage.hpp

@@ -29,6 +29,8 @@ struct ExternalBlockStorageBackend
     [[nodiscard]] virtual auto externalStorageLocation() const
         -> nlohmann::json = 0;
 
+    virtual void sync();
+
     virtual ~ExternalBlockStorageBackend();
 };
 } // namespace openPMD::internal
@@ -105,6 +107,8 @@ class ExternalBlockStorage
         nlohmann::json::json_pointer const &path,
         T *data);
 
+    void sync();
+
     [[nodiscard]] auto externalStorageLocation() const -> nlohmann::json;
 
     static void sanitizeString(std::string &s);

src/IO/JSON/JSONIOHandlerImpl.cpp

@@ -689,6 +689,13 @@ std::future<void> JSONIOHandlerImpl::flush()
         putJsonContents(file, false);
     }
     m_dirty.clear();
+    std::visit(
+        auxiliary::overloaded{
+            [](DatasetMode::External_t &externalStorage) {
+                externalStorage->sync();
+            },
+            [](auto &&) {}},
+        this->m_datasetMode.m_mode.as_base());
     return std::future<void>();
 }
 

src/toolkit/Aws.cpp

@@ -1,10 +1,12 @@
 #include "openPMD/toolkit/Aws.hpp"
 
+#include <aws/s3/S3Client.h>
 #include <aws/s3/model/CreateBucketRequest.h>
 #include <aws/s3/model/GetObjectRequest.h>
 #include <aws/s3/model/PutObjectRequest.h>
 
 #include <iostream>
+#include <mutex>
 #include <stdexcept>
 
 namespace
@@ -31,13 +33,44 @@ struct imemstream : std::iostream
 
 namespace openPMD::internal
 {
+void AwsAsyncHandler::wait()
+{
+    std::cerr << "Waiting for remaining tasks. Have " << completion_counter
+              << " of " << request_counter << std::endl;
+    size_t target = this->request_counter;
+    std::unique_lock lk(this->mutex);
+    this->event.wait(
+        lk, [this, target]() { return this->completion_counter >= target; });
+    std::cerr << "Finished waiting for remaining tasks" << std::endl;
+}
+
+void AwsAsyncHandler::add_task()
+{
+    this->request_counter++;
+}
+
+void AwsAsyncHandler::add_and_notify_result()
+{
+    std::unique_lock lk(this->mutex);
+    this->completion_counter++;
+    lk.unlock();
+    this->event.notify_all();
+}
+
+AwsAsyncHandler::~AwsAsyncHandler()
+{
+    this->wait();
+}
+
 ExternalBlockStorageAws::ExternalBlockStorageAws(
     Aws::S3::S3Client client,
     std::string bucketName,
-    std::optional<std::string> endpoint)
+    std::optional<std::string> endpoint,
+    bool async)
     : m_client{std::move(client)}
     , m_bucketName(std::move(bucketName))
     , m_endpoint(std::move(endpoint))
+    , m_async(async ? std::make_optional<AwsAsyncHandler>() : std::nullopt)
 {
     Aws::S3::Model::CreateBucketRequest create_request;
     create_request.SetBucket(m_bucketName);
@@ -71,16 +104,49 @@ auto ExternalBlockStorageAws::put(
     put_request.SetBody(input_data);
     put_request.SetContentLength(static_cast<long long>(len));
 
-    auto put_outcome = m_client.PutObject(put_request);
-
-    if (put_outcome.IsSuccess())
+    if (!m_async.has_value())
     {
-        std::cout << "File uploaded successfully to S3!" << std::endl;
+        auto put_outcome = m_client.PutObject(put_request);
+
+        if (put_outcome.IsSuccess())
+        {
+            std::cout << "File synchronously uploaded successfully to S3!"
+                      << std::endl;
+        }
+        else
+        {
+            std::cerr << "Synchronous upload failed: "
+                      << put_outcome.GetError().GetMessage() << std::endl;
+        }
     }
     else
     {
-        std::cerr << "Upload failed: " << put_outcome.GetError().GetMessage()
-                  << std::endl;
+        auto &async_handler = *m_async;
+        auto responseReceivedHandler =
+            [&async_handler](
+                const Aws::S3::S3Client *,
+                const Aws::S3::Model::PutObjectRequest &,
+                const Aws::S3::Model::PutObjectOutcome &put_outcome,
+                const std::shared_ptr<const Aws::Client::AsyncCallerContext>
+                    &) {
+                if (put_outcome.IsSuccess())
+                {
+                    std::cout
+                        << "File asynchronously uploaded successfully to S3!"
+                        << std::endl;
+                }
+                else
+                {
+                    std::cerr << "Asynchronous upload failed: "
+                              << put_outcome.GetError().GetMessage()
+                              << std::endl;
+                }
+                async_handler.add_and_notify_result();
+            };
+        async_handler.add_task();
+        m_client.PutObjectAsync(put_request, responseReceivedHandler);
+        // todo replace this
+        async_handler.wait();
     }
     return sanitized;
 }
@@ -117,6 +183,15 @@ void ExternalBlockStorageAws::get(
     }
 }
 
+void ExternalBlockStorageAws::sync()
+{
+    if (!this->m_async.has_value())
+    {
+        return;
+    }
+    this->m_async->wait();
+}
+
 [[nodiscard]] auto ExternalBlockStorageAws::externalStorageLocation() const
     -> nlohmann::json
 {

src/toolkit/AwsBuilder.cpp

@@ -119,7 +119,9 @@ AwsBuilder::operator ExternalBlockStorage()
     return ExternalBlockStorage{std::make_unique<ExternalBlockStorageAws>(
         std::move(s3_client),
         std::move(m_bucketName),
-        std::move(m_endpointOverride))};
+        std::move(m_endpointOverride),
+        // TODO: Add config option for this
+        /* async = */ true)};
 }
 
 auto AwsBuilder::build() -> ExternalBlockStorage

src/toolkit/ExternalBlockStorage.cpp

@@ -12,9 +12,14 @@
 
 namespace openPMD::internal
 {
-ExternalBlockStorageBackend::~ExternalBlockStorageBackend() = default;
+void ExternalBlockStorageBackend::sync()
+{
+    // default for non-async backends: no-op
 }
 
+ExternalBlockStorageBackend::~ExternalBlockStorageBackend() = default;
+} // namespace openPMD::internal
+
 namespace openPMD
 {
 
@@ -161,7 +166,9 @@ void ExternalBlockStorage::read(
     [[maybe_unused]] nlohmann::json const &fullJsonDataset,
     [[maybe_unused]] nlohmann::json::json_pointer const &path,
     [[maybe_unused]] T *data)
-{}
+{
+    throw std::runtime_error("Unimplemented!");
+}
 
 template <typename DatatypeHandling, typename T>
 void ExternalBlockStorage::read(
@@ -211,6 +218,11 @@ void ExternalBlockStorage::read(
     }
 }
 
+void ExternalBlockStorage::sync()
+{
+    this->m_worker->sync();
+}
+
 [[nodiscard]] auto ExternalBlockStorage::externalStorageLocation() const
     -> nlohmann::json
 {