Implement and test S3 Sink, with usage in runtime.

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Added
+
+- Support for writing to S3 buckets.
+
 ### Changed
 
 - Calling `acquire_get_configuration` with a Zarr storage device now returns a URI of the storage device, with file://

src/CMakeLists.txt

@@ -19,6 +19,8 @@ add_library(${tgt} MODULE
         writers/sink.creator.cpp
         writers/file.sink.hh
         writers/file.sink.cpp
+        writers/s3.sink.hh
+        writers/s3.sink.cpp
         writers/writer.hh
         writers/writer.cpp
         writers/zarrv2.writer.hh

src/common/utilities.cpp

@@ -201,9 +201,16 @@ common::split_uri(const std::string& uri)
     auto end = uri.find_first_of(delim);
 
     std::vector<std::string> out;
-    while (end != std::string::npos) {
+    while (end <= std::string::npos) {
         if (end > begin) {
-            out.emplace_back(uri.substr(begin, end - begin));
+            std::string part = uri.substr(begin, end - begin);
+            if (!part.empty()) {
+                out.emplace_back(part);
+            }
+        }
+
+        if (end == std::string::npos) {
+            break;
         }
 
         begin = end + 1;
@@ -480,6 +487,76 @@ extern "C"
 
         return retval;
     }
-}
 
-#endif
+    acquire_export int unit_test__split_uri()
+    {
+        try {
+            auto parts = common::split_uri("s3://bucket/key");
+            CHECK(parts.size() == 3);
+            CHECK(parts[0] == "s3:");
+            CHECK(parts[1] == "bucket");
+            CHECK(parts[2] == "key");
+
+            parts = common::split_uri("s3://bucket/key/");
+            CHECK(parts.size() == 3);
+            CHECK(parts[0] == "s3:");
+            CHECK(parts[1] == "bucket");
+            CHECK(parts[2] == "key");
+
+            parts = common::split_uri("s3://bucket/key/with/slashes");
+            CHECK(parts.size() == 5);
+            CHECK(parts[0] == "s3:");
+            CHECK(parts[1] == "bucket");
+            CHECK(parts[2] == "key");
+            CHECK(parts[3] == "with");
+            CHECK(parts[4] == "slashes");
+
+            parts = common::split_uri("s3://bucket");
+            CHECK(parts.size() == 2);
+            CHECK(parts[0] == "s3:");
+            CHECK(parts[1] == "bucket");
+
+            parts = common::split_uri("s3://");
+            CHECK(parts.size() == 1);
+            CHECK(parts[0] == "s3:");
+
+            parts = common::split_uri("s3:///");
+            CHECK(parts.size() == 1);
+            CHECK(parts[0] == "s3:");
+
+            parts = common::split_uri("s3://bucket/");
+            CHECK(parts.size() == 2);
+            CHECK(parts[0] == "s3:");
+            CHECK(parts[1] == "bucket");
+
+            parts = common::split_uri("s3://bucket/");
+            CHECK(parts.size() == 2);
+            CHECK(parts[0] == "s3:");
+            CHECK(parts[1] == "bucket");
+
+            parts = common::split_uri("s3://bucket/key/with/slashes/");
+            CHECK(parts.size() == 5);
+            CHECK(parts[0] == "s3:");
+            CHECK(parts[1] == "bucket");
+            CHECK(parts[2] == "key");
+            CHECK(parts[3] == "with");
+            CHECK(parts[4] == "slashes");
+
+            parts = common::split_uri("s3://bucket/key/with/slashes//");
+            CHECK(parts.size() == 5);
+            CHECK(parts[0] == "s3:");
+            CHECK(parts[1] == "bucket");
+            CHECK(parts[2] == "key");
+            CHECK(parts[3] == "with");
+            CHECK(parts[4] == "slashes");
+            return 1;
+        } catch (const std::exception& exc) {
+            LOGE("Exception: %s\n", exc.what());
+        } catch (...) {
+            LOGE("Exception: (unknown)");
+        }
+
+        return 0;
+    }
+}
+#endif

src/common/utilities.hh

@@ -22,8 +22,13 @@ namespace acquire::sink::zarr {
 
 struct Zarr;
 
-namespace common {
+enum class ZarrVersion
+{
+    V2 = 2,
+    V3
+};
 
+namespace common {
 /// @brief Get the number of chunks along a dimension.
 /// @param dimension A dimension.
 /// @return The number of, possibly ragged, chunks along the dimension, given

src/writers/file.sink.cpp

@@ -16,9 +16,8 @@ zarr::FileSink::FileSink(const std::string& uri)
 bool
 zarr::FileSink::write(size_t offset, const uint8_t* buf, size_t bytes_of_buf)
 {
-    if (!file_) {
-        return false;
-    }
+    CHECK(buf);
+    CHECK(bytes_of_buf);
 
     return file_write(file_.get(), offset, buf, buf + bytes_of_buf);
 }

src/writers/s3.sink.cpp

@@ -0,0 +1,160 @@
+#include "s3.sink.hh"
+
+#include "common/utilities.hh"
+#include "logger.h"
+
+#include <miniocpp/client.h>
+
+namespace zarr = acquire::sink::zarr;
+
+zarr::S3Sink::S3Sink(const std::string& bucket_name,
+                     const std::string& object_key,
+                     std::shared_ptr<common::S3ConnectionPool> connection_pool)
+  : bucket_name_{ bucket_name }
+  , object_key_{ object_key }
+  , connection_pool_{ connection_pool }
+  , buf_(5 << 20, 0) // 5 MiB is the minimum multipart upload size
+{
+    CHECK(!bucket_name_.empty());
+    CHECK(!object_key_.empty());
+    CHECK(connection_pool_);
+}
+
+zarr::S3Sink::~S3Sink()
+{
+    try {
+        // upload_id_ is only populated after successfully uploading a part
+        if (upload_id_.empty() && buf_size_ > 0) {
+            CHECK(put_object_());
+        } else if (!upload_id_.empty()) {
+            if (buf_size_ > 0) {
+                CHECK(flush_part_());
+            }
+            CHECK(finalize_multipart_upload_());
+        }
+    } catch (const std::exception& exc) {
+        LOGE("Error: %s", exc.what());
+    } catch (...) {
+        LOGE("Error: (unknown)");
+    }
+}
+
+bool
+zarr::S3Sink::write(size_t offset, const uint8_t* buf, size_t bytes_of_buf)
+{
+    CHECK(buf);
+    CHECK(bytes_of_buf);
+
+    while (bytes_of_buf > 0) {
+        const auto n = std::min(bytes_of_buf, buf_.size() - buf_size_);
+        if (n) {
+            std::copy(buf, buf + n, buf_.begin() + buf_size_);
+            buf_size_ += n;
+            buf += n;
+            bytes_of_buf -= n;
+        }
+
+        if (buf_size_ == buf_.size()) {
+            CHECK(flush_part_());
+        }
+    }
+
+    return true;
+}
+
+bool
+zarr::S3Sink::put_object_() noexcept
+{
+    if (buf_size_ == 0) {
+        return false;
+    }
+
+    auto connection = connection_pool_->get_connection();
+
+    bool retval = false;
+    try {
+        std::span<uint8_t> data(buf_.data(), buf_size_);
+        std::string etag =
+          connection->put_object(bucket_name_, object_key_, data);
+        EXPECT(
+          !etag.empty(), "Failed to upload object: %s", object_key_.c_str());
+
+        retval = true;
+    } catch (const std::exception& exc) {
+        LOGE("Error: %s", exc.what());
+    } catch (...) {
+        LOGE("Error: (unknown)");
+    }
+
+    // cleanup
+    connection_pool_->return_connection(std::move(connection));
+    buf_size_ = 0;
+
+    return retval;
+}
+
+bool
+zarr::S3Sink::flush_part_() noexcept
+{
+    if (buf_size_ == 0) {
+        return false;
+    }
+
+    auto connection = connection_pool_->get_connection();
+
+    bool retval = false;
+    try {
+        std::string upload_id = upload_id_;
+        if (upload_id.empty()) {
+            upload_id =
+              connection->create_multipart_object(bucket_name_, object_key_);
+            EXPECT(!upload_id.empty(),
+                   "Failed to create multipart object: %s",
+                   object_key_.c_str());
+        }
+
+        minio::s3::Part part;
+        part.number = static_cast<unsigned int>(parts_.size()) + 1;
+
+        std::span<uint8_t> data(buf_.data(), buf_size_);
+        part.etag = connection->upload_multipart_object_part(
+          bucket_name_, object_key_, upload_id, data, part.number);
+        EXPECT(!part.etag.empty(),
+               "Failed to upload part %u of object %s",
+               part.number,
+               object_key_.c_str());
+
+        // set these only when the part is successfully uploaded
+        parts_.push_back(part);
+        upload_id_ = upload_id;
+
+        retval = true;
+    } catch (const std::exception& exc) {
+        LOGE("Error: %s", exc.what());
+    } catch (...) {
+        LOGE("Error: (unknown)");
+    }
+
+    // cleanup
+    connection_pool_->return_connection(std::move(connection));
+    buf_size_ = 0;
+
+    return retval;
+}
+
+bool
+zarr::S3Sink::finalize_multipart_upload_() noexcept
+{
+    if (upload_id_.empty()) {
+        return false;
+    }
+
+    auto connection = connection_pool_->get_connection();
+
+    bool retval = connection->complete_multipart_object(
+      bucket_name_, object_key_, upload_id_, parts_);
+
+    connection_pool_->return_connection(std::move(connection));
+
+    return retval;
+}

src/writers/s3.sink.hh

@@ -0,0 +1,49 @@
+#pragma once
+
+#include "sink.hh"
+#include "platform.h"
+#include "common/s3.connection.hh"
+
+#include <miniocpp/types.h>
+
+#include <future>
+#include <string>
+
+namespace acquire::sink::zarr {
+struct S3Sink final : public Sink
+{
+    S3Sink() = delete;
+    S3Sink(const std::string& bucket_name,
+           const std::string& object_key,
+           std::shared_ptr<common::S3ConnectionPool> connection_pool);
+    ~S3Sink() override;
+
+    bool write(size_t offset, const uint8_t* buf, size_t bytes_of_buf) override;
+
+  private:
+    std::string bucket_name_;
+    std::string object_key_;
+
+    std::shared_ptr<common::S3ConnectionPool> connection_pool_;
+
+    // multipart upload
+    std::vector<uint8_t> buf_; // temporary 5MiB buffer for multipart upload
+    size_t buf_size_ = 0;
+
+    std::string upload_id_;
+    std::list<minio::s3::Part> parts_;
+
+    // single-part upload
+    /// @brief Upload the object to S3.
+    /// @returns True if the object was successfully uploaded, otherwise false.
+    [[nodiscard]] bool put_object_() noexcept;
+
+    // multipart upload
+    /// @brief Flush the current part to S3.
+    /// @returns True if the part was successfully flushed, otherwise false.
+    [[nodiscard]] bool flush_part_() noexcept;
+    /// @brief Finalize the multipart upload.
+    /// @returns True if a multipart upload was successfully finalized, otherwise false.
+    [[nodiscard]] bool finalize_multipart_upload_() noexcept;
+};
+} // namespace acquire::sink::zarr