char/char-rtc-obs

git clone https://git.t4t.associates/char/char-rtc-obs

videogame hackerInitial commiteeaf4aa

main
branchesmainall refs →
16.0 KiB443 linesraw
1use crate::whip;
2use anyhow::{anyhow, Result};
3use bytes::Bytes;
4use log::{debug, error, info, trace};
5use reqwest::Url;
6use std::{
7    boxed::Box,
8    sync::{Arc, Mutex, RwLock},
9    thread::JoinHandle,
10    time::{Duration, Instant},
11};
12use tokio::{
13    select,
14    sync::mpsc::{self, UnboundedReceiver, UnboundedSender},
15    time::interval,
16};
17use webrtc::{
18    api::{
19        interceptor_registry::register_default_interceptors,
20        media_engine::{MediaEngine, MIME_TYPE_H264, MIME_TYPE_OPUS},
21        setting_engine::SettingEngine,
22        APIBuilder,
23    },
24    ice_transport::ice_connection_state::RTCIceConnectionState,
25    interceptor::registry::Registry,
26    media::Sample,
27    peer_connection::{
28        configuration::RTCConfiguration, peer_connection_state::RTCPeerConnectionState,
29        RTCPeerConnection,
30    },
31    rtp_transceiver::{
32        rtp_codec::{RTCRtpCodecCapability, RTCRtpCodecParameters, RTPCodecType},
33        rtp_transceiver_direction::RTCRtpTransceiverDirection,
34        RTCPFeedback, RTCRtpTransceiverInit,
35    },
36    stats::StatsReportType,
37    track::track_local::track_local_static_sample::TrackLocalStaticSample,
38};
39
40pub type ErrorCallback = Box<dyn FnMut(OutputStreamError) + Send + Sync>;
41#[derive(Debug)]
42pub struct EncodedPacket {
43    pub data: Vec<u8>,
44    pub duration: Duration,
45    pub typ: EncodedPacketType,
46}
47
48#[derive(Debug)]
49pub enum EncodedPacketType {
50    Audio,
51    Video,
52}
53
54/// Errors for [`OutputStream`] worker thread
55#[derive(Debug)]
56pub enum OutputStreamError {
57    ConnectFailed,
58    NetworkError,
59    WriteError(webrtc::Error),
60}
61
62#[derive(Debug)]
63pub enum Message {
64    Packet(EncodedPacket),
65    Error(OutputStreamError),
66    Close,
67}
68
69pub enum WorkerResult {
70    Error(OutputStreamError),
71    Close,
72}
73
74#[derive(Default)]
75struct OutputStreamStats {
76    bytes_sent: u64,
77    congestion: f64,
78    dropped_frames: i32,
79    connect_time: Option<Instant>,
80    connect_duration: Duration,
81}
82
83pub struct OutputStream {
84    video_track: Arc<TrackLocalStaticSample>,
85    audio_track: Arc<TrackLocalStaticSample>,
86    peer_connection: Arc<RTCPeerConnection>,
87    stats: Arc<RwLock<OutputStreamStats>>,
88    worker_tx: UnboundedSender<Message>,
89    worker_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
90    whip_resource: Arc<Mutex<Option<Url>>>,
91    error_callback: Arc<Mutex<Option<ErrorCallback>>>,
92}
93
94impl OutputStream {
95    fn start_worker(&mut self, mut worker_rx: UnboundedReceiver<Message>) {
96        let worker_handle = std::thread::spawn({
97            let audio_track = self.audio_track.clone();
98            let video_track = self.video_track.clone();
99            let peer_connection = self.peer_connection.clone();
100            let stats = self.stats.clone();
101            let error_callback = self.error_callback.clone();
102            move || {
103                let runtime = tokio::runtime::Builder::new_multi_thread()
104                    .enable_all()
105                    .thread_name("webrtc_worker")
106                    .worker_threads(4)
107                    .build()
108                    .unwrap();
109                let runtime = runtime;
110                let result = runtime.block_on(
111                    async move {
112                        let mut interval = interval(Duration::from_millis(500));
113                        'worker_loop: loop {
114                            let result = select! {
115                                message = worker_rx.recv() => {
116                                    match message {
117                                        Some(Message::Packet(packet)) => {
118                                            let sample = Sample {
119                                                data: Bytes::from(packet.data),
120                                                duration: packet.duration,
121                                                ..Default::default()
122                                            };
123                                            match (packet.typ, peer_connection.connection_state()) {
124                                                (_, RTCPeerConnectionState::Failed | RTCPeerConnectionState::Closed) => Err(WorkerResult::Error(OutputStreamError::NetworkError)),
125                                                (EncodedPacketType::Audio, _) => audio_track.write_sample(&sample).await.map_err(|e| WorkerResult::Error(OutputStreamError::WriteError(e))),
126                                                (EncodedPacketType::Video, _) => video_track.write_sample(&sample).await.map_err(|e| WorkerResult::Error(OutputStreamError::WriteError(e)))
127                                            }
128                                        },
129                                        Some(Message::Error(e)) => Err(WorkerResult::Error(e)),
130                                        Some(Message::Close) | None => Err(WorkerResult::Close),
131                                    }
132                                }
133                                _ = interval.tick() => {
134                                    let pc_stats = peer_connection.get_stats().await;
135
136                                    let stats = &mut stats.write().unwrap();
137                                    if let Some(StatsReportType::Transport(transport_stats)) = pc_stats.reports.get("ice_transport") {
138                                        stats.bytes_sent = transport_stats.bytes_sent as u64;
139                                    }
140                                    Ok(())
141                                }
142                            };
143
144                           if let Err(e) = result {
145                                break 'worker_loop e;
146                           }
147                        }
148                    }
149                );
150
151                match result {
152                    WorkerResult::Close => {}
153                    WorkerResult::Error(e) => {
154                        if let Some(callback) = &mut *error_callback.lock().unwrap() {
155                            error!("Worker encountered error: {e:?}");
156                            callback(e);
157                        }
158                    }
159                }
160
161                info!("Exiting worker thread");
162            }
163        });
164
165        *self.worker_handle.lock().unwrap() = Some(worker_handle);
166    }
167
168    pub async fn new() -> Result<Self> {
169        let video_track = Arc::new(TrackLocalStaticSample::new(
170            RTCRtpCodecCapability {
171                mime_type: MIME_TYPE_H264.to_owned(),
172                clock_rate: 90000,
173                sdp_fmtp_line: "profile-level-id=42e01f; max-fs=3600; max-mbps=108000; max-br=1400"
174                    .to_string(),
175                ..Default::default()
176            },
177            "video".to_owned(),
178            "webrtc-rs".to_owned(),
179        ));
180
181        let audio_track = Arc::new(TrackLocalStaticSample::new(
182            RTCRtpCodecCapability {
183                mime_type: MIME_TYPE_OPUS.to_owned(),
184                ..Default::default()
185            },
186            "audio".to_owned(),
187            "webrtc-rs".to_owned(),
188        ));
189
190        let mut m = MediaEngine::default();
191        m.register_codec(
192            RTCRtpCodecParameters {
193                capability: RTCRtpCodecCapability {
194                    mime_type: MIME_TYPE_OPUS.to_owned(),
195                    clock_rate: 48000,
196                    channels: 2,
197                    sdp_fmtp_line: "minptime=10;useinbandfec=1".to_owned(),
198                    rtcp_feedback: vec![],
199                },
200                payload_type: 111,
201                ..Default::default()
202            },
203            RTPCodecType::Audio,
204        )?;
205
206        m.register_codec(
207            RTCRtpCodecParameters {
208                capability: RTCRtpCodecCapability {
209                    mime_type: MIME_TYPE_H264.to_owned(),
210                    clock_rate: 90000,
211                    channels: 0,
212                    sdp_fmtp_line:
213                        "level-asymmetry-allowed=1;packetization-mode=1;profile-level-id=42e01f"
214                            .to_owned(),
215                    rtcp_feedback: vec![
216                        RTCPFeedback {
217                            typ: "goog-remb".to_owned(),
218                            parameter: "".to_owned(),
219                        },
220                        RTCPFeedback {
221                            typ: "ccm".to_owned(),
222                            parameter: "fir".to_owned(),
223                        },
224                        RTCPFeedback {
225                            typ: "nack".to_owned(),
226                            parameter: "".to_owned(),
227                        },
228                        RTCPFeedback {
229                            typ: "nack".to_owned(),
230                            parameter: "pli".to_owned(),
231                        },
232                    ],
233                },
234                payload_type: 125,
235                ..Default::default()
236            },
237            RTPCodecType::Video,
238        )?;
239
240        let mut registry = Registry::new();
241        registry = register_default_interceptors(registry, &mut m)?;
242
243        let gather_ips: Vec<_> = if_addrs::get_if_addrs()
244            .unwrap_or_else(|_| Vec::new())
245            .into_iter()
246            .filter(|addr| {
247                if !addr.is_loopback() {
248                    trace!(
249                        "Valid gather address for interface {}: {:#?}",
250                        addr.name,
251                        addr.ip()
252                    );
253                    true
254                } else {
255                    trace!(
256                        "Loopback address ignored for interface {}: {:#?}",
257                        addr.name,
258                        addr.ip()
259                    );
260                    false
261                }
262            })
263            .map(|addr| addr.ip())
264            .collect();
265
266        let mut setting_engine = SettingEngine::default();
267
268        // Only attempt to limit if we found valid interfaces, otherwise do not attempt
269        // to limit as we may be on a platform that doesn't expose enough information
270        if !gather_ips.is_empty() {
271            debug!("Gather addresses: {gather_ips:#?}");
272            setting_engine.set_ip_filter(Box::new(move |ip: std::net::IpAddr| -> bool {
273                gather_ips.contains(&ip)
274            }));
275        }
276
277        let api = APIBuilder::new()
278            .with_media_engine(m)
279            .with_interceptor_registry(registry)
280            .with_setting_engine(setting_engine)
281            .build();
282
283        // Prepare the configuration
284        let config = RTCConfiguration {
285            ..Default::default()
286        };
287
288        let peer_connection = Arc::new(api.new_peer_connection(config).await?);
289        let stats: Arc<RwLock<OutputStreamStats>> = Arc::new(RwLock::new(OutputStreamStats {
290            connect_time: None,
291            ..Default::default()
292        }));
293
294        let (worker_tx, worker_rx): (UnboundedSender<Message>, UnboundedReceiver<Message>) =
295            mpsc::unbounded_channel();
296        let error_callback: Arc<Mutex<Option<ErrorCallback>>> = Arc::new(Mutex::new(None));
297
298        let mut output_stream = Self {
299            audio_track,
300            video_track,
301            peer_connection,
302            stats,
303            worker_tx,
304            worker_handle: Arc::new(Mutex::new(None)),
305            whip_resource: Arc::new(Mutex::new(None)),
306            error_callback,
307        };
308
309        output_stream.start_worker(worker_rx);
310
311        Ok(output_stream)
312    }
313
314    pub async fn connect(&self, url: &str, bearer_token: Option<&str>) {
315        self.connect_internal(url, bearer_token)
316            .await
317            .unwrap_or_else(|e| {
318                error!("Failed connecting to: {e}");
319                let _ = self
320                    .worker_tx
321                    .send(Message::Error(OutputStreamError::ConnectFailed));
322            })
323    }
324
325    async fn connect_internal(&self, url: &str, bearer_token: Option<&str>) -> Result<()> {
326        println!("Setting up webrtc!");
327
328        self.peer_connection
329            .add_transceiver_from_track(
330                self.video_track.clone(),
331                &[RTCRtpTransceiverInit {
332                    direction: RTCRtpTransceiverDirection::Sendonly,
333                    send_encodings: Vec::new(),
334                }],
335            )
336            .await?;
337
338        self.peer_connection
339            .add_transceiver_from_track(
340                self.audio_track.clone(),
341                &[RTCRtpTransceiverInit {
342                    direction: RTCRtpTransceiverDirection::Sendonly,
343                    send_encodings: Vec::new(),
344                }],
345            )
346            .await?;
347
348        self.peer_connection
349            .on_ice_connection_state_change(Box::new(
350                move |connection_state: RTCIceConnectionState| {
351                    info!("Connection State has changed {}", connection_state);
352                    if connection_state == RTCIceConnectionState::Connected {
353                        // on_success();
354                    }
355                    Box::pin(async {})
356                },
357            ));
358
359        self.peer_connection
360            .on_peer_connection_state_change(Box::new(move |s: RTCPeerConnectionState| {
361                debug!("Peer Connection State has changed: {}", s);
362
363                if s == RTCPeerConnectionState::Failed {
364                    error!("Peer connection state went to failed")
365                }
366
367                Box::pin(async {})
368            }));
369
370        let offer = self.peer_connection.create_offer(None).await?;
371        let mut gather_complete = self.peer_connection.gathering_complete_promise().await;
372        self.peer_connection.set_local_description(offer).await?;
373
374        // Block until gathering complete
375        let _ = gather_complete.recv().await;
376
377        let offer = self
378            .peer_connection
379            .local_description()
380            .await
381            .ok_or_else(|| anyhow!("No local description available"))?;
382        let (answer, whip_resource) = whip::offer(url, bearer_token, offer).await?;
383        self.peer_connection.set_remote_description(answer).await?;
384
385        *self.whip_resource.lock().unwrap() = Some(whip_resource);
386
387        Ok(())
388    }
389
390    pub async fn close(&self) -> Result<()> {
391        let close_result = self.worker_tx.send(Message::Close);
392
393        // Take worker handle so it's dropped
394        let worker_future = self.worker_handle.lock().unwrap().take();
395        // If close was a success (worker could receive messages), join on thread
396        // Otherwise, worker thread is already dead or currently closing due to error callback
397        if close_result.is_ok() {
398            if let Some(worker_future) = worker_future {
399                worker_future
400                    .join()
401                    .unwrap_or_else(|e| error!("Failed joining worker thread: {e:?}"));
402            }
403        }
404
405        let whip_resource = self.whip_resource.lock().unwrap().take();
406        if let Some(whip_resource) = whip_resource {
407            whip::delete(&whip_resource).await?;
408        }
409        Ok(self.peer_connection.close().await?)
410    }
411
412    pub fn bytes_sent(&self) -> u64 {
413        self.stats.read().unwrap().bytes_sent
414    }
415
416    pub fn congestion(&self) -> f64 {
417        self.stats.read().unwrap().congestion
418    }
419
420    pub fn connect_time(&self) -> Duration {
421        let mut stats = self.stats.write().unwrap();
422
423        if let Some(connect_time) = stats.connect_time {
424            stats.connect_duration = Instant::now() - connect_time;
425        }
426
427        stats.connect_duration
428    }
429
430    pub fn dropped_frames(&self) -> i32 {
431        self.stats.read().unwrap().dropped_frames
432    }
433
434    pub fn write(&self, packet: EncodedPacket) -> Result<()> {
435        self.worker_tx
436            .send(Message::Packet(packet))
437            .map_err(|e| e.into())
438    }
439
440    pub fn set_error_callback(&self, callback: ErrorCallback) {
441        *self.error_callback.lock().unwrap() = Some(callback);
442    }
443}