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 18 use ylong_http::response::Response; 19 20 use crate::async_impl::HttpBody; 21 use crate::{ErrorKind, HttpClientError, 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> Future for TimeoutFuture<T> 29 where 30 T: Future<Output = Result<Response<HttpBody>, HttpClientError>> + Unpin, 31 { 32 type Output = Result<Response<HttpBody>, HttpClientError>; 33 poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output>34 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { 35 let this = self.get_mut(); 36 37 if let Some(delay) = this.timeout.as_mut() { 38 if let Poll::Ready(()) = delay.as_mut().poll(cx) { 39 return Poll::Ready(Err(HttpClientError::new_with_message( 40 ErrorKind::Timeout, 41 "Request timeout", 42 ))); 43 } 44 } 45 match Pin::new(&mut this.future).poll(cx) { 46 Poll::Ready(Ok(mut response)) => { 47 response.body_mut().set_sleep(this.timeout.take()); 48 Poll::Ready(Ok(response)) 49 } 50 Poll::Ready(Err(e)) => Poll::Ready(Err(e)), 51 Poll::Pending => Poll::Pending, 52 } 53 } 54 } 55 56 #[cfg(all(test, feature = "ylong_base"))] 57 mod ut_timeout { 58 use ylong_http::response::status::StatusCode; 59 use ylong_http::response::{Response, ResponsePart}; 60 use ylong_http::version::Version; 61 62 use crate::async_impl::timeout::TimeoutFuture; 63 use crate::async_impl::HttpBody; 64 use crate::util::normalizer::BodyLength; 65 use crate::HttpClientError; 66 67 /// UT test cases for `TimeoutFuture`. 68 /// 69 /// # Brief 70 /// 1. Creates a `Future`. 71 /// 2. Calls `ylong_runtime::block_on` to run the future. 72 /// 3. Checks if result is correct. 73 #[test] ut_timeout_future()74 fn ut_timeout_future() { 75 let future1 = Box::pin(async { 76 let part = ResponsePart { 77 version: Version::HTTP1_1, 78 status: StatusCode::OK, 79 headers: Default::default(), 80 }; 81 let body = HttpBody::new(BodyLength::Empty, Box::new([].as_slice()), &[]).unwrap(); 82 Ok(Response::from_raw_parts(part, body)) 83 }); 84 85 let future2 = Box::pin(async { 86 Result::<Response<HttpBody>, HttpClientError>::Err(HttpClientError::user_aborted()) 87 }); 88 89 let time_future1 = TimeoutFuture { 90 timeout: None, 91 future: future1, 92 }; 93 94 let time_future2 = TimeoutFuture { 95 timeout: None, 96 future: future2, 97 }; 98 99 assert!(ylong_runtime::block_on(time_future1).is_ok()); 100 assert!(ylong_runtime::block_on(time_future2).is_err()); 101 } 102 } 103