coapwithudp/src/bin/stream.rs

//! Stream a media file from the CoAP server and play it in an ffplay window.
//!
//!   cargo run --release --bin stream -- image
//!   cargo run --release --bin stream -- audio
//!   cargo run --release --bin stream -- audio --no-play --out copy.wav
//!
//! The file arrives as 1024-byte Block2 blocks over UDP. Each block is written
//! to ffplay's stdin as soon as it lands; when the player's pipe is full the
//! write waits, so blocks are requested at playback speed instead of flooded.

use std::io::{self, IsTerminal, Write};
use std::process::Stdio;
use std::time::Duration;

use coap_hello_udp::{BlockFetcher, Error, FetchOptions, TransferStats};
use tokio::io::AsyncWriteExt;
use tokio::process::{Child, Command};

const USAGE: &str = "usage: stream audio|image [--out FILE] [--no-play] [--timeout-ms N]";
const IMAGE_BYTES: u64 = 320 * 240 * 2;

#[derive(Clone, Copy, PartialEq)]
enum Media {
    Audio,
    Image,
}

impl Media {
    fn path(self) -> &'static str {
        match self {
            Media::Audio => "/audio",
            Media::Image => "/image",
        }
    }

    /// ffplay arguments for reading this media from stdin.
    fn player_args(self) -> &'static [&'static str] {
        match self {
            Media::Audio => &["-autoexit", "-showmode", "waves", "-window_title", "CoAP audio", "-f", "wav"],
            Media::Image => &[
                "-f", "rawvideo", "-pixel_format", "rgb565le", "-video_size", "320x240",
                "-window_title", "CoAP image",
            ],
        }
    }
}

struct Args {
    media: Media,
    out: Option<String>,
    play: bool,
    timeout: Duration,
}

fn parse_args() -> Result<Args, String> {
    let mut args = std::env::args().skip(1);
    let media = match args.next().as_deref() {
        Some("audio") => Media::Audio,
        Some("image") => Media::Image,
        _ => return Err(USAGE.into()),
    };

    let mut parsed = Args { media, out: None, play: true, timeout: FetchOptions::default().timeout };
    while let Some(arg) = args.next() {
        match arg.as_str() {
            "--out" => parsed.out = Some(args.next().ok_or(USAGE)?),
            "--no-play" => parsed.play = false,
            "--timeout-ms" => {
                let ms = args.next().and_then(|v| v.parse().ok()).ok_or(USAGE)?;
                parsed.timeout = Duration::from_millis(ms);
            }
            _ => return Err(USAGE.into()),
        }
    }
    Ok(parsed)
}

#[tokio::main]
async fn main() -> Result<(), Error> {
    let args = match parse_args() {
        Ok(args) => args,
        Err(usage) => {
            eprintln!("{usage}");
            std::process::exit(2);
        }
    };
    let host = std::env::var("COAP_HOST").unwrap_or_else(|_| "127.0.0.1".into());
    let port = std::env::var("COAP_PORT").unwrap_or_else(|_| "5683".into());
    let path = args.media.path();

    let mut player = if args.play { Some(spawn_player(args.media)?) } else { None };
    let mut player_stdin = player.as_mut().and_then(|p| p.stdin.take());
    let mut out = match &args.out {
        Some(file) => Some(io::BufWriter::new(std::fs::File::create(file)?)),
        None => None,
    };

    let options = FetchOptions { timeout: args.timeout, ..FetchOptions::default() };
    let mut fetcher = BlockFetcher::connect(&host, &port, path, options).await?;
    eprintln!("streaming coap://{host}:{port}{path}");

    let show_progress = io::stderr().is_terminal();
    let mut stopped_by_player = false;

    while let Some(block) = fetcher.next_block().await? {
        if let Some(file) = out.as_mut() {
            file.write_all(&block)?;
        }
        if let Some(stdin) = player_stdin.as_mut()
            && stdin.write_all(&block).await.is_err()
        {
            // The player window was closed.
            stopped_by_player = true;
            break;
        }
        if show_progress {
            let stats = fetcher.stats();
            let total = stats.total.map(|t| format!(" / {t}")).unwrap_or_default();
            eprint!("\r  {}{total} B, {} blocks, {} retransmits", stats.bytes, stats.blocks, stats.retransmits);
        }
    }

    if let Some(mut file) = out {
        file.flush()?;
    }
    let stats = fetcher.stats();
    if show_progress {
        eprintln!();
    }
    print_summary(&stats, stopped_by_player);

    if args.media == Media::Image
        && let Some(total) = stats.total
        && total != IMAGE_BYTES
    {
        eprintln!("warning: server sent {total} B, a 320x240 RGB565 image is {IMAGE_BYTES} B");
    }

    // Closing stdin tells ffplay the stream has ended.
    drop(player_stdin);
    if let Some(player) = player.as_mut() {
        if !stopped_by_player && args.media == Media::Image {
            eprintln!("close the image window to exit");
        }
        player.wait().await?;
    }
    Ok(())
}

fn spawn_player(media: Media) -> Result<Child, String> {
    Command::new("ffplay")
        .args(["-hide_banner", "-loglevel", "error"])
        .args(media.player_args())
        .args(["-i", "pipe:0"])
        .stdin(Stdio::piped())
        .kill_on_drop(true)
        .spawn()
        .map_err(|e| match e.kind() {
            io::ErrorKind::NotFound => {
                "ffplay not found: install ffmpeg, or use --no-play --out FILE".to_string()
            }
            _ => format!("cannot start ffplay: {e}"),
        })
}

fn print_summary(stats: &TransferStats, stopped_by_player: bool) {
    let secs = stats.elapsed.as_secs_f64();
    let rate = if secs > 0.0 { stats.bytes as f64 / 1024.0 / secs } else { 0.0 };
    let etag = match &stats.etag {
        Some(tag) => tag.iter().map(|b| format!("{b:02x}")).collect(),
        None => "none".to_string(),
    };
    let what = if stopped_by_player { "stopped by player after" } else { "received" };
    eprintln!(
        "{what} {} B in {} blocks, {secs:.2} s ({rate:.0} KiB/s), {} retransmits, ETag {etag}",
        stats.bytes, stats.blocks, stats.retransmits
    );
}