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 pub async fn is_closed(&self) -> bool {
74 self.stream.lock().await.is_closed()
75 }
76}
77
78#[derive(Clone)]
79pub struct TextStreamWriter {
81 info: Arc<TextStreamInfo>,
82 stream: Arc<Mutex<RawStream>>,
83}
84
85impl TextStreamWriter {
86 pub(crate) fn new(info: Arc<TextStreamInfo>, stream: Arc<Mutex<RawStream>>) -> Self {
87 Self { info, stream }
88 }
89
90 pub async fn is_closed(&self) -> bool {
92 self.stream.lock().await.is_closed()
93 }
94}
95
96impl<'a> StreamWriter<'a> for ByteStreamWriter {
97 type Input = &'a [u8];
98 type Info = ByteStreamInfo;
99
100 fn info(&self) -> &Self::Info {
101 &self.info
102 }
103
104 async fn write(&self, bytes: &'a [u8]) -> StreamResult<()> {
105 let mut stream = self.stream.lock().await;
106 for chunk in bytes.chunks(STREAM_CHUNK_SIZE_BYTES) {
107 stream.write_chunk(chunk).await?;
108 }
109 Ok(())
110 }
111
112 async fn close(self) -> StreamResult<()> {
113 self.stream.lock().await.close(None, None).await
114 }
115
116 async fn close_with_reason(self, reason: &str) -> StreamResult<()> {
117 self.stream.lock().await.close(Some(reason), None).await
118 }
119
120 async fn close_with_options(
121 self,
122 reason: Option<&str>,
123 attributes: Option<HashMap<String, String>>,
124 ) -> StreamResult<()> {
125 self.stream.lock().await.close(reason, attributes).await
126 }
127}
128
129impl<'a> StreamWriter<'a> for TextStreamWriter {
130 type Input = &'a str;
131 type Info = TextStreamInfo;
132
133 fn info(&self) -> &Self::Info {
134 &self.info
135 }
136
137 async fn write(&self, text: &'a str) -> StreamResult<()> {
138 let mut stream = self.stream.lock().await;
139 for chunk in text.as_bytes().utf8_aware_chunks(STREAM_CHUNK_SIZE_BYTES) {
140 stream.write_chunk(chunk).await?;
141 }
142 Ok(())
143 }
144
145 async fn close(self) -> StreamResult<()> {
146 self.stream.lock().await.close(None, None).await
147 }
148
149 async fn close_with_reason(self, reason: &str) -> StreamResult<()> {
150 self.stream.lock().await.close(Some(reason), None).await
151 }
152
153 async fn close_with_options(
154 self,
155 reason: Option<&str>,
156 attributes: Option<HashMap<String, String>>,
157 ) -> StreamResult<()> {
158 self.stream.lock().await.close(reason, attributes).await
159 }
160}