1use std::any::Any;
16use std::future::Future;
17use std::io;
18use std::panic::AssertUnwindSafe;
19use std::sync::{Arc, Mutex, PoisonError};
20use std::time::Duration;
21
22use bao_tree::io::EncodeError;
23use bao_tree::io::fsm::encode_ranges_validated;
24use bao_tree::io::outboard::PreOrderMemOutboard;
25use bao_tree::{BaoTree, ChunkRanges};
26use bytes::{Bytes, BytesMut};
27use futures_util::FutureExt;
28use iroh_blobs::Hash;
29use iroh_io::{AsyncSliceReader, AsyncStreamWriter};
30use tokio::sync::{OwnedSemaphorePermit, mpsc};
31use tokio_util::task::AbortOnDropHandle;
32
33use crate::error::{CacheError, CacheResult};
34use crate::metrics::CacheMetrics;
35use crate::origin::{Origin, OriginKind, OriginRangeFetch, OriginRangeRequest};
36use crate::outboard_cache::OutboardCache;
37
38#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43pub(crate) struct OriginReadBudget {
44 pub(crate) head_start: Duration,
46 pub(crate) min_bps: u64,
49}
50
51impl OriginReadBudget {
52 pub(crate) fn for_len(self, len: u64) -> Duration {
54 let nanos = (u128::from(len) * 1_000_000_000).div_ceil(u128::from(self.min_bps.max(1)));
55 self.head_start.saturating_add(Duration::from_nanos(
56 u64::try_from(nanos).unwrap_or(u64::MAX),
57 ))
58 }
59}
60use crate::range_pull::{AlignedRange, IROH_BLOCK_SIZE};
61
62pub const RANGE_PULL_WINDOW_BYTES: u64 = 4 * 1024 * 1024;
68
69pub const MAX_CONCURRENT_RANGE_PULLS: usize = 4;
73
74pub(crate) const RANGE_PULL_PERMIT_WARN_AFTER: Duration = Duration::from_secs(10);
77
78const WIRE_CHANNEL_CAP: usize = 8;
82
83pub(crate) async fn within_origin_timeout<F: Future>(
93 budget: Option<OriginReadBudget>,
94 len: u64,
95 hash: Hash,
96 what: &'static str,
97 metrics: Option<&CacheMetrics>,
98 fut: F,
99) -> CacheResult<F::Output> {
100 let Some(budget) = budget else {
101 return Ok(fut.await);
102 };
103 let limit = budget.for_len(len);
104 tokio::time::timeout(limit, fut).await.map_err(|_| {
105 if let Some(m) = metrics {
106 m.origin_range_timeouts.inc();
107 }
108 CacheError::OriginError {
109 hash,
110 source: anyhow::anyhow!(
111 "origin {what} of {len} bytes took longer than its {limit:?} budget \
112 (cache.node_pull_stall_window_sec plus the size at \
113 cache.node_pull_min_throughput_bps)"
114 ),
115 }
116 })
117}
118
119#[derive(Debug, Clone)]
124pub(crate) struct WindowSpans {
125 next: u64,
126 end: u64,
127 done: bool,
128}
129
130impl WindowSpans {
131 pub(crate) const fn new(aligned: &AlignedRange) -> Self {
132 Self {
133 next: aligned.fetch_start(),
134 end: aligned.fetch_end(),
135 done: false,
136 }
137 }
138}
139
140impl Iterator for WindowSpans {
141 type Item = (u64, u64);
142
143 fn next(&mut self) -> Option<(u64, u64)> {
144 if self.done {
145 return None;
146 }
147 let start = self.next;
148 let end = start.saturating_add(RANGE_PULL_WINDOW_BYTES).min(self.end);
149 self.next = end;
150 self.done = end >= self.end;
151 Some((start, end))
152 }
153}
154
155#[derive(Debug)]
158pub(crate) struct OriginWindow {
159 pub(crate) start: u64,
160 pub(crate) data: Bytes,
161}
162
163#[derive(Debug)]
167pub(crate) enum WindowFetch {
168 Window(OriginWindow),
170 Declined { start: u64, end: u64 },
172 WrongLength { start: u64, end: u64, got: usize },
174}
175
176pub(crate) struct OriginRangeCursor {
180 origin: Arc<dyn Origin>,
181 hash: Hash,
182 outboard: Bytes,
183 budget: Option<OriginReadBudget>,
185 metrics: Option<Arc<CacheMetrics>>,
186 spans: WindowSpans,
187 pending: Option<WindowFetch>,
189}
190
191impl OriginRangeCursor {
192 pub(crate) async fn open(
207 origin: Arc<dyn Origin>,
208 hash: Hash,
209 aligned: &AlignedRange,
210 outboard: Bytes,
211 budget: Option<OriginReadBudget>,
212 metrics: Option<Arc<CacheMetrics>>,
213 ) -> CacheResult<Option<Self>> {
214 let mut cursor = Self {
215 origin,
216 hash,
217 outboard,
218 budget,
219 metrics,
220 spans: WindowSpans::new(aligned),
221 pending: None,
222 };
223 match cursor.fetch_next().await? {
224 None | Some(WindowFetch::Declined { .. }) => Ok(None),
225 Some(first) => {
226 cursor.pending = Some(first);
227 Ok(Some(cursor))
228 }
229 }
230 }
231
232 pub(crate) fn outboard(&self) -> Bytes {
234 self.outboard.clone()
235 }
236
237 pub(crate) fn origin_kind(&self) -> OriginKind {
239 self.origin.kind()
240 }
241
242 pub(crate) const fn first_is_wrong_length(&self) -> bool {
244 matches!(self.pending, Some(WindowFetch::WrongLength { .. }))
245 }
246
247 pub(crate) async fn next_window(&mut self) -> CacheResult<Option<WindowFetch>> {
254 if let Some(first) = self.pending.take() {
255 return Ok(Some(first));
256 }
257 self.fetch_next().await
258 }
259
260 async fn fetch_next(&mut self) -> CacheResult<Option<WindowFetch>> {
261 let Some((start, end)) = self.spans.next() else {
262 return Ok(None);
263 };
264 let req = OriginRangeRequest {
265 fetch_start: start,
266 fetch_end: end,
267 };
268 let data = match within_origin_timeout(
269 self.budget,
270 end - start,
271 self.hash,
272 "range window fetch",
273 self.metrics.as_deref(),
274 self.origin.fetch_range_data(self.hash, req),
275 )
276 .await?
277 .map_err(|e| CacheError::OriginError {
278 hash: self.hash,
279 source: e.into_inner(),
280 })? {
281 OriginRangeFetch::Ranged { data } => data,
282 OriginRangeFetch::Unsupported | OriginRangeFetch::NotFound => {
283 return Ok(Some(WindowFetch::Declined { start, end }));
284 }
285 };
286 if let Some(m) = &self.metrics {
287 m.pull_through_bytes
288 .inc_by(u64::try_from(data.len()).unwrap_or(u64::MAX));
289 }
290 if u64::try_from(data.len()).ok() != Some(end - start) {
291 return Ok(Some(WindowFetch::WrongLength {
292 start,
293 end,
294 got: data.len(),
295 }));
296 }
297 Ok(Some(WindowFetch::Window(OriginWindow { start, data })))
298 }
299}
300
301type FaultSlot = Arc<Mutex<Option<CacheError>>>;
305
306fn park_fault(slot: &FaultSlot, fault: CacheError) {
310 let mut guard = slot.lock().unwrap_or_else(PoisonError::into_inner);
311 if guard.is_none() {
312 *guard = Some(fault);
313 }
314}
315
316fn park_panic(slot: &FaultSlot, hash: Hash, kind: OriginKind, payload: &(dyn Any + Send)) {
320 let msg = payload
321 .downcast_ref::<&str>()
322 .copied()
323 .or_else(|| payload.downcast_ref::<String>().map(String::as_str))
324 .unwrap_or("non-string panic payload");
325 tracing::error!(
326 %hash,
327 ?kind,
328 panic = msg,
329 "origin range encode panicked; internal fault (code bug), not an origin fault",
330 );
331 park_fault(
332 slot,
333 CacheError::Internal(anyhow::anyhow!(
334 "origin range encode for {hash} panicked: {msg}"
335 )),
336 );
337}
338
339fn internal_fault(hash: Hash, what: &str) -> CacheError {
342 tracing::error!(%hash, what, "origin range reader invariant broken");
343 CacheError::Internal(anyhow::anyhow!(
344 "origin range reader invariant broken for {hash}: {what}"
345 ))
346}
347
348pub(crate) struct OriginWindowReader {
359 cursor: OriginRangeCursor,
360 blob_size: u64,
361 fetch_end: u64,
362 window_start: u64,
363 window: Bytes,
364 fault: FaultSlot,
365}
366
367impl OriginWindowReader {
368 const fn new(cursor: OriginRangeCursor, aligned: &AlignedRange, fault: FaultSlot) -> Self {
369 Self {
370 cursor,
371 blob_size: aligned.blob_size(),
372 fetch_end: aligned.fetch_end(),
373 window_start: aligned.fetch_start(),
374 window: Bytes::new(),
375 fault,
376 }
377 }
378
379 fn window_end(&self) -> u64 {
380 self.window_start
381 .saturating_add(u64::try_from(self.window.len()).unwrap_or(u64::MAX))
382 }
383
384 fn fail(&self, fault: CacheError) -> io::Error {
385 let msg = fault.display_chain().to_string();
386 park_fault(&self.fault, fault);
387 io::Error::other(msg)
388 }
389
390 async fn load_next(&mut self) -> io::Result<()> {
392 let hash = self.cursor.hash;
393 self.window = Bytes::new();
395 match self.cursor.next_window().await {
396 Ok(Some(WindowFetch::Window(w))) if !w.data.is_empty() => {
397 self.window_start = w.start;
398 self.window = w.data;
399 Ok(())
400 }
401 Ok(Some(WindowFetch::Window(_)) | None) => Err(self.fail(internal_fault(
402 hash,
403 "read past the last window of the span",
404 ))),
405 Ok(Some(WindowFetch::Declined { start, end })) => {
406 tracing::warn!(
407 %hash,
408 kind = ?self.cursor.origin_kind(),
409 window_start = start,
410 window_end = end,
411 "own origin stopped serving the range mid-stream",
412 );
413 Err(self.fail(CacheError::OriginError {
414 hash,
415 source: anyhow::anyhow!(
416 "origin stopped serving range [{start}, {end}) of {hash} mid-stream"
417 ),
418 }))
419 }
420 Ok(Some(WindowFetch::WrongLength { start, end, got })) => {
421 tracing::warn!(
422 %hash,
423 kind = ?self.cursor.origin_kind(),
424 window_start = start,
425 window_end = end,
426 got,
427 "own origin served a wrong-length range window; hard local-origin fault \
428 (no degrade — committed to serving under H)",
429 );
430 Err(self.fail(CacheError::VerifyFailed { expected: hash }))
431 }
432 Err(e) => {
433 tracing::warn!(
434 %hash,
435 kind = ?self.cursor.origin_kind(),
436 error = %e.display_chain(),
437 "own origin range window fetch failed mid-stream",
438 );
439 Err(self.fail(e))
440 }
441 }
442 }
443}
444
445impl AsyncSliceReader for OriginWindowReader {
446 async fn read_at(&mut self, offset: u64, len: usize) -> io::Result<Bytes> {
447 let end = offset
448 .saturating_add(u64::try_from(len).unwrap_or(u64::MAX))
449 .min(self.fetch_end);
450 if offset < self.window_start {
451 return Err(self.fail(internal_fault(
452 self.cursor.hash,
453 "backward read before the held window",
454 )));
455 }
456 let mut out: Option<BytesMut> = None;
457 let mut pos = offset;
458 while pos < end {
459 while pos >= self.window_end() {
460 self.load_next().await?;
461 }
462 let rel = usize::try_from(pos - self.window_start).unwrap_or(usize::MAX);
463 let take_end = end.min(self.window_end());
464 let take = usize::try_from(take_end - pos).unwrap_or(usize::MAX);
465 let Some(slice) = self.window.get(rel..rel.saturating_add(take)) else {
466 return Err(self.fail(internal_fault(
467 self.cursor.hash,
468 "held window shorter than its span",
469 )));
470 };
471 let piece = self.window.slice_ref(slice);
472 pos = take_end;
473 if pos >= end && out.is_none() {
477 return Ok(piece);
478 }
479 out.get_or_insert_with(BytesMut::new)
480 .extend_from_slice(&piece);
481 }
482 Ok(out.map(BytesMut::freeze).unwrap_or_default())
483 }
484
485 async fn size(&mut self) -> io::Result<u64> {
486 Ok(self.blob_size)
487 }
488}
489
490struct ChannelWriter {
493 tx: mpsc::Sender<Bytes>,
494}
495
496impl AsyncStreamWriter for ChannelWriter {
497 async fn write(&mut self, data: &[u8]) -> io::Result<()> {
498 self.write_bytes(Bytes::copy_from_slice(data)).await
499 }
500
501 async fn write_bytes(&mut self, data: Bytes) -> io::Result<()> {
502 if data.is_empty() {
503 return Ok(());
504 }
505 self.tx
506 .send(data)
507 .await
508 .map_err(|_| io::Error::other("origin range wire consumer dropped"))
509 }
510
511 async fn sync(&mut self) -> io::Result<()> {
512 Ok(())
513 }
514}
515
516struct UnfinishedGuard {
523 slot: FaultSlot,
524 hash: Hash,
525 finished: bool,
526}
527
528impl UnfinishedGuard {
529 const fn disarm(&mut self) {
532 self.finished = true;
533 }
534}
535
536impl Drop for UnfinishedGuard {
537 fn drop(&mut self) {
538 if !self.finished {
539 park_fault(
540 &self.slot,
541 CacheError::OriginError {
542 hash: self.hash,
543 source: anyhow::anyhow!(
544 "origin range encode for {} was cancelled, or panicked outside the encode, before completing",
545 self.hash
546 ),
547 },
548 );
549 }
550 }
551}
552
553pub struct OriginRangeWire {
570 rx: mpsc::Receiver<Bytes>,
571 fault: FaultSlot,
572 ended: bool,
573 _task: AbortOnDropHandle<()>,
574}
575
576impl std::fmt::Debug for OriginRangeWire {
577 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
578 f.debug_struct("OriginRangeWire")
579 .field("ended", &self.ended)
580 .finish_non_exhaustive()
581 }
582}
583
584impl OriginRangeWire {
585 pub(crate) fn spawn(
593 cursor: OriginRangeCursor,
594 aligned: &AlignedRange,
595 permit: OwnedSemaphorePermit,
596 outboards: OutboardCache,
597 ) -> CacheResult<Self> {
598 let hash = cursor.hash;
599 let kind = cursor.origin_kind();
600 let tree = BaoTree::new(aligned.blob_size(), IROH_BLOCK_SIZE);
601 let outboard = cursor.outboard();
602 if u64::try_from(outboard.len()).ok() != Some(tree.outboard_size()) {
603 tracing::warn!(
604 %hash,
605 ?kind,
606 expected = tree.outboard_size(),
607 got = outboard.len(),
608 "outboard length does not match the blob size; hard local-origin fault \
609 (no degrade — committed to serving under H)",
610 );
611 outboards.evict_if_same(hash, &outboard);
612 return Err(CacheError::VerifyFailed { expected: hash });
613 }
614 debug_assert!(!cursor.first_is_wrong_length());
617 let used_outboard = outboard.clone();
620 let ob = PreOrderMemOutboard {
621 root: bao_tree::blake3::Hash::from(*hash.as_bytes()),
622 tree,
623 data: outboard,
624 };
625 let ranges: ChunkRanges = aligned.chunk_ranges().clone();
626 let fault: FaultSlot = Arc::new(Mutex::new(None));
627 let reader = OriginWindowReader::new(cursor, aligned, Arc::clone(&fault));
628 let (tx, rx) = mpsc::channel(WIRE_CHANNEL_CAP);
629 let mut guard = UnfinishedGuard {
630 slot: Arc::clone(&fault),
631 hash,
632 finished: false,
633 };
634 let task = tokio::spawn(async move {
635 let _permit = permit;
636 let mut writer = ChannelWriter { tx };
641 let mut reader = reader;
642 let mut ob = ob;
643 let encode = async {
647 if ranges.is_empty() {
648 Ok(())
649 } else {
650 encode_ranges_validated(&mut reader, &mut ob, ranges.as_ref(), &mut writer)
651 .await
652 }
653 };
654 let result = match AssertUnwindSafe(encode).catch_unwind().await {
658 Ok(result) => result,
659 Err(payload) => {
660 park_panic(&guard.slot, hash, kind, payload.as_ref());
661 Ok(())
663 }
664 };
665 match result {
666 Ok(()) => {}
667 Err(EncodeError::ParentHashMismatch(_) | EncodeError::LeafHashMismatch(_)) => {
668 tracing::warn!(
669 %hash,
670 ?kind,
671 "own origin served a range that failed bao verification against H; \
672 hard local-origin fault (no degrade — committed to serving under H)",
673 );
674 outboards.evict_if_same(hash, &used_outboard);
675 park_fault(&guard.slot, CacheError::VerifyFailed { expected: hash });
676 }
677 Err(EncodeError::Io(e)) => park_fault(
681 &guard.slot,
682 CacheError::OriginError {
683 hash,
684 source: anyhow::Error::new(e).context("origin range encode failed"),
685 },
686 ),
687 Err(other) => park_fault(
688 &guard.slot,
689 CacheError::OriginError {
690 hash,
691 source: anyhow::anyhow!("origin range encode failed: {other}"),
692 },
693 ),
694 }
695 guard.disarm();
696 drop(writer);
697 });
698 Ok(Self {
699 rx,
700 fault,
701 ended: false,
702 _task: AbortOnDropHandle::new(task),
703 })
704 }
705
706 pub async fn next_chunk(&mut self) -> Option<CacheResult<Bytes>> {
710 if self.ended {
711 return None;
712 }
713 if let Some(chunk) = self.rx.recv().await {
714 return Some(Ok(chunk));
715 }
716 self.ended = true;
717 self.fault
718 .lock()
719 .unwrap_or_else(PoisonError::into_inner)
720 .take()
721 .map(Err)
722 }
723}
724
725#[cfg(test)]
726#[allow(
727 clippy::unwrap_used,
728 clippy::expect_used,
729 clippy::indexing_slicing,
730 clippy::cast_possible_truncation,
731 reason = "tests"
732)]
733mod tests;