Antalya: Cache the list objects operation on object storage using a TTL + prefix matching cache implementation

programs/server/Server.cpp

@@ -83,6 +83,7 @@
 #include <Storages/System/attachInformationSchemaTables.h>
 #include <Storages/Cache/ExternalDataSourceCache.h>
 #include <Storages/Cache/registerRemoteFileMetadatas.h>
+#include <Storages/Cache/ObjectStorageListObjectsCache.h>
 #include <AggregateFunctions/registerAggregateFunctions.h>
 #include <Functions/UserDefined/IUserDefinedSQLObjectsStorage.h>
 #include <Functions/registerFunctions.h>
@@ -318,6 +319,9 @@ namespace ServerSetting
     extern const ServerSettingsUInt64 max_prefixes_deserialization_thread_pool_free_size;
     extern const ServerSettingsUInt64 prefixes_deserialization_thread_pool_thread_pool_queue_size;
     extern const ServerSettingsUInt64 input_format_parquet_metadata_cache_max_size;
+    extern const ServerSettingsUInt64 object_storage_list_objects_cache_size;
+    extern const ServerSettingsUInt64 object_storage_list_objects_cache_max_entries;
+    extern const ServerSettingsUInt64 object_storage_list_objects_cache_ttl;
 }
 
 }
@@ -2327,6 +2331,10 @@ try
     if (dns_cache_updater)
         dns_cache_updater->start();
 
+    ObjectStorageListObjectsCache::instance().setMaxSizeInBytes(server_settings[ServerSetting::object_storage_list_objects_cache_size]);
+    ObjectStorageListObjectsCache::instance().setMaxCount(server_settings[ServerSetting::object_storage_list_objects_cache_max_entries]);
+    ObjectStorageListObjectsCache::instance().setTTL(server_settings[ServerSetting::object_storage_list_objects_cache_ttl]);
+
     auto replicas_reconnector = ReplicasReconnector::init(global_context);
     ParquetFileMetaDataCache::instance()->setMaxSizeInBytes(server_settings[ServerSetting::input_format_parquet_metadata_cache_max_size]);
 

src/Common/TTLCachePolicy.h

@@ -243,10 +243,10 @@ class TTLCachePolicy : public ICachePolicy<Key, Mapped, HashFunction, WeightFunc
         return res;
     }
 
-private:
+protected:
     using Cache = std::unordered_map<Key, MappedPtr, HashFunction>;
     Cache cache;
-
+private:
     /// TODO To speed up removal of stale entries, we could also add another container sorted on expiry times which maps keys to iterators
     /// into the cache. To insert an entry, add it to the cache + add the iterator to the sorted container. To remove stale entries, do a
     /// binary search on the sorted container and erase all left of the found key.

src/Core/ServerSettings.cpp

@@ -1004,8 +1004,10 @@ namespace DB
     ```
     )", 0) \
     DECLARE(Bool, storage_shared_set_join_use_inner_uuid, false, "If enabled, an inner UUID is generated during the creation of SharedSet and SharedJoin. ClickHouse Cloud only", 0) \
-    DECLARE(UInt64, input_format_parquet_metadata_cache_max_size, 500000000, "Maximum size of parquet file metadata cache", 0) \
-
+    DECLARE(UInt64, input_format_parquet_metadata_cache_max_size, 500000000, "Maximum size of parquet file metadata cache", 0)                                                                                \
+    DECLARE(UInt64, object_storage_list_objects_cache_size, 5000000, "Maximum size of ObjectStorage list objects cache in bytes. Zero means disabled.", 0)                            \
+    DECLARE(UInt64, object_storage_list_objects_cache_max_entries, 1000, "Maximum size of ObjectStorage list objects cache in entries. Zero means disabled.", 0)                      \
+    DECLARE(UInt64, object_storage_list_objects_cache_ttl, 3600, "Time to live of records in ObjectStorage list objects cache in seconds. Zero means unlimited", 0) \
 // clang-format on
 
 /// If you add a setting which can be updated at runtime, please update 'changeable_settings' map in dumpToSystemServerSettingsColumns below

src/Core/Settings.cpp

@@ -6108,6 +6108,9 @@ Limit for hosts used for request in object storage cluster table functions - azu
 Possible values:
 - Positive integer.
 - 0 — All hosts in cluster.
+)", EXPERIMENTAL)                       \
+    DECLARE(Bool, use_object_storage_list_objects_cache, true, R"(
+Cache the list of objects returned by list objects calls in object storage
 )", EXPERIMENTAL) \
     \
     /* ####################################################### */ \

src/Storages/Cache/ObjectStorageListObjectsCache.cpp

@@ -0,0 +1,213 @@
+#include <Storages/Cache/ObjectStorageListObjectsCache.h>
+#include <boost/functional/hash.hpp>
+
+namespace DB
+{
+
+template <typename Key, typename Mapped, typename HashFunction, typename WeightFunction, typename IsStaleFunction>
+class ObjectStorageListObjectsCachePolicy : public TTLCachePolicy<Key, Mapped, HashFunction, WeightFunction, IsStaleFunction>
+{
+public:
+    using BasePolicy = TTLCachePolicy<Key, Mapped, HashFunction, WeightFunction, IsStaleFunction>;
+    using typename BasePolicy::MappedPtr;
+    using typename BasePolicy::KeyMapped;
+    using BasePolicy::cache;
+
+    ObjectStorageListObjectsCachePolicy()
+        : BasePolicy(std::make_unique<NoCachePolicyUserQuota>())
+    {
+    }
+
+    MappedPtr get(const Key & key) override
+    {
+        if (const auto it = cache.find(key); it != cache.end())
+        {
+            if (IsStaleFunction()(it->first))
+            {
+                BasePolicy::remove(it->first);
+                return {};
+            }
+            return it->second;
+        }
+
+        if (const auto it = findBestMatchingPrefix(key); it != cache.end())
+        {
+            if (IsStaleFunction()(it->first))
+            {
+                BasePolicy::remove(it->first);
+                return {};
+            }
+            return it->second;
+        }
+
+        return {};
+    }
+
+    std::optional<KeyMapped> getWithKey(const Key & key) override
+    {
+        if (const auto it = cache.find(key); it != cache.end())
+        {
+            if (IsStaleFunction()(it->first))
+            {
+                BasePolicy::remove(it->first);
+                return std::nullopt;
+            }
+            return std::make_optional<KeyMapped>({it->first, it->second});
+        }
+
+        if (const auto it = findBestMatchingPrefix(key); it != cache.end())
+        {
+            if (IsStaleFunction()(it->first))
+            {
+                BasePolicy::remove(it->first);
+                return std::nullopt;
+            }
+            return std::make_optional<KeyMapped>({it->first, it->second});
+        }
+
+        return std::nullopt;
+    }
+
+private:
+    auto findBestMatchingPrefix(const Key & key)
+    {
+        const auto & prefix = key.prefix;
+
+        auto best_match = cache.end();
+        size_t best_length = 0;
+
+        std::vector<Key> to_remove;
+
+        for (auto it = cache.begin(); it != cache.end(); ++it)
+        {
+            const auto & candidate_bucket = it->first.bucket;
+            const auto & candidate_prefix = it->first.prefix;
+
+            if (candidate_bucket == key.bucket && prefix.starts_with(candidate_prefix))
+            {
+                if (IsStaleFunction()(it->first))
+                {
+                    to_remove.push_back(it->first);
+                    continue;
+                }
+
+                if (candidate_prefix.size() > best_length)
+                {
+                    best_match = it;
+                    best_length = candidate_prefix.size();
+                }
+            }
+        }
+
+        for (const auto & k : to_remove)
+            BasePolicy::remove(k);
+
+        return best_match;
+    }
+};
+
+ObjectStorageListObjectsCache::Key::Key(
+    const String & bucket_,
+    const String & prefix_,
+    const std::chrono::system_clock::time_point & expires_at_,
+    std::optional<UUID> user_id_)
+    : bucket(bucket_), prefix(prefix_), expires_at(expires_at_), user_id(user_id_) {}
+
+bool ObjectStorageListObjectsCache::Key::operator==(const Key & other) const
+{
+    return bucket == other.bucket && prefix == other.prefix;
+}
+
+size_t ObjectStorageListObjectsCache::KeyHasher::operator()(const Key & key) const
+{
+    std::size_t seed = 0;
+
+    boost::hash_combine(seed, std::hash<String>()(key.bucket));
+    boost::hash_combine(seed, std::hash<String>()(key.prefix));
+
+    return seed;
+}
+
+bool ObjectStorageListObjectsCache::IsStale::operator()(const Key & key) const
+{
+    return key.expires_at < std::chrono::system_clock::now();
+}
+
+size_t ObjectStorageListObjectsCache::WeightFunction::operator()(const Value & value) const
+{
+    std::size_t weight = 0;
+
+    for (const auto & object : value)
+    {
+        weight += object->relative_path.capacity() + sizeof(ObjectMetadata);
+    }
+
+    return weight;
+}
+
+ObjectStorageListObjectsCache::ObjectStorageListObjectsCache()
+    : cache(std::make_unique<ObjectStorageListObjectsCachePolicy<Key, Value, KeyHasher, WeightFunction, IsStale>>())
+{
+}
+
+void ObjectStorageListObjectsCache::set(
+    const std::string & bucket,
+    const std::string & prefix,
+    const std::shared_ptr<Value> & value)
+{
+    const auto key = Key{bucket, prefix, std::chrono::system_clock::now() + std::chrono::seconds(ttl)};
+
+    cache.set(key, value);
+}
+
+ObjectStorageListObjectsCache::Cache::MappedPtr ObjectStorageListObjectsCache::get(const String & bucket, const String & prefix, bool filter_by_prefix)
+{
+    const auto input_key = Key{bucket, prefix};
+    auto pair = cache.getWithKey(input_key);
+
+    if (!pair)
+    {
+        return {};
+    }
+
+    if (pair->key == input_key || filter_by_prefix)
+    {
+        return pair->mapped;
+    }
+
+    auto filtered_objects = std::make_shared<std::vector<ObjectInfoPtr>>();
+    filtered_objects->reserve(pair->mapped->size());
+
+    for (const auto & object : *pair->mapped)
+    {
+        if (object->relative_path.starts_with(input_key.prefix))
+        {
+            filtered_objects->push_back(object);
+        }
+    }
+
+    return filtered_objects;
+}
+
+void ObjectStorageListObjectsCache::setMaxSizeInBytes(std::size_t size_in_bytes_)
+{
+    cache.setMaxSizeInBytes(size_in_bytes_);
+}
+
+void ObjectStorageListObjectsCache::setMaxCount(std::size_t count)
+{
+    cache.setMaxCount(count);
+}
+
+void ObjectStorageListObjectsCache::setTTL(std::size_t ttl_)
+{
+    ttl = ttl_;
+}
+
+ObjectStorageListObjectsCache & ObjectStorageListObjectsCache::instance()
+{
+    static ObjectStorageListObjectsCache instance;
+    return instance;
+}
+
+}

src/Storages/Cache/ObjectStorageListObjectsCache.h

@@ -0,0 +1,69 @@
+#pragma once
+
+#include <chrono>
+#include <Disks/ObjectStorages/IObjectStorage.h>
+#include <Storages/ObjectStorage/StorageObjectStorage.h>
+#include <Common/TTLCachePolicy.h>
+#include <Common/CacheBase.h>
+
+namespace DB
+{
+
+class ObjectStorageListObjectsCache
+{
+public:
+    static ObjectStorageListObjectsCache & instance();
+
+    struct Key
+    {
+        Key(
+            const String & bucket_,
+            const String & prefix_,
+            const std::chrono::system_clock::time_point & expires_at_ = std::chrono::system_clock::now(),
+            std::optional<UUID> user_id_ = std::nullopt);
+
+        std::string bucket;
+        std::string prefix;
+        std::chrono::system_clock::time_point expires_at;
+        std::optional<UUID> user_id;
+
+        bool operator==(const Key & other) const;
+    };
+
+    using Value = StorageObjectStorage::ObjectInfos;
+    struct KeyHasher
+    {
+        size_t operator()(const Key & key) const;
+    };
+
+    struct IsStale
+    {
+        bool operator()(const Key & key) const;
+    };
+
+    struct WeightFunction
+    {
+        size_t operator()(const Value & value) const;
+    };
+
+    using Cache = CacheBase<Key, Value, KeyHasher, WeightFunction>;
+
+    void set(
+        const std::string & bucket,
+        const std::string & prefix,
+        const std::shared_ptr<Value> & value);
+
+    Cache::MappedPtr get(const String & bucket, const String & prefix, bool filter_by_prefix = true);
+
+    void setMaxSizeInBytes(std::size_t size_in_bytes_);
+    void setMaxCount(std::size_t count);
+    void setTTL(std::size_t ttl_);
+
+private:
+    ObjectStorageListObjectsCache();
+
+    Cache cache;
+    size_t ttl;
+};
+
+}