Skip to main content

livekit_data_stream/outgoing/
stream_writer.rs

1// Copyright 2026 LiveKit, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use 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
25/// Writer for an open data stream.
26pub trait StreamWriter<'a> {
27    /// Type of input this writer accepts.
28    type Input: 'a;
29
30    /// Information about the underlying data stream.
31    type Info;
32
33    /// Returns a reference to the stream info.
34    fn info(&self) -> &Self::Info;
35
36    /// Writes to the stream.
37    fn write(
38        &self,
39        input: Self::Input,
40    ) -> impl std::future::Future<Output = StreamResult<()>> + Send;
41
42    /// Closes the stream normally.
43    fn close(self) -> impl std::future::Future<Output = StreamResult<()>> + Send;
44
45    /// Closes the stream abnormally, specifying the reason for closure.
46    fn close_with_reason(
47        self,
48        reason: &str,
49    ) -> impl std::future::Future<Output = StreamResult<()>> + Send;
50
51    /// Closes the stream, optionally specifying a closure reason (abnormal
52    /// closure) and attributes to attach to the stream trailer.
53    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)]
61/// Writer for an open byte data stream.
62pub 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    /// Whether the stream has been closed — either locally (via `close`) or because a send failed.
73    pub async fn is_closed(&self) -> bool {
74        self.stream.lock().await.is_closed()
75    }
76}
77
78#[derive(Clone)]
79/// Writer for an open text data stream.
80pub 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    /// Whether the stream has been closed — either locally (via `close`) or because a send failed.
91    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}