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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
use pin_project_lite::pin_project;
use std::pin::Pin;
use std::task::{Context, Poll};
use super::{Frame, SizeHint, StreamingBody};
pin_project! {
/// A [`StreamingBody`] wrapper that owns an opaque `guard` and releases it
/// when the body is **fully consumed on the wire** ([`StreamingBody::poll_frame`]
/// returns `Poll::Ready(None)`), or on drop if the body is abandoned early.
///
/// This ties the guard's lifetime to the body it accompanies rather than to
/// the response headers. The motivating use is a pooled connection handle: the
/// connector returns the response at headers and immediately drops its handle,
/// but the connection is logically still in use until the body has been read.
/// By moving the handle in here, its release (e.g. freeing a stream slot or
/// returning the connection to the pool) happens when the body is done.
///
/// The guard is released exactly once: either at end-of-stream or, for an
/// early drop (e.g. a cancelled/aborted read), when this body is dropped.
#[must_use = "GuardedBody does nothing unless polled as a response body"]
pub struct GuardedBody<B, G> {
#[pin]
inner: B,
guard: Option<G>,
}
}
impl<B, G> GuardedBody<B, G> {
/// Wrap `inner`, holding `guard` until the body reaches end-of-stream or is dropped.
pub fn new(inner: B, guard: G) -> Self {
Self {
inner,
guard: Some(guard),
}
}
}
impl<B, G> StreamingBody for GuardedBody<B, G>
where
B: StreamingBody,
{
type Data = B::Data;
type Error = B::Error;
fn poll_frame(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
let this = self.project();
let result = this.inner.poll_frame(cx);
if matches!(result, Poll::Ready(None | Some(Err(_)))) {
// Fully consumed (or errored) on the wire: release the guard now,
// rather than waiting for this body to be dropped.
*this.guard = None;
}
result
}
#[inline(always)]
fn is_end_stream(&self) -> bool {
self.inner.is_end_stream()
}
#[inline(always)]
fn size_hint(&self) -> SizeHint {
self.inner.size_hint()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::body::util::{BodyExt, Empty, Full};
use rama_core::bytes::Bytes;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
/// A guard whose drop flips a shared flag, so tests can observe release timing.
struct ReleaseFlag(Arc<AtomicBool>);
impl Drop for ReleaseFlag {
fn drop(&mut self) {
self.0.store(true, Ordering::Relaxed);
}
}
#[tokio::test]
async fn releases_at_end_of_stream() {
let released = Arc::new(AtomicBool::new(false));
let body = GuardedBody::new(
Full::<Bytes>::from(Bytes::from("hello")),
ReleaseFlag(released.clone()),
);
// Not released until the stream is actually exhausted.
assert!(!released.load(Ordering::Relaxed));
body.collect().await.unwrap();
assert!(released.load(Ordering::Relaxed));
}
#[tokio::test]
async fn releases_on_early_drop() {
let released = Arc::new(AtomicBool::new(false));
// Consume only the single data frame, then drop before end-of-stream.
let mut body = GuardedBody::new(
Full::<Bytes>::from(Bytes::from("hello")),
ReleaseFlag(released.clone()),
);
let _ = body.frame().await;
assert!(!released.load(Ordering::Relaxed));
drop(body);
assert!(released.load(Ordering::Relaxed));
}
#[tokio::test]
async fn releases_once_on_empty_body() {
let released = Arc::new(AtomicBool::new(false));
GuardedBody::new(Empty::<Bytes>::new(), ReleaseFlag(released.clone()))
.collect()
.await
.unwrap();
assert!(released.load(Ordering::Relaxed));
}
}