use std::sync::Arc;
use std::time::Duration;
use affinidi_messaging_sdk::errors::ATMError;
use affinidi_messaging_sdk::{EvictionPolicy, PersistentRelationshipStore, RelationshipKv};
use tracing::{info, warn};
use crate::store::KeyspaceHandle;
pub struct KeyspaceRelationshipKv {
keyspace: KeyspaceHandle,
}
impl KeyspaceRelationshipKv {
pub fn new(keyspace: KeyspaceHandle) -> Self {
Self { keyspace }
}
}
#[async_trait::async_trait]
impl RelationshipKv for KeyspaceRelationshipKv {
async fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, ATMError> {
self.keyspace
.get_raw(key.to_vec())
.await
.map_err(|e| ATMError::SDKError(format!("relationships keyspace get: {e}")))
}
async fn put(&self, key: &[u8], value: &[u8]) -> Result<(), ATMError> {
self.keyspace
.insert_raw(key.to_vec(), value.to_vec())
.await
.map_err(|e| ATMError::SDKError(format!("relationships keyspace put: {e}")))
}
async fn delete(&self, key: &[u8]) -> Result<(), ATMError> {
self.keyspace
.remove(key.to_vec())
.await
.map_err(|e| ATMError::SDKError(format!("relationships keyspace delete: {e}")))
}
async fn scan_prefix(&self, prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, ATMError> {
self.keyspace
.prefix_iter_raw(prefix.to_vec())
.await
.map_err(|e| ATMError::SDKError(format!("relationships keyspace scan: {e}")))
}
}
pub type KeyspaceRelationshipStore = PersistentRelationshipStore<KeyspaceRelationshipKv>;
pub fn build_relationship_store(keyspace: KeyspaceHandle) -> Arc<KeyspaceRelationshipStore> {
Arc::new(PersistentRelationshipStore::new(
KeyspaceRelationshipKv::new(keyspace),
))
}
const SWEEP_INTERVAL: Duration = Duration::from_secs(6 * 60 * 60);
pub async fn maintenance_loop(store: Arc<KeyspaceRelationshipStore>) {
match store.established_relationships().await {
Ok(established) => info!(
count = established.len(),
"TSP relationships restored from the durable store"
),
Err(e) => warn!(error = %e, "could not enumerate restored TSP relationships"),
}
let policy = EvictionPolicy::default();
let mut ticker = tokio::time::interval(SWEEP_INTERVAL);
loop {
ticker.tick().await;
let Some(now_ms) = unix_millis() else {
continue;
};
match store.evict_idle(now_ms, &policy).await {
Ok(evicted) if !evicted.is_empty() => {
info!(count = evicted.len(), "evicted idle TSP relationships")
}
Ok(_) => {}
Err(e) => warn!(error = %e, "TSP relationship eviction sweep failed"),
}
}
}
fn unix_millis() -> Option<u64> {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.ok()
.map(|d| d.as_millis() as u64)
}