Files
lmcache--lmcache/csrc/storage_manager/ttl_lock.cpp
T
2026-07-13 12:24:33 +08:00

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