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
|
//! Wrapper around [`tokio::net::TcpStream`]. This allows in-flight requests to be cancelled
//! immediately instead of after the socket times out.
use futures::channel::oneshot;
use hyper::client::connect::{Connected, Connection};
use std::{
io,
net::Shutdown,
pin::Pin,
sync::{Arc, Mutex, Weak},
task::{Context, Poll},
};
use tokio::{
io::{AsyncRead, AsyncWrite, ReadBuf},
net::TcpStream as TokioTcpStream,
};
#[derive(Debug)]
pub struct TcpStreamHandle {
inner: Weak<Mutex<Option<StreamInner>>>,
}
impl TcpStreamHandle {
pub fn close(self) {
if let Some(inner_lock) = self.inner.upgrade() {
if let Ok(Some(inner)) = inner_lock.lock().map(|mut inner| inner.take()) {
if let Err(err) = flatten_result(
inner
.stream
.into_std()
.map(|stream| stream.shutdown(Shutdown::Both)),
) {
log::error!("Failed to shut down TCP socket: {}", err);
}
}
}
}
}
pub struct TcpStream {
inner: Arc<Mutex<Option<StreamInner>>>,
}
impl TcpStream {
pub fn new(
stream: TokioTcpStream,
shutdown_tx: Option<oneshot::Sender<()>>,
) -> (Self, TcpStreamHandle) {
let inner = Arc::new(Mutex::new(Some(StreamInner {
stream,
shutdown_tx,
})));
let stream_handle = TcpStreamHandle {
inner: Arc::downgrade(&inner),
};
(Self { inner }, stream_handle)
}
fn do_stream<T>(
&self,
mut stream_fn: impl FnMut(&mut TokioTcpStream) -> T,
closed_value: T,
) -> T {
let mut inner = self.inner.lock().expect("TCP lock poisoned");
if let Some(inner) = &mut *inner {
stream_fn(&mut inner.stream)
} else {
closed_value
}
}
}
impl Drop for TcpStream {
fn drop(&mut self) {
if let Ok(Some(mut inner)) = self.inner.lock().map(|mut inner| inner.take()) {
if let Some(tx) = inner.shutdown_tx.take() {
let _ = tx.send(());
}
}
}
}
impl AsyncWrite for TcpStream {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
self.do_stream(
|stream| Pin::new(stream).poll_write(cx, buf),
Poll::Ready(Err(io::Error::new(
io::ErrorKind::ConnectionReset,
"socket is closed",
))),
)
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
self.do_stream(
|stream| Pin::new(stream).poll_flush(cx),
Poll::Ready(Ok(())),
)
}
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
self.do_stream(
|stream| Pin::new(stream).poll_shutdown(cx),
Poll::Ready(Ok(())),
)
}
}
impl AsyncRead for TcpStream {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
self.do_stream(
|stream| Pin::new(stream).poll_read(cx, buf),
Poll::Ready(Err(io::Error::new(
io::ErrorKind::ConnectionReset,
"socket is closed",
))),
)
}
}
impl Connection for TcpStream {
fn connected(&self) -> Connected {
Connected::new()
}
}
#[derive(Debug)]
struct StreamInner {
stream: TokioTcpStream,
shutdown_tx: Option<oneshot::Sender<()>>,
}
fn flatten_result<T, E>(result: Result<Result<T, E>, E>) -> Result<T, E> {
match result {
Ok(value) => value,
Err(err) => Err(err),
}
}
|
