@@ -5,6 +5,7 @@ use std::fmt;
5
5
use std:: io:: { self , Read , Write } ;
6
6
use std:: net:: { SocketAddr , Shutdown } ;
7
7
use std:: sync:: { Arc , Mutex } ;
8
+ use std:: sync:: atomic:: { AtomicBool , Ordering } ;
8
9
9
10
use std:: time:: Duration ;
10
11
@@ -130,7 +131,7 @@ impl<C: NetworkConnector<Stream=S>, S: NetworkStream + Send> NetworkConnector fo
130
131
} ;
131
132
Ok ( PooledStream {
132
133
inner : Some ( inner) ,
133
- is_closed : false ,
134
+ is_closed : AtomicBool :: new ( false ) ,
134
135
pool : self . inner . clone ( ) ,
135
136
} )
136
137
}
@@ -139,7 +140,7 @@ impl<C: NetworkConnector<Stream=S>, S: NetworkStream + Send> NetworkConnector fo
139
140
/// A Stream that will try to be returned to the Pool when dropped.
140
141
pub struct PooledStream < S > {
141
142
inner : Option < PooledStreamInner < S > > ,
142
- is_closed : bool ,
143
+ is_closed : AtomicBool ,
143
144
pool : Arc < Mutex < PoolImpl < S > > > ,
144
145
}
145
146
@@ -148,7 +149,7 @@ impl<S> fmt::Debug for PooledStream<S> where S: fmt::Debug + 'static {
148
149
fn fmt ( & self , fmt : & mut fmt:: Formatter ) -> fmt:: Result {
149
150
fmt. debug_struct ( "PooledStream" )
150
151
. field ( "inner" , & self . inner )
151
- . field ( "is_closed" , & self . is_closed )
152
+ . field ( "is_closed" , & self . is_closed . load ( Ordering :: Relaxed ) )
152
153
. field ( "pool" , & self . pool )
153
154
. finish ( )
154
155
}
@@ -176,7 +177,7 @@ impl<S: NetworkStream> Read for PooledStream<S> {
176
177
// if the wrapped stream returns EOF (Ok(0)), that means the
177
178
// server has closed the stream. we must be sure this stream
178
179
// is dropped and not put back into the pool.
179
- self . is_closed = true ;
180
+ self . is_closed . store ( true , Ordering :: Relaxed ) ;
180
181
Ok ( 0 )
181
182
} ,
182
183
r => r
@@ -200,21 +201,33 @@ impl<S: NetworkStream> NetworkStream for PooledStream<S> {
200
201
#[ inline]
201
202
fn peer_addr ( & mut self ) -> io:: Result < SocketAddr > {
202
203
self . inner . as_mut ( ) . unwrap ( ) . stream . peer_addr ( )
204
+ . map_err ( |e| {
205
+ self . is_closed . store ( true , Ordering :: Relaxed ) ;
206
+ e
207
+ } )
203
208
}
204
209
205
210
#[ inline]
206
211
fn set_read_timeout ( & self , dur : Option < Duration > ) -> io:: Result < ( ) > {
207
212
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
+ } )
208
217
}
209
218
210
219
#[ inline]
211
220
fn set_write_timeout ( & self , dur : Option < Duration > ) -> io:: Result < ( ) > {
212
221
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
+ } )
213
226
}
214
227
215
228
#[ inline]
216
229
fn close ( & mut self , how : Shutdown ) -> io:: Result < ( ) > {
217
- self . is_closed = true ;
230
+ self . is_closed . store ( true , Ordering :: Relaxed ) ;
218
231
self . inner . as_mut ( ) . unwrap ( ) . stream . close ( how)
219
232
}
220
233
@@ -234,8 +247,9 @@ impl<S: NetworkStream> NetworkStream for PooledStream<S> {
234
247
235
248
impl < S > Drop for PooledStream < S > {
236
249
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 {
239
253
self . inner . take ( ) . map ( |inner| {
240
254
if let Ok ( mut pool) = self . pool . lock ( ) {
241
255
pool. reuse ( inner. key . clone ( ) , inner) ;
0 commit comments