use bytes::Bytes;
use std::fmt::Debug;
use std::sync::Arc;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RetryKind {
Free,
ViaFactory,
Impossible,
}
pub type RewindFactory = Arc<dyn Fn() -> RequestBody + Send + Sync>;
#[derive(Default)]
pub enum RequestBody {
#[default]
Empty,
Full(Bytes),
Rewindable(RewindFactory),
Streaming(Box<dyn http_body::Body<Data = Bytes, Error = crate::Error> + Unpin + Send>), }
impl Debug for RequestBody {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
RequestBody::Empty => f.write_str("Empty"),
RequestBody::Full(b) => write!(f, "Full({} bytes)", b.len()),
RequestBody::Rewindable(_) => f.write_str("Rewindable(..)"),
RequestBody::Streaming(_) => f.write_str("Streaming(..)"),
}
}
}
impl RequestBody {
pub fn rewindable<F>(f: F) -> Self
where
F: Fn() -> RequestBody + Send + Sync + 'static, {
RequestBody::Rewindable(Arc::new(f))
}
pub fn retry_kind(&self) -> RetryKind {
match self {
RequestBody::Empty | RequestBody::Full(_) => RetryKind::Free,
RequestBody::Rewindable(_) => RetryKind::ViaFactory,
RequestBody::Streaming(_) => RetryKind::Impossible,
}
}
pub fn rewind(&self) -> Option<RequestBody> {
match self {
RequestBody::Empty => Some(RequestBody::Empty),
RequestBody::Full(b) => Some(RequestBody::Full(b.clone())),
RequestBody::Rewindable(f) => Some(f()),
RequestBody::Streaming(_) => None,
}
}
pub fn size_hint(&self) -> Option<u64> {
match self {
RequestBody::Empty => Some(0),
RequestBody::Full(b) => Some(b.len() as u64),
_ => None,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use bytes::Bytes;
use std::pin::Pin;
use std::task::Context;
use std::task::Poll;
#[test]
fn replayability_is_knowable_before_sending() {
assert_eq!(RequestBody::Empty.retry_kind(), RetryKind::Free);
assert_eq!(
RequestBody::Full(Bytes::from_static(b"x")).retry_kind(),
RetryKind::Free
);
}
#[test]
fn rewindable_replays_through_factory() {
let b = RequestBody::rewindable(|| RequestBody::Full(Bytes::from_static(b"same")));
assert_eq!(b.retry_kind(), RetryKind::ViaFactory);
let again = b.rewind().expect("rewindable must rewind");
assert!(matches!(again, RequestBody::Full(ref x) if &x[..] == b"same"));
}
#[test]
fn full_rewind_preserves_the_payload() {
let b = RequestBody::Full(Bytes::from_static(b"abc"));
match b.rewind().expect("Full replays") {
RequestBody::Full(x) => assert_eq!(&x[..], b"abc"),
other => panic!("expected Full, got {other:?}"),
}
}
#[test]
fn a_factory_survives_repeated_replays() {
use std::sync::atomic::{AtomicUsize, Ordering};
let calls = Arc::new(AtomicUsize::new(0));
let c = calls.clone();
let b = RequestBody::rewindable(move || {
c.fetch_add(1, Ordering::SeqCst);
RequestBody::Full(Bytes::from_static(b"same"))
});
for _ in 0..3 {
let again = b.rewind().expect("rewindable replays");
assert!(matches!(again, RequestBody::Full(ref x) if &x[..] == b"same"));
assert_eq!(
b.retry_kind(),
RetryKind::ViaFactory,
"kind doesn't change across replays"
);
}
assert_eq!(calls.load(Ordering::SeqCst), 3);
}
#[test]
fn size_hint_is_known_for_empty_and_full_bodies() {
assert_eq!(RequestBody::Empty.size_hint(), Some(0));
assert_eq!(
RequestBody::Full(Bytes::from_static(b"abcd")).size_hint(),
Some(4)
);
}
struct EmptyStream;
impl http_body::Body for EmptyStream {
type Data = Bytes;
type Error = crate::Error;
fn poll_frame(
self: Pin<&mut Self>,
_: &mut Context<'_>,
) -> Poll<Option<Result<http_body::Frame<Bytes>, Self::Error>>> {
Poll::Ready(None)
}
}
#[test]
fn streaming_is_honest_about_being_unreplayable() {
let b = RequestBody::Streaming(Box::new(EmptyStream));
assert_eq!(b.retry_kind(), RetryKind::Impossible);
assert!(b.rewind().is_none(), "must return None, not panic");
assert_eq!(b.size_hint(), None);
}
}