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
73#[derive(Clone)]
74/// Writer for an open text data stream.
75pub 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}