1#![allow(clippy::disallowed_types)]
9
10use std::{
11 collections::HashMap,
12 future::Future,
13 pin::Pin,
14 sync::{Arc, Mutex},
15 task::{Context, Poll},
16 time::{Duration, Instant},
17};
18
19use bytes::Bytes;
20use futures::sink::Sink;
21use http_body::Frame;
22use slotmap::{KeyData, SlotMap};
23use tokio::sync::mpsc;
24
25use crate::{http::body::BoxError, CoreError};
26
27pub const DEFAULT_CHANNEL_CAPACITY: usize = 32;
30pub const DEFAULT_MAX_CHUNK_SIZE: usize = 64 * 1024; pub const DEFAULT_DRAIN_TIMEOUT_MS: u64 = 30_000; pub const DEFAULT_SLAB_TTL_MS: u64 = 300_000; pub const DEFAULT_SWEEP_INTERVAL_MS: u64 = 60_000; pub const DEFAULT_MAX_HANDLES: usize = 65_536;
35
36pub struct SessionEntry {
39 pub conn: iroh::endpoint::Connection,
40}
41
42pub struct ResponseHeadEntry {
43 pub status: u16,
44 pub headers: Vec<(String, String)>,
45}
46
47slotmap::new_key_type! { pub(crate) struct ReaderKey; }
50slotmap::new_key_type! { pub(crate) struct WriterKey; }
51slotmap::new_key_type! { pub(crate) struct FetchCancelKey; }
52slotmap::new_key_type! { pub(crate) struct SessionKey; }
53slotmap::new_key_type! { pub(crate) struct RequestHeadKey; }
54
55fn key_to_handle<K: slotmap::Key>(k: K) -> u64 {
58 k.data().as_ffi()
59}
60
61macro_rules! handle_to_key {
62 ($fn_name:ident, $key_type:ty) => {
63 fn $fn_name(h: u64) -> $key_type {
64 <$key_type>::from(KeyData::from_ffi(h))
65 }
66 };
67}
68
69handle_to_key!(handle_to_reader_key, ReaderKey);
70handle_to_key!(handle_to_writer_key, WriterKey);
71handle_to_key!(handle_to_session_key, SessionKey);
72handle_to_key!(handle_to_request_head_key, RequestHeadKey);
73handle_to_key!(handle_to_fetch_cancel_key, FetchCancelKey);
74
75pub struct BodyReader {
91 pub(crate) rx: Arc<tokio::sync::Mutex<mpsc::Receiver<Bytes>>>,
92 pub(crate) terminal_error: Arc<Mutex<Option<CoreError>>>,
93 pub(crate) cancel: Arc<tokio::sync::Notify>,
96 pending: Option<Pin<Box<dyn Future<Output = Option<Bytes>> + Send + Sync>>>,
102}
103
104pub struct BodyWriter {
120 pub(crate) tx: mpsc::Sender<Bytes>,
121 pub(crate) terminal_error: Arc<Mutex<Option<CoreError>>>,
122 pub(crate) drain_timeout: Duration,
124 sending: Option<BodyWriterSendFuture>,
128}
129
130type BodyWriterSendFuture = Pin<Box<dyn Future<Output = Result<(), BoxError>> + Send + Sync>>;
133
134pub fn make_body_channel() -> (BodyWriter, BodyReader) {
142 make_body_channel_with(
143 DEFAULT_CHANNEL_CAPACITY,
144 Duration::from_millis(DEFAULT_DRAIN_TIMEOUT_MS),
145 )
146}
147
148fn make_body_channel_with(capacity: usize, drain_timeout: Duration) -> (BodyWriter, BodyReader) {
149 let (tx, rx) = mpsc::channel(capacity);
150 let terminal_error = Arc::new(Mutex::new(None));
151 (
152 BodyWriter {
153 tx,
154 terminal_error: terminal_error.clone(),
155 drain_timeout,
156 sending: None,
157 },
158 BodyReader {
159 rx: Arc::new(tokio::sync::Mutex::new(rx)),
160 terminal_error,
161 cancel: Arc::new(tokio::sync::Notify::new()),
162 pending: None,
163 },
164 )
165}
166
167async fn recv_with_cancel(
176 rx: Arc<tokio::sync::Mutex<mpsc::Receiver<Bytes>>>,
177 cancel: Arc<tokio::sync::Notify>,
178) -> Option<Bytes> {
179 tokio::select! {
180 biased;
181 _ = cancel.notified() => None,
182 chunk = async { rx.lock().await.recv().await } => chunk,
183 }
184}
185
186impl BodyReader {
187 pub async fn next_chunk(&self) -> Option<Bytes> {
190 recv_with_cancel(self.rx.clone(), self.cancel.clone()).await
191 }
192}
193
194impl http_body::Body for BodyReader {
197 type Data = Bytes;
198 type Error = std::convert::Infallible;
199
200 fn poll_frame(
201 self: Pin<&mut Self>,
202 cx: &mut Context<'_>,
203 ) -> Poll<Option<Result<Frame<Bytes>, Self::Error>>> {
204 let this = self.get_mut();
205 if this.pending.is_none() {
206 if let Ok(mut rx) = this.rx.try_lock() {
207 match rx.try_recv() {
208 Ok(data) => return Poll::Ready(Some(Ok(Frame::data(data)))),
209 Err(mpsc::error::TryRecvError::Disconnected) => return Poll::Ready(None),
210 Err(mpsc::error::TryRecvError::Empty) => {}
211 }
212 }
213 }
214
215 let fut = this.pending.get_or_insert_with(|| {
216 Box::pin(recv_with_cancel(this.rx.clone(), this.cancel.clone()))
217 });
218 match fut.as_mut().poll(cx) {
219 Poll::Pending => Poll::Pending,
220 Poll::Ready(opt) => {
221 this.pending = None;
222 Poll::Ready(opt.map(|data| Ok(Frame::data(data))))
223 }
224 }
225 }
226}
227
228impl BodyWriter {
229 pub async fn send_chunk(&self, chunk: Bytes) -> Result<(), String> {
232 match self.tx.try_send(chunk) {
233 Ok(()) => Ok(()),
234 Err(mpsc::error::TrySendError::Closed(_)) => Err("body reader dropped".to_string()),
235 Err(mpsc::error::TrySendError::Full(chunk)) => {
236 tokio::time::timeout(self.drain_timeout, self.tx.send(chunk))
237 .await
238 .map_err(|_| "drain timeout: body reader is too slow".to_string())?
239 .map_err(|_| "body reader dropped".to_string())
240 }
241 }
242 }
243
244 pub(crate) fn abort(&self, error: CoreError) {
247 *self
248 .terminal_error
249 .lock()
250 .unwrap_or_else(|e| e.into_inner()) = Some(error);
251 }
252}
253
254async fn send_chunk_with_fast_path(
255 tx: &mpsc::Sender<Bytes>,
256 chunk: Bytes,
257 drain_timeout: Duration,
258) -> Result<(), CoreError> {
259 match tx.try_send(chunk) {
260 Ok(()) => Ok(()),
261 Err(mpsc::error::TrySendError::Closed(_)) => {
262 Err(CoreError::internal("body reader dropped"))
263 }
264 Err(mpsc::error::TrySendError::Full(chunk)) => {
265 tokio::time::timeout(drain_timeout, tx.send(chunk))
266 .await
267 .map_err(|_| CoreError::timeout("drain timeout: body reader is too slow"))?
268 .map_err(|_| CoreError::internal("body reader dropped"))
269 }
270 }
271}
272
273impl Sink<Bytes> for BodyWriter {
278 type Error = BoxError;
279
280 fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
281 let this = self.get_mut();
282 match this.sending.as_mut() {
283 None => Poll::Ready(Ok(())),
284 Some(fut) => match fut.as_mut().poll(cx) {
285 Poll::Pending => Poll::Pending,
286 Poll::Ready(res) => {
287 this.sending = None;
288 Poll::Ready(res)
289 }
290 },
291 }
292 }
293
294 fn start_send(self: Pin<&mut Self>, item: Bytes) -> Result<(), Self::Error> {
295 debug_assert!(
299 self.sending.is_none(),
300 "Sink contract: poll_ready must return Ready(Ok) before start_send"
301 );
302 let tx = self.tx.clone();
303 let drain_timeout = self.drain_timeout;
304 self.get_mut().sending = Some(Box::pin(async move {
305 tokio::time::timeout(drain_timeout, tx.send(item))
306 .await
307 .map_err(|_| -> BoxError { "drain timeout: body reader is too slow".into() })?
308 .map_err(|_| -> BoxError { "body reader dropped".into() })
309 }));
310 Ok(())
311 }
312
313 fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
314 self.poll_ready(cx)
316 }
317
318 fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
319 self.poll_flush(cx)
323 }
324}
325
326#[derive(Debug, Clone)]
330pub struct StoreConfig {
331 pub channel_capacity: usize,
333 pub max_chunk_size: usize,
335 pub drain_timeout: Duration,
337 pub max_handles: usize,
339 pub ttl: Duration,
342}
343
344impl Default for StoreConfig {
345 fn default() -> Self {
346 Self {
347 channel_capacity: DEFAULT_CHANNEL_CAPACITY,
348 max_chunk_size: DEFAULT_MAX_CHUNK_SIZE,
349 drain_timeout: Duration::from_millis(DEFAULT_DRAIN_TIMEOUT_MS),
350 max_handles: DEFAULT_MAX_HANDLES,
351 ttl: Duration::from_millis(DEFAULT_SLAB_TTL_MS),
352 }
353 }
354}
355
356struct Timed<T> {
359 value: T,
360 last_accessed: Instant,
363}
364
365impl<T> Timed<T> {
366 fn new(value: T) -> Self {
367 Self {
368 value,
369 last_accessed: Instant::now(),
370 }
371 }
372
373 fn touch(&mut self) {
375 self.last_accessed = Instant::now();
376 }
377
378 fn is_expired(&self, ttl: Duration) -> bool {
379 self.last_accessed.elapsed() > ttl
380 }
381}
382
383struct PendingReaderEntry {
385 reader: BodyReader,
386 created: Instant,
387}
388
389pub(crate) struct InsertGuard<'a> {
395 store: &'a HandleStore,
396 tracked: Vec<TrackedHandle>,
397 committed: bool,
398}
399
400enum TrackedHandle {
415 Reader(u64),
416 Writer(u64),
417 ReqHead(u64),
418}
419
420impl<'a> InsertGuard<'a> {
421 fn new(store: &'a HandleStore) -> Self {
422 Self {
423 store,
424 tracked: Vec::new(),
425 committed: false,
426 }
427 }
428
429 pub fn insert_reader(&mut self, reader: BodyReader) -> Result<u64, CoreError> {
430 let h = self.store.insert_reader(reader)?;
431 self.tracked.push(TrackedHandle::Reader(h));
432 Ok(h)
433 }
434
435 pub fn insert_writer(&mut self, writer: BodyWriter) -> Result<u64, CoreError> {
436 let h = self.store.insert_writer(writer)?;
437 self.tracked.push(TrackedHandle::Writer(h));
438 Ok(h)
439 }
440
441 pub fn allocate_req_handle(
442 &mut self,
443 sender: tokio::sync::oneshot::Sender<ResponseHeadEntry>,
444 ) -> Result<u64, CoreError> {
445 let h = self.store.allocate_req_handle(sender)?;
446 self.tracked.push(TrackedHandle::ReqHead(h));
447 Ok(h)
448 }
449
450 pub fn commit(mut self) {
452 self.committed = true;
453 }
454}
455
456impl Drop for InsertGuard<'_> {
457 fn drop(&mut self) {
458 if self.committed {
459 return;
460 }
461 for handle in &self.tracked {
462 match handle {
463 TrackedHandle::Reader(h) => self.store.cancel_reader(*h),
464 TrackedHandle::Writer(h) => {
465 let _ = self.store.finish_body(*h);
466 }
467 TrackedHandle::ReqHead(h) => {
468 self.store
469 .request_heads
470 .lock()
471 .unwrap_or_else(|e| e.into_inner())
472 .remove(handle_to_request_head_key(*h));
473 }
474 }
475 }
476 }
477}
478
479pub struct HandleStore {
486 readers: Mutex<SlotMap<ReaderKey, Timed<BodyReader>>>,
487 writers: Mutex<SlotMap<WriterKey, Timed<BodyWriter>>>,
488 sessions: Mutex<SlotMap<SessionKey, Timed<Arc<SessionEntry>>>>,
489 request_heads:
490 Mutex<SlotMap<RequestHeadKey, Timed<tokio::sync::oneshot::Sender<ResponseHeadEntry>>>>,
491 fetch_cancels: Mutex<SlotMap<FetchCancelKey, Timed<Arc<tokio::sync::Notify>>>>,
492 pending_readers: Mutex<HashMap<u64, PendingReaderEntry>>,
493 pub(crate) config: StoreConfig,
494}
495
496impl HandleStore {
497 pub fn new(config: StoreConfig) -> Self {
499 Self {
500 readers: Mutex::new(SlotMap::with_key()),
501 writers: Mutex::new(SlotMap::with_key()),
502 sessions: Mutex::new(SlotMap::with_key()),
503 request_heads: Mutex::new(SlotMap::with_key()),
504 fetch_cancels: Mutex::new(SlotMap::with_key()),
505 pending_readers: Mutex::new(HashMap::new()),
506 config,
507 }
508 }
509
510 pub(crate) fn insert_guard(&self) -> InsertGuard<'_> {
514 InsertGuard::new(self)
515 }
516
517 pub fn drain_timeout(&self) -> Duration {
519 self.config.drain_timeout
520 }
521
522 pub fn max_chunk_size(&self) -> usize {
524 self.config.max_chunk_size
525 }
526
527 pub fn count_handles(&self) -> (usize, usize, usize, usize) {
531 let readers = self.readers.lock().unwrap_or_else(|e| e.into_inner()).len();
532 let writers = self.writers.lock().unwrap_or_else(|e| e.into_inner()).len();
533 let sessions = self
534 .sessions
535 .lock()
536 .unwrap_or_else(|e| e.into_inner())
537 .len();
538 let total = readers
539 .saturating_add(writers)
540 .saturating_add(sessions)
541 .saturating_add(
542 self.request_heads
543 .lock()
544 .unwrap_or_else(|e| e.into_inner())
545 .len(),
546 )
547 .saturating_add(
548 self.fetch_cancels
549 .lock()
550 .unwrap_or_else(|e| e.into_inner())
551 .len(),
552 );
553 (readers, writers, sessions, total)
554 }
555
556 pub fn make_body_channel(&self) -> (BodyWriter, BodyReader) {
560 make_body_channel_with(self.config.channel_capacity, self.config.drain_timeout)
561 }
562
563 fn insert_checked<K: slotmap::Key, T>(
566 registry: &Mutex<SlotMap<K, Timed<T>>>,
567 value: T,
568 max: usize,
569 ) -> Result<u64, CoreError> {
570 let mut reg = registry.lock().unwrap_or_else(|e| e.into_inner());
571 if reg.len() >= max {
572 return Err(CoreError::internal("handle registry at capacity"));
573 }
574 let key = reg.insert(Timed::new(value));
575 Ok(key_to_handle(key))
576 }
577
578 pub fn insert_reader(&self, reader: BodyReader) -> Result<u64, CoreError> {
582 Self::insert_checked(&self.readers, reader, self.config.max_handles)
583 }
584
585 pub fn insert_writer(&self, writer: BodyWriter) -> Result<u64, CoreError> {
587 Self::insert_checked(&self.writers, writer, self.config.max_handles)
588 }
589
590 pub fn alloc_body_writer(&self) -> Result<(u64, BodyReader), CoreError> {
596 let (writer, reader) = self.make_body_channel();
597 let handle = self.insert_writer(writer)?;
598 Ok((handle, reader))
599 }
600
601 pub fn store_pending_reader(&self, writer_handle: u64, reader: BodyReader) {
604 self.pending_readers
605 .lock()
606 .unwrap_or_else(|e| e.into_inner())
607 .insert(
608 writer_handle,
609 PendingReaderEntry {
610 reader,
611 created: Instant::now(),
612 },
613 );
614 }
615
616 pub fn claim_pending_reader(&self, writer_handle: u64) -> Option<BodyReader> {
619 self.pending_readers
620 .lock()
621 .unwrap_or_else(|e| e.into_inner())
622 .remove(&writer_handle)
623 .map(|e| e.reader)
624 }
625
626 pub async fn next_chunk(&self, handle: u64) -> Result<Option<Bytes>, CoreError> {
633 let (rx_arc, cancel, terminal_error) = {
635 let mut reg = self.readers.lock().unwrap_or_else(|e| e.into_inner());
636 let entry = reg
637 .get_mut(handle_to_reader_key(handle))
638 .ok_or_else(|| CoreError::invalid_handle(handle))?;
639 entry.touch();
640 (
641 entry.value.rx.clone(),
642 entry.value.cancel.clone(),
643 entry.value.terminal_error.clone(),
644 )
645 };
646
647 let chunk = recv_with_cancel(rx_arc, cancel).await;
648
649 if chunk.is_none() {
651 self.readers
652 .lock()
653 .unwrap_or_else(|e| e.into_inner())
654 .remove(handle_to_reader_key(handle));
655 if let Some(error) = terminal_error
656 .lock()
657 .unwrap_or_else(|e| e.into_inner())
658 .take()
659 {
660 return Err(error);
661 }
662 }
663
664 Ok(chunk)
665 }
666
667 pub fn try_next_chunk(&self, handle: u64) -> Result<Option<Bytes>, CoreError> {
679 let (rx_arc, terminal_error) = {
680 let mut reg = self.readers.lock().unwrap_or_else(|e| e.into_inner());
681 let entry = reg
682 .get_mut(handle_to_reader_key(handle))
683 .ok_or_else(|| CoreError::invalid_handle(handle))?;
684 entry.touch();
685 (entry.value.rx.clone(), entry.value.terminal_error.clone())
686 };
687
688 let mut rx_guard = match rx_arc.try_lock() {
690 Ok(g) => g,
691 Err(_) => return Err(CoreError::internal("try_next_chunk: lock contended")),
692 };
693
694 match rx_guard.try_recv() {
695 Ok(chunk) => Ok(Some(chunk)),
696 Err(tokio::sync::mpsc::error::TryRecvError::Empty) => {
697 Err(CoreError::internal("try_next_chunk: channel empty"))
698 }
699 Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => {
700 drop(rx_guard);
702 self.readers
703 .lock()
704 .unwrap_or_else(|e| e.into_inner())
705 .remove(handle_to_reader_key(handle));
706 if let Some(error) = terminal_error
707 .lock()
708 .unwrap_or_else(|e| e.into_inner())
709 .take()
710 {
711 return Err(error);
712 }
713 Ok(None)
714 }
715 }
716 }
717
718 pub async fn send_chunk(&self, handle: u64, chunk: Bytes) -> Result<(), CoreError> {
723 let (tx, timeout) = {
725 let mut reg = self.writers.lock().unwrap_or_else(|e| e.into_inner());
726 let entry = reg
727 .get_mut(handle_to_writer_key(handle))
728 .ok_or_else(|| CoreError::invalid_handle(handle))?;
729 entry.touch();
730 (entry.value.tx.clone(), entry.value.drain_timeout)
731 };
732 let max = self.config.max_chunk_size;
733 if chunk.len() <= max {
734 send_chunk_with_fast_path(&tx, chunk, timeout).await
735 } else {
736 let mut offset = 0;
738 while offset < chunk.len() {
739 let end = offset.saturating_add(max).min(chunk.len());
740 send_chunk_with_fast_path(&tx, chunk.slice(offset..end), timeout).await?;
741 offset = end;
742 }
743 Ok(())
744 }
745 }
746
747 pub fn finish_body(&self, handle: u64) -> Result<(), CoreError> {
749 self.writers
750 .lock()
751 .unwrap_or_else(|e| e.into_inner())
752 .remove(handle_to_writer_key(handle))
753 .ok_or_else(|| CoreError::invalid_handle(handle))?;
754 Ok(())
755 }
756
757 pub fn cancel_reader(&self, handle: u64) {
759 let entry = self
760 .readers
761 .lock()
762 .unwrap_or_else(|e| e.into_inner())
763 .remove(handle_to_reader_key(handle));
764 if let Some(e) = entry {
765 e.value.cancel.notify_waiters();
766 }
767 }
768
769 pub fn insert_session(&self, entry: SessionEntry) -> Result<u64, CoreError> {
773 Self::insert_checked(&self.sessions, Arc::new(entry), self.config.max_handles)
774 }
775
776 pub fn lookup_session(&self, handle: u64) -> Option<Arc<SessionEntry>> {
778 self.sessions
779 .lock()
780 .unwrap_or_else(|e| e.into_inner())
781 .get(handle_to_session_key(handle))
782 .map(|e| e.value.clone())
783 }
784
785 pub fn remove_session(&self, handle: u64) -> Option<Arc<SessionEntry>> {
787 self.sessions
788 .lock()
789 .unwrap_or_else(|e| e.into_inner())
790 .remove(handle_to_session_key(handle))
791 .map(|e| e.value)
792 }
793
794 pub fn allocate_req_handle(
798 &self,
799 sender: tokio::sync::oneshot::Sender<ResponseHeadEntry>,
800 ) -> Result<u64, CoreError> {
801 Self::insert_checked(&self.request_heads, sender, self.config.max_handles)
802 }
803
804 pub fn take_req_sender(
806 &self,
807 handle: u64,
808 ) -> Option<tokio::sync::oneshot::Sender<ResponseHeadEntry>> {
809 self.request_heads
810 .lock()
811 .unwrap_or_else(|e| e.into_inner())
812 .remove(handle_to_request_head_key(handle))
813 .map(|e| e.value)
814 }
815
816 pub fn alloc_fetch_token(&self) -> Result<u64, CoreError> {
820 let notify = Arc::new(tokio::sync::Notify::new());
821 Self::insert_checked(&self.fetch_cancels, notify, self.config.max_handles)
822 }
823
824 pub fn cancel_in_flight(&self, token: u64) {
826 if let Some(entry) = self
827 .fetch_cancels
828 .lock()
829 .unwrap_or_else(|e| e.into_inner())
830 .get(handle_to_fetch_cancel_key(token))
831 {
832 entry.value.notify_one();
833 }
834 }
835
836 pub fn get_fetch_cancel_notify(&self, token: u64) -> Option<Arc<tokio::sync::Notify>> {
838 self.fetch_cancels
839 .lock()
840 .unwrap_or_else(|e| e.into_inner())
841 .get(handle_to_fetch_cancel_key(token))
842 .map(|e| e.value.clone())
843 }
844
845 pub fn remove_fetch_token(&self, token: u64) {
847 self.fetch_cancels
848 .lock()
849 .unwrap_or_else(|e| e.into_inner())
850 .remove(handle_to_fetch_cancel_key(token));
851 }
852
853 pub fn sweep(&self, ttl: Duration) {
859 Self::sweep_readers(&self.readers, ttl);
860 Self::sweep_registry(&self.writers, ttl);
861 Self::sweep_registry(&self.request_heads, ttl);
862 Self::sweep_registry(&self.sessions, ttl);
863 Self::sweep_registry(&self.fetch_cancels, ttl);
864 self.sweep_pending_readers(ttl);
865 }
866
867 fn sweep_readers(registry: &Mutex<SlotMap<ReaderKey, Timed<BodyReader>>>, ttl: Duration) {
870 let mut reg = registry.lock().unwrap_or_else(|e| e.into_inner());
871 let expired: Vec<ReaderKey> = reg
872 .iter()
873 .filter(|(_, e)| e.is_expired(ttl))
874 .map(|(k, _)| k)
875 .collect();
876
877 if expired.is_empty() {
878 return;
879 }
880
881 for key in &expired {
882 if let Some(entry) = reg.remove(*key) {
883 entry.value.cancel.notify_waiters();
884 }
885 }
886 tracing::debug!(
887 "[iroh-http] swept {} expired reader entries (ttl={ttl:?})",
888 expired.len()
889 );
890 if reg.is_empty() && reg.capacity() > 128 {
891 *reg = SlotMap::with_key();
892 }
893 }
894
895 fn sweep_registry<K: slotmap::Key, T>(registry: &Mutex<SlotMap<K, Timed<T>>>, ttl: Duration) {
896 let mut reg = registry.lock().unwrap_or_else(|e| e.into_inner());
897 let expired: Vec<K> = reg
898 .iter()
899 .filter(|(_, e)| e.is_expired(ttl))
900 .map(|(k, _)| k)
901 .collect();
902
903 if expired.is_empty() {
904 return;
905 }
906
907 for key in &expired {
908 reg.remove(*key);
909 }
910 tracing::debug!(
911 "[iroh-http] swept {} expired registry entries (ttl={ttl:?})",
912 expired.len()
913 );
914 if reg.is_empty() && reg.capacity() > 128 {
916 *reg = SlotMap::with_key();
917 }
918 }
919
920 fn sweep_pending_readers(&self, ttl: Duration) {
921 let mut map = self
922 .pending_readers
923 .lock()
924 .unwrap_or_else(|e| e.into_inner());
925 let before = map.len();
926 map.retain(|_, e| e.created.elapsed() < ttl);
927 let removed = before.saturating_sub(map.len());
928 if removed > 0 {
929 tracing::debug!("[iroh-http] swept {removed} stale pending readers (ttl={ttl:?})");
930 }
931 }
932}
933
934#[cfg(test)]
935mod tests {
936 use super::*;
937
938 fn test_store() -> HandleStore {
939 HandleStore::new(StoreConfig::default())
940 }
941
942 #[tokio::test]
945 async fn body_channel_send_recv() {
946 let (writer, reader) = make_body_channel();
947 writer.send_chunk(Bytes::from("hello")).await.unwrap();
948 drop(writer); let chunk = reader.next_chunk().await;
950 assert_eq!(chunk, Some(Bytes::from("hello")));
951 let eof = reader.next_chunk().await;
952 assert!(eof.is_none());
953 }
954
955 #[tokio::test]
956 async fn body_channel_multiple_chunks() {
957 let (writer, reader) = make_body_channel();
958 writer.send_chunk(Bytes::from("a")).await.unwrap();
959 writer.send_chunk(Bytes::from("b")).await.unwrap();
960 writer.send_chunk(Bytes::from("c")).await.unwrap();
961 drop(writer);
962
963 let mut collected = Vec::new();
964 while let Some(chunk) = reader.next_chunk().await {
965 collected.push(chunk);
966 }
967 assert_eq!(
968 collected,
969 vec![Bytes::from("a"), Bytes::from("b"), Bytes::from("c"),]
970 );
971 }
972
973 #[tokio::test]
974 async fn body_channel_reader_dropped_returns_error() {
975 let (writer, reader) = make_body_channel();
976 drop(reader);
977 let result = writer.send_chunk(Bytes::from("data")).await;
978 assert!(result.is_err());
979 }
980
981 #[tokio::test]
984 async fn body_writer_sink_forward_from_stream() {
985 use futures::{stream, StreamExt};
986 let (writer, reader) = make_body_channel();
987
988 let chunks = vec![
989 Ok::<_, BoxError>(Bytes::from("one")),
990 Ok(Bytes::from("two")),
991 Ok(Bytes::from("three")),
992 ];
993 let producer = tokio::spawn(async move {
994 stream::iter(chunks).forward(writer).await.unwrap();
995 });
997
998 let mut collected = Vec::new();
999 while let Some(chunk) = reader.next_chunk().await {
1000 collected.push(chunk);
1001 }
1002 producer.await.unwrap();
1003
1004 assert_eq!(
1005 collected,
1006 vec![Bytes::from("one"), Bytes::from("two"), Bytes::from("three")]
1007 );
1008 }
1009
1010 #[tokio::test]
1011 async fn body_writer_sink_send_via_sinkext() {
1012 use futures::SinkExt;
1013 let (mut writer, reader) = make_body_channel();
1014 writer.send(Bytes::from("a")).await.unwrap();
1015 writer.send(Bytes::from("b")).await.unwrap();
1016 writer.close().await.unwrap();
1017 drop(writer);
1018
1019 let mut collected = Vec::new();
1020 while let Some(chunk) = reader.next_chunk().await {
1021 collected.push(chunk);
1022 }
1023 assert_eq!(collected, vec![Bytes::from("a"), Bytes::from("b")]);
1024 }
1025
1026 #[tokio::test]
1027 async fn body_writer_sink_propagates_reader_dropped() {
1028 use futures::SinkExt;
1029 let (mut writer, reader) = make_body_channel();
1030 drop(reader);
1031 let mut err = None;
1035 for _ in 0..(DEFAULT_CHANNEL_CAPACITY + 1) {
1036 if let Err(e) = writer.send(Bytes::from("x")).await {
1037 err = Some(e);
1038 break;
1039 }
1040 }
1041 assert!(err.is_some(), "expected reader-dropped error from Sink");
1042 }
1043
1044 #[tokio::test]
1047 async fn insert_reader_and_next_chunk() {
1048 let store = test_store();
1049 let (writer, reader) = store.make_body_channel();
1050 let handle = store.insert_reader(reader).unwrap();
1051
1052 writer.send_chunk(Bytes::from("slab-data")).await.unwrap();
1053 drop(writer);
1054
1055 let chunk = store.next_chunk(handle).await.unwrap();
1056 assert_eq!(chunk, Some(Bytes::from("slab-data")));
1057
1058 let eof = store.next_chunk(handle).await.unwrap();
1060 assert!(eof.is_none());
1061 }
1062
1063 #[tokio::test]
1064 async fn next_chunk_invalid_handle() {
1065 let store = test_store();
1066 let result = store.next_chunk(999999).await;
1067 assert!(result.is_err());
1068 assert_eq!(result.unwrap_err().code, crate::ErrorCode::InvalidInput);
1069 }
1070
1071 #[tokio::test]
1072 async fn send_chunk_via_handle() {
1073 let store = test_store();
1074 let (writer, reader) = store.make_body_channel();
1075 let handle = store.insert_writer(writer).unwrap();
1076
1077 store
1078 .send_chunk(handle, Bytes::from("via-slab"))
1079 .await
1080 .unwrap();
1081 store.finish_body(handle).unwrap();
1082
1083 let chunk = reader.next_chunk().await;
1084 assert_eq!(chunk, Some(Bytes::from("via-slab")));
1085 let eof = reader.next_chunk().await;
1086 assert!(eof.is_none());
1087 }
1088
1089 #[tokio::test]
1090 async fn capacity_cap_rejects_overflow() {
1091 let store = HandleStore::new(StoreConfig {
1092 max_handles: 2,
1093 ..StoreConfig::default()
1094 });
1095 let (_, r1) = store.make_body_channel();
1096 let (_, r2) = store.make_body_channel();
1097 let (_, r3) = store.make_body_channel();
1098 store.insert_reader(r1).unwrap();
1099 store.insert_reader(r2).unwrap();
1100 let err = store.insert_reader(r3).unwrap_err();
1101 assert!(err.message.contains("capacity"));
1102 }
1103
1104 #[tokio::test]
1107 async fn recv_with_cancel_returns_none_on_cancel() {
1108 let (_tx, rx) = mpsc::channel::<Bytes>(4);
1109 let rx = Arc::new(tokio::sync::Mutex::new(rx));
1110 let cancel = Arc::new(tokio::sync::Notify::new());
1111 cancel.notify_one();
1113 let result = recv_with_cancel(rx, cancel).await;
1114 assert!(result.is_none());
1115 }
1116
1117 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1122 async fn concurrent_alloc_dealloc_leaves_no_leak() {
1123 let store = Arc::new(test_store());
1124 let mut tasks = Vec::new();
1125 for _ in 0..16 {
1126 let s = store.clone();
1127 tasks.push(tokio::spawn(async move {
1128 for _ in 0..100 {
1129 let (writer, reader) = s.make_body_channel();
1130 let reader_handle = s.insert_reader(reader).unwrap();
1131 let writer_handle = s.insert_writer(writer).unwrap();
1132 s.cancel_reader(reader_handle);
1133 s.finish_body(writer_handle).unwrap();
1134 }
1135 }));
1136 }
1137 for t in tasks {
1138 t.await.unwrap();
1139 }
1140
1141 let (readers, writers, _sessions, total) = store.count_handles();
1142 assert_eq!(readers, 0, "all reader slots should be freed");
1143 assert_eq!(writers, 0, "all writer slots should be freed");
1144 assert_eq!(total, 0, "no handles should leak");
1145 }
1146
1147 #[tokio::test]
1150 async fn sweep_removes_expired_keeps_fresh() {
1151 let store = test_store();
1152 let (_old_writer, old_reader) = make_body_channel();
1154 let old_handle = store.insert_reader(old_reader).unwrap();
1155 tokio::time::sleep(Duration::from_millis(60)).await;
1156 let (_new_writer, new_reader) = make_body_channel();
1158 store.insert_reader(new_reader).unwrap();
1159
1160 store.sweep(Duration::from_millis(30));
1162
1163 let (readers, _, _, _) = store.count_handles();
1164 assert_eq!(readers, 1, "only the fresh reader should remain");
1165 assert!(
1166 store.next_chunk(old_handle).await.is_err(),
1167 "the swept handle must be invalid",
1168 );
1169 }
1170
1171 #[tokio::test]
1174 async fn freed_handle_is_invalid_and_slot_reusable() {
1175 let store = test_store();
1176 let (_writer_a, reader_a) = make_body_channel();
1177 let handle_a = store.insert_reader(reader_a).unwrap();
1178 store.cancel_reader(handle_a); assert!(
1181 store.next_chunk(handle_a).await.is_err(),
1182 "freed handle must be invalid",
1183 );
1184
1185 let (writer_b, reader_b) = store.make_body_channel();
1187 let handle_b = store.insert_reader(reader_b).unwrap();
1188 assert_ne!(
1189 handle_a, handle_b,
1190 "reused slot must yield a new versioned handle",
1191 );
1192
1193 writer_b.send_chunk(Bytes::from("ok")).await.unwrap();
1195 drop(writer_b);
1196 assert_eq!(
1197 store.next_chunk(handle_b).await.unwrap(),
1198 Some(Bytes::from("ok")),
1199 );
1200 }
1201}