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 //! Yields the current task and wakes it for a reschedule.
15 use std::future::Future;
16 use std::pin::Pin;
17 use std::task::{Context, Poll};
18
19 cfg_not_ffrt!(
20 use crate::executor::worker;
21 use crate::executor::worker::WorkerContext;
22 );
23
24 /// Yields the current task and wakes it for a reschedule.
25 /// # Examples
26 ///
27 /// ```
28 /// use ylong_runtime::task::*;
29 ///
30 /// let res = ylong_runtime::block_on(ylong_runtime::spawn(async {
31 /// yield_now().await;
32 /// }))
33 /// .unwrap();
34 /// assert_eq!(res, ());
35 /// ```
yield_now()36 pub async fn yield_now() {
37 YieldTask(false).await
38 }
39
40 struct YieldTask(bool);
41
42 impl Future for YieldTask {
43 type Output = ();
44
45 #[cfg(not(feature = "ffrt"))]
poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output>46 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
47 if !self.0 {
48 self.0 = true;
49 let ctx = worker::get_current_ctx();
50
51 // Under worker context, we push the waker into the yielded list owned by the
52 // worker to avoid waking the waker immediately. This is because
53 // waking the waker in a worker context will put the task in the
54 // lifo slot, we don't want that.
55 if let Some(ctx) = ctx {
56 let mut yielded = ctx.worker.yielded.borrow_mut();
57 yielded.push(cx.waker().clone());
58 } else {
59 cx.waker().wake_by_ref();
60 }
61 Poll::Pending
62 } else {
63 Poll::Ready(())
64 }
65 }
66
67 #[cfg(feature = "ffrt")]
poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output>68 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
69 if !self.0 {
70 self.0 = true;
71 cx.waker().wake_by_ref();
72 Poll::Pending
73 } else {
74 Poll::Ready(())
75 }
76 }
77 }
78
79 #[cfg(not(feature = "ffrt"))]
wake_yielded_tasks(worker_ctx: &WorkerContext)80 pub(crate) fn wake_yielded_tasks(worker_ctx: &WorkerContext) {
81 let mut yielded = worker_ctx.worker.yielded.borrow_mut();
82 if yielded.is_empty() {
83 return;
84 }
85 for waker in yielded.drain(..) {
86 waker.wake();
87 }
88 }
89
90 #[cfg(test)]
91 mod test {
92 use crate::task::yield_now;
93
94 /// UT test cases for yield.
95 ///
96 /// # Brief
97 /// 1. Create two tasks that adds a number to 1000
98 /// 2. Make the tasks yield during adding
99 /// 3. Check if value is correct
100 #[test]
ut_yield_now()101 fn ut_yield_now() {
102 let handle = crate::spawn(async move {
103 let mut i = 0;
104 for _ in 0..1000 {
105 i += 1;
106 yield_now().await;
107 }
108 i
109 });
110 let handle2 = crate::spawn(async move {
111 let mut i = 0;
112 for _ in 0..1000 {
113 i += 1;
114 yield_now().await;
115 }
116 i
117 });
118 let ret = crate::block_on(handle).unwrap();
119 assert_eq!(ret, 1000);
120 let ret = crate::block_on(handle2).unwrap();
121 assert_eq!(ret, 1000);
122 }
123
124 /// UT test cases for calling yield inside block_on.
125 ///
126 /// # Brief
127 /// 1. Blocks on an async function that adds a number to 1000
128 /// 2. Checks if the returned value is correct
129 #[test]
ut_yield_block_on()130 fn ut_yield_block_on() {
131 let ret = crate::block_on(async move {
132 let mut i = 0;
133 for _ in 0..1000 {
134 i += 1;
135 yield_now().await;
136 }
137 i
138 });
139 assert_eq!(ret, 1000);
140 }
141 }
142