moq-transport 0.12.2

//! A broadcast is a collection of tracks, split into two handles: [Writer] and [Reader].
//!
//! The [Writer] can create tracks, either manually or on request.
//! It receives all requests by a [Reader] for a tracks that don't exist.
//! The simplest implementation is to close every unknown track with [ServeError::NotFound].
//!
//! A [Reader] can request tracks by name.
//! If the track already exists, it will be returned.
//! If the track doesn't exist, it will be sent to [Unknown] to be handled.
//! A [Reader] can be cloned to create multiple subscriptions.
//!
//! The broadcast is automatically closed with [ServeError::Done] when [Writer] is dropped, or all [Reader]s are dropped.
use std::{collections::HashMap, ops::Deref, sync::Arc};

use super::{ServeError, Track, TrackReader, TrackWriter};
use crate::coding::TrackNamespace;
use crate::watch::{Queue, State};

/// Full track identifier: namespace + track name
#[derive(Hash, Eq, PartialEq, Clone, Debug)]
pub struct FullTrackName {
    pub namespace: TrackNamespace,
    pub name: String,
}

/// Static information about a broadcast.
#[derive(Debug)]
pub struct Tracks {
    pub namespace: TrackNamespace,
}

impl Tracks {
    pub fn new(namespace: TrackNamespace) -> Self {
        Self { namespace }
    }

    pub fn produce(self) -> (TracksWriter, TracksRequest, TracksReader) {
        let info = Arc::new(self);
        let state = State::default().split();
        let queue = Queue::default().split();

        let writer = TracksWriter::new(state.0.clone(), info.clone());
        let request = TracksRequest::new(state.0, queue.0, info.clone());
        let reader = TracksReader::new(state.1, queue.1, info);

        (writer, request, reader)
    }
}

#[derive(Default)]
pub struct TracksState {
    tracks: HashMap<FullTrackName, TrackReader>,
}

/// Publish new tracks for a broadcast by name.
pub struct TracksWriter {
    state: State<TracksState>,
    pub info: Arc<Tracks>,
}

impl TracksWriter {
    fn new(state: State<TracksState>, info: Arc<Tracks>) -> Self {
        Self { state, info }
    }

    /// Create a new track with the given name, inserting it into the broadcast.
    /// The track will use this writer's namespace.
    /// None is returned if all [TracksReader]s have been dropped.
    pub fn create(&mut self, track: &str) -> Option<TrackWriter> {
        let (writer, reader) = Track {
            namespace: self.namespace.clone(),
            name: track.to_owned(),
        }
        .produce();

        // NOTE: We overwrite the track if it already exists.
        let full_name = FullTrackName {
            namespace: self.namespace.clone(),
            name: track.to_owned(),
        };
        self.state.lock_mut()?.tracks.insert(full_name, reader);

        Some(writer)
    }

    /// Remove a track from the broadcast by full name.
    pub fn remove(&mut self, namespace: &TrackNamespace, track_name: &str) -> Option<TrackReader> {
        let full_name = FullTrackName {
            namespace: namespace.clone(),
            name: track_name.to_owned(),
        };
        self.state.lock_mut()?.tracks.remove(&full_name)
    }
}

impl Deref for TracksWriter {
    type Target = Tracks;

    fn deref(&self) -> &Self::Target {
        &self.info
    }
}

pub struct TracksRequest {
    #[allow(dead_code)] // Avoid dropping the write side
    state: State<TracksState>,
    incoming: Option<Queue<TrackWriter>>,
    pub info: Arc<Tracks>,
}

impl TracksRequest {
    fn new(state: State<TracksState>, incoming: Queue<TrackWriter>, info: Arc<Tracks>) -> Self {
        Self {
            state,
            incoming: Some(incoming),
            info,
        }
    }

    /// Wait for a request to create a new track.
    /// None is returned if all [TracksReader]s have been dropped.
    pub async fn next(&mut self) -> Option<TrackWriter> {
        self.incoming.as_mut()?.pop().await
    }
}

impl Deref for TracksRequest {
    type Target = Tracks;

    fn deref(&self) -> &Self::Target {
        &self.info
    }
}

impl Drop for TracksRequest {
    fn drop(&mut self) {
        // Close any tracks still in the Queue
        for track in self.incoming.take().unwrap().close() {
            let _ = track.close(ServeError::not_found_ctx(
                "tracks request dropped before track handled",
            ));
        }
    }
}

/// Subscribe to a broadcast by requesting tracks.
///
/// This can be cloned to create handles.
#[derive(Clone)]
pub struct TracksReader {
    state: State<TracksState>,
    queue: Queue<TrackWriter>,
    pub info: Arc<Tracks>,
}

impl TracksReader {
    fn new(state: State<TracksState>, queue: Queue<TrackWriter>, info: Arc<Tracks>) -> Self {
        Self { state, queue, info }
    }

    /// Get a track from the broadcast by full name, if it exists and is still alive.
    /// Returns None if the track doesn't exist or has been closed.
    pub fn get_track_reader(
        &mut self,
        namespace: &TrackNamespace,
        track_name: &str,
    ) -> Option<TrackReader> {
        let state = self.state.lock();
        let full_name = FullTrackName {
            namespace: namespace.clone(),
            name: track_name.to_owned(),
        };

        if let Some(track_reader) = state.tracks.get(&full_name) {
            if !track_reader.is_closed() {
                return Some(track_reader.clone());
            }
            // Track exists but is closed/stale - don't return it
        }
        None
    }

    /// Get or request a track from the broadcast by full name.
    /// The namespace parameter should be the full requested namespace, not just the announced prefix.
    /// None is returned if [TracksWriter] or [TracksRequest] cannot fufill the request.
    pub fn subscribe(
        &mut self,
        namespace: TrackNamespace,
        track_name: &str,
    ) -> Option<TrackReader> {
        let state = self.state.lock();
        let full_name = FullTrackName {
            namespace: namespace.clone(),
            name: track_name.to_owned(),
        };

        // Check if we have a cached track that is still alive
        if let Some(track_reader) = state.tracks.get(&full_name) {
            if !track_reader.is_closed() {
                // Track is still active, return the cached reader
                return Some(track_reader.clone());
            }
            // Track is closed/stale, fall through to create a new one
            // We'll remove the stale entry and request a fresh track from the publisher
        }

        let mut state = state.into_mut()?;

        // Remove the stale track if it exists (it was closed)
        state.tracks.remove(&full_name);
        // Use the full requested namespace, not self.namespace
        let track_writer_reader = Track {
            namespace: namespace.clone(),
            name: track_name.to_owned(),
        }
        .produce();

        if self.queue.push(track_writer_reader.0).is_err() {
            return None;
        }

        // We requested the track successfully so we can deduplicate it by full name.
        state
            .tracks
            .insert(full_name, track_writer_reader.1.clone());

        Some(track_writer_reader.1)
    }
}

impl Deref for TracksReader {
    type Target = Tracks;

    fn deref(&self) -> &Self::Target {
        &self.info
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Regression test for the stale track caching bug.
    ///
    /// Scenario:
    /// 1. Subscriber requests a track via subscribe()
    /// 2. Publisher receives TrackWriter, closes it with an error (simulating failure)
    /// 3. Subscriber requests the same track again
    /// 4. Publisher should receive a new TrackWriter (previously didn't due to stale cache)
    ///
    /// This test verifies the fix for an issue seen in production where a track became
    /// "stale" after a connection timeout, and subsequent subscribers never received
    /// data because the publisher was never notified of new subscriptions.
    #[tokio::test]
    async fn test_stale_track_cache_bug() {
        let namespace = TrackNamespace::from_utf8_path("test/namespace");
        let track_name = "test-track";

        // Create the Tracks producer (simulates what the relay does)
        let (_writer, mut request, mut reader) = Tracks::new(namespace.clone()).produce();

        // First subscription: subscriber requests the track
        let track_reader_1 = reader
            .subscribe(namespace.clone(), track_name)
            .expect("first subscribe should succeed");

        // Publisher receives the request and gets a TrackWriter
        let track_writer_1 = request
            .next()
            .await
            .expect("publisher should receive first track request");

        assert_eq!(track_writer_1.name, track_name);

        // Publisher closes the track with an error (simulates connection failure)
        track_writer_1
            .close(ServeError::Cancel)
            .expect("close should succeed");

        // Verify the first track reader is now closed
        // (This is what makes subsequent reads fail immediately)
        let closed_result = tokio::time::timeout(
            std::time::Duration::from_millis(100),
            track_reader_1.closed(),
        )
        .await;
        assert!(
            closed_result.is_ok(),
            "track_reader_1 should be closed after writer closes"
        );

        // Second subscription: subscriber requests the SAME track again
        let track_reader_2 = reader
            .subscribe(namespace.clone(), track_name)
            .expect("second subscribe should succeed");

        // With the fix, the stale cached TrackReader is detected and evicted,
        // so the publisher receives a new TrackWriter for the second subscription.
        let maybe_track_writer_2 =
            tokio::time::timeout(std::time::Duration::from_millis(100), request.next()).await;

        // Publisher should receive a new TrackWriter (stale cache entry was evicted)
        assert!(
            maybe_track_writer_2.is_ok(),
            "Publisher should receive a new track request after the first one was closed"
        );

        let track_writer_2 = maybe_track_writer_2
            .unwrap()
            .expect("publisher should receive second track request");

        assert_eq!(track_writer_2.name, track_name);

        // Verify that track_reader_2 is NOT already closed
        // (It should be a fresh, working track)
        let closed_result_2 = tokio::time::timeout(
            std::time::Duration::from_millis(100),
            track_reader_2.closed(),
        )
        .await;
        assert!(
            closed_result_2.is_err(),
            "track_reader_2 should NOT be immediately closed - it should be a fresh track"
        );
    }

    /// Test that normal track caching works correctly when tracks are still alive.
    ///
    /// Multiple subscribers to the same track should share the same TrackReader
    /// (deduplication), and the publisher should only receive one request.
    #[tokio::test]
    async fn test_track_deduplication_while_alive() {
        let namespace = TrackNamespace::from_utf8_path("test/namespace");
        let track_name = "test-track";

        let (_writer, mut request, mut reader) = Tracks::new(namespace.clone()).produce();

        // First subscription
        let track_reader_1 = reader
            .subscribe(namespace.clone(), track_name)
            .expect("first subscribe should succeed");

        // Publisher receives request
        let _track_writer = request
            .next()
            .await
            .expect("publisher should receive track request");

        // Second subscription to the SAME track (while it's still alive)
        let track_reader_2 = reader
            .subscribe(namespace.clone(), track_name)
            .expect("second subscribe should succeed");

        // Publisher should NOT receive another request (track is cached and alive)
        let maybe_second_request =
            tokio::time::timeout(std::time::Duration::from_millis(100), request.next()).await;

        assert!(
            maybe_second_request.is_err(),
            "Publisher should NOT receive a second request - track is cached and alive"
        );

        // Both readers should refer to the same track
        assert_eq!(track_reader_1.name, track_reader_2.name);
        assert_eq!(track_reader_1.namespace, track_reader_2.namespace);
    }
}