blob: cdc0aff9d2ff66f87e7eec06b822983410c5ebd4 [file] [log] [blame]
Janis Danisevskis93927dd2020-12-23 12:23:08 -08001// Copyright 2020, The Android Open Source Project
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! This module implements the handling of async tasks.
16//! The worker thread has a high priority and a low priority queue. Adding a job to either
17//! will cause one thread to be spawned if none exists. As a compromise between performance
18//! and resource consumption, the thread will linger for about 30 seconds after it has
19//! processed all tasks before it terminates.
20//! Note that low priority tasks are processed only when the high priority queue is empty.
21
Janis Danisevskis40f0e6b2021-02-10 15:48:44 -080022use std::{any::Any, any::TypeId, time::Duration};
Janis Danisevskis93927dd2020-12-23 12:23:08 -080023use std::{
Janis Danisevskis40f0e6b2021-02-10 15:48:44 -080024 collections::{HashMap, VecDeque},
Janis Danisevskis93927dd2020-12-23 12:23:08 -080025 sync::Arc,
26 sync::{Condvar, Mutex, MutexGuard},
27 thread,
28};
29
30#[derive(Debug, PartialEq, Eq)]
31enum State {
32 Exiting,
33 Running,
34}
35
Janis Danisevskis40f0e6b2021-02-10 15:48:44 -080036/// The Shelf allows async tasks to store state across invocations.
37/// Note: Store elves at your own peril ;-).
38#[derive(Debug, Default)]
39pub struct Shelf(HashMap<TypeId, Box<dyn Any + Send>>);
40
41impl Shelf {
42 /// Get a reference to the shelved data of type T. Returns Some if the data exists.
43 pub fn get_downcast_ref<T: Any + Send>(&self) -> Option<&T> {
44 self.0.get(&TypeId::of::<T>()).and_then(|v| v.downcast_ref::<T>())
45 }
46
47 /// Get a mutable reference to the shelved data of type T. If a T was inserted using put,
48 /// get_mut, or get_or_put_with.
49 pub fn get_downcast_mut<T: Any + Send>(&mut self) -> Option<&mut T> {
50 self.0.get_mut(&TypeId::of::<T>()).and_then(|v| v.downcast_mut::<T>())
51 }
52
53 /// Remove the entry of the given type and returns the stored data if it existed.
54 pub fn remove_downcast_ref<T: Any + Send>(&mut self) -> Option<T> {
55 self.0.remove(&TypeId::of::<T>()).and_then(|v| v.downcast::<T>().ok().map(|b| *b))
56 }
57
58 /// Puts data `v` on the shelf. If there already was an entry of type T it is returned.
59 pub fn put<T: Any + Send>(&mut self, v: T) -> Option<T> {
60 self.0
61 .insert(TypeId::of::<T>(), Box::new(v) as Box<dyn Any + Send>)
62 .and_then(|v| v.downcast::<T>().ok().map(|b| *b))
63 }
64
65 /// Gets a mutable reference to the entry of the given type and default creates it if necessary.
66 /// The type must implement Default.
67 pub fn get_mut<T: Any + Send + Default>(&mut self) -> &mut T {
68 self.0
69 .entry(TypeId::of::<T>())
70 .or_insert_with(|| Box::new(T::default()) as Box<dyn Any + Send>)
71 .downcast_mut::<T>()
72 .unwrap()
73 }
Janis Danisevskis7e8b4622021-02-13 10:01:59 -080074
75 /// Gets a mutable reference to the entry of the given type or creates it using the init
76 /// function. Init is not executed if the entry already existed.
77 pub fn get_or_put_with<T: Any + Send, F>(&mut self, init: F) -> &mut T
78 where
79 F: FnOnce() -> T,
80 {
81 self.0
82 .entry(TypeId::of::<T>())
83 .or_insert_with(|| Box::new(init()) as Box<dyn Any + Send>)
84 .downcast_mut::<T>()
85 .unwrap()
86 }
Janis Danisevskis40f0e6b2021-02-10 15:48:44 -080087}
88
Janis Danisevskis93927dd2020-12-23 12:23:08 -080089struct AsyncTaskState {
90 state: State,
91 thread: Option<thread::JoinHandle<()>>,
Janis Danisevskis40f0e6b2021-02-10 15:48:44 -080092 hi_prio_req: VecDeque<Box<dyn FnOnce(&mut Shelf) + Send>>,
93 lo_prio_req: VecDeque<Box<dyn FnOnce(&mut Shelf) + Send>>,
94 /// The store allows tasks to store state across invocations. It is passed to each invocation
95 /// of each task. Tasks need to cooperate on the ids they use for storing state.
96 shelf: Option<Shelf>,
Janis Danisevskis93927dd2020-12-23 12:23:08 -080097}
98
99/// AsyncTask spawns one worker thread on demand to process jobs inserted into
Janis Danisevskis7e8b4622021-02-13 10:01:59 -0800100/// a low and a high priority work queue. The queues are processed FIFO, and low
101/// priority queue is processed if the high priority queue is empty.
102/// Note: Because there is only one worker thread at a time for a given AsyncTask instance,
103/// all scheduled requests are guaranteed to be serialized with respect to one another.
Janis Danisevskis93927dd2020-12-23 12:23:08 -0800104pub struct AsyncTask {
105 state: Arc<(Condvar, Mutex<AsyncTaskState>)>,
106}
107
108impl Default for AsyncTask {
109 fn default() -> Self {
110 Self {
111 state: Arc::new((
112 Condvar::new(),
113 Mutex::new(AsyncTaskState {
114 state: State::Exiting,
115 thread: None,
116 hi_prio_req: VecDeque::new(),
117 lo_prio_req: VecDeque::new(),
Janis Danisevskis40f0e6b2021-02-10 15:48:44 -0800118 shelf: None,
Janis Danisevskis93927dd2020-12-23 12:23:08 -0800119 }),
120 )),
121 }
122 }
123}
124
125impl AsyncTask {
126 /// Adds a job to the high priority queue. High priority jobs are completed before
127 /// low priority jobs and can also overtake low priority jobs. But they cannot
128 /// preempt them.
129 pub fn queue_hi<F>(&self, f: F)
130 where
Janis Danisevskis40f0e6b2021-02-10 15:48:44 -0800131 F: for<'r> FnOnce(&'r mut Shelf) + Send + 'static,
Janis Danisevskis93927dd2020-12-23 12:23:08 -0800132 {
133 self.queue(f, true)
134 }
135
136 /// Adds a job to the low priority queue. Low priority jobs are completed after
137 /// high priority. And they are not executed as long as high priority jobs are
138 /// present. Jobs always run to completion and are never preempted by high
139 /// priority jobs.
140 pub fn queue_lo<F>(&self, f: F)
141 where
Janis Danisevskis40f0e6b2021-02-10 15:48:44 -0800142 F: FnOnce(&mut Shelf) + Send + 'static,
Janis Danisevskis93927dd2020-12-23 12:23:08 -0800143 {
144 self.queue(f, false)
145 }
146
147 fn queue<F>(&self, f: F, hi_prio: bool)
148 where
Janis Danisevskis40f0e6b2021-02-10 15:48:44 -0800149 F: for<'r> FnOnce(&'r mut Shelf) + Send + 'static,
Janis Danisevskis93927dd2020-12-23 12:23:08 -0800150 {
151 let (ref condvar, ref state) = *self.state;
152 let mut state = state.lock().unwrap();
153 if hi_prio {
154 state.hi_prio_req.push_back(Box::new(f));
155 } else {
156 state.lo_prio_req.push_back(Box::new(f));
157 }
158
159 if state.state != State::Running {
160 self.spawn_thread(&mut state);
161 }
162 drop(state);
163 condvar.notify_all();
164 }
165
166 fn spawn_thread(&self, state: &mut MutexGuard<AsyncTaskState>) {
167 if let Some(t) = state.thread.take() {
168 t.join().expect("AsyncTask panicked.");
169 }
170
171 let cloned_state = self.state.clone();
172
173 state.thread = Some(thread::spawn(move || {
174 let (ref condvar, ref state) = *cloned_state;
Janis Danisevskis40f0e6b2021-02-10 15:48:44 -0800175 // When the worker starts, it takes the shelf and puts it on the stack.
176 let mut shelf = state.lock().unwrap().shelf.take().unwrap_or_default();
Janis Danisevskis93927dd2020-12-23 12:23:08 -0800177 loop {
178 if let Some(f) = {
179 let (mut state, timeout) = condvar
180 .wait_timeout_while(
181 state.lock().unwrap(),
182 Duration::from_secs(30),
183 |state| state.hi_prio_req.is_empty() && state.lo_prio_req.is_empty(),
184 )
185 .unwrap();
186 match (
187 state.hi_prio_req.pop_front(),
188 state.lo_prio_req.is_empty(),
189 timeout.timed_out(),
190 ) {
191 (Some(f), _, _) => Some(f),
192 (None, false, _) => state.lo_prio_req.pop_front(),
193 (None, true, true) => {
Janis Danisevskis40f0e6b2021-02-10 15:48:44 -0800194 // When the worker exits it puts the shelf back into the shared
195 // state for the next worker to use. So state is preserved not
196 // only across invocations but also across worker thread shut down.
197 state.shelf = Some(shelf);
Janis Danisevskis93927dd2020-12-23 12:23:08 -0800198 state.state = State::Exiting;
199 break;
200 }
201 (None, true, false) => None,
202 }
203 } {
Janis Danisevskis40f0e6b2021-02-10 15:48:44 -0800204 f(&mut shelf)
Janis Danisevskis93927dd2020-12-23 12:23:08 -0800205 }
206 }
207 }));
208 state.state = State::Running;
209 }
210}
David Drysdale2d3e5012021-02-23 12:30:27 +0000211
212#[cfg(test)]
213mod tests {
214 use super::Shelf;
215
216 #[test]
217 fn test_shelf() {
218 let mut shelf = Shelf::default();
219
220 let s = "A string".to_string();
221 assert_eq!(shelf.put(s), None);
222
223 let s2 = "Another string".to_string();
224 assert_eq!(shelf.put(s2), Some("A string".to_string()));
225
226 // Put something of a different type on the shelf.
227 #[derive(Debug, PartialEq, Eq)]
228 struct Elf {
229 pub name: String,
230 }
231 let e1 = Elf { name: "Glorfindel".to_string() };
232 assert_eq!(shelf.put(e1), None);
233
234 // The String value is still on the shelf.
235 let s3 = shelf.get_downcast_ref::<String>().unwrap();
236 assert_eq!(s3, "Another string");
237
238 // As is the Elf.
239 {
240 let e2 = shelf.get_downcast_mut::<Elf>().unwrap();
241 assert_eq!(e2.name, "Glorfindel");
242 e2.name = "Celeborn".to_string();
243 }
244
245 // Take the Elf off the shelf.
246 let e3 = shelf.remove_downcast_ref::<Elf>().unwrap();
247 assert_eq!(e3.name, "Celeborn");
248
249 assert_eq!(shelf.remove_downcast_ref::<Elf>(), None);
250
251 // No u64 value has been put on the shelf, so getting one gives the default value.
252 {
253 let i = shelf.get_mut::<u64>();
254 assert_eq!(*i, 0);
255 *i = 42;
256 }
257 let i2 = shelf.get_downcast_ref::<u64>().unwrap();
258 assert_eq!(*i2, 42);
259
260 // No i32 value has ever been seen near the shelf.
261 assert_eq!(shelf.get_downcast_ref::<i32>(), None);
262 assert_eq!(shelf.get_downcast_mut::<i32>(), None);
263 assert_eq!(shelf.remove_downcast_ref::<i32>(), None);
264 }
265}