arcly_stream/protocol/
udp.rs1use std::net::SocketAddr;
25
26use async_trait::async_trait;
27use tokio_util::sync::CancellationToken;
28use tracing::{info, warn};
29
30use crate::inbound::{InboundProtocol, IngestContext, PublishSession};
31use crate::protocol::tsdemux::{TsDemuxer, TsTrackKind};
32use crate::{CodecId, MediaFrame, Result};
33
34pub struct UdpTsHandler {
37 bind: SocketAddr,
38 key: crate::StreamKey,
39 recv_buf: usize,
42}
43
44impl UdpTsHandler {
45 pub fn new(bind: SocketAddr, key: crate::StreamKey) -> Self {
47 Self {
48 bind,
49 key,
50 recv_buf: 2048,
51 }
52 }
53
54 pub fn recv_buffer(mut self, bytes: usize) -> Self {
56 self.recv_buf = bytes.max(188);
57 self
58 }
59}
60
61#[async_trait]
62impl InboundProtocol for UdpTsHandler {
63 fn name(&self) -> &'static str {
64 "udp"
65 }
66
67 async fn serve(&self, ctx: IngestContext, shutdown: CancellationToken) -> Result<()> {
68 use tokio::net::UdpSocket;
69
70 let socket = UdpSocket::bind(self.bind).await?;
71 info!(bind = %self.bind, "udp ts listener bound");
72 let mut buf = vec![0u8; self.recv_buf];
73 let mut demux = TsDemuxer::new();
74 let mut session: Option<PublishSession> = None;
75
76 loop {
77 let n = tokio::select! {
78 _ = shutdown.cancelled() => break,
79 r = socket.recv_from(&mut buf) => match r {
80 Ok((n, _from)) => n,
81 Err(e) => {
82 warn!(error = %e, "udp recv failed");
83 continue;
84 }
85 }
86 };
87
88 for au in demux.push(&buf[..n]) {
89 if au.codec == CodecId::Unknown {
90 continue;
91 }
92 if session.is_none() {
94 session = Some(ctx.open_publish(self.key.clone()).await?);
95 }
96 let sess = session.as_ref().unwrap();
97 let pts = au.pts_ms;
98 let mut frame = match au.kind {
99 TsTrackKind::Video => {
100 MediaFrame::new_video(pts, pts, au.data, au.codec, au.keyframe)
101 }
102 TsTrackKind::Audio => MediaFrame::new_audio(pts, au.data, au.codec),
103 };
104 if au.is_config {
105 frame.flags |= crate::FrameFlags::CONFIG;
106 }
107 let _ = sess.publish_frame(frame)?;
108 }
109 }
110
111 if let Some(sess) = session {
112 sess.finish().await?;
113 }
114 Ok(())
115 }
116}
117
118#[cfg(test)]
119mod tests {
120 use super::*;
121 use crate::Engine;
122 use std::sync::Arc;
123
124 #[tokio::test]
125 async fn binds_and_shuts_down_cleanly() {
126 let engine: Arc<Engine> = Engine::builder()
127 .application(crate::AppSpec::new("live"))
128 .build();
129 let ctx = IngestContext::new(engine);
130 let handler = UdpTsHandler::new(
131 "127.0.0.1:0".parse().unwrap(),
132 crate::StreamKey::new("live", "feed"),
133 );
134 let shutdown = CancellationToken::new();
135
136 let token = shutdown.clone();
139 let task = tokio::spawn(async move { handler.serve(ctx, token).await });
140 tokio::task::yield_now().await;
141 shutdown.cancel();
142 let res = tokio::time::timeout(std::time::Duration::from_secs(5), task)
143 .await
144 .expect("serve returned after cancel")
145 .expect("task joined");
146 assert!(res.is_ok(), "clean shutdown");
147 }
148}