96 lines
2.8 KiB
C++
96 lines
2.8 KiB
C++
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
#include "ttl_lock.h"
|
|
|
|
namespace lmcache {
|
|
namespace storage_manager {
|
|
|
|
TTLLock::TTLLock(uint32_t ttl_sec)
|
|
: counter_(0),
|
|
expiration_ms_(0),
|
|
ttl_ms_(static_cast<int64_t>(ttl_sec) * 1000) {}
|
|
|
|
void TTLLock::lock() {
|
|
int64_t current_time = now_ms();
|
|
int64_t new_expiration = current_time + ttl_ms_;
|
|
|
|
// Use compare-and-swap loop to handle the TTL expiration case
|
|
while (true) {
|
|
int64_t old_expiration = expiration_ms_.load(std::memory_order_acquire);
|
|
int64_t old_counter = counter_.load(std::memory_order_acquire);
|
|
|
|
// Check if TTL has expired
|
|
bool expired = (current_time >= old_expiration);
|
|
|
|
if (expired) {
|
|
// TTL expired, try to reset counter to 1 and set new expiration
|
|
// First, try to update the expiration
|
|
if (expiration_ms_.compare_exchange_strong(old_expiration, new_expiration,
|
|
std::memory_order_seq_cst)) {
|
|
// Successfully updated expiration, now set counter to 1
|
|
counter_.store(1, std::memory_order_seq_cst);
|
|
return;
|
|
}
|
|
// Another thread updated expiration, retry
|
|
continue;
|
|
} else {
|
|
// TTL not expired, try to increment counter
|
|
if (counter_.compare_exchange_strong(old_counter, old_counter + 1,
|
|
std::memory_order_seq_cst)) {
|
|
// Successfully incremented counter, update expiration
|
|
expiration_ms_.store(new_expiration, std::memory_order_seq_cst);
|
|
return;
|
|
}
|
|
// Another thread modified counter, retry
|
|
continue;
|
|
}
|
|
}
|
|
}
|
|
|
|
void TTLLock::unlock() {
|
|
// Use compare-and-swap loop to ensure we don't go below 0
|
|
while (true) {
|
|
int64_t old_counter = counter_.load(std::memory_order_acquire);
|
|
|
|
if (old_counter <= 0) {
|
|
// Already at 0, nothing to do
|
|
return;
|
|
}
|
|
|
|
if (counter_.compare_exchange_strong(old_counter, old_counter - 1,
|
|
std::memory_order_seq_cst)) {
|
|
return;
|
|
}
|
|
// Another thread modified counter, retry
|
|
}
|
|
}
|
|
|
|
bool TTLLock::is_locked() const {
|
|
int64_t current_time = now_ms();
|
|
int64_t expiration = expiration_ms_.load(std::memory_order_acquire);
|
|
int64_t counter = counter_.load(std::memory_order_acquire);
|
|
|
|
// Lock is held if counter > 0 AND TTL not expired
|
|
return (counter > 0) && (current_time < expiration);
|
|
}
|
|
|
|
void TTLLock::reset() {
|
|
counter_.store(0, std::memory_order_seq_cst);
|
|
expiration_ms_.store(0, std::memory_order_seq_cst);
|
|
}
|
|
|
|
int64_t TTLLock::now_ms() {
|
|
auto now = Clock::now();
|
|
return to_ms(now);
|
|
}
|
|
|
|
int64_t TTLLock::to_ms(const TimePoint& tp) {
|
|
return static_cast<int64_t>(
|
|
std::chrono::duration_cast<std::chrono::milliseconds>(
|
|
tp.time_since_epoch())
|
|
.count());
|
|
}
|
|
|
|
} // namespace storage_manager
|
|
} // namespace lmcache
|