Files
lxc-streamutils/apps/server-rust/src/live/supervisor.rs
T

135 lines
4.5 KiB
Rust

//! Per-account source task ownership, restart and cancellation.
//!
//! The supervisor guarantees at most one provider generation for a source ID.
//! Reconfiguration cancels the old task before a replacement starts, preventing
//! duplicate Bilibili listeners and duplicate downstream events.
use std::{collections::HashMap, sync::Arc};
use async_trait::async_trait;
use tokio::sync::{RwLock, mpsc, watch};
use tokio_util::sync::CancellationToken;
use tracing::{error, info};
use uuid::Uuid;
use crate::{
domain::LiveEvent,
live::{LiveProvider, SourceContext, SourceStatus},
};
#[async_trait]
pub trait ProviderFactory: Send + Sync {
async fn build(&self, source: &SourceContext) -> Result<Arc<dyn LiveProvider>, String>;
}
struct RunningSource {
cancel: CancellationToken,
status: watch::Receiver<SourceStatus>,
}
/// Owns exactly one provider task per account. Restart always cancels
/// the prior generation before starting another, which prevents the duplicate
/// listeners produced by the legacy `/reconnect` handler.
#[derive(Clone)]
pub struct SourceSupervisor {
factory: Arc<dyn ProviderFactory>,
events: mpsc::Sender<Arc<LiveEvent>>,
running: Arc<RwLock<HashMap<Uuid, RunningSource>>>,
}
impl SourceSupervisor {
pub fn new(factory: Arc<dyn ProviderFactory>, events: mpsc::Sender<Arc<LiveEvent>>) -> Self {
Self {
factory,
events,
running: Arc::new(RwLock::new(HashMap::new())),
}
}
pub async fn start(&self, source: SourceContext) -> Result<(), String> {
self.stop(source.source_id).await;
let cancel = CancellationToken::new();
let provider = match self.factory.build(&source).await {
Ok(provider) => provider,
Err(error) => {
let (_, status) = watch::channel(SourceStatus {
source_id: source.source_id,
room_id: source.room_id.clone(),
connected: false,
cookie_cloud: false,
detail: error.clone(),
});
self.running
.write()
.await
.insert(source.source_id, RunningSource { cancel, status });
return Err(error);
}
};
let initial = SourceStatus::starting(source.source_id, source.room_id.clone());
let (status_tx, status_rx) = watch::channel(initial);
self.running.write().await.insert(
source.source_id,
RunningSource {
cancel: cancel.clone(),
status: status_rx,
},
);
let events = self.events.clone();
let source_id = source.source_id;
let terminal_room_id = source.room_id.clone();
let terminal_status = status_tx.clone();
tokio::spawn(async move {
info!(%source_id, room_id = %source.room_id, provider = provider.provider_name(), "starting live source");
if let Err(error) = provider.run(source, events, status_tx, cancel).await {
error!(%source_id, %error, "live source stopped with error");
let _ = terminal_status.send(SourceStatus {
source_id,
room_id: terminal_room_id,
connected: false,
cookie_cloud: true,
detail: error,
});
}
// Keep the terminal status visible. A deliberate stop or restart
// removes the receiver first, so sends from an older task cannot
// affect a newer generation.
});
Ok(())
}
pub async fn restart(&self, source: SourceContext) -> Result<(), String> {
self.start(source).await
}
pub async fn stop(&self, source_id: Uuid) {
if let Some(entry) = self.running.write().await.remove(&source_id) {
entry.cancel.cancel();
}
}
pub async fn stop_all(&self) {
let mut running = self.running.write().await;
for (_, entry) in running.drain() {
entry.cancel.cancel();
}
}
pub async fn status(&self, source_id: Uuid) -> Option<SourceStatus> {
self.running
.read()
.await
.get(&source_id)
.map(|entry| entry.status.borrow().clone())
}
pub async fn statuses(&self) -> Vec<SourceStatus> {
self.running
.read()
.await
.values()
.map(|entry| entry.status.borrow().clone())
.collect()
}
}