#![allow(dead_code)]
use std::collections::{BTreeMap, BTreeSet};
use std::pin::pin;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::SystemTime;
use parking_lot::Mutex;
use tokio::sync::Notify;
use crate::backend::Backend;
use crate::facade::{
Batch, LivenessSource, Pin, PinState, ProtectCallback, Registration, RegistrationKind,
};
use crate::keys::Key;
use crate::Error;
pub(crate) struct GcHome {
core: Arc<RegistryCore>,
}
impl Default for GcHome {
fn default() -> GcHome {
GcHome {
core: Arc::new(RegistryCore::new()),
}
}
}
impl GcHome {
pub(crate) async fn acquire(&self, key: &Key) -> Pin {
self.core.acquire(key).await
}
pub(crate) fn new_batch(&self) -> Batch {
Batch::created(Arc::clone(&self.core))
}
pub(crate) fn core(&self) -> &Arc<RegistryCore> {
&self.core
}
pub(crate) fn convert<F: FnOnce()>(&self, pin: Pin, register: F) {
self.core.convert(pin, register)
}
pub(crate) fn registered_source_count(&self) -> usize {
self.core.registered_source_count()
}
pub(crate) fn register_source(&self, source: Arc<dyn LivenessSource>) -> Registration {
self.core.register_source(source)
}
pub(crate) fn register_fleet_source(&self, reference: String) -> Registration {
self.core.register_fleet_source(reference)
}
pub(crate) fn unregister_fleet_source(&self, registration: Registration) -> Option<String> {
self.core.unregister_fleet_source(registration)
}
pub(crate) fn fleet_source_count(&self) -> usize {
self.core.fleet_source_count()
}
pub(crate) fn registered_liveness_count(&self) -> usize {
self.registered_source_count() + self.fleet_source_count()
}
pub(crate) fn unregister_source(&self, registration: Registration) -> bool {
self.core.unregister_source(registration)
}
pub(crate) fn register_protect(&self, callback: Arc<dyn ProtectCallback>) -> Registration {
self.core.register_protect(callback)
}
pub(crate) fn unregister_protect(&self, registration: Registration) -> bool {
self.core.unregister_protect(registration)
}
pub(crate) fn registered_sources(&self) -> Vec<Arc<dyn LivenessSource>> {
self.core.registered_sources()
}
pub(crate) fn registered_protects(&self) -> Vec<Arc<dyn ProtectCallback>> {
self.core.registered_protects()
}
pub(crate) fn pinned_keys(&self) -> Vec<Key> {
self.core.pinned_keys()
}
pub(crate) fn pin_count(&self, key: &Key) -> u64 {
self.core.pin_count(key)
}
pub(crate) fn try_begin_sweep(&self) -> bool {
self.core.try_begin_sweep()
}
pub(crate) fn end_sweep(&self) {
self.core.end_sweep()
}
pub(crate) fn begin_sweep_guarded(&self) -> Option<SweepLeaseGuard<'_>> {
if self.try_begin_sweep() {
Some(SweepLeaseGuard { home: self })
} else {
None
}
}
pub(crate) async fn arbitrate(
&self,
key: &Key,
consult_protect: bool,
) -> Result<Arbitration, Error> {
self.core.arbitrate(key, consult_protect).await
}
pub(crate) fn conclude_delete(&self, key: &Key) {
self.core.conclude_delete(key)
}
}
pub(crate) enum Arbitration {
Cancelled,
Proceed,
}
pub(crate) enum DeleteOutcome {
Cancelled,
Deleted(bool),
}
pub(crate) struct SweepLeaseGuard<'a> {
home: &'a GcHome,
}
impl Drop for SweepLeaseGuard<'_> {
fn drop(&mut self) {
self.home.end_sweep();
}
}
pub(crate) struct RegistryCore {
state: Mutex<RegistryState>,
delete_done: Notify,
sweeping: AtomicBool,
next_registration: AtomicU64,
}
struct RegistryState {
pins: BTreeMap<Key, u64>,
deleting: BTreeSet<Key>,
sources: Vec<(u64, Arc<dyn LivenessSource>)>,
protects: Vec<(u64, Arc<dyn ProtectCallback>)>,
fleet_sources: Vec<(u64, String)>,
}
impl RegistryCore {
fn new() -> RegistryCore {
RegistryCore {
state: Mutex::new(RegistryState {
pins: BTreeMap::new(),
deleting: BTreeSet::new(),
sources: Vec::new(),
protects: Vec::new(),
fleet_sources: Vec::new(),
}),
delete_done: Notify::new(),
sweeping: AtomicBool::new(false),
next_registration: AtomicU64::new(1),
}
}
fn allocate_registration(&self) -> u64 {
self.next_registration.fetch_add(1, Ordering::Relaxed)
}
pub(crate) fn register_source(&self, source: Arc<dyn LivenessSource>) -> Registration {
let id = self.allocate_registration();
self.state.lock().sources.push((id, source));
Registration {
id,
kind: RegistrationKind::Source,
}
}
pub(crate) fn unregister_source(&self, registration: Registration) -> bool {
self.unregister(®istration, RegistrationKind::Source)
}
pub(crate) fn register_protect(&self, callback: Arc<dyn ProtectCallback>) -> Registration {
let id = self.allocate_registration();
self.state.lock().protects.push((id, callback));
Registration {
id,
kind: RegistrationKind::Protect,
}
}
pub(crate) fn unregister_protect(&self, registration: Registration) -> bool {
self.unregister(®istration, RegistrationKind::Protect)
}
pub(crate) fn register_fleet_source(&self, reference: String) -> Registration {
let mut state = self.state.lock();
if let Some((id, _)) = state
.fleet_sources
.iter()
.find(|(_, held)| *held == reference)
{
return Registration {
id: *id,
kind: RegistrationKind::FleetSource,
};
}
let id = self.allocate_registration();
state.fleet_sources.push((id, reference));
Registration {
id,
kind: RegistrationKind::FleetSource,
}
}
pub(crate) fn unregister_fleet_source(&self, registration: Registration) -> Option<String> {
let mut state = self.state.lock();
if registration.kind != RegistrationKind::FleetSource {
return None;
}
let position = state
.fleet_sources
.iter()
.position(|(id, _)| *id == registration.id)?;
Some(state.fleet_sources.remove(position).1)
}
pub(crate) fn fleet_source_count(&self) -> usize {
self.state.lock().fleet_sources.len()
}
fn unregister(&self, registration: &Registration, kind: RegistrationKind) -> bool {
let mut state = self.state.lock();
match kind {
RegistrationKind::Source => {
let before = state.sources.len();
state.sources.retain(|(id, _)| *id != registration.id);
state.sources.len() != before
}
RegistrationKind::Protect => {
let before = state.protects.len();
state.protects.retain(|(id, _)| *id != registration.id);
state.protects.len() != before
}
RegistrationKind::FleetSource => {
let _ = registration;
false
}
}
}
pub(crate) fn registered_source_count(&self) -> usize {
self.state.lock().sources.len()
}
fn registered_sources(&self) -> Vec<Arc<dyn LivenessSource>> {
self.state
.lock()
.sources
.iter()
.map(|(_, s)| Arc::clone(s))
.collect()
}
fn registered_protects(&self) -> Vec<Arc<dyn ProtectCallback>> {
self.state
.lock()
.protects
.iter()
.map(|(_, c)| Arc::clone(c))
.collect()
}
pub(crate) fn pinned_keys(&self) -> Vec<Key> {
self.state.lock().pins.keys().cloned().collect()
}
pub(crate) fn pin_count(&self, key: &Key) -> u64 {
self.state.lock().pins.get(key).copied().unwrap_or(0)
}
pub(crate) async fn acquire(self: &Arc<Self>, key: &Key) -> Pin {
loop {
let notified = self.delete_done.notified();
let mut notified = pin!(notified);
notified.as_mut().enable();
{
let mut state = self.state.lock();
if !state.deleting.contains(key) {
*state.pins.entry(key.clone()).or_insert(0) += 1;
return Pin::minted(Arc::clone(self), key.clone());
}
}
notified.as_mut().await;
}
}
pub(crate) fn release_pin(&self, key: &Key) {
let mut state = self.state.lock();
Self::release_in(&mut state, key);
}
pub(crate) fn release_pins(&self, mut keys: Vec<Key>) {
keys.sort();
let mut state = self.state.lock();
for key in &keys {
Self::release_in(&mut state, key);
}
}
fn release_in(state: &mut RegistryState, key: &Key) {
if let Some(count) = state.pins.get_mut(key) {
*count -= 1;
if *count == 0 {
state.pins.remove(key);
}
}
}
pub(crate) fn convert<F: FnOnce()>(&self, pin: Pin, register: F) {
#[cfg(feature = "postgres")]
let pin = {
let mut pin = pin;
pin.stop_renewal();
pin
};
let pin = std::mem::ManuallyDrop::new(pin);
let key = pin.key().clone();
let mut state = self.state.lock();
Self::release_in(&mut state, &key);
register();
}
pub(crate) fn try_begin_sweep(&self) -> bool {
self.sweeping
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
.is_ok()
}
pub(crate) fn end_sweep(&self) {
self.sweeping.store(false, Ordering::Release);
}
pub(crate) async fn arbitrate(
&self,
key: &Key,
consult_protect: bool,
) -> Result<Arbitration, Error> {
loop {
let notified = self.delete_done.notified();
let mut notified = pin!(notified);
notified.as_mut().enable();
enum Plan {
Cancel,
Wait,
Proceed,
}
let plan = {
let mut state = self.state.lock();
if state.pins.contains_key(key) {
Plan::Cancel
} else if state.deleting.contains(key) {
Plan::Wait
} else {
let mut protected = false;
if consult_protect {
for (_, callback) in &state.protects {
match callback.protect() {
Some(keys) if keys.iter().any(|live| live == key) => {
protected = true;
break;
}
None => {
protected = true;
break;
}
Some(_) => {}
}
}
}
if !protected {
for (_, source) in &state.sources {
if source.contains(key) {
protected = true;
break;
}
}
}
if protected {
Plan::Cancel
} else {
state.deleting.insert(key.clone());
Plan::Proceed
}
}
};
match plan {
Plan::Cancel => return Ok(Arbitration::Cancelled),
Plan::Proceed => return Ok(Arbitration::Proceed),
Plan::Wait => notified.as_mut().await,
}
}
}
pub(crate) fn conclude_delete(&self, key: &Key) {
let removed = self.state.lock().deleting.remove(key);
if removed {
self.delete_done.notify_waiters();
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct WindowCandidate {
pub key: Vec<u8>,
pub visible_after: SystemTime,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct PinRow {
pub key: Vec<u8>,
pub owner: Vec<u8>,
pub expiry: SystemTime,
}
pub(crate) const WINDOW_TAKE_BATCH: usize = 100;
pub(crate) trait EngineWindowState: Backend {
async fn ensure_gc_schema(&self) -> Result<(), Error>;
async fn window_stage(&self, candidates: &[WindowCandidate]) -> Result<(), Error>;
async fn window_take(&self, now: SystemTime) -> Result<Vec<Vec<u8>>, Error>;
}
pub(crate) trait EngineFleetState: EngineWindowState {
type JointTx: Send;
async fn begin_joint_tx(&self) -> Result<Self::JointTx, Error>;
async fn joint_put(
&self,
tx: &mut Self::JointTx,
key: &[u8],
value: &[u8],
) -> Result<(), Error>;
async fn pin_commit(&self, tx: Self::JointTx, entries: &[PinRow]) -> Result<(), Error>;
async fn pin_renew(&self, key: &[u8], owner: &[u8]) -> Result<bool, Error>;
async fn pin_check(&self, key: &[u8]) -> Result<PinState, Error>;
async fn pin_reap(&self, now: SystemTime) -> Result<u64, Error>;
async fn arbitrate(&self, key: &[u8]) -> Result<bool, Error>;
async fn sweep_lock_acquire(&self) -> Result<bool, Error>;
async fn sweep_lock_release(&self) -> Result<(), Error>;
}