1 // Copyright (c) 2023 Huawei Device Co., Ltd.
2 // Licensed under the Apache License, Version 2.0 (the "License");
3 // you may not use this file except in compliance with the License.
4 // You may obtain a copy of the License at
5 //
6 //     http://www.apache.org/licenses/LICENSE-2.0
7 //
8 // Unless required by applicable law or agreed to in writing, software
9 // distributed under the License is distributed on an "AS IS" BASIS,
10 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11 // See the License for the specific language governing permissions and
12 // limitations under the License.
13 
14 use std::future::Future;
15 use std::pin::Pin;
16 use std::task::{Context, Poll};
17 use std::time::Duration;
18 
19 use super::Response;
20 use crate::error::HttpClientError;
21 use crate::runtime::{sleep, Sleep};
22 
23 pub(crate) struct TimeoutFuture<T> {
24     pub(crate) timeout: Option<Pin<Box<Sleep>>>,
25     pub(crate) future: T,
26 }
27 
28 impl<T> TimeoutFuture<Pin<Box<T>>> {
new(future: T, timeout: Duration) -> Self29     pub(crate) fn new(future: T, timeout: Duration) -> Self {
30         Self {
31             timeout: Some(Box::pin(sleep(timeout))),
32             future: Box::pin(future),
33         }
34     }
35 }
36 
37 impl<T> Future for TimeoutFuture<T>
38 where
39     T: Future<Output = Result<Response, HttpClientError>> + Unpin,
40 {
41     type Output = Result<Response, HttpClientError>;
42 
poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output>43     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
44         let this = self.get_mut();
45 
46         if let Some(delay) = this.timeout.as_mut() {
47             if let Poll::Ready(()) = delay.as_mut().poll(cx) {
48                 return Poll::Ready(err_from_io!(Timeout, std::io::ErrorKind::TimedOut.into()));
49             }
50         }
51         match Pin::new(&mut this.future).poll(cx) {
52             Poll::Ready(Ok(mut response)) => {
53                 response.body_mut().set_sleep(this.timeout.take());
54                 Poll::Ready(Ok(response))
55             }
56             Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
57             Poll::Pending => Poll::Pending,
58         }
59     }
60 }
61 
62 #[cfg(all(test, feature = "ylong_base"))]
63 mod ut_timeout {
64     use std::sync::Arc;
65 
66     use ylong_http::response::status::StatusCode;
67     use ylong_http::response::{Response, ResponsePart};
68     use ylong_http::version::Version;
69 
70     use crate::async_impl::interceptor::IdleInterceptor;
71     use crate::async_impl::timeout::TimeoutFuture;
72     use crate::async_impl::HttpBody;
73     use crate::util::normalizer::BodyLength;
74     use crate::HttpClientError;
75 
76     /// UT test cases for `TimeoutFuture`.
77     ///
78     /// # Brief
79     /// 1. Creates a `Future`.
80     /// 2. Calls `ylong_runtime::block_on` to run the future.
81     /// 3. Checks if result is correct.
82     #[test]
ut_timeout_future()83     fn ut_timeout_future() {
84         let future1 = Box::pin(async {
85             let part = ResponsePart {
86                 version: Version::HTTP1_1,
87                 status: StatusCode::OK,
88                 headers: Default::default(),
89             };
90             let body = HttpBody::new(
91                 Arc::new(IdleInterceptor),
92                 BodyLength::Empty,
93                 Box::new([].as_slice()),
94                 &[],
95             )
96             .unwrap();
97             Ok(crate::async_impl::Response::new(Response::from_raw_parts(
98                 part, body,
99             )))
100         });
101 
102         let future2 = Box::pin(async {
103             Result::<crate::async_impl::Response, HttpClientError>::Err(
104                 HttpClientError::user_aborted(),
105             )
106         });
107 
108         let time_future1 = TimeoutFuture {
109             timeout: None,
110             future: future1,
111         };
112 
113         let time_future2 = TimeoutFuture {
114             timeout: None,
115             future: future2,
116         };
117 
118         assert!(ylong_runtime::block_on(time_future1).is_ok());
119         assert!(ylong_runtime::block_on(time_future2).is_err());
120     }
121 }
122