@@ -5,6 +5,7 @@ use std::fmt;
55use std:: io:: { self , Read , Write } ;
66use std:: net:: { SocketAddr , Shutdown } ;
77use std:: sync:: { Arc , Mutex } ;
8+ use std:: sync:: atomic:: { AtomicBool , Ordering } ;
89
910use 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.
140141pub 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
235248impl < 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