Skip to content

Commit d5ffee2

Browse files
committed
fix(client): close Pooled streams on sockopt error
1 parent 4bc4f1f commit d5ffee2

File tree

1 file changed

+21
-7
lines changed

1 file changed

+21
-7
lines changed

src/client/pool.rs

+21-7
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ use std::fmt;
55
use std::io::{self, Read, Write};
66
use std::net::{SocketAddr, Shutdown};
77
use std::sync::{Arc, Mutex};
8+
use std::sync::atomic::{AtomicBool, Ordering};
89

910
use std::time::Duration;
1011

@@ -130,7 +131,7 @@ impl<C: NetworkConnector<Stream=S>, S: NetworkStream + Send> NetworkConnector fo
130131
};
131132
Ok(PooledStream {
132133
inner: Some(inner),
133-
is_closed: false,
134+
is_closed: AtomicBool::new(false),
134135
pool: self.inner.clone(),
135136
})
136137
}
@@ -139,7 +140,7 @@ impl<C: NetworkConnector<Stream=S>, S: NetworkStream + Send> NetworkConnector fo
139140
/// A Stream that will try to be returned to the Pool when dropped.
140141
pub struct PooledStream<S> {
141142
inner: Option<PooledStreamInner<S>>,
142-
is_closed: bool,
143+
is_closed: AtomicBool,
143144
pool: Arc<Mutex<PoolImpl<S>>>,
144145
}
145146

@@ -148,7 +149,7 @@ impl<S> fmt::Debug for PooledStream<S> where S: fmt::Debug + 'static {
148149
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
149150
fmt.debug_struct("PooledStream")
150151
.field("inner", &self.inner)
151-
.field("is_closed", &self.is_closed)
152+
.field("is_closed", &self.is_closed.load(Ordering::Relaxed))
152153
.field("pool", &self.pool)
153154
.finish()
154155
}
@@ -176,7 +177,7 @@ impl<S: NetworkStream> Read for PooledStream<S> {
176177
// if the wrapped stream returns EOF (Ok(0)), that means the
177178
// server has closed the stream. we must be sure this stream
178179
// is dropped and not put back into the pool.
179-
self.is_closed = true;
180+
self.is_closed.store(true, Ordering::Relaxed);
180181
Ok(0)
181182
},
182183
r => r
@@ -200,21 +201,33 @@ impl<S: NetworkStream> NetworkStream for PooledStream<S> {
200201
#[inline]
201202
fn peer_addr(&mut self) -> io::Result<SocketAddr> {
202203
self.inner.as_mut().unwrap().stream.peer_addr()
204+
.map_err(|e| {
205+
self.is_closed.store(true, Ordering::Relaxed);
206+
e
207+
})
203208
}
204209

205210
#[inline]
206211
fn set_read_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
207212
self.inner.as_ref().unwrap().stream.set_read_timeout(dur)
213+
.map_err(|e| {
214+
self.is_closed.store(true, Ordering::Relaxed);
215+
e
216+
})
208217
}
209218

210219
#[inline]
211220
fn set_write_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
212221
self.inner.as_ref().unwrap().stream.set_write_timeout(dur)
222+
.map_err(|e| {
223+
self.is_closed.store(true, Ordering::Relaxed);
224+
e
225+
})
213226
}
214227

215228
#[inline]
216229
fn close(&mut self, how: Shutdown) -> io::Result<()> {
217-
self.is_closed = true;
230+
self.is_closed.store(true, Ordering::Relaxed);
218231
self.inner.as_mut().unwrap().stream.close(how)
219232
}
220233

@@ -234,8 +247,9 @@ impl<S: NetworkStream> NetworkStream for PooledStream<S> {
234247

235248
impl<S> Drop for PooledStream<S> {
236249
fn drop(&mut self) {
237-
trace!("PooledStream.drop, is_closed={}", self.is_closed);
238-
if !self.is_closed {
250+
let is_closed = self.is_closed.load(Ordering::Relaxed);
251+
trace!("PooledStream.drop, is_closed={}", is_closed);
252+
if !is_closed {
239253
self.inner.take().map(|inner| {
240254
if let Ok(mut pool) = self.pool.lock() {
241255
pool.reuse(inner.key.clone(), inner);

0 commit comments

Comments
 (0)