cat / goq
clone

turns a linux box into a private steam streaming appliance: plug in, tap your key, play from another computer

fix/drop-hero-notarization-note
goq / crates / sigil-host / src / server / generation.rs
generation.rs 18.2K
1use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};2use std::sync::{Arc, Mutex, Weak};3use std::time::Duration;45use anyhow::{Context, Result};6use moq_net::{Broadcast, BroadcastConsumer, BroadcastProducer, Error as MoqError, Track};7use sigil_protocol::{8 MOQ_AUDIO_OPUS_TRACK, MOQ_AUDIO_TRACK_PRIORITY, MOQ_VIDEO_H264_TRACK, MOQ_VIDEO_TRACK_PRIORITY,9 MediaGenerationSigningKey, SignedMediaGenerationCertificate,10 media_generation_moq_broadcast_name,11};12use tokio::sync::{Notify, watch};13use tracing::{debug, info, warn};1415use super::adaptive::GenerationAdaptiveCoordinator;16use super::moq::{run_generation_audio_publisher, run_generation_video_publisher};17use super::session::{MediaV3Telemetry, SourceTaskGuard};18use crate::audio::spawn_pipewire_audio;19use crate::authorization::unix_timestamp_now;20use crate::clock::SessionClock;21use crate::config::{HostConfig, VideoSource};22use crate::moq_catalog::{GoqCatalogProducer, publish_goq_catalog_v2};23use crate::source::{24 EncodedFrame, EncodedSource, EncoderControl, spawn_gamescope_pipewire_after_static_preflight,25 spawn_test_pattern,26};2728const GENERATION_STOP_TIMEOUT: Duration = Duration::from_secs(2);29// Producer resources are generation-global. Viewer growth is represented only30// by at most MAX_VIEWERS leases and consumers; no capture, encoder, publisher,31// actuator, frame watch, or retained GOP allocation scales with viewer count.3233#[derive(Clone, Copy, Debug, PartialEq, Eq)]34pub(super) enum MediaGenerationLifecycle {35 Idle,36 Starting,37 Active,38 Stopping,39}4041#[derive(Debug)]42struct GenerationState {43 lifecycle: MediaGenerationLifecycle,44 active: Option<ActiveGeneration>,45 leases: usize,46}4748impl Default for GenerationState {49 fn default() -> Self {50 Self {51 lifecycle: MediaGenerationLifecycle::Idle,52 active: None,53 leases: 0,54 }55 }56}5758#[derive(Debug)]59pub(crate) struct MediaGenerationManager {60 config: HostConfig,61 host_secret: [u8; 32],62 next_generation_id: AtomicU64,63 state: tokio::sync::Mutex<GenerationState>,64 state_changed: Notify,65 starts: AtomicU64,66 stops: AtomicU64,67 recovering_viewers: Arc<AtomicUsize>,68}6970impl MediaGenerationManager {71 pub(crate) fn new(config: HostConfig, host_secret: [u8; 32]) -> Arc<Self> {72 Arc::new(Self {73 config,74 host_secret,75 next_generation_id: AtomicU64::new(0),76 state: tokio::sync::Mutex::new(GenerationState::default()),77 state_changed: Notify::new(),78 starts: AtomicU64::new(0),79 stops: AtomicU64::new(0),80 recovering_viewers: Arc::new(AtomicUsize::new(0)),81 })82 }8384 pub(super) async fn acquire(self: &Arc<Self>) -> Result<MediaGenerationLease> {85 loop {86 let mut state = self.state.lock().await;87 match state.lifecycle {88 MediaGenerationLifecycle::Active => {89 state.leases = state90 .leases91 .checked_add(1)92 .context("media generation lease count overflowed")?;93 let shared = state94 .active95 .as_ref()96 .context("active generation has no resources")?97 .shared98 .clone();99 return Ok(MediaGenerationLease::new(self, shared));100 }101 MediaGenerationLifecycle::Idle => {102 state.lifecycle = MediaGenerationLifecycle::Starting;103 drop(state);104 let generation_id = self105 .next_generation_id106 .fetch_add(1, Ordering::Relaxed)107 .checked_add(1)108 .context("media generation id overflowed")?;109 match ActiveGeneration::start(110 &self.config,111 self.host_secret,112 generation_id,113 Arc::clone(&self.recovering_viewers),114 )115 .await116 {117 Ok(active) => {118 let shared = active.shared.clone();119 let mut state = self.state.lock().await;120 state.active = Some(active);121 state.leases = 1;122 state.lifecycle = MediaGenerationLifecycle::Active;123 self.starts.fetch_add(1, Ordering::Relaxed);124 self.state_changed.notify_waiters();125 return Ok(MediaGenerationLease::new(self, shared));126 }127 Err(error) => {128 let mut state = self.state.lock().await;129 state.lifecycle = MediaGenerationLifecycle::Idle;130 self.state_changed.notify_waiters();131 return Err(error);132 }133 }134 }135 MediaGenerationLifecycle::Starting | MediaGenerationLifecycle::Stopping => {136 let changed = self.state_changed.notified();137 drop(state);138 changed.await;139 }140 }141 }142 }143144 async fn release(&self, generation_id: u64) {145 let active = {146 let mut state = self.state.lock().await;147 if state.lifecycle != MediaGenerationLifecycle::Active148 || state149 .active150 .as_ref()151 .map(|active| active.shared.generation_id)152 != Some(generation_id)153 {154 return;155 }156 state.leases = state.leases.saturating_sub(1);157 if state.leases != 0 {158 return;159 }160 state.lifecycle = MediaGenerationLifecycle::Stopping;161 state.active.take()162 };163 if let Some(active) = active {164 active.stop().await;165 self.stops.fetch_add(1, Ordering::Relaxed);166 }167 self.recovering_viewers.store(0, Ordering::Relaxed);168 let mut state = self.state.lock().await;169 state.lifecycle = MediaGenerationLifecycle::Idle;170 self.state_changed.notify_waiters();171 }172173 pub(super) async fn adaptive(174 &self,175 generation_id: u64,176 ) -> Result<Arc<GenerationAdaptiveCoordinator>> {177 let state = self.state.lock().await;178 state179 .active180 .as_ref()181 .filter(|active| {182 state.lifecycle == MediaGenerationLifecycle::Active183 && active.shared.generation_id == generation_id184 })185 .map(|active| Arc::clone(&active.shared.adaptive))186 .context("feedback does not match the active media generation")187 }188189 pub(crate) fn recovering_viewer_count(&self) -> usize {190 self.recovering_viewers.load(Ordering::Relaxed)191 }192193 #[cfg(test)]194 async fn lifecycle(&self) -> MediaGenerationLifecycle {195 self.state.lock().await.lifecycle196 }197}198199pub(super) struct GenerationShared {200 pub(super) generation_id: u64,201 pub(super) broadcast_name: String,202 broadcast: Mutex<Option<BroadcastConsumer>>,203 pub(super) certificate: SignedMediaGenerationCertificate,204 pub(super) session_clock: SessionClock,205 pub(super) pointer_surface_dimensions: Option<sigil_protocol::PointerSurfaceDimensions>,206 pub(super) encoder_control: Option<EncoderControl>,207 pub(super) telemetry: Arc<MediaV3Telemetry>,208 pub(super) audio_enabled: bool,209 pub(super) keyframe_requests: watch::Sender<Option<sigil_protocol::MediaControlRequestV3>>,210 pub(super) adaptive: Arc<GenerationAdaptiveCoordinator>,211}212213impl std::fmt::Debug for GenerationShared {214 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {215 formatter216 .debug_struct("GenerationShared")217 .field("generation_id", &self.generation_id)218 .field("broadcast_name", &self.broadcast_name)219 .field("audio_enabled", &self.audio_enabled)220 .finish_non_exhaustive()221 }222}223224impl GenerationShared {225 pub(super) fn consumer(&self) -> Result<BroadcastConsumer> {226 self.broadcast227 .lock()228 .map_err(|_| anyhow::anyhow!("media generation broadcast lock poisoned"))?229 .as_ref()230 .cloned()231 .context("media generation broadcast is stopping")232 }233}234235#[derive(Debug)]236pub(super) struct MediaGenerationLease {237 manager: Weak<MediaGenerationManager>,238 pub(super) shared: Arc<GenerationShared>,239 released: bool,240}241242impl MediaGenerationLease {243 fn new(manager: &Arc<MediaGenerationManager>, shared: Arc<GenerationShared>) -> Self {244 Self {245 manager: Arc::downgrade(manager),246 shared,247 released: false,248 }249 }250251 pub(super) async fn release(mut self) {252 self.release_inner().await;253 self.released = true;254 }255256 async fn release_inner(&self) {257 if let Some(manager) = self.manager.upgrade() {258 manager.release(self.shared.generation_id).await;259 }260 }261}262263impl Drop for MediaGenerationLease {264 fn drop(&mut self) {265 if self.released {266 return;267 }268 let Some(manager) = self.manager.upgrade() else {269 return;270 };271 let generation_id = self.shared.generation_id;272 if let Ok(runtime) = tokio::runtime::Handle::try_current() {273 runtime.spawn(async move { manager.release(generation_id).await });274 }275 }276}277278struct ActiveGeneration {279 shared: Arc<GenerationShared>,280 shutdown: watch::Sender<bool>,281 video_publisher: tokio::task::JoinHandle<Result<()>>,282 audio_publisher: Option<tokio::task::JoinHandle<Result<()>>>,283 adaptive: tokio::task::JoinHandle<Result<()>>,284 video_source: SourceTaskGuard,285 audio_source: Option<SourceTaskGuard>,286 _frame_receiver: watch::Receiver<Option<EncodedFrame>>,287 catalog: Option<GoqCatalogProducer>,288 broadcast: Option<BroadcastProducer>,289}290291impl std::fmt::Debug for ActiveGeneration {292 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {293 formatter294 .debug_struct("ActiveGeneration")295 .field("generation_id", &self.shared.generation_id)296 .field("audio_enabled", &self.shared.audio_enabled)297 .finish_non_exhaustive()298 }299}300301impl ActiveGeneration {302 async fn start(303 config: &HostConfig,304 host_secret: [u8; 32],305 generation_id: u64,306 recovering_viewers: Arc<AtomicUsize>,307 ) -> Result<Self> {308 let session_clock = SessionClock::start();309 let source = match config.source {310 VideoSource::TestPattern => Ok(spawn_test_pattern(config.clone(), session_clock)),311 VideoSource::GamescopePipewire => {312 let primary =313 spawn_gamescope_pipewire_after_static_preflight(config.clone(), session_clock)314 .await?;315 super::select_gamescope_startup_source(config.clone(), session_clock, primary).await316 }317 }?;318 let EncodedSource {319 frames,320 current_gop,321 task,322 pointer_surface_dimensions,323 encoder_control,324 } = source;325 let video_source = SourceTaskGuard::new(task);326327 let mut generation_secret = [0_u8; 32];328 getrandom::fill(&mut generation_secret)329 .context("generating shared media authentication key")?;330 let signing_key = MediaGenerationSigningKey::from_bytes(&generation_secret);331 generation_secret.fill(0);332 let issued_at_unix = unix_timestamp_now()?;333 let certificate = signing_key.certify(334 *iroh::SecretKey::from_bytes(&host_secret)335 .public()336 .as_bytes(),337 &host_secret,338 generation_id,339 issued_at_unix,340 issued_at_unix.saturating_add(60 * 60),341 )?;342343 let mut broadcast = Broadcast::new().produce();344 let video_track = broadcast345 .create_track(Track {346 name: MOQ_VIDEO_H264_TRACK.to_owned(),347 priority: MOQ_VIDEO_TRACK_PRIORITY,348 })349 .context("creating generation H.264 track")?;350 let audio_enabled = config.audio.is_some();351 let audio_track = if audio_enabled {352 Some(353 broadcast354 .create_track(Track {355 name: MOQ_AUDIO_OPUS_TRACK.to_owned(),356 priority: MOQ_AUDIO_TRACK_PRIORITY,357 })358 .context("creating generation Opus track")?,359 )360 } else {361 None362 };363 let catalog =364 publish_goq_catalog_v2(&mut broadcast, generation_id, &certificate, audio_enabled)?;365 let broadcast_consumer = broadcast.consume();366 let broadcast_name = media_generation_moq_broadcast_name(generation_id)?;367 let telemetry = Arc::new(MediaV3Telemetry::default());368 let (keyframe_requests, keyframes) = watch::channel(None);369 let (shutdown, shutdown_rx) = watch::channel(false);370 let (adaptive, adaptive_task) = GenerationAdaptiveCoordinator::start(371 config.clone(),372 generation_id,373 encoder_control.clone(),374 Arc::clone(&telemetry),375 keyframe_requests.clone(),376 shutdown_rx.clone(),377 recovering_viewers,378 )?;379 let video_publisher = tokio::spawn(run_generation_video_publisher(380 config.clone(),381 current_gop,382 keyframes,383 video_track,384 generation_id,385 signing_key.clone(),386 encoder_control.clone(),387 Arc::clone(&telemetry),388 shutdown_rx.clone(),389 ));390391 let (audio_source, audio_publisher) = if let Some(audio_track) = audio_track {392 let (audio_packets, task) =393 match spawn_pipewire_audio(config.clone(), session_clock).await {394 Ok(source) => source,395 Err(error) => {396 video_publisher.abort();397 video_source.abort_and_wait().await;398 return Err(error).context("starting generation Opus source");399 }400 };401 (402 Some(SourceTaskGuard::new(task)),403 Some(tokio::spawn(run_generation_audio_publisher(404 audio_packets,405 audio_track,406 generation_id,407 signing_key,408 shutdown_rx,409 ))),410 )411 } else {412 (None, None)413 };414415 let shared = Arc::new(GenerationShared {416 generation_id,417 broadcast_name,418 broadcast: Mutex::new(Some(broadcast_consumer)),419 certificate,420 session_clock,421 pointer_surface_dimensions,422 encoder_control,423 telemetry,424 audio_enabled,425 keyframe_requests,426 adaptive,427 });428 info!(429 generation_id,430 video_sources = 1,431 video_encoders = 1,432 audio_sources = usize::from(audio_enabled),433 audio_encoders = usize::from(audio_enabled),434 publishers = 1,435 adaptive_actuators = 1,436 "shared media generation started"437 );438 Ok(Self {439 shared,440 shutdown,441 video_publisher,442 audio_publisher,443 adaptive: adaptive_task,444 video_source,445 audio_source,446 _frame_receiver: frames,447 catalog: Some(catalog),448 broadcast: Some(broadcast),449 })450 }451452 async fn stop(mut self) {453 let generation_id = self.shared.generation_id;454 if let Ok(mut broadcast) = self.shared.broadcast.lock() {455 *broadcast = None;456 }457 self.shutdown.send_replace(true);458 wait_or_abort_task(&mut self.adaptive, "adaptive coordinator").await;459 wait_or_abort_task(&mut self.video_publisher, "video publisher").await;460 if let Some(mut audio_publisher) = self.audio_publisher.take() {461 wait_or_abort_task(&mut audio_publisher, "audio publisher").await;462 }463 if let Some(catalog) = self.catalog.take()464 && let Err(error) = catalog.finish()465 {466 warn!(generation_id, %error, "finishing shared media catalog failed");467 }468 if let Some(mut broadcast) = self.broadcast.take() {469 let _ = broadcast.abort(MoqError::Cancel);470 }471 if let Some(audio_source) = self.audio_source.take() {472 audio_source.abort_and_wait().await;473 }474 self.video_source475 .wait_or_abort(super::SOURCE_REAP_GRACE_TIMEOUT)476 .await;477 info!(478 generation_id,479 "shared media generation stopped with all sources cleaned up"480 );481 }482}483484async fn wait_or_abort_task(task: &mut tokio::task::JoinHandle<Result<()>>, name: &'static str) {485 match tokio::time::timeout(GENERATION_STOP_TIMEOUT, &mut *task).await {486 Ok(Ok(Ok(()))) => {}487 Ok(Ok(Err(error))) => warn!(%error, task = name, "media generation task ended with error"),488 Ok(Err(error)) => warn!(%error, task = name, "media generation task failed"),489 Err(_) => {490 debug!(task = name, "aborting stalled media generation task");491 task.abort();492 let _ = task.await;493 }494 }495}496497#[cfg(test)]498mod tests {499 use super::*;500501 #[test]502 fn lifecycle_is_explicit_and_bounded() {503 let state = GenerationState::default();504 assert_eq!(state.lifecycle, MediaGenerationLifecycle::Idle);505 assert_eq!(state.leases, 0);506 assert!(state.active.is_none());507 }508509 #[tokio::test]510 async fn manager_begins_idle_without_starting_capture() {511 let manager = MediaGenerationManager::new(super::super::moq_test_config(), [7; 32]);512 assert_eq!(manager.lifecycle().await, MediaGenerationLifecycle::Idle);513 assert_eq!(manager.starts.load(Ordering::Relaxed), 0);514 assert_eq!(manager.stops.load(Ordering::Relaxed), 0);515 }516}