livekit_data_stream/outgoing/
stream_writer.rs1use std::{collections::HashMap, sync::Arc};
16use tokio::sync::Mutex;
17
18use crate::{
19 info::{ByteStreamInfo, TextStreamInfo},
20 outgoing::{constants::STREAM_CHUNK_SIZE_BYTES, raw_stream::RawStream},
21 utf8_chunk::Utf8AwareChunkExt,
22 utils::StreamResult,
23};
24
25pub trait StreamWriter<'a> {
27 type Input: 'a;
29
30 type Info;
32
33 fn info(&self) -> &Self::Info;
35
36 fn write(
38 &self,
39 input: Self::Input,
40 ) -> impl std::future::Future<Output = StreamResult<()>> + Send;
41
42 fn close(self) -> impl std::future::Future<Output = StreamResult<()>> + Send;
44
45 fn close_with_reason(
47 self,
48 reason: &str,
49 ) -> impl std::future::Future<Output = StreamResult<()>> + Send;
50
51 fn close_with_options(
54 self,
55 reason: Option<&str>,
56 attributes: Option<HashMap<String, String>>,
57 ) -> impl std::future::Future<Output = StreamResult<()>> + Send;
58}
59
60#[derive(Clone)]
61pub struct ByteStreamWriter {
63 info: Arc<ByteStreamInfo>,
64 stream: Arc<Mutex<RawStream>>,
65}
66
67impl ByteStreamWriter {
68 pub(crate) fn new(info: Arc<ByteStreamInfo>, stream: Arc<Mutex<RawStream>>) -> Self {
69 Self { info, stream }
70 }
71}
72
73#[derive(Clone)]
74pub struct TextStreamWriter {
76 info: Arc<TextStreamInfo>,
77 stream: Arc<Mutex<RawStream>>,
78}
79
80impl TextStreamWriter {
81 pub(crate) fn new(info: Arc<TextStreamInfo>, stream: Arc<Mutex<RawStream>>) -> Self {
82 Self { info, stream }
83 }
84}
85
86impl<'a> StreamWriter<'a> for ByteStreamWriter {
87 type Input = &'a [u8];
88 type Info = ByteStreamInfo;
89
90 fn info(&self) -> &Self::Info {
91 &self.info
92 }
93
94 async fn write(&self, bytes: &'a [u8]) -> StreamResult<()> {
95 let mut stream = self.stream.lock().await;
96 for chunk in bytes.chunks(STREAM_CHUNK_SIZE_BYTES) {
97 stream.write_chunk(chunk).await?;
98 }
99 Ok(())
100 }
101
102 async fn close(self) -> StreamResult<()> {
103 self.stream.lock().await.close(None, None).await
104 }
105
106 async fn close_with_reason(self, reason: &str) -> StreamResult<()> {
107 self.stream.lock().await.close(Some(reason), None).await
108 }
109
110 async fn close_with_options(
111 self,
112 reason: Option<&str>,
113 attributes: Option<HashMap<String, String>>,
114 ) -> StreamResult<()> {
115 self.stream.lock().await.close(reason, attributes).await
116 }
117}
118
119impl<'a> StreamWriter<'a> for TextStreamWriter {
120 type Input = &'a str;
121 type Info = TextStreamInfo;
122
123 fn info(&self) -> &Self::Info {
124 &self.info
125 }
126
127 async fn write(&self, text: &'a str) -> StreamResult<()> {
128 let mut stream = self.stream.lock().await;
129 for chunk in text.as_bytes().utf8_aware_chunks(STREAM_CHUNK_SIZE_BYTES) {
130 stream.write_chunk(chunk).await?;
131 }
132 Ok(())
133 }
134
135 async fn close(self) -> StreamResult<()> {
136 self.stream.lock().await.close(None, None).await
137 }
138
139 async fn close_with_reason(self, reason: &str) -> StreamResult<()> {
140 self.stream.lock().await.close(Some(reason), None).await
141 }
142
143 async fn close_with_options(
144 self,
145 reason: Option<&str>,
146 attributes: Option<HashMap<String, String>>,
147 ) -> StreamResult<()> {
148 self.stream.lock().await.close(reason, attributes).await
149 }
150}