camber 0.2.0

Opinionated async Rust for IO-bound services on top of Tokio
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100

use std::io::{self, Read};
use std::net::TcpStream;
use std::time::{Duration, Instant};

#[derive(Debug, thiserror::Error)]
pub enum BoundedReadError {
    #[error("stream read timed out after {timeout:?} with {bytes_read} bytes read")]
    Timeout {
        timeout: Duration,
        bytes_read: usize,
    },
    #[error("stream exceeded the {limit}-byte read limit")]
    LimitExceeded { limit: usize },
    #[error("stream read failed: {0}")]
    Io(#[source] io::Error),
    #[error("supporting thread did not finish before {timeout:?}")]
    JoinTimeout { timeout: Duration },
    #[error("supporting thread panicked")]
    ThreadPanicked,
}

pub fn join_thread_bounded<T>(
    thread: &mut Option<std::thread::JoinHandle<T>>,
    timeout: Duration,
) -> Result<T, BoundedReadError> {
    let deadline = Instant::now() + timeout;
    loop {
        match thread.as_ref() {
            Some(handle) if handle.is_finished() => break,
            Some(_) if Instant::now() < deadline => std::thread::yield_now(),
            Some(_) => return Err(BoundedReadError::JoinTimeout { timeout }),
            None => return Err(BoundedReadError::ThreadPanicked),
        }
    }
    let handle = thread.take().ok_or(BoundedReadError::ThreadPanicked)?;
    handle.join().map_err(|_| BoundedReadError::ThreadPanicked)
}

pub fn bounded_read(
    stream: &mut TcpStream,
    timeout: Duration,
    limit: usize,
) -> Result<Box<[u8]>, BoundedReadError> {
    let prior_timeout = stream.read_timeout().map_err(BoundedReadError::Io)?;
    let result = read_with_limit(stream, timeout, limit);
    let restore = stream
        .set_read_timeout(prior_timeout)
        .map_err(BoundedReadError::Io);
    match (result, restore) {
        (Ok(bytes), Ok(())) => Ok(bytes),
        (Err(error), _) => Err(error),
        (Ok(_), Err(error)) => Err(error),
    }
}

fn read_with_limit(
    stream: &mut TcpStream,
    timeout: Duration,
    limit: usize,
) -> Result<Box<[u8]>, BoundedReadError> {
    let deadline = Instant::now() + timeout;
    let mut bytes = Vec::new();
    let mut chunk = [0_u8; 4096];
    loop {
        let remaining = deadline.saturating_duration_since(Instant::now());
        if remaining.is_zero() {
            return Err(BoundedReadError::Timeout {
                timeout,
                bytes_read: bytes.len(),
            });
        }
        stream
            .set_read_timeout(Some(remaining))
            .map_err(BoundedReadError::Io)?;
        match stream.read(&mut chunk) {
            Ok(0) => return Ok(bytes.into_boxed_slice()),
            Ok(count)
                if bytes
                    .len()
                    .checked_add(count)
                    .is_some_and(|length| length <= limit) =>
            {
                bytes.extend_from_slice(&chunk[..count]);
            }
            Ok(_) => return Err(BoundedReadError::LimitExceeded { limit }),
            Err(error)
                if matches!(
                    error.kind(),
                    io::ErrorKind::TimedOut | io::ErrorKind::WouldBlock
                ) =>
            {
                return Err(BoundedReadError::Timeout {
                    timeout,
                    bytes_read: bytes.len(),
                });
            }
            Err(error) => return Err(BoundedReadError::Io(error)),
        }
    }
}