Skip to main content

strand_http_video_streaming/
lib.rs

1// Copyright (C) The Strand-Braid Authors
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4use std::{
5    collections::HashMap,
6    sync::{Arc, Mutex},
7};
8use strand_http_video_streaming_types::StrokeStyle;
9
10use event_stream_types::{ConnectionEvent, ConnectionEventType, EventChunkSender};
11use strand_bui_backend_session_types::ConnectionKey;
12use strand_dynamic_frame::DynamicFrameOwned;
13
14pub use strand_http_video_streaming_types::{CircleParams, DrawableShape, Point, Shape, ToClient};
15
16type Result<T> = std::result::Result<T, Error>;
17
18#[derive(Debug, thiserror::Error)]
19pub enum Error {
20    #[error("unknown path")]
21    UnknownPath(),
22    #[error(transparent)]
23    ConvertImageError(#[from] convert_image::Error),
24}
25
26// future: use MediaSource API? https://w3c.github.io/media-source
27
28#[derive(Debug)]
29pub struct AnnotatedFrame {
30    pub frame: Arc<DynamicFrameOwned>,
31    pub found_points: Vec<Point>,
32    pub valid_display: Option<Shape>,
33    pub annotations: Vec<DrawableShape>,
34}
35
36fn _test_annotated_frame_is_send() {
37    // Compile-time test to ensure AnnotatedFrame implements Send trait.
38    fn implements<T: Send>() {}
39    implements::<AnnotatedFrame>();
40}
41
42struct PerSender {
43    /// The Sender to the specific HTTP connection.
44    conn_tx: EventChunkSender,
45    frame_lifo: Option<Arc<Mutex<AnnotatedFrame>>>,
46    ready_to_send: bool,
47    conn_key: ConnectionKey,
48    fno: u64,
49    green_stroke: StrokeStyle,
50}
51
52fn _test_per_sender_is_send() {
53    // Compile-time test to ensure PerSender implements Send trait.
54    fn implements<T: Send>() {}
55    implements::<PerSender>();
56}
57
58#[derive(Debug)]
59pub enum NameSelector {
60    All,
61    None,
62    Name(String),
63}
64
65impl PerSender {
66    fn new(
67        conn_tx: EventChunkSender,
68        conn_key: ConnectionKey,
69        frame: Arc<Mutex<AnnotatedFrame>>,
70    ) -> PerSender {
71        PerSender {
72            conn_tx,
73            frame_lifo: Some(frame),
74            ready_to_send: true,
75            conn_key,
76            fno: 0,
77            green_stroke: StrokeStyle::from_rgb(0x7F, 0xFF, 0x7F),
78        }
79    }
80    fn push(&mut self, frame: Arc<Mutex<AnnotatedFrame>>) {
81        self.fno += 1;
82        self.frame_lifo = Some(frame);
83    }
84    fn got_callback(&mut self, _msg: ConnectionKey) {
85        self.ready_to_send = true;
86    }
87    async fn service(&mut self) -> Result<()> {
88        // check if we should send frame(s) and send if so.
89
90        // should we send it?
91        // TODO cache the converted frame.
92        // TODO allow client to throttle?
93        // TODO make algorithm smarter to have more in-flight frames?
94        // TODO include sent time in message to clients so we don't maintain that
95
96        if let Some(ref most_recent_frame_data) = self.frame_lifo
97            && self.ready_to_send
98        {
99            // sent_time computed early so that latency includes duration to encode, etc.
100            let sent_time = chrono::Local::now();
101            let tc = {
102                let most_recent_frame_data = most_recent_frame_data.lock().unwrap();
103                let bytes = &most_recent_frame_data
104                    .frame
105                    .borrow()
106                    .to_encoded_buffer(convert_image::EncoderOptions::Jpeg(80))?;
107                use base64::Engine;
108                let firehose_frame_base64 = base64::engine::general_purpose::STANDARD.encode(bytes);
109                let data_url = format!("data:image/jpeg;base64,{}", firehose_frame_base64);
110                // most_recent_frame_data.data_url = Some(data_url.clone()); // todo: cache like this
111                let mut annotations = most_recent_frame_data.annotations.clone();
112                // Convert found points into normal annotations. (This should perhaps be done earlier.)
113                for found_point in most_recent_frame_data.found_points.iter() {
114                    let line_width = 5.0;
115                    let shape = Shape::Circle(CircleParams {
116                        center_x: found_point.x.round() as i16,
117                        center_y: found_point.y.round() as i16,
118                        radius: 10,
119                    });
120                    let green_shape = strand_http_video_streaming_types::DrawableShape::from_shape(
121                        &shape,
122                        &self.green_stroke,
123                        line_width,
124                    );
125                    annotations.push(green_shape);
126                }
127                ToClient {
128                    firehose_frame_data_url: data_url,
129                    valid_display: most_recent_frame_data.valid_display.clone(),
130                    annotations,
131                    fno: self.fno,
132                    ts_rfc3339: sent_time.to_rfc3339(),
133                    ck: self.conn_key,
134                }
135            };
136            let buf = serde_json::to_string(&tc).expect("encode");
137            let buf = format!(
138                "event: {}\ndata: {}\n\n",
139                strand_http_video_streaming_types::VIDEO_STREAM_EVENT_NAME,
140                buf
141            );
142            let chunk = http_body::Frame::data(bytes::Bytes::from(buf));
143
144            match self.conn_tx.send(chunk).await {
145                Ok(()) => {}
146                Err(_) => {
147                    tracing::info!("failed to send data to connection. dropping.");
148                    // Failed to send data to event stream key.
149                    // TODO: drop this sender.
150                }
151            }
152            self.ready_to_send = false;
153        }
154
155        self.frame_lifo = None;
156
157        Ok(())
158    }
159}
160
161struct TaskState {
162    /// cache of senders
163    per_sender_map: HashMap<ConnectionKey, PerSender>,
164    /// most recent image frame, with annotations
165    frame: Arc<Mutex<AnnotatedFrame>>,
166}
167
168fn _test_task_state_is_send() {
169    // Compile-time test to ensure PerSender implements Send trait.
170    fn implements<T: Send>() {}
171    implements::<TaskState>();
172}
173
174impl TaskState {
175    async fn service(&mut self) -> Result<()> {
176        // TODO: make sending concurrent on all listeners and set a timeout.
177        for ps in self.per_sender_map.values_mut() {
178            ps.service().await?;
179        }
180        Ok(())
181    }
182    /// Event handler when a new connection or disconnection is made.
183    fn handle_connection(&mut self, conn_evt: ConnectionEvent) -> Result<()> {
184        match conn_evt.typ {
185            ConnectionEventType::Connect(conn_tx) => {
186                // sender was added.
187                let ps = PerSender::new(conn_tx, conn_evt.connection_key, self.frame.clone());
188                self.per_sender_map.insert(conn_evt.connection_key, ps);
189            }
190            ConnectionEventType::Disconnect => {
191                self.per_sender_map.remove(&conn_evt.connection_key);
192            }
193        }
194        Ok(())
195    }
196    fn handle_frame(&mut self, new_frame: AnnotatedFrame) -> Result<()> {
197        // Move the frame into a reference-counted pointer.
198        self.frame = Arc::new(Mutex::new(new_frame));
199        for ps in self.per_sender_map.values_mut() {
200            // Clone the pointer and move the pointer into each sender.
201            ps.push(self.frame.clone());
202        }
203        Ok(())
204    }
205    fn handle_callback(&mut self, ck: ConnectionKey) -> Result<()> {
206        if let Some(ps) = self.per_sender_map.get_mut(&ck) {
207            ps.got_callback(ck)
208        } else {
209            tracing::debug!(
210                "Got firehose_callback for non-existant connection key. \
211                            Did connection disconnect?"
212            );
213        }
214        Ok(())
215    }
216}
217
218/// Option unwrap macro for use in loop.
219///
220/// Logs message and breaks out of loop if None.
221macro_rules! otry {
222    ($e:expr_2021, $msg:expr_2021) => {
223        match $e {
224            Some(v) => v,
225            None => {
226                tracing::debug!($msg);
227                break;
228            }
229        }
230    };
231}
232
233pub async fn firehose_task(
234    mut connection_callback_rx: tokio::sync::mpsc::Receiver<ConnectionEvent>,
235    mut firehose_rx: tokio::sync::mpsc::Receiver<AnnotatedFrame>,
236    mut firehose_callback_rx: tokio::sync::mpsc::Receiver<ConnectionKey>,
237) -> Result<()> {
238    // Wait for the first frame so we don't need to deal with an Option<>.
239    let first_frame = firehose_rx.recv().await.unwrap();
240    let frame = Arc::new(Mutex::new(first_frame));
241
242    let mut task_state = TaskState {
243        per_sender_map: HashMap::new(),
244        frame,
245    };
246
247    let mut interval = tokio::time::interval(std::time::Duration::from_millis(100));
248    loop {
249        tokio::select! {
250            opt_new_connection = connection_callback_rx.recv() => {
251                task_state.handle_connection(otry!(opt_new_connection, "new connection senders done."))?;
252            }
253            opt_new_frame = firehose_rx.recv() => {
254                task_state.handle_frame(otry!(opt_new_frame, "new frame senders done."))?;
255            },
256            opt_callback = firehose_callback_rx.recv() => {
257                task_state.handle_callback(otry!(opt_callback, "new callback senders done."))?;
258            },
259            _ = interval.tick() => {
260                task_state.service().await?;
261            }
262        }
263    }
264    tracing::debug!("firehose task done.");
265    Ok(())
266}