1use std::cmp::min;
2use std::io::{Error as IoError, ErrorKind, Read, Result as IoResult, Seek, SeekFrom};
3use std::sync::Arc;
4use std::time::Duration;
5use tokio::sync::{Mutex, Notify};
6use tracing::Instrument;
7
8const MIN_PREBUFFER_BYTES: usize = 256 * 1024; const MIN_BUFFER_AHEAD: usize = 128 * 1024; const MAX_BUFFER_SIZE: usize = 1024 * 1024 * 1024; struct SharedState {
31 buffer: Vec<u8>,
32 done: bool,
33 error: Option<String>,
34 total_size: Option<u64>,
35 prebuffer_ready: bool,
36}
37
38pub struct StreamBuffer {
39 state: Arc<(Mutex<SharedState>, Notify)>,
40 pos: u64,
41}
42
43impl StreamBuffer {
44 pub fn new(url: String, is_radio: bool) -> Self {
45 Self::with_user_agent(url, is_radio, None)
46 }
47
48 pub fn with_user_agent(url: String, is_radio: bool, user_agent: Option<String>) -> Self {
49 let prebuffer_size = if is_radio {
50 16 * 1024
51 } else {
52 MIN_PREBUFFER_BYTES
53 };
54
55 let state = Arc::new((
56 Mutex::new(SharedState {
57 buffer: Vec::with_capacity(prebuffer_size * 4),
58 done: false,
59 error: None,
60 total_size: None,
61 prebuffer_ready: false,
62 }),
63 Notify::new(),
64 ));
65
66 let state_clone = state.clone();
67
68 let handle = tokio::runtime::Handle::current();
73 handle.spawn(
74 async move {
75 let ua = user_agent
76 .unwrap_or_else(|| concat!("Kopuz/", env!("CARGO_PKG_VERSION")).to_string());
77 let client = reqwest::Client::builder()
78 .tcp_nodelay(true)
79 .user_agent(ua)
80 .build()
81 .unwrap_or_else(|_| reqwest::Client::new());
82
83 match client.get(&url).send().await {
84 Ok(mut response) => {
85 tracing::trace!(
86 status = %response.status(),
87 content_length = ?response.content_length(),
88 content_type = ?response
89 .headers()
90 .get("content-type")
91 .and_then(|v| v.to_str().ok()),
92 "stream buffer HTTP response",
93 );
94 if !response.status().is_success() {
95 let (lock, notify) = &*state_clone;
96 let mut state = lock.lock().await;
97 state.error = Some(format!("HTTP {}", response.status()));
98 state.done = true;
99 state.prebuffer_ready = true;
100 notify.notify_waiters();
101 return;
102 }
103
104 let total_size = response.content_length();
105 {
106 let (lock, notify) = &*state_clone;
107 let mut state = lock.lock().await;
108 state.total_size = total_size;
109 notify.notify_waiters();
110 }
111
112 let mut total_buffered = 0usize;
113
114 while let Ok(Some(chunk)) = response.chunk().await {
115 if Arc::strong_count(&state_clone) == 1 {
116 break;
117 }
118 let chunk_len = chunk.len();
119
120 if total_buffered + chunk_len > MAX_BUFFER_SIZE {
121 let (lock, notify) = &*state_clone;
122 let mut state = lock.lock().await;
123 state.error = Some("Buffer limit exceeded (1GB)".to_string());
124 state.done = true;
125 state.prebuffer_ready = true;
126 notify.notify_waiters();
127 break;
128 }
129
130 {
131 let (lock, notify) = &*state_clone;
132 let mut state = lock.lock().await;
133 state.buffer.extend_from_slice(&chunk);
134 total_buffered += chunk_len;
135
136 if !state.prebuffer_ready {
137 let is_small_file = state
138 .total_size
139 .map(|s| s <= prebuffer_size as u64)
140 .unwrap_or(false);
141
142 if total_buffered >= prebuffer_size || is_small_file {
143 state.prebuffer_ready = true;
144 }
145 }
146
147 notify.notify_waiters();
148 }
149 }
150
151 let (lock, notify) = &*state_clone;
152 let mut state = lock.lock().await;
153 state.done = true;
154 state.prebuffer_ready = true;
155 notify.notify_waiters();
156 }
157 Err(e) => {
158 let (lock, notify) = &*state_clone;
159 let mut state = lock.lock().await;
160 state.error = Some(e.to_string());
161 state.done = true;
162 state.prebuffer_ready = true;
163 notify.notify_waiters();
164 }
165 }
166 }
167 .instrument(tracing::info_span!("player.stream_buffer")),
168 );
169
170 Self { state, pos: 0 }
171 }
172
173 fn wait_for_prebuffer(&self) {
181 let (lock, _notify) = &*self.state;
182 loop {
183 let state = lock.blocking_lock();
184 if state.prebuffer_ready || state.done {
185 return;
186 }
187 drop(state);
188 std::thread::sleep(Duration::from_millis(5));
189 }
190 }
191
192 pub fn wait_for_total_size(&self) {
193 let (lock, _notify) = &*self.state;
194 loop {
195 let state = lock.blocking_lock();
196 if state.total_size.is_some() || state.done {
197 return;
198 }
199 drop(state);
200 std::thread::sleep(Duration::from_millis(5));
201 }
202 }
203
204 pub fn known_total_size(&self) -> Option<u64> {
205 let (lock, _) = &*self.state;
206 let state = lock.blocking_lock();
207 state.total_size.or(Some(state.buffer.len() as u64))
208 }
209
210 fn wait_for_buffer_ahead(&self, min_ahead: usize) {
211 let (lock, _notify) = &*self.state;
212 loop {
213 let state = lock.blocking_lock();
214 let buffer_len = state.buffer.len() as u64;
215 let buffered_ahead = buffer_len.saturating_sub(self.pos) as usize;
216
217 if buffered_ahead >= min_ahead || state.done || state.error.is_some() {
218 return;
219 }
220 drop(state);
221 std::thread::sleep(Duration::from_millis(5));
222 }
223 }
224}
225
226impl Read for StreamBuffer {
227 fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
228 if self.pos == 0 {
229 self.wait_for_prebuffer();
230 }
231
232 {
233 let (lock, _) = &*self.state;
234 let state = lock.blocking_lock();
235 let buffer_len = state.buffer.len() as u64;
236 let buffered_ahead = buffer_len.saturating_sub(self.pos) as usize;
237
238 if buffered_ahead < MIN_BUFFER_AHEAD && !state.done && state.error.is_none() {
239 drop(state);
240 self.wait_for_buffer_ahead(MIN_BUFFER_AHEAD);
241 }
242 }
243
244 let (lock, _notify) = &*self.state;
249 loop {
250 let state = lock.blocking_lock();
251
252 if let Some(err) = &state.error {
253 return Err(IoError::other(err.clone()));
254 }
255
256 let current_len = state.buffer.len() as u64;
257
258 if self.pos < current_len {
259 let available = (current_len - self.pos) as usize;
260 let to_read = min(buf.len(), available);
261 buf[0..to_read].copy_from_slice(
262 &state.buffer[self.pos as usize..(self.pos as usize + to_read)],
263 );
264 self.pos += to_read as u64;
265 return Ok(to_read);
266 }
267
268 if state.done {
269 if let Some(err) = &state.error {
270 return Err(IoError::other(err.clone()));
271 }
272 return Ok(0);
273 }
274
275 drop(state);
276 std::thread::sleep(Duration::from_millis(5));
277 }
278 }
279}
280
281impl Seek for StreamBuffer {
282 fn seek(&mut self, pos: SeekFrom) -> IoResult<u64> {
283 let (lock, _notify) = &*self.state;
284 let state = lock.blocking_lock();
285
286 let len = state.buffer.len() as u64;
287 let total = state.total_size.unwrap_or(len);
288
289 let new_pos = match pos {
290 SeekFrom::Start(p) => p as i64,
291 SeekFrom::Current(p) => self.pos as i64 + p,
292 SeekFrom::End(p) => total as i64 + p,
293 };
294
295 if new_pos < 0 {
296 return Err(IoError::new(
297 ErrorKind::InvalidInput,
298 "Seek to negative position",
299 ));
300 }
301
302 self.pos = new_pos as u64;
303 Ok(self.pos)
304 }
305}