1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
|
#include <Interpreters/Cache/QueryLimit.h>
#include <Interpreters/Cache/Metadata.h>
namespace DB
{
namespace ErrorCodes
{
extern const int LOGICAL_ERROR;
}
static bool isQueryInitialized()
{
return CurrentThread::isInitialized()
&& CurrentThread::get().getQueryContext()
&& !CurrentThread::getQueryId().empty();
}
FileCacheQueryLimit::QueryContextPtr FileCacheQueryLimit::tryGetQueryContext(const CacheGuard::Lock &)
{
if (!isQueryInitialized())
return nullptr;
auto query_iter = query_map.find(std::string(CurrentThread::getQueryId()));
return (query_iter == query_map.end()) ? nullptr : query_iter->second;
}
void FileCacheQueryLimit::removeQueryContext(const std::string & query_id, const CacheGuard::Lock &)
{
auto query_iter = query_map.find(query_id);
if (query_iter == query_map.end())
{
throw Exception(
ErrorCodes::LOGICAL_ERROR,
"Attempt to release query context that does not exist (query_id: {})",
query_id);
}
query_map.erase(query_iter);
}
FileCacheQueryLimit::QueryContextPtr FileCacheQueryLimit::getOrSetQueryContext(
const std::string & query_id,
const ReadSettings & settings,
const CacheGuard::Lock &)
{
if (query_id.empty())
return nullptr;
auto [it, inserted] = query_map.emplace(query_id, nullptr);
if (inserted)
{
it->second = std::make_shared<QueryContext>(
settings.filesystem_cache_max_download_size,
!settings.skip_download_if_exceeds_query_cache);
}
return it->second;
}
FileCacheQueryLimit::QueryContext::QueryContext(
size_t query_cache_size,
bool recache_on_query_limit_exceeded_)
: priority(LRUFileCachePriority(query_cache_size, 0))
, recache_on_query_limit_exceeded(recache_on_query_limit_exceeded_)
{
}
void FileCacheQueryLimit::QueryContext::add(
KeyMetadataPtr key_metadata,
size_t offset,
size_t size,
const CacheGuard::Lock & lock)
{
auto it = getPriority().add(key_metadata, offset, size, lock);
auto [_, inserted] = records.emplace(FileCacheKeyAndOffset{key_metadata->key, offset}, it);
if (!inserted)
{
it->remove(lock);
throw Exception(
ErrorCodes::LOGICAL_ERROR,
"Cannot add offset {} to query context under key {}, it already exists",
offset, key_metadata->key);
}
}
void FileCacheQueryLimit::QueryContext::remove(
const Key & key,
size_t offset,
const CacheGuard::Lock & lock)
{
auto record = records.find({key, offset});
if (record == records.end())
throw Exception(ErrorCodes::LOGICAL_ERROR, "There is no {}:{} in query context", key, offset);
record->second->remove(lock);
records.erase({key, offset});
}
IFileCachePriority::Iterator FileCacheQueryLimit::QueryContext::tryGet(
const Key & key,
size_t offset,
const CacheGuard::Lock &)
{
auto it = records.find({key, offset});
if (it == records.end())
return nullptr;
return it->second;
}
}
|
