Skip to content

Commit

Permalink
QueryCache: Fix odd bugs with user quota
Browse files Browse the repository at this point in the history
- SYSTEM DROP QUERY CACHE did not clear the user policy statistics, this
  made the cache block all inserts afterwards

- Fixes two other oddities in the approval logic
  • Loading branch information
rschu1ze committed Jan 12, 2024
1 parent 3954036 commit 5357e8f
Show file tree
Hide file tree
Showing 4 changed files with 68 additions and 4 deletions.
4 changes: 4 additions & 0 deletions src/Common/ICachePolicyUserQuota.h
Expand Up @@ -25,6 +25,9 @@ class ICachePolicyUserQuota
/// 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;
};

Expand All @@ -38,6 +41,7 @@ class NoCachePolicyUserQuota : public ICachePolicyUserQuota
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 {}
};


Expand Down
14 changes: 10 additions & 4 deletions src/Common/TTLCachePolicy.h
Expand Up @@ -38,12 +38,12 @@ class PerUserTTLCachePolicyUserQuota : public ICachePolicyUserQuota
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}; /// default if no user is found 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()}; /// default if no user is found is no threshold
if (it_quota != quotas.end())
quota_for_user = it_quota->second;

Expand All @@ -54,16 +54,21 @@ class PerUserTTLCachePolicyUserQuota : public ICachePolicyUserQuota
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;
Expand Down Expand Up @@ -125,6 +130,7 @@ class TTLCachePolicy : public ICachePolicy<Key, Mapped, HashFunction, WeightFunc
void clear() override
{
cache.clear();
Base::user_quotas->clear();
}

void remove(const Key & key) override
Expand Down
@@ -0,0 +1,13 @@
a
b
1
c
d
3
--
a
b
1
c
d
3
@@ -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 '---';

0 comments on commit 5357e8f

Please sign in to comment.