@@ -2,6 +2,7 @@ extern crate futures;
2
2
extern crate tokio_core;
3
3
extern crate tokio_io;
4
4
extern crate tokio_service;
5
+ extern crate tokio_io_timeout;
5
6
extern crate hyper;
6
7
7
8
use std:: time:: Duration ;
@@ -12,6 +13,7 @@ use futures::future::{Either, Future};
12
13
use tokio_core:: reactor:: { Handle , Timeout } ;
13
14
use tokio_io:: { AsyncRead , AsyncWrite } ;
14
15
use tokio_service:: Service ;
16
+ use tokio_io_timeout:: TimeoutStream ;
15
17
16
18
use hyper:: client:: Connect ;
17
19
@@ -23,46 +25,94 @@ pub struct TimeoutConnector<T> {
23
25
/// Handle to be used to set the timeout within tokio's core
24
26
handle : Handle ,
25
27
/// Amount of time to wait connecting
26
- connect_timeout : Duration ,
28
+ connect_timeout : Option < Duration > ,
29
+ /// Amount of time to wait reading response
30
+ read_timeout : Option < Duration > ,
31
+ /// Amount of time to wait writing request
32
+ write_timeout : Option < Duration > ,
27
33
}
28
34
29
35
impl < T : Connect > TimeoutConnector < T > {
30
36
/// Construct a new TimeoutConnector with a given connector implementing the `Connect` trait
31
- pub fn new ( connector : T , handle : & Handle , timeout : Duration ) -> Self {
37
+ pub fn new ( connector : T , handle : & Handle ) -> Self {
32
38
TimeoutConnector {
33
39
connector : connector,
34
40
handle : handle. clone ( ) ,
35
- connect_timeout : timeout,
41
+ connect_timeout : None ,
42
+ read_timeout : None ,
43
+ write_timeout : None ,
36
44
}
37
45
}
46
+
47
+ /// Set the timeout for connecting to a URL.
48
+ ///
49
+ /// Default is no timeout.
50
+ #[ inline]
51
+ pub fn set_connect_timeout ( & mut self , val : Option < Duration > ) {
52
+ self . connect_timeout = val;
53
+ }
54
+
55
+ /// Set the timeout for the response.
56
+ ///
57
+ /// Default is no timeout.
58
+ #[ inline]
59
+ pub fn set_read_timeout ( & mut self , val : Option < Duration > ) {
60
+ self . read_timeout = val;
61
+ }
62
+
63
+ /// Set the timeout for the request.
64
+ ///
65
+ /// Default is no timeout.
66
+ #[ inline]
67
+ pub fn set_write_timeout ( & mut self , val : Option < Duration > ) {
68
+ self . write_timeout = val;
69
+ }
38
70
}
39
71
40
72
impl < T > Service for TimeoutConnector < T >
41
- where T : Service < Error =io:: Error > + ' static ,
42
- T :: Response : AsyncRead + AsyncWrite ,
43
- T :: Future : Future < Error =io:: Error > ,
73
+ where
74
+ T : Service < Error = io:: Error > + ' static ,
75
+ T :: Response : AsyncRead + AsyncWrite ,
76
+ T :: Future : Future < Error = io:: Error > ,
44
77
{
45
78
type Request = T :: Request ;
46
- type Response = T :: Response ;
79
+ type Response = TimeoutStream < T :: Response > ;
47
80
type Error = T :: Error ;
48
- type Future = Box < Future < Item = Self :: Response , Error = Self :: Error > > ;
81
+ type Future = Box < Future < Item = Self :: Response , Error = Self :: Error > > ;
49
82
50
83
fn call ( & self , req : Self :: Request ) -> Self :: Future {
84
+ let handle = self . handle . clone ( ) ;
85
+ let read_timeout = self . read_timeout . clone ( ) ;
86
+ let write_timeout = self . write_timeout . clone ( ) ;
51
87
let connecting = self . connector . call ( req) ;
52
- let timeout = Timeout :: new ( self . connect_timeout , & self . handle ) . unwrap ( ) ;
53
-
54
- Box :: new ( connecting. select2 ( timeout) . then ( |res| {
55
- match res {
56
- Ok ( Either :: A ( ( io, _) ) ) => Ok ( io) ,
57
- Ok ( Either :: B ( ( _, _) ) ) => {
58
- Err ( io:: Error :: new (
59
- io:: ErrorKind :: TimedOut ,
60
- "Client timed out while connecting"
61
- ) )
62
- }
63
- Err ( Either :: A ( ( e, _) ) ) => Err ( e) ,
64
- Err ( Either :: B ( ( e, _) ) ) => Err ( e) ,
88
+
89
+ if self . connect_timeout . is_none ( ) {
90
+ return Box :: new ( connecting. map ( move |io| {
91
+ let mut tm = TimeoutStream :: new ( io, & handle) ;
92
+ tm. set_read_timeout ( read_timeout) ;
93
+ tm. set_write_timeout ( write_timeout) ;
94
+ tm
95
+ } ) ) ;
96
+ }
97
+
98
+ let connect_timeout = self . connect_timeout . expect ( "Connect timeout should be set" ) ;
99
+ let timeout = Timeout :: new ( connect_timeout, & self . handle ) . unwrap ( ) ;
100
+
101
+ Box :: new ( connecting. select2 ( timeout) . then ( move |res| match res {
102
+ Ok ( Either :: A ( ( io, _) ) ) => {
103
+ let mut tm = TimeoutStream :: new ( io, & handle) ;
104
+ tm. set_read_timeout ( read_timeout) ;
105
+ tm. set_write_timeout ( write_timeout) ;
106
+ Ok ( tm)
65
107
}
108
+ Ok ( Either :: B ( ( _, _) ) ) => {
109
+ Err ( io:: Error :: new (
110
+ io:: ErrorKind :: TimedOut ,
111
+ "Client timed out while connecting" ,
112
+ ) )
113
+ }
114
+ Err ( Either :: A ( ( e, _) ) ) => Err ( e) ,
115
+ Err ( Either :: B ( ( e, _) ) ) => Err ( e) ,
66
116
} ) )
67
117
}
68
118
}
@@ -80,12 +130,15 @@ mod tests {
80
130
let mut core = Core :: new ( ) . unwrap ( ) ;
81
131
// 10.255.255.1 is a not a routable IP address
82
132
let url = "http://10.255.255.1" . parse ( ) . unwrap ( ) ;
83
- let connector = TimeoutConnector :: new (
84
- HttpConnector :: new ( 1 , & core. handle ( ) ) ,
85
- & core. handle ( ) ,
86
- Duration :: from_millis ( 1 )
87
- ) ;
133
+ let mut connector =
134
+ TimeoutConnector :: new ( HttpConnector :: new ( 1 , & core. handle ( ) ) , & core. handle ( ) ) ;
135
+ connector. set_connect_timeout ( Some ( Duration :: from_millis ( 1 ) ) ) ;
88
136
89
- assert_eq ! ( core. run( connector. connect( url) ) . unwrap_err( ) . kind( ) , io:: ErrorKind :: TimedOut ) ;
137
+ match core. run ( connector. connect ( url) ) {
138
+ Err ( e) => {
139
+ assert_eq ! ( e. kind( ) , io:: ErrorKind :: TimedOut ) ;
140
+ }
141
+ _ => panic ! ( "Expected timeout error" ) ,
142
+ }
90
143
}
91
144
}
0 commit comments