Merge pull request #58731 from rschu1ze/qc-quota

Query cache: Fix per-user quota

src/Common/ICachePolicyUserQuota.h

@@ -25,6 +25,9 @@ public:
     /// Is the user allowed to write a new entry into the cache?
     virtual bool approveWrite(const UUID & user_id, size_t entry_size_in_bytes) const = 0;
 
+    /// Clears the policy contents
+    virtual void clear() = 0;
+
     virtual ~ICachePolicyUserQuota() = default;
 };
 
@@ -38,6 +41,7 @@ public:
     void increaseActual(const UUID & /*user_id*/, size_t /*entry_size_in_bytes*/) override {}
     void decreaseActual(const UUID & /*user_id*/, size_t /*entry_size_in_bytes*/) override {}
     bool approveWrite(const UUID & /*user_id*/, size_t /*entry_size_in_bytes*/) const override { return true; }
+    void clear() override {}
 };
 
 

src/Common/TTLCachePolicy.h

@@ -38,12 +38,12 @@ public:
     bool approveWrite(const UUID & user_id, size_t entry_size_in_bytes) const override
     {
         auto it_actual = actual.find(user_id);
-        Resources actual_for_user{.size_in_bytes = 0, .num_items = 0}; /// assume zero actual resource consumption is user isn't found
+        Resources actual_for_user{.size_in_bytes = 0, .num_items = 0}; /// if no user is found, the default is no resource consumption
         if (it_actual != actual.end())
             actual_for_user = it_actual->second;
 
         auto it_quota = quotas.find(user_id);
-        Resources quota_for_user{.size_in_bytes = std::numeric_limits<size_t>::max(), .num_items = std::numeric_limits<size_t>::max()}; /// assume no threshold if no quota is found
+        Resources quota_for_user{.size_in_bytes = std::numeric_limits<size_t>::max(), .num_items = std::numeric_limits<size_t>::max()}; /// if no user is found, the default is no threshold
         if (it_quota != quotas.end())
             quota_for_user = it_quota->second;
 
@@ -54,16 +54,21 @@ public:
             quota_for_user.num_items = std::numeric_limits<UInt64>::max();
 
         /// Check size quota
-        if (actual_for_user.size_in_bytes + entry_size_in_bytes >= quota_for_user.size_in_bytes)
+        if (actual_for_user.size_in_bytes + entry_size_in_bytes > quota_for_user.size_in_bytes)
             return false;
 
         /// Check items quota
-        if (quota_for_user.num_items + 1 >= quota_for_user.num_items)
+        if (actual_for_user.num_items + 1 > quota_for_user.num_items)
             return false;
 
         return true;
     }
 
+    void clear() override
+    {
+        actual.clear();
+    }
+
     struct Resources
     {
         size_t size_in_bytes = 0;
@@ -125,6 +130,7 @@ public:
     void clear() override
     {
         cache.clear();
+        Base::user_quotas->clear();
     }
 
     void remove(const Key & key) override

tests/queries/0_stateless/02494_query_cache_user_quotas_after_drop.reference

@@ -0,0 +1,13 @@
+a
+b
+1
+c
+d
+3
+--
+a
+b
+1
+c
+d
+3

tests/queries/0_stateless/02494_query_cache_user_quotas_after_drop.sql

@@ -0,0 +1,41 @@
+-- Tags: no-parallel
+-- Tag no-parallel: Messes with internal cache
+
+-- Tests per-user quotas of the query cache. Settings 'query_cache_max_size_in_bytes' and 'query_cache_max_entries' are actually supposed to
+-- be used in a settings profile, together with a readonly constraint. For simplicity, test both settings stand-alone in a stateless test
+-- instead of an integration test - the relevant logic will still be covered by that.
+
+SYSTEM DROP QUERY CACHE;
+
+-- Run SELECT with quota that current user may write only 1 entry in the query cache
+SET query_cache_max_entries = 1;
+SELECT 'a' SETTINGS use_query_cache = true;
+SELECT 'b' SETTINGS use_query_cache = true;
+SELECT count(*) FROM system.query_cache; -- expect 1 entry
+
+-- Run SELECTs again but w/o quota
+SET query_cache_max_entries = DEFAULT;
+SELECT 'c' SETTINGS use_query_cache = true;
+SELECT 'd' SETTINGS use_query_cache = true;
+SELECT count(*) FROM system.query_cache; -- expect 3 entries
+
+SYSTEM DROP QUERY CACHE;
+
+-- Run the same as above after a DROP QUERY CACHE.
+SELECT '--';
+
+SET query_cache_max_entries = 1;
+SELECT 'a' SETTINGS use_query_cache = true;
+SELECT 'b' SETTINGS use_query_cache = true;
+SELECT count(*) FROM system.query_cache; -- expect 1 entry
+
+-- Run SELECTs again but w/o quota
+SET query_cache_max_entries = DEFAULT;
+SELECT 'c' SETTINGS use_query_cache = true;
+SELECT 'd' SETTINGS use_query_cache = true;
+SELECT count(*) FROM system.query_cache; -- expect 3 entries
+
+SYSTEM DROP QUERY CACHE;
+
+-- SELECT '---';
+