use std::collections::HashMap;
use std::sync::Arc;
use miette::{IntoDiagnostic, Result};
use serde::Serialize;
use smol_str::SmolStr;
use tracing::debug;
use crate::config::RateTier;
use crate::db::keys::pds_account_count_key;
use crate::db::pds_meta as db_pds;
use crate::pds_discovery::canonical_pds_host;
use crate::pds_meta::{HostDesc, HostStatus, PdsMeta};
use crate::state::AppState;
/// a single PDS-to-tier assignment.
#[derive(Debug, Clone, Serialize)]
pub struct PdsTierAssignment {
pub host: String,
pub tier: String,
}
/// a rate tier definition, as returned by the API.
#[derive(Debug, Clone, Serialize)]
pub struct PdsTierDefinition {
pub per_second_base: u64,
pub per_second_account_mul: f64,
pub per_hour: u64,
pub per_day: u64,
pub account_limit: Option<u64>,
}
impl From<RateTier> for PdsTierDefinition {
fn from(t: RateTier) -> Self {
Self {
per_second_base: t.per_second_base,
per_second_account_mul: t.per_second_account_mul,
per_hour: t.per_hour,
per_day: t.per_day,
account_limit: t.account_limit,
}
}
}
/// runtime control over pds related behaviour (eg. ratelimits).
#[derive(Clone)]
pub struct PdsControl(pub(super) Arc<AppState>);
impl PdsControl {
async fn update<F, G>(&self, db_op: F, mem_op: G) -> Result<()>
where
F: FnOnce(&mut fjall::OwnedWriteBatch, &fjall::Keyspace) + Send + 'static,
G: FnOnce(&mut PdsMeta) + Send + 'static,
{
let _admission = self.0.pds_admission.lock().await;
self.update_locked(db_op, mem_op).await
}
async fn update_locked<F, G>(&self, db_op: F, mem_op: G) -> Result<()>
where
F: FnOnce(&mut fjall::OwnedWriteBatch, &fjall::Keyspace) + Send + 'static,
G: FnOnce(&mut PdsMeta) + Send + 'static,
{
let status_write = self.0.pds_status_write.clone();
let pds_meta = self.0.pds_meta.clone();
self.0
.db
.run(move |db| {
let _status_write = status_write.lock();
let mut batch = db.inner.batch();
db_op(&mut batch, &db.filter);
batch.commit().into_diagnostic()?;
db.persist()?;
let mut snapshot = (**pds_meta.load()).clone();
mem_op(&mut snapshot);
pds_meta.store(Arc::new(snapshot));
Ok(())
})
.await
}
/// list all current per-PDS tier assignments (explicit api-assigned overrides only).
pub async fn list_tiers(&self) -> HashMap<String, String> {
let snapshot = self.0.pds_meta.load();
snapshot
.hosts
.iter()
.filter_map(|(host, desc): (&String, &HostDesc)| {
desc.tier
.as_ref()
.map(|t: &smol_str::SmolStr| (host.clone(), t.to_string()))
})
.collect()
}
/// returns the assigned tier for `host`, or \"default\" if none is assigned.
pub fn get_tier(&self, host: impl AsRef<str>) -> String {
let snapshot = self.0.pds_meta.load();
snapshot
.hosts
.get(host.as_ref())
.and_then(|h| h.tier.as_ref())
.map(|t| t.to_string())
.unwrap_or_else(|| "default".to_string())
}
/// returns true if `host` is currently banned.
pub fn is_banned(&self, host: impl AsRef<str>) -> bool {
self.0.pds_meta.load().is_banned(host.as_ref())
}
/// list all currently banned PDS hosts.
pub async fn list_banned(&self) -> Vec<String> {
let snapshot = self.0.pds_meta.load();
snapshot
.hosts
.iter()
.filter(|(_, desc)| desc.operator_ban || desc.status == HostStatus::Banned)
.map(|(host, _)| host.clone())
.collect()
}
/// list all configured rate tier definitions.
pub fn list_rate_tiers(&self) -> HashMap<String, PdsTierDefinition> {
self.0
.tier_policy
.tiers
.iter()
.map(|(name, tier): (&smol_str::SmolStr, &RateTier)| {
(name.to_string(), PdsTierDefinition::from(*tier))
})
.collect()
}
/// assign `host` to `tier`.
/// returns an error if `tier` is not a known tier name.
pub async fn set_tier(&self, host: impl AsRef<str>, tier: String) -> Result<()> {
if !self.0.tier_policy.tiers.contains_key(tier.as_str()) {
miette::bail!(
"unknown tier '{tier}'; known tiers: {:?}",
self.0.tier_policy.tiers.keys().collect::<Vec<_>>()
);
}
let _admission = self.0.pds_admission.lock().await;
let host = host.as_ref().to_string();
let host_clone = host.clone();
let tier_clone = tier.clone();
// read the new tier's account limit and check for a status transition,
// now that the override is about to change.
let new_tier_limit = self
.0
.tier_policy
.tiers
.get(tier.as_str())
.unwrap()
.account_limit;
let count = self.0.db.get_count_sync(&pds_account_count_key(&host));
let current_status = self.0.pds_meta.load().status(&host);
let maybe_status = current_status.check_limit_transition(count, new_tier_limit);
self.update_locked(
move |batch, ks| {
db_pds::set_tier(batch, ks, &host_clone, &tier_clone);
if let Some(status) = maybe_status {
let _ = db_pds::set_status(batch, ks, &host_clone, status);
}
},
move |meta| {
meta.update_host_entry(&host, |entry| {
entry.tier = Some(SmolStr::new(&tier));
if let Some(status) = maybe_status {
entry.status = status;
}
});
},
)
.await
}
/// remove any explicit tier assignment for `host`, reverting it to the matched rule or default.
pub async fn remove_tier(&self, host: impl AsRef<str>) -> Result<()> {
let _admission = self.0.pds_admission.lock().await;
let host = host.as_ref().to_string();
let host_clone = host.clone();
// after removing the override, the effective tier is determined by glob rules.
// resolve it without the override to get the correct limit.
let effective_limit = self.0.tier_policy.resolve(&host, None).account_limit;
let count = self.0.db.get_count_sync(&pds_account_count_key(&host));
let current_status = self.0.pds_meta.load().status(&host);
let maybe_status = current_status.check_limit_transition(count, effective_limit);
debug!(
host,
?current_status,
?effective_limit,
count,
?maybe_status,
"remove_tier: computed status transition"
);
self.update_locked(
move |batch, ks| {
db_pds::remove_tier(batch, ks, &host_clone);
if let Some(status) = maybe_status {
let _ = db_pds::set_status(batch, ks, &host_clone, status);
}
},
move |meta| {
meta.update_host_entry(&host, |desc| {
desc.tier = None;
if let Some(status) = maybe_status {
desc.status = status;
}
});
},
)
.await
}
/// ban `host`
pub async fn ban(&self, host: impl AsRef<str>) -> Result<()> {
let host = canonical_pds_host(host.as_ref())
.map_err(|error| miette::miette!("invalid PDS host: {error}"))?
.as_str()
.to_string();
let host_clone = host.clone();
self.update(
move |batch, ks| db_pds::set_ban(batch, ks, &host_clone),
move |meta| {
meta.update_host_entry(&host, |desc| {
desc.operator_ban = true;
});
},
)
.await
}
/// unban `host`
pub async fn unban(&self, host: impl AsRef<str>) -> Result<()> {
let host = canonical_pds_host(host.as_ref())
.map_err(|error| miette::miette!("invalid PDS host: {error}"))?
.as_str()
.to_string();
let host_clone = host.clone();
self.update(
move |batch, ks| db_pds::remove_ban(batch, ks, &host_clone),
move |meta| {
meta.update_host_entry(&host, |desc| {
desc.operator_ban = false;
});
},
)
.await
}
}