thevar1able/ClickHouse
1
0
Fork 0
mirror of https://github.com/ClickHouse/ClickHouse.git synced 2024-11-22 23:52:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

Incorporate review feedback

Browse Source
This commit is contained in:
Robert Schulze 2024-01-11 13:33:16 +00:00
parent 9031c671f6
commit bd9e38ff47
No known key found for this signature in database
GPG Key ID: 26703B55FB13728A
6 changed files with 41 additions and 45 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/Common/CacheBase.h
View File

@ -227,7 +227,7 @@ public:
cache_policy->setMaxSizeInBytes(max_size_in_bytes);
}
void setQuotaForUser(std::optional<UUID> user_id, size_t max_size_in_bytes, size_t max_entries)
void setQuotaForUser(const UUID & user_id, size_t max_size_in_bytes, size_t max_entries)
{
std::lock_guard lock(mutex);
cache_policy->setQuotaForUser(user_id, max_size_in_bytes, max_entries);

2
src/Common/ICachePolicy.h
View File

@ -44,7 +44,7 @@ public:
virtual void setMaxCount(size_t /*max_count*/) = 0;
virtual void setMaxSizeInBytes(size_t /*max_size_in_bytes*/) = 0;
virtual void setQuotaForUser(std::optional<UUID> user_id, size_t max_size_in_bytes, size_t max_entries) { user_quotas->setQuotaForUser(user_id, max_size_in_bytes, max_entries); }
virtual void setQuotaForUser(const UUID & user_id, size_t max_size_in_bytes, size_t max_entries) { user_quotas->setQuotaForUser(user_id, max_size_in_bytes, max_entries); }
/// HashFunction usually hashes the entire key and the found key will be equal the provided key. In such cases, use get(). It is also
/// possible to store other, non-hashed data in the key. In that case, the found key is potentially different from the provided key.

18
src/Common/ICachePolicyUserQuota.h
View File

@ -3,8 +3,6 @@
#include <base/UUID.h>
#include <base/types.h>
#include <optional>
namespace DB
{
@ -18,14 +16,14 @@ class ICachePolicyUserQuota
{
public:
/// Register or update the user's quota for the given resource.
virtual void setQuotaForUser(std::optional<UUID> user_id, size_t max_size_in_bytes, size_t max_entries) = 0;
virtual void setQuotaForUser(const UUID & user_id, size_t max_size_in_bytes, size_t max_entries) = 0;
/// Update the actual resource usage for the given user.
virtual void increaseActual(std::optional<UUID> user_id, size_t entry_size_in_bytes) = 0;
virtual void decreaseActual(std::optional<UUID> user_id, size_t entry_size_in_bytes) = 0;
virtual void increaseActual(const UUID & user_id, size_t entry_size_in_bytes) = 0;
virtual void decreaseActual(const UUID & user_id, size_t entry_size_in_bytes) = 0;
/// Is the user allowed to write a new entry into the cache?
virtual bool approveWrite(std::optional<UUID> user_id, size_t entry_size_in_bytes) const = 0;
virtual bool approveWrite(const UUID & user_id, size_t entry_size_in_bytes) const = 0;
virtual ~ICachePolicyUserQuota() = default;
};
@ -36,10 +34,10 @@ using CachePolicyUserQuotaPtr = std::unique_ptr<ICachePolicyUserQuota>;
class NoCachePolicyUserQuota : public ICachePolicyUserQuota
{
public:
void setQuotaForUser(std::optional<UUID> /*user_id*/, size_t /*max_size_in_bytes*/, size_t /*max_entries*/) override {}
void increaseActual(std::optional<UUID> /*user_id*/, size_t /*entry_size_in_bytes*/) override {}
void decreaseActual(std::optional<UUID> /*user_id*/, size_t /*entry_size_in_bytes*/) override {}
bool approveWrite(std::optional<UUID> /*user_id*/, size_t /*entry_size_in_bytes*/) const override { return true; }
void setQuotaForUser(const UUID & /*user_id*/, size_t /*max_size_in_bytes*/, size_t /*max_entries*/) override {}
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; }
};

53
src/Common/TTLCachePolicy.h
View File

@ -4,7 +4,6 @@
#include <base/UUID.h>
#include <limits>
#include <optional>
#include <unordered_map>
namespace DB
@ -13,47 +12,37 @@ namespace DB
class PerUserTTLCachePolicyUserQuota : public ICachePolicyUserQuota
{
public:
void setQuotaForUser(std::optional<UUID> user_id, size_t max_size_in_bytes, size_t max_entries) override
void setQuotaForUser(const UUID & user_id, size_t max_size_in_bytes, size_t max_entries) override
{
if (!user_id.has_value())
return;
quotas[*user_id] = {max_size_in_bytes, max_entries};
quotas[user_id] = {max_size_in_bytes, max_entries};
}
void increaseActual(std::optional<UUID> user_id, size_t entry_size_in_bytes) override
void increaseActual(const UUID & user_id, size_t entry_size_in_bytes) override
{
if (!user_id.has_value())
return;
auto & actual_for_user = actual[*user_id];
auto & actual_for_user = actual[user_id];
actual_for_user.size_in_bytes += entry_size_in_bytes;
actual_for_user.num_items += 1;
}
void decreaseActual(std::optional<UUID> user_id, size_t entry_size_in_bytes) override
void decreaseActual(const UUID & user_id, size_t entry_size_in_bytes) override
{
if (!user_id.has_value())
return;
chassert(actual.contains(user_id));
chassert(actual.contains(*user_id));
chassert(actual[user_id].size_in_bytes >= entry_size_in_bytes);
actual[user_id].size_in_bytes -= entry_size_in_bytes;
chassert(actual[*user_id].size_in_bytes >= entry_size_in_bytes);
actual[*user_id].size_in_bytes -= entry_size_in_bytes;
chassert(actual[*user_id].num_items >= 1);
actual[*user_id].num_items -= 1;
chassert(actual[user_id].num_items >= 1);
actual[user_id].num_items -= 1;
}
bool approveWrite(std::optional<UUID> user_id, size_t entry_size_in_bytes) const override
bool approveWrite(const UUID & user_id, size_t entry_size_in_bytes) const override
{
if (!user_id.has_value())
return true;
auto it_actual = actual.find(*user_id);
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
if (it_actual != actual.end())
actual_for_user = it_actual->second;
auto it_quota = quotas.find(*user_id);
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
if (it_quota != quotas.end())
quota_for_user = it_quota->second;
@ -144,7 +133,8 @@ public:
if (it == cache.end())
return;
size_t sz = weight_function(*it->second);
Base::user_quotas->decreaseActual(it->first.user_id, sz);
if (it->first.user_id.has_value())
Base::user_quotas->decreaseActual(*it->first.user_id, sz);
cache.erase(it);
size_in_bytes -= sz;
}
@ -181,7 +171,9 @@ public:
/// Checks against per-user limits
auto sufficient_space_in_cache_for_user = [&]()
{
return Base::user_quotas->approveWrite(key.user_id, entry_size_in_bytes);
if (key.user_id.has_value())
return Base::user_quotas->approveWrite(*key.user_id, entry_size_in_bytes);
return true;
};
if (!sufficient_space_in_cache() || !sufficient_space_in_cache_for_user())
@ -191,7 +183,8 @@ public:
if (is_stale_function(it->first))
{
size_t sz = weight_function(*it->second);
Base::user_quotas->decreaseActual(it->first.user_id, sz);
if (it->first.user_id.has_value())
Base::user_quotas->decreaseActual(*it->first.user_id, sz);
it = cache.erase(it);
size_in_bytes -= sz;
}
@ -205,14 +198,16 @@ public:
if (auto it = cache.find(key); it != cache.end())
{
size_t sz = weight_function(*it->second);
Base::user_quotas->decreaseActual(it->first.user_id, sz);
if (it->first.user_id.has_value())
Base::user_quotas->decreaseActual(*it->first.user_id, sz);
cache.erase(it); // stupid bug: (*) doesn't replace existing entries (likely due to custom hash function), need to erase explicitly
size_in_bytes -= sz;
}
cache[key] = std::move(mapped); // (*)
size_in_bytes += entry_size_in_bytes;
Base::user_quotas->increaseActual(key.user_id, entry_size_in_bytes);
if (key.user_id.has_value())
Base::user_quotas->increaseActual(*key.user_id, entry_size_in_bytes);
}
}

4
src/Interpreters/Cache/QueryCache.cpp
View File

@ -507,7 +507,9 @@ QueryCache::Writer QueryCache::createWriter(const Key & key, std::chrono::millis
/// Update the per-user cache quotas with the values stored in the query context. This happens per query which writes into the query
/// cache. Obviously, this is overkill but I could find the good place to hook into which is called when the settings profiles in
/// users.xml change.
cache.setQuotaForUser(key.user_id, max_query_cache_size_in_bytes_quota, max_query_cache_entries_quota);
/// user_id == std::nullopt is the internal user for which no quota can be configured
if (key.user_id.has_value())
cache.setQuotaForUser(*key.user_id, max_query_cache_size_in_bytes_quota, max_query_cache_entries_quota);
std::lock_guard lock(mutex);
return Writer(cache, key, max_entry_size_in_bytes, max_entry_size_in_rows, min_query_runtime, squash_partial_results, max_block_size);

7
tests/queries/0_stateless/02494_query_cache_user_isolation.sh
View File

@ -1,7 +1,8 @@
#!/usr/bin/env bash
# Tags: no-parallel, no-fasttest
# Tags: no-parallel, no-fasttest, long
# Tag no-parallel: Messes with internal cache
# no-fasttest: Produces wrong results in fasttest, unclear why, didn't reproduce locally.
# long: Sloooow ...
CURDIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
# shellcheck source=../shell_config.sh
@ -13,7 +14,7 @@ CURDIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
# - drop the user, recreate it with the same name
# - test that the cache entry is inaccessible
${CLICKHOUSE_CLIENT} --query "SELECT 'Attack 1'"
echo "Attack 1"
rnd=`tr -dc 1-9 </dev/urandom | head -c 5` # disambiguates the specific query in system.query_log below
# echo $rnd
@ -57,7 +58,7 @@ ${CLICKHOUSE_CLIENT} --query "SYSTEM DROP QUERY CACHE"
# - create two roles, each with different row policies
# - cached query result in the context of the 1st role must must not be visible in the context of the 2nd role
${CLICKHOUSE_CLIENT} --query "SELECT 'Attack 2'"
echo "Attack 2"
# Start with empty query cache (QC).
${CLICKHOUSE_CLIENT} --query "SYSTEM DROP QUERY CACHE"