Get object sizes based on S3's ListObjects output (just with scan_object_sizes() to start with)

8560764974

Author: poodlewars

cpp/arcticdb/async/async_store.hpp

@@ -223,6 +223,30 @@ void iterate_type(
     library_->iterate_type(type, func, prefix);
 }
 
+folly::Future<storage::ObjectSizes> get_object_sizes(KeyType type, const std::string& prefix) override {
+    if (library_->supports_object_size_calculation()) {
+        // The library has native support for some kind of clever size calculation, so let it take over
+        return async::submit_io_task(ObjectSizesTask{type, prefix, library_});
+    }
+
+    // No native support for a clever size calculation, so just read keys and sum their sizes
+    auto counter = std::make_shared<std::atomic_uint64_t>(0);
+    auto bytes = std::make_shared<std::atomic_uint64_t>(0);
+    KeySizeCalculators key_size_calculators;
+    iterate_type(type, [&key_size_calculators, &counter, &bytes](const VariantKey&& k) {
+        key_size_calculators.emplace_back(std::forward<const VariantKey>(k), [counter, bytes] (auto&& ks) {
+            counter->fetch_add(1);
+            auto key_seg = std::forward<decltype(ks)>(ks);
+            auto compressed_size = key_seg.segment().size();
+            bytes->fetch_add(compressed_size);
+            return key_seg.variant_key();
+        });
+    }, prefix);
+
+    read_ignoring_key_not_found(std::move(key_size_calculators));
+    return folly::makeFuture(storage::ObjectSizes{type, *counter, *bytes});
+}
+
 bool scan_for_matching_key(
     KeyType key_type, const IterateTypePredicate& predicate) override {
     return library_->scan_for_matching_key(key_type, predicate);

cpp/arcticdb/async/tasks.hpp

@@ -665,4 +665,28 @@ struct RemoveBatchTask : BaseTask {
     }
 };
 
+struct ObjectSizesTask : BaseTask {
+    KeyType type_;
+    std::string prefix_;
+    std::shared_ptr<storage::Library> lib_;
+
+    ObjectSizesTask(
+        KeyType type,
+        std::string prefix,
+        std::shared_ptr<storage::Library> lib
+        ) :
+        type_(type),
+        prefix_(std::move(prefix)),
+        lib_(std::move(lib)) {
+        ARCTICDB_DEBUG(log::storage(), "Creating object sizes task for key type {} prefix {}", type_, prefix_);
+    }
+
+    ARCTICDB_MOVE_ONLY_DEFAULT(ObjectSizesTask)
+
+    folly::Future<storage::ObjectSizes> operator()() {
+        util::check(lib_->supports_object_size_calculation(), "ObjectSizesBytesTask should only be used with storages"
+                                                              " that natively support size calculation");
+        return lib_->get_object_sizes(type_, prefix_);
+    }
+};
 }

cpp/arcticdb/storage/library.hpp

@@ -70,6 +70,15 @@ class Library {
         storages_->iterate_type(key_type, visitor, prefix);
     }
 
+    bool supports_object_size_calculation() {
+        return storages_->supports_object_size_calculation();
+    }
+
+    ObjectSizes get_object_sizes(KeyType type, const std::string& prefix) {
+        ARCTICDB_SAMPLE(GetObjectSizes, 0)
+        return storages_->get_object_sizes(type, prefix);
+    }
+
     /**
      * Scan through every key of the given type until one matches the predicate.
      *

cpp/arcticdb/storage/s3/detail-inl.hpp

@@ -405,23 +405,30 @@ void do_update_impl(
     do_write_impl(std::move(kvs), root_folder, bucket_name, s3_client, std::forward<KeyBucketizer>(bucketizer));
 }
 
-inline auto default_prefix_handler() {
+inline PrefixHandler default_prefix_handler() {
     return [](const std::string& prefix, const std::string& key_type_dir, const KeyDescriptor& key_descriptor, KeyType) {
         return !prefix.empty() ? fmt::format("{}/{}*{}", key_type_dir, key_descriptor, prefix) : key_type_dir;
     };
 }
 
-template<class KeyBucketizer, class PrefixHandler>
-bool do_iterate_type_impl(
-    KeyType key_type,
-    const IterateTypePredicate& visitor,
+struct PathInfo {
+    PathInfo(std::string prefix, std::string key_type_dir, size_t path_to_key_size) :
+        key_prefix_(std::move(prefix)), key_type_dir_(std::move(key_type_dir)), path_to_key_size_(path_to_key_size) {
+
+    }
+
+    std::string key_prefix_;
+    std::string key_type_dir_;
+    size_t path_to_key_size_;
+};
+
+template<class KeyBucketizer>
+PathInfo calculate_path_info(
     const std::string& root_folder,
-    const std::string& bucket_name,
-    const S3ClientInterface& s3_client,
-    KeyBucketizer&& bucketizer,
-    PrefixHandler&& prefix_handler = default_prefix_handler(),
-    const std::string& prefix = std::string{}) {
-    ARCTICDB_SAMPLE(S3StorageIterateType, 0)
+    KeyType key_type,
+    const PrefixHandler& prefix_handler,
+    const std::string& prefix,
+    KeyBucketizer&& bucketizer) {
     auto key_type_dir = key_type_folder(root_folder, key_type);
     const auto path_to_key_size = key_type_dir.size() + 1 + bucketizer.bucketize_length(key_type);
     // if prefix is empty, add / to avoid matching both 'log' and 'logc' when key_type_dir is {root_folder}/log
@@ -438,19 +445,36 @@ bool do_iterate_type_impl(
                                                             : IndexDescriptorImpl::Type::TIMESTAMP,
                                  FormatType::TOKENIZED);
     auto key_prefix = prefix_handler(prefix, key_type_dir, key_descriptor, key_type);
-    ARCTICDB_RUNTIME_DEBUG(log::storage(), "Searching for objects in bucket {} with prefix {}", bucket_name,
-                           key_prefix);
+
+    return {key_prefix, key_type_dir, path_to_key_size};
+}
+
+template<class KeyBucketizer>
+bool do_iterate_type_impl(
+    KeyType key_type,
+    const IterateTypePredicate& visitor,
+    const std::string& root_folder,
+    const std::string& bucket_name,
+    const S3ClientInterface& s3_client,
+    KeyBucketizer&& bucketizer,
+    const PrefixHandler& prefix_handler = default_prefix_handler(),
+    const std::string& prefix = std::string{}) {
+    ARCTICDB_SAMPLE(S3StorageIterateType, 0)
+
+    auto path_info = calculate_path_info(root_folder, key_type, prefix_handler, prefix, std::move(bucketizer));
+    ARCTICDB_RUNTIME_DEBUG(log::storage(), "Iterating over objects in bucket {} with prefix {}", bucket_name,
+                           path_info.key_prefix_);
 
     auto continuation_token = std::optional<std::string>();
     do {
-        auto list_objects_result = s3_client.list_objects(key_prefix, bucket_name, continuation_token);
+        auto list_objects_result = s3_client.list_objects(path_info.key_prefix_, bucket_name, continuation_token);
         if (list_objects_result.is_success()) {
             auto& output = list_objects_result.get_output();
 
             ARCTICDB_RUNTIME_DEBUG(log::storage(), "Received object list");
 
             for (auto& s3_object_name : output.s3_object_names) {
-                auto key = s3_object_name.substr(path_to_key_size);
+                auto key = s3_object_name.substr(path_info.path_to_key_size_);
                 ARCTICDB_TRACE(log::version(), "Got object_list: {}, key: {}", s3_object_name, key);
                 auto k = variant_key_from_bytes(
                     reinterpret_cast<uint8_t *>(key.data()),
@@ -474,13 +498,55 @@ bool do_iterate_type_impl(
                                 error.GetMessage().c_str());
             // We don't raise on expected errors like NoSuchKey because we want to return an empty list
             // instead of raising.
-            raise_if_unexpected_error(error, key_prefix);
+            raise_if_unexpected_error(error, path_info.key_prefix_);
             return false;
         }
     } while (continuation_token.has_value());
     return false;
 }
 
+template<class KeyBucketizer>
+ObjectSizes do_calculate_sizes_for_type_impl(
+    KeyType key_type,
+    const std::string& root_folder,
+    const std::string& bucket_name,
+    const S3ClientInterface& s3_client,
+    KeyBucketizer&& bucketizer,
+    const PrefixHandler& prefix_handler = default_prefix_handler(),
+    const std::string& prefix = std::string{}) {
+    ARCTICDB_SAMPLE(S3StorageCalculateSizesForType, 0)
+
+    auto path_info = calculate_path_info(root_folder, key_type, prefix_handler, prefix, std::move(bucketizer));
+    ARCTICDB_RUNTIME_DEBUG(log::storage(), "Calculating sizes for objects in bucket {} with prefix {}", bucket_name,
+                           path_info.key_prefix_);
+
+    auto continuation_token = std::optional<std::string>();
+    ObjectSizes res{key_type};
+    do {
+        auto list_objects_result = s3_client.list_objects(path_info.key_prefix_, bucket_name, continuation_token);
+        if (list_objects_result.is_success()) {
+            auto& output = list_objects_result.get_output();
+
+            ARCTICDB_RUNTIME_DEBUG(log::storage(), "Received object list");
+
+            for (auto& s3_object_size : output.s3_object_sizes) {
+                res.count_ += 1;
+                res.compressed_size_bytes_ += s3_object_size;
+            }
+            continuation_token = output.next_continuation_token;
+        } else {
+            const auto& error = list_objects_result.get_error();
+            log::storage().warn("Failed to iterate key type with key '{}' {}: {}",
+                                key_type,
+                                error.GetExceptionName().c_str(),
+                                error.GetMessage().c_str());
+            raise_if_unexpected_error(error, path_info.key_prefix_);
+        }
+    } while (continuation_token.has_value());
+
+    return res;
+}
+
 template<class KeyBucketizer>
 bool do_key_exists_impl(
     const VariantKey& key,

cpp/arcticdb/storage/s3/nfs_backed_storage.cpp

@@ -197,6 +197,16 @@ void NfsBackedStorage::do_remove(std::span<VariantKey> variant_keys, RemoveOpts)
     s3::detail::do_remove_impl(std::span(enc), root_folder_, bucket_name_, *s3_client_, NfsBucketizer{});
 }
 
+static std::string prefix_handler(const std::string& prefix, const std::string& key_type_dir, const KeyDescriptor&, KeyType key_type) {
+    std::string new_prefix;
+    if(!prefix.empty()) {
+        uint32_t id = get_id_bucket(encode_item<StreamId, StringId, NumericId>(StringId{prefix}, is_ref_key_class(key_type)));
+        new_prefix = fmt::format("{:03}", id);
+    }
+
+    return !prefix.empty() ? fmt::format("{}/{}", key_type_dir, new_prefix) : key_type_dir;
+}
+
 bool NfsBackedStorage::do_iterate_type_until_match(KeyType key_type, const IterateTypePredicate& visitor, const std::string& prefix) {
     const IterateTypePredicate func = [&v = visitor, prefix=prefix] (VariantKey&& k) {
         auto key = unencode_object_id(k);
@@ -207,16 +217,6 @@ bool NfsBackedStorage::do_iterate_type_until_match(KeyType key_type, const Itera
         }
     };
 
-    auto prefix_handler = [] (const std::string& prefix, const std::string& key_type_dir, const KeyDescriptor&, KeyType key_type) {
-        std::string new_prefix;
-        if(!prefix.empty()) {
-            uint32_t id = get_id_bucket(encode_item<StreamId, StringId, NumericId>(StringId{prefix}, is_ref_key_class(key_type)));
-            new_prefix = fmt::format("{:03}", id);
-        }
-
-        return !prefix.empty() ? fmt::format("{}/{}", key_type_dir, new_prefix) : key_type_dir;
-    };
-
     return s3::detail::do_iterate_type_impl(key_type, func, root_folder_, bucket_name_, *s3_client_, NfsBucketizer{}, prefix_handler, prefix);
 }
 
@@ -225,4 +225,12 @@ bool NfsBackedStorage::do_key_exists(const VariantKey& key) {
     return s3::detail::do_key_exists_impl(encoded_key, root_folder_, bucket_name_, *s3_client_, NfsBucketizer{});
 }
 
+bool NfsBackedStorage::supports_object_size_calculation() const {
+    return true;
+}
+
+ObjectSizes NfsBackedStorage::do_get_object_sizes(KeyType key_type, const std::string& prefix) {
+    return s3::detail::do_calculate_sizes_for_type_impl(key_type, root_folder_, bucket_name_, *s3_client_, NfsBucketizer{}, prefix_handler, prefix);
+}
+
 } //namespace arcticdb::storage::nfs_backed

cpp/arcticdb/storage/s3/nfs_backed_storage.hpp

@@ -32,6 +32,8 @@ class NfsBackedStorage final : public Storage {
 
     std::string name() const final;
 
+    bool supports_object_size_calculation() const final override;
+
 private:
     void do_write(KeySegmentPair& key_seg) final;
 
@@ -51,6 +53,8 @@ class NfsBackedStorage final : public Storage {
 
     bool do_iterate_type_until_match(KeyType key_type, const IterateTypePredicate& visitor, const std::string &prefix) final;
 
+    ObjectSizes do_get_object_sizes(KeyType key_type, const std::string& prefix) final;
+
     bool do_key_exists(const VariantKey& key) final;
 
     bool do_supports_prefix_matching() const final {

cpp/arcticdb/storage/s3/s3_client_impl.cpp

@@ -299,11 +299,13 @@ S3Result<ListObjectsOutput> S3ClientImpl::list_objects(
         next_continuation_token = {result.GetNextContinuationToken()};
 
     auto s3_object_names = std::vector<std::string>();
+    auto s3_object_sizes = std::vector<uint64_t>();
     for (const auto &s3_object: result.GetContents()) {
         s3_object_names.emplace_back(s3_object.GetKey());
+        s3_object_sizes.emplace_back(s3_object.GetSize());
     }
 
-    ListObjectsOutput output = {s3_object_names, next_continuation_token};
+    ListObjectsOutput output = {s3_object_names, s3_object_sizes, next_continuation_token};
     return {output};
 }
 

cpp/arcticdb/storage/s3/s3_client_interface.hpp

@@ -49,6 +49,7 @@ using S3Result = StorageResult<Output, Aws::S3::S3Error>;
 
 struct ListObjectsOutput{
     std::vector<std::string> s3_object_names;
+    std::vector<uint64_t> s3_object_sizes;
     // next_continuation_token indicates there are more s3_objects to be listed because they didn't fit in one response.
     // If set can be used to get the remaining s3_objects.
     std::optional<std::string> next_continuation_token;

cpp/arcticdb/storage/s3/s3_storage.cpp

@@ -105,7 +105,15 @@ bool S3Storage::do_iterate_type_until_match(KeyType key_type, const IterateTypeP
         return !prefix.empty() ? fmt::format("{}/{}*{}", key_type_dir, key_descriptor, prefix) : key_type_dir;
     };
 
-    return detail::do_iterate_type_impl(key_type, visitor, root_folder_, bucket_name_, client(), FlatBucketizer{}, std::move(prefix_handler), prefix);
+    return detail::do_iterate_type_impl(key_type, visitor, root_folder_, bucket_name_, client(), FlatBucketizer{}, prefix_handler, prefix);
+}
+
+ObjectSizes S3Storage::do_get_object_sizes(KeyType key_type, const std::string& prefix) {
+    auto prefix_handler = [] (const std::string& prefix, const std::string& key_type_dir, const KeyDescriptor& key_descriptor, KeyType) {
+        return !prefix.empty() ? fmt::format("{}/{}*{}", key_type_dir, key_descriptor, prefix) : key_type_dir;
+    };
+
+    return detail::do_calculate_sizes_for_type_impl(key_type, root_folder_, bucket_name_, client(), FlatBucketizer{}, prefix_handler, prefix);
 }
 
 bool S3Storage::do_key_exists(const VariantKey& key) {
@@ -183,6 +191,10 @@ S3Storage::S3Storage(const LibraryPath &library_path, OpenMode mode, const S3Set
     ARCTICDB_DEBUG(log::storage(), "Opened S3 backed storage at {}", root_folder_);
 }
 
+bool S3Storage::supports_object_size_calculation() const {
+    return true;
+}
+
 GCPXMLStorage::GCPXMLStorage(const arcticdb::storage::LibraryPath& lib,
                              arcticdb::storage::OpenMode mode,
                              const arcticdb::storage::s3::GCPXMLSettings& conf) :

cpp/arcticdb/storage/s3/s3_storage.hpp

@@ -29,6 +29,8 @@
 
 namespace arcticdb::storage::s3 {
 
+using PrefixHandler = std::function<std::string(const std::string&, const std::string&, const KeyDescriptor&, KeyType)>;
+
 const std::string USE_AWS_CRED_PROVIDERS_TOKEN = "_RBAC_";
 
 class S3Storage : public Storage, AsyncStorage {
@@ -48,6 +50,8 @@ class S3Storage : public Storage, AsyncStorage {
         return dynamic_cast<AsyncStorage*>(this);
     }
 
+    bool supports_object_size_calculation() const final override;
+
   protected:
     void do_write(KeySegmentPair& key_seg) final;
 
@@ -67,6 +71,8 @@ class S3Storage : public Storage, AsyncStorage {
 
     void do_remove(std::span<VariantKey> variant_keys, RemoveOpts opts);
 
+    ObjectSizes do_get_object_sizes(KeyType key_type, const std::string& prefix) override final;
+
     bool do_iterate_type_until_match(KeyType key_type, const IterateTypePredicate& visitor, const std::string &prefix) final;
 
     bool do_key_exists(const VariantKey& key) final;