add MappedHolder to get cache values

This commit is contained in:
lgbo-ustc 2021-12-29 18:18:38 +08:00 committed by liangjiabiao
parent 5a3b215f24
commit 6b6a82f3b9
2 changed files with 242 additions and 201 deletions

View File

@ -32,102 +32,59 @@ template <
class LRUResourceCache
{
public:
LRUResourceCache(size_t max_weight_, size_t max_element_size_ = 0) : max_weight(max_weight_), max_element_size(max_element_size_) { }
~LRUResourceCache() = default;
using Key = TKey;
using Mapped = TMapped;
using MappedPtr = std::shared_ptr<Mapped>;
// - load_func : when key is not exists in cache, load_func is called to generate a new key
// - return: is null when there is no more space for the new value or the old value is in used.
template <typename LoadFunc>
MappedPtr acquire(const Key & key, LoadFunc && load_func)
class MappedHolder
{
InsertToken * insert_token = nullptr;
public:
~MappedHolder()
{
std::lock_guard lock(mutex);
auto it = cells.find(key);
if (it != cells.end())
{
hits++;
it->second.reference_count += 1;
queue.splice(queue.end(), queue, it->second.queue_iterator);
return it->second.value;
}
misses++;
insert_token = acquireInsertToken(key);
cache->release(key);
}
Cell * cell_ptr = nullptr;
Mapped & value()
{
std::lock_guard lock(insert_token->mutex);
if (!insert_token->value)
{
insert_token->value = load_func();
std::lock_guard cell_lock(mutex);
cell_ptr = insert_value(key, insert_token->value);
if (cell_ptr)
{
cell_ptr->reference_count += 1;
}
else
{
insert_token->value = nullptr;
}
}
return *(val.get());
}
static bool tryRemove(std::unique_ptr<MappedHolder> * holder_ptr)
{
auto & holder = *holder_ptr;
auto cache = holder->cache;
auto key = holder->key;
*holder_ptr = nullptr;
return cache->tryRemove(key);
}
std::lock_guard lock(mutex);
releaseInsertToken(key);
if (cell_ptr)
MappedHolder(LRUResourceCache * cache_, const Key & key_, MappedPtr value_) : cache(cache_), key(key_), val(value_)
{
return cell_ptr->value;
}
return nullptr;
}
protected:
LRUResourceCache * cache;
Key key;
MappedPtr val;
};
using MappedHolderPtr = std::unique_ptr<MappedHolder>;
MappedPtr acquire(const Key & key)
// use get() or getOrSet() to access the elements
MappedHolderPtr get(const Key & key)
{
std::lock_guard lock(mutex);
auto it = cells.find(key);
if (it == cells.end())
{
misses++;
auto mappedptr = getImpl(key);
if (!mappedptr)
return nullptr;
}
hits++;
it->second.reference_count += 1;
queue.splice(queue.end(), queue, it->second.queue_iterator);
return it->second.value;
return std::make_unique<MappedHolder>(this, key, mappedptr);
}
template<typename LoadFunc>
MappedHolderPtr getOrSet(const Key & key, LoadFunc && load_func)
{
auto mappedptr = getImpl(key, load_func);
if (!mappedptr)
return nullptr;
return std::make_unique<MappedHolder>(this, key, mappedptr);
}
// mark a reference is released
void release(const Key & key)
{
std::lock_guard lock(mutex);
auto it = cells.find(key);
if (it == cells.end() || it->second.reference_count == 0)
{
LOG_ERROR(&Poco::Logger::get("LRUResourceCache"), "try to release an invalid element");
abort();
}
it->second.reference_count -= 1;
}
// If you want to update a value, call tryRemove() at first and then call acquire() with load_func.
bool tryRemove(const Key & key)
{
std::lock_guard guard(mutex);
auto it = cells.find(key);
if (it == cells.end())
return true;
auto & cell = it->second;
if (cell.reference_count)
return false;
queue.erase(cell.queue_iterator);
current_weight -= cell.weight;
cells.erase(it);
return true;
}
LRUResourceCache(size_t max_weight_, size_t max_element_size_ = 0) : max_weight(max_weight_), max_element_size(max_element_size_) { }
~LRUResourceCache() = default;
size_t weight()
{
@ -181,6 +138,81 @@ private:
std::atomic<size_t> hits{0};
std::atomic<size_t> misses{0};
std::atomic<size_t> evict_count{0};
// - load_func : when key is not exists in cache, load_func is called to generate a new key
// - return: is null when there is no more space for the new value or the old value is in used.
template <typename LoadFunc>
MappedPtr getImpl(const Key & key, LoadFunc && load_func)
{
InsertToken * insert_token = nullptr;
{
std::lock_guard lock(mutex);
auto it = cells.find(key);
if (it != cells.end())
{
hits++;
it->second.reference_count += 1;
queue.splice(queue.end(), queue, it->second.queue_iterator);
return it->second.value;
}
misses++;
insert_token = acquireInsertToken(key);
}
Cell * cell_ptr = nullptr;
{
std::lock_guard lock(insert_token->mutex);
if (!insert_token->value)
{
insert_token->value = load_func();
std::lock_guard cell_lock(mutex);
cell_ptr = set(key, insert_token->value);
if (cell_ptr)
{
cell_ptr->reference_count += 1;
}
else
{
insert_token->value = nullptr;
}
}
}
std::lock_guard lock(mutex);
releaseInsertToken(key);
if (cell_ptr)
{
return cell_ptr->value;
}
return nullptr;
}
MappedPtr getImpl(const Key & key)
{
std::lock_guard lock(mutex);
auto it = cells.find(key);
if (it == cells.end())
{
misses++;
return nullptr;
}
hits++;
it->second.reference_count += 1;
queue.splice(queue.end(), queue, it->second.queue_iterator);
return it->second.value;
}
// mark a reference is released
void release(const Key & key)
{
std::lock_guard lock(mutex);
auto it = cells.find(key);
if (it == cells.end() || it->second.reference_count == 0)
{
LOG_ERROR(&Poco::Logger::get("LRUResourceCache"), "try to release an invalid element");
abort();
}
it->second.reference_count -= 1;
}
InsertToken * acquireInsertToken(const Key & key)
{
@ -201,7 +233,7 @@ private:
}
// key mustn't be in the cache
Cell * insert_value(const Key & insert_key, MappedPtr value)
Cell * set(const Key & insert_key, MappedPtr value)
{
auto weight = value ? weight_function(*value) : 0;
auto queue_size = cells.size() + 1;
@ -251,5 +283,21 @@ private:
new_cell.queue_iterator = queue.insert(queue.end(), insert_key);
return &new_cell;
}
// If you want to update a value, call tryRemove() at first and then call acquire() with load_func.
bool tryRemove(const Key & key)
{
std::lock_guard guard(mutex);
auto it = cells.find(key);
if (it == cells.end())
return true;
auto & cell = it->second;
if (cell.reference_count)
return false;
queue.erase(cell.queue_iterator);
current_weight -= cell.weight;
cells.erase(it);
return true;
}
};
}

View File

@ -3,21 +3,21 @@
#include <gtest/gtest.h>
#include <Common/LRUResourceCache.h>
TEST(LRUResourceCache, acquire)
TEST(LRUResourceCache, get)
{
using MyCache = DB::LRUResourceCache<int, int>;
auto mcache = MyCache(10, 10);
int x = 10;
auto load_int = [&] { return std::make_shared<int>(x); };
auto val = mcache.acquire(1, load_int);
auto holder1 = mcache.getOrSet(1, load_int);
x = 11;
val = mcache.acquire(2, load_int);
ASSERT_TRUE(val != nullptr);
ASSERT_TRUE(*val == 11);
auto holder2 = mcache.getOrSet(2, load_int);
ASSERT_TRUE(holder2 != nullptr);
ASSERT_TRUE(holder2->value() == 11);
val = mcache.acquire(1);
ASSERT_TRUE(val != nullptr);
ASSERT_TRUE(*val == 10);
auto holder3 = mcache.get(1);
ASSERT_TRUE(holder3 != nullptr);
ASSERT_TRUE(holder3->value() == 10);
}
TEST(LRUResourceCache, remove)
@ -26,27 +26,20 @@ TEST(LRUResourceCache, remove)
auto mcache = MyCache(10, 10);
int x = 10;
auto load_int = [&] { return std::make_shared<int>(x); };
auto val = mcache.acquire(1, load_int);
x = 11;
val = mcache.acquire(2, load_int);
auto holder0 = mcache.getOrSet(1, load_int);
auto holder1 = mcache.getOrSet(1, load_int);
auto succ = mcache.tryRemove(3);
ASSERT_TRUE(succ);
succ = mcache.tryRemove(1);
auto succ = MyCache::MappedHolder::tryRemove(&holder0);
ASSERT_TRUE(!succ);
val = mcache.acquire(1);
ASSERT_TRUE(val != nullptr);
ASSERT_TRUE(*val == 10);
holder0 = mcache.get(1);
ASSERT_TRUE(holder0 != nullptr);
ASSERT_TRUE(holder0->value() == 10);
mcache.release(1);
succ = mcache.tryRemove(1);
ASSERT_TRUE(!succ);
mcache.release(1);
succ = mcache.tryRemove(1);
holder0 = nullptr;
succ = MyCache::MappedHolder::tryRemove(&holder1);
ASSERT_TRUE(succ);
val = mcache.acquire(1);
ASSERT_TRUE(val == nullptr);
holder1 = mcache.get(1);
ASSERT_TRUE(holder1 == nullptr);
}
struct MyWeight
@ -60,27 +53,27 @@ TEST(LRUResourceCache, evict_on_weight)
auto mcache = MyCache(5, 10);
int x = 2;
auto load_int = [&] { return std::make_shared<int>(x); };
auto val = mcache.acquire(1, load_int);
mcache.release(1);
auto holder1 = mcache.getOrSet(1, load_int);
holder1 = nullptr;
val = mcache.acquire(2, load_int);
mcache.release(2);
auto holder2 = mcache.getOrSet(2, load_int);
holder2 = nullptr;
x = 3;
val = mcache.acquire(3, load_int);
ASSERT_TRUE(val != nullptr);
auto holder3 = mcache.getOrSet(3, load_int);
ASSERT_TRUE(holder3 != nullptr);
auto w = mcache.weight();
ASSERT_EQ(w, 5);
auto n = mcache.size();
ASSERT_EQ(n, 2);
val = mcache.acquire(1);
ASSERT_TRUE(val == nullptr);
val = mcache.acquire(2);
ASSERT_TRUE(val != nullptr);
val = mcache.acquire(3);
ASSERT_TRUE(val != nullptr);
holder1 = mcache.get(1);
ASSERT_TRUE(holder1 == nullptr);
holder2 = mcache.get(2);
ASSERT_TRUE(holder2 != nullptr);
holder3 = mcache.get(3);
ASSERT_TRUE(holder3 != nullptr);
}
TEST(LRUResourceCache, evict_on_weight_v2)
@ -89,30 +82,30 @@ TEST(LRUResourceCache, evict_on_weight_v2)
auto mcache = MyCache(5, 10);
int x = 2;
auto load_int = [&] { return std::make_shared<int>(x); };
auto val = mcache.acquire(1, load_int);
mcache.release(1);
auto holder1 = mcache.getOrSet(1, load_int);
holder1 = nullptr;
val = mcache.acquire(2, load_int);
mcache.release(2);
auto holder2 = mcache.getOrSet(2, load_int);
holder2 = nullptr;
val = mcache.acquire(1);
mcache.release(1);
holder1 = mcache.get(1);
holder1 = nullptr;
x = 3;
val = mcache.acquire(3, load_int);
ASSERT_TRUE(val != nullptr);
auto holder3 = mcache.getOrSet(3, load_int);
ASSERT_TRUE(holder3 != nullptr);
auto w = mcache.weight();
ASSERT_EQ(w, 5);
auto n = mcache.size();
ASSERT_EQ(n, 2);
val = mcache.acquire(1);
ASSERT_TRUE(val != nullptr);
val = mcache.acquire(2);
ASSERT_TRUE(val == nullptr);
val = mcache.acquire(3);
ASSERT_TRUE(val != nullptr);
holder1 = mcache.get(1);
ASSERT_TRUE(holder1 != nullptr);
holder2 = mcache.get(2);
ASSERT_TRUE(holder2 == nullptr);
holder3 = mcache.get(3);
ASSERT_TRUE(holder3 != nullptr);
}
TEST(LRUResourceCache, evict_on_weight_v3)
@ -121,30 +114,30 @@ TEST(LRUResourceCache, evict_on_weight_v3)
auto mcache = MyCache(5, 10);
int x = 2;
auto load_int = [&] { return std::make_shared<int>(x); };
auto val = mcache.acquire(1, load_int);
mcache.release(1);
auto holder1 = mcache.getOrSet(1, load_int);
holder1 = nullptr;
val = mcache.acquire(2, load_int);
mcache.release(2);
auto holder2 = mcache.getOrSet(2, load_int);
holder2 = nullptr;
val = mcache.acquire(1, load_int);
mcache.release(1);
holder1 = mcache.getOrSet(1, load_int);
holder1 = nullptr;
x = 3;
val = mcache.acquire(3, load_int);
ASSERT_TRUE(val != nullptr);
auto holder3 = mcache.getOrSet(3, load_int);
ASSERT_TRUE(holder3 != nullptr);
auto w = mcache.weight();
ASSERT_EQ(w, 5);
auto n = mcache.size();
ASSERT_EQ(n, 2);
val = mcache.acquire(1);
ASSERT_TRUE(val != nullptr);
val = mcache.acquire(2);
ASSERT_TRUE(val == nullptr);
val = mcache.acquire(3);
ASSERT_TRUE(val != nullptr);
holder1 = mcache.get(1);
ASSERT_TRUE(holder1 != nullptr);
holder2 = mcache.get(2);
ASSERT_TRUE(holder2 == nullptr);
holder3 = mcache.get(3);
ASSERT_TRUE(holder3 != nullptr);
}
TEST(LRUResourceCache, evict_on_size)
@ -153,27 +146,27 @@ TEST(LRUResourceCache, evict_on_size)
auto mcache = MyCache(5, 2);
int x = 2;
auto load_int = [&] { return std::make_shared<int>(x); };
auto val = mcache.acquire(1, load_int);
mcache.release(1);
auto holder1 = mcache.getOrSet(1, load_int);
holder1 = nullptr;
val = mcache.acquire(2, load_int);
mcache.release(2);
auto holder2 = mcache.getOrSet(2, load_int);
holder2 = nullptr;
x = 3;
val = mcache.acquire(3, load_int);
ASSERT_TRUE(val != nullptr);
auto holder3 = mcache.getOrSet(3, load_int);
ASSERT_TRUE(holder3 != nullptr);
auto n = mcache.size();
ASSERT_EQ(n, 2);
auto w = mcache.weight();
ASSERT_EQ(w, 2);
val = mcache.acquire(1);
ASSERT_TRUE(val == nullptr);
val = mcache.acquire(2);
ASSERT_TRUE(val != nullptr);
val = mcache.acquire(3);
ASSERT_TRUE(val != nullptr);
holder1 = mcache.get(1);
ASSERT_TRUE(holder1 == nullptr);
holder2 = mcache.get(2);
ASSERT_TRUE(holder2 != nullptr);
holder3 = mcache.get(3);
ASSERT_TRUE(holder3 != nullptr);
}
TEST(LRUResourceCache, not_evict_used_element)
@ -182,95 +175,95 @@ TEST(LRUResourceCache, not_evict_used_element)
auto mcache = MyCache(7, 10);
int x = 2;
auto load_int = [&] { return std::make_shared<int>(x); };
auto val = mcache.acquire(1, load_int);
auto holder1 = mcache.getOrSet(1, load_int);
val = mcache.acquire(2, load_int);
mcache.release(2);
auto holder2 = mcache.getOrSet(2, load_int);
holder2 = nullptr;
val = mcache.acquire(3, load_int);
mcache.release(3);
auto holder3 = mcache.getOrSet(3, load_int);
holder3 = nullptr;
x = 3;
val = mcache.acquire(4, load_int);
ASSERT_TRUE(val != nullptr);
auto holder4 = mcache.getOrSet(4, load_int);
ASSERT_TRUE(holder4 != nullptr);
auto n = mcache.size();
ASSERT_EQ(n, 3);
auto w = mcache.weight();
ASSERT_EQ(w, 7);
val = mcache.acquire(1);
ASSERT_TRUE(val != nullptr);
val = mcache.acquire(2);
ASSERT_TRUE(val == nullptr);
val = mcache.acquire(3);
ASSERT_TRUE(val != nullptr);
val = mcache.acquire(4);
ASSERT_TRUE(val != nullptr);
holder1 = mcache.get(1);
ASSERT_TRUE(holder1 != nullptr);
holder2 = mcache.get(2);
ASSERT_TRUE(holder2 == nullptr);
holder3 = mcache.get(3);
ASSERT_TRUE(holder3 != nullptr);
holder4 = mcache.get(4);
ASSERT_TRUE(holder4 != nullptr);
}
TEST(LRUResourceCache, acquire_fail)
TEST(LRUResourceCache, get_fail)
{
using MyCache = DB::LRUResourceCache<int, int, MyWeight>;
auto mcache = MyCache(5, 10);
int x = 2;
auto load_int = [&] { return std::make_shared<int>(x); };
auto val = mcache.acquire(1, load_int);
val = mcache.acquire(2, load_int);
val = mcache.acquire(3, load_int);
ASSERT_TRUE(val == nullptr);
auto holder1 = mcache.getOrSet(1, load_int);
auto holder2 = mcache.getOrSet(2, load_int);
auto holder3 = mcache.getOrSet(3, load_int);
ASSERT_TRUE(holder3 == nullptr);
auto n = mcache.size();
ASSERT_EQ(n, 2);
auto w = mcache.weight();
ASSERT_EQ(w, 4);
val = mcache.acquire(1);
ASSERT_TRUE(val != nullptr);
val = mcache.acquire(2);
ASSERT_TRUE(val != nullptr);
val = mcache.acquire(3);
ASSERT_TRUE(val == nullptr);
holder1 = mcache.get(1);
ASSERT_TRUE(holder1 != nullptr);
holder2 = mcache.get(2);
ASSERT_TRUE(holder2 != nullptr);
holder3 = mcache.get(3);
ASSERT_TRUE(holder3 == nullptr);
}
TEST(LRUResourceCache, dup_acquire)
TEST(LRUResourceCache, dup_get)
{
using MyCache = DB::LRUResourceCache<int, int, MyWeight>;
auto mcache = MyCache(20, 10);
int x = 2;
auto load_int = [&] { return std::make_shared<int>(x); };
auto val = mcache.acquire(1, load_int);
mcache.release(1);
auto holder1 = mcache.getOrSet(1, load_int);
holder1 = nullptr;
x = 11;
val = mcache.acquire(1, load_int);
ASSERT_TRUE(val != nullptr);
holder1 = mcache.getOrSet(1, load_int);
ASSERT_TRUE(holder1 != nullptr);
auto n = mcache.size();
ASSERT_EQ(n, 1);
auto w = mcache.weight();
ASSERT_EQ(w, 2);
val = mcache.acquire(1);
ASSERT_TRUE(val != nullptr);
ASSERT_TRUE(*val == 2);
holder1 = mcache.get(1);
ASSERT_TRUE(holder1 != nullptr);
ASSERT_TRUE(holder1->value() == 2);
}
TEST(LRUResourceCache, re_acquire)
TEST(LRUResourceCache, re_get)
{
using MyCache = DB::LRUResourceCache<int, int, MyWeight>;
auto mcache = MyCache(20, 10);
int x = 2;
auto load_int = [&] { return std::make_shared<int>(x); };
auto val = mcache.acquire(1, load_int);
mcache.release(1);
mcache.tryRemove(1);
auto holder1 = mcache.getOrSet(1, load_int);
MyCache::MappedHolder::tryRemove(&holder1);
x = 11;
val = mcache.acquire(1, load_int);
ASSERT_TRUE(val != nullptr);
holder1 = mcache.getOrSet(1, load_int);
ASSERT_TRUE(holder1 != nullptr);
auto n = mcache.size();
ASSERT_EQ(n, 1);
auto w = mcache.weight();
ASSERT_EQ(w, 11);
val = mcache.acquire(1);
ASSERT_TRUE(val != nullptr);
ASSERT_TRUE(*val == 11);
holder1 = mcache.get(1);
ASSERT_TRUE(holder1 != nullptr);
ASSERT_TRUE(holder1->value() == 11);
}