135 lines
4.5 KiB
Rust
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()
|
|
}
|
|
}
|