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
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
|
// Copyright 2018-2024 the Deno authors. All rights reserved. MIT license.
use std::pin::Pin;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::task::Context;
use std::task::Poll;
use bytes::Bytes;
use deno_core::futures::Stream;
use pin_project::pin_project;
use tokio::io::AsyncRead;
use tokio_util::io::ReaderStream;
/// [ExternallyAbortableByteStream] adapts a [tokio::AsyncRead] into a [Stream].
/// It is used to bridge between the HTTP response body resource, and
/// `hyper::Body`. The stream has the special property that it errors if the
/// underlying reader is closed before an explicit EOF is sent (in the form of
/// setting the `shutdown` flag to true).
#[pin_project]
pub struct ExternallyAbortableReaderStream<R: AsyncRead> {
#[pin]
inner: ReaderStream<R>,
done: Arc<AtomicBool>,
}
pub struct ShutdownHandle(Arc<AtomicBool>);
impl ShutdownHandle {
pub fn shutdown(&self) {
self.0.store(true, std::sync::atomic::Ordering::SeqCst);
}
}
impl<R: AsyncRead> ExternallyAbortableReaderStream<R> {
pub fn new(reader: R) -> (Self, ShutdownHandle) {
let done = Arc::new(AtomicBool::new(false));
let this = Self {
inner: ReaderStream::new(reader),
done: done.clone(),
};
(this, ShutdownHandle(done))
}
}
impl<R: AsyncRead> Stream for ExternallyAbortableReaderStream<R> {
type Item = std::io::Result<Bytes>;
fn poll_next(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<Self::Item>> {
let this = self.project();
let val = std::task::ready!(this.inner.poll_next(cx));
match val {
None if this.done.load(Ordering::SeqCst) => Poll::Ready(None),
None => Poll::Ready(Some(Err(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
"stream reader has shut down",
)))),
Some(val) => Poll::Ready(Some(val)),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use bytes::Bytes;
use deno_core::futures::StreamExt;
use tokio::io::AsyncWriteExt;
#[tokio::test]
async fn success() {
let (a, b) = tokio::io::duplex(64 * 1024);
let (reader, _) = tokio::io::split(a);
let (_, mut writer) = tokio::io::split(b);
let (mut stream, shutdown_handle) =
ExternallyAbortableReaderStream::new(reader);
writer.write_all(b"hello").await.unwrap();
assert_eq!(stream.next().await.unwrap().unwrap(), Bytes::from("hello"));
writer.write_all(b"world").await.unwrap();
assert_eq!(stream.next().await.unwrap().unwrap(), Bytes::from("world"));
shutdown_handle.shutdown();
writer.shutdown().await.unwrap();
drop(writer);
assert!(stream.next().await.is_none());
}
#[tokio::test]
async fn error() {
let (a, b) = tokio::io::duplex(64 * 1024);
let (reader, _) = tokio::io::split(a);
let (_, mut writer) = tokio::io::split(b);
let (mut stream, _shutdown_handle) =
ExternallyAbortableReaderStream::new(reader);
writer.write_all(b"hello").await.unwrap();
assert_eq!(stream.next().await.unwrap().unwrap(), Bytes::from("hello"));
drop(writer);
assert_eq!(
stream.next().await.unwrap().unwrap_err().kind(),
std::io::ErrorKind::UnexpectedEof
);
}
#[tokio::test]
async fn error2() {
let (a, b) = tokio::io::duplex(64 * 1024);
let (reader, _) = tokio::io::split(a);
let (_, mut writer) = tokio::io::split(b);
let (mut stream, _shutdown_handle) =
ExternallyAbortableReaderStream::new(reader);
writer.write_all(b"hello").await.unwrap();
assert_eq!(stream.next().await.unwrap().unwrap(), Bytes::from("hello"));
writer.shutdown().await.unwrap();
drop(writer);
assert_eq!(
stream.next().await.unwrap().unwrap_err().kind(),
std::io::ErrorKind::UnexpectedEof
);
}
#[tokio::test]
async fn write_after_shutdown() {
let (a, b) = tokio::io::duplex(64 * 1024);
let (reader, _) = tokio::io::split(a);
let (_, mut writer) = tokio::io::split(b);
let (mut stream, shutdown_handle) =
ExternallyAbortableReaderStream::new(reader);
writer.write_all(b"hello").await.unwrap();
assert_eq!(stream.next().await.unwrap().unwrap(), Bytes::from("hello"));
writer.write_all(b"world").await.unwrap();
assert_eq!(stream.next().await.unwrap().unwrap(), Bytes::from("world"));
shutdown_handle.shutdown();
writer.shutdown().await.unwrap();
assert!(writer.write_all(b"!").await.is_err());
drop(writer);
assert!(stream.next().await.is_none());
}
}
|
