use std::{
collections::HashMap,
hash::{Hash, Hasher},
sync::{Arc, Mutex, OnceLock},
};
use tokio::{sync::mpsc, time::Duration};
use tracing::{debug, warn};
const DEFAULT_QUEUE_CAPACITY: usize = 64;
const CLEANUP_YIELD: Duration = Duration::from_millis(100);
const IN_FLIGHT_SUFFIX: &str = "~new";
#[must_use]
pub fn in_flight_pin_name(name: &str) -> String {
format!("{name}{IN_FLIGHT_SUFFIX}")
}
fn stale_pins(pins: Vec<String>, protected_cid: &str) -> Vec<String> {
pins.into_iter()
.filter(|cid| cid != protected_cid)
.collect()
}
#[derive(Clone, Debug)]
pub struct PinCleanupRequest {
pub kubo_url: String,
pub name: String,
pub protected_cid: String,
pub cleanup_local: bool,
pub remote_service: Option<String>,
}
#[derive(Clone, Debug, Eq)]
struct PinCleanupKey {
kubo_url: String,
name: String,
cleanup_local: bool,
remote_service: Option<String>,
}
impl PartialEq for PinCleanupKey {
fn eq(&self, other: &Self) -> bool {
self.kubo_url == other.kubo_url
&& self.name == other.name
&& self.cleanup_local == other.cleanup_local
&& self.remote_service == other.remote_service
}
}
impl Hash for PinCleanupKey {
fn hash<H: Hasher>(&self, state: &mut H) {
self.kubo_url.hash(state);
self.name.hash(state);
self.cleanup_local.hash(state);
self.remote_service.hash(state);
}
}
impl From<&PinCleanupRequest> for PinCleanupKey {
fn from(request: &PinCleanupRequest) -> Self {
Self {
kubo_url: request.kubo_url.clone(),
name: request.name.clone(),
cleanup_local: request.cleanup_local,
remote_service: request.remote_service.clone(),
}
}
}
#[derive(Clone, Debug)]
pub struct PinCleanupScheduler {
pending: Arc<Mutex<HashMap<PinCleanupKey, PinCleanupRequest>>>,
wake: mpsc::Sender<PinCleanupKey>,
}
impl PinCleanupScheduler {
#[must_use]
pub fn global() -> &'static Self {
static SCHEDULER: OnceLock<PinCleanupScheduler> = OnceLock::new();
SCHEDULER.get_or_init(Self::new)
}
#[must_use]
pub fn new() -> Self {
Self::with_capacity(DEFAULT_QUEUE_CAPACITY)
}
#[must_use]
pub fn with_capacity(capacity: usize) -> Self {
let capacity = capacity.max(1);
let pending = Arc::new(Mutex::new(HashMap::new()));
let (wake, receiver) = mpsc::channel(capacity);
tokio::spawn(run_cleanup_worker(
Arc::clone(&pending),
wake.clone(),
receiver,
));
Self { pending, wake }
}
pub fn schedule(&self, request: PinCleanupRequest) -> bool {
let key = PinCleanupKey::from(&request);
let mut pending = self.pending.lock().expect("pin cleanup scheduler poisoned");
if let std::collections::hash_map::Entry::Occupied(mut entry) = pending.entry(key.clone()) {
entry.insert(request);
return true;
}
match self.wake.try_send(key.clone()) {
Ok(()) => {
pending.insert(key, request);
true
}
Err(error) => {
debug!(name = %request.name, error = %error, "dropping bounded pin cleanup job");
false
}
}
}
}
impl Default for PinCleanupScheduler {
fn default() -> Self {
Self::new()
}
}
pub fn delete_local_pins_named_in_background(
kubo_url: impl Into<String>,
name: impl Into<String>,
protected_cid: impl Into<String>,
) -> bool {
PinCleanupScheduler::global().schedule(PinCleanupRequest {
kubo_url: kubo_url.into(),
name: name.into(),
protected_cid: protected_cid.into(),
cleanup_local: true,
remote_service: None,
})
}
pub fn delete_remote_pins_named_in_background(
kubo_url: impl Into<String>,
service: impl Into<String>,
name: impl Into<String>,
protected_cid: impl Into<String>,
) -> bool {
PinCleanupScheduler::global().schedule(PinCleanupRequest {
kubo_url: kubo_url.into(),
name: name.into(),
protected_cid: protected_cid.into(),
cleanup_local: false,
remote_service: Some(service.into()),
})
}
pub async fn remote_pin_replace_named(
kubo_url: &str,
service: &str,
name: &str,
cid: &str,
overwrite: bool,
) -> anyhow::Result<bool> {
let add_name = if overwrite {
in_flight_pin_name(name)
} else {
name.to_string()
};
crate::kubo::kubo::remote_pin_add_named(kubo_url, service, cid, &add_name).await?;
Ok(overwrite && delete_remote_pins_named_in_background(kubo_url, service, name, cid))
}
async fn run_cleanup_worker(
pending: Arc<Mutex<HashMap<PinCleanupKey, PinCleanupRequest>>>,
wake: mpsc::Sender<PinCleanupKey>,
mut receiver: mpsc::Receiver<PinCleanupKey>,
) {
while let Some(key) = receiver.recv().await {
let request = pending
.lock()
.expect("pin cleanup scheduler poisoned")
.remove(&key);
let Some(request) = request else {
continue;
};
if cleanup_one_batch(&request).await {
let should_requeue = {
let mut pending = pending.lock().expect("pin cleanup scheduler poisoned");
if pending.contains_key(&key) {
false
} else {
pending.insert(key.clone(), request);
true
}
};
if should_requeue {
tokio::time::sleep(CLEANUP_YIELD).await;
if let Err(error) = wake.try_send(key) {
debug!(error = %error, "dropping delayed pin cleanup batch");
}
}
}
}
}
async fn cleanup_one_batch(request: &PinCleanupRequest) -> bool {
let local_more = if request.cleanup_local {
cleanup_local_pass(request).await
} else {
false
};
let remote_more = if request.remote_service.is_some() {
cleanup_remote_pass(request).await
} else {
false
};
local_more || remote_more
}
async fn cleanup_local_pass(request: &PinCleanupRequest) -> bool {
let temp_name = in_flight_pin_name(&request.name);
let final_pins = match crate::kubo::kubo::list_named_recursive_pins(
&request.kubo_url,
&request.name,
)
.await
{
Ok(pins) => pins,
Err(error) => {
warn!(name = %request.name, error = %error, "local old-pin lookup failed");
return false;
}
};
let temp_pins =
match crate::kubo::kubo::list_named_recursive_pins(&request.kubo_url, &temp_name).await {
Ok(pins) => pins,
Err(error) => {
warn!(name = %temp_name, error = %error, "local in-flight pin lookup failed");
return false;
}
};
let mut all_pins = final_pins.clone();
all_pins.extend(temp_pins);
let stale = stale_pins(all_pins, &request.protected_cid);
if !stale.is_empty() {
let mut removed_any = false;
for cid in stale {
match crate::kubo::kubo::pin_rm(&request.kubo_url, &cid).await {
Ok(()) => removed_any = true,
Err(error) => {
warn!(name = %request.name, cid = %cid, error = %error, "local old-pin cleanup failed");
}
}
}
return removed_any;
}
if !final_pins.iter().any(|cid| cid == &request.protected_cid) {
if let Err(error) = crate::kubo::kubo::pin_add_named(
&request.kubo_url,
&request.protected_cid,
&request.name,
)
.await
{
warn!(name = %request.name, cid = %request.protected_cid, error = %error, "local pin finalisation failed");
}
}
false
}
async fn cleanup_remote_pass(request: &PinCleanupRequest) -> bool {
let Some(service) = request.remote_service.as_deref() else {
return false;
};
let temp_name = in_flight_pin_name(&request.name);
let final_pins = match crate::kubo::kubo::list_named_remote_pins(
&request.kubo_url,
service,
&request.name,
)
.await
{
Ok(pins) => pins,
Err(error) => {
warn!(name = %request.name, service, error = %error, "remote old-pin lookup failed");
return false;
}
};
let temp_pins = match crate::kubo::kubo::list_named_remote_pins(
&request.kubo_url,
service,
&temp_name,
)
.await
{
Ok(pins) => pins,
Err(error) => {
warn!(name = %temp_name, service, error = %error, "remote in-flight pin lookup failed");
return false;
}
};
let stale_final = stale_pins(final_pins.clone(), &request.protected_cid);
let stale_temp = stale_pins(temp_pins.clone(), &request.protected_cid);
if !stale_final.is_empty() || !stale_temp.is_empty() {
let mut removed_any = false;
for (cid, pin_name) in stale_final
.iter()
.map(|cid| (cid, request.name.as_str()))
.chain(stale_temp.iter().map(|cid| (cid, temp_name.as_str())))
{
match crate::kubo::kubo::remote_pin_rm_named(&request.kubo_url, service, cid, pin_name)
.await
{
Ok(()) => removed_any = true,
Err(error) => {
warn!(name = %pin_name, service, cid = %cid, error = %error, "remote old-pin cleanup failed");
}
}
}
return removed_any;
}
if temp_pins.iter().any(|cid| cid == &request.protected_cid) {
if !final_pins.iter().any(|cid| cid == &request.protected_cid) {
if let Err(error) = crate::kubo::kubo::remote_pin_add_named(
&request.kubo_url,
service,
&request.protected_cid,
&request.name,
)
.await
{
warn!(name = %request.name, service, cid = %request.protected_cid, error = %error, "remote pin finalisation failed");
return false;
}
}
if let Err(error) = crate::kubo::kubo::remote_pin_rm_named(
&request.kubo_url,
service,
&request.protected_cid,
&temp_name,
)
.await
{
warn!(name = %temp_name, service, cid = %request.protected_cid, error = %error, "remote in-flight pin removal failed");
}
}
false
}
#[cfg(test)]
mod tests {
use super::stale_pins;
#[test]
fn stale_pins_protects_current_cid() {
let stale = stale_pins(
vec![
"old-a".to_string(),
"current".to_string(),
"old-b".to_string(),
"old-c".to_string(),
],
"current",
);
assert_eq!(stale, ["old-a", "old-b", "old-c"]);
}
#[test]
fn stale_pins_returns_empty_when_only_current_remains() {
let stale = stale_pins(vec!["current".to_string()], "current");
assert!(stale.is_empty());
}
}