mito2/schedule/
scheduler.rs

1// Copyright 2023 Greptime Team
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
15use std::future::Future;
16use std::pin::Pin;
17use std::sync::atomic::{AtomicU8, Ordering};
18use std::sync::{Arc, RwLock};
19
20use common_telemetry::warn;
21use snafu::{OptionExt, ResultExt, ensure};
22use tokio::sync::Mutex;
23use tokio::task::JoinHandle;
24use tokio_util::sync::CancellationToken;
25
26use crate::error::{InvalidSchedulerStateSnafu, InvalidSenderSnafu, Result, StopSchedulerSnafu};
27
28pub type Job = Pin<Box<dyn Future<Output = ()> + Send>>;
29
30///The state of scheduler
31const STATE_RUNNING: u8 = 0;
32const STATE_STOP: u8 = 1;
33const STATE_AWAIT_TERMINATION: u8 = 2;
34
35/// [Scheduler] defines a set of API to schedule Jobs
36#[async_trait::async_trait]
37pub trait Scheduler: Send + Sync {
38    /// Schedules a Job
39    fn schedule(&self, job: Job) -> Result<()>;
40
41    /// Runs completion only after the job and its owned resources have been dropped.
42    fn schedule_with_completion(&self, job: Job, completion: Job) -> Result<()> {
43        self.schedule(Box::pin(async move {
44            job.await;
45            completion.await;
46        }))
47    }
48
49    /// Stops scheduler. If `await_termination` is set to true, the scheduler will wait until all tasks are processed.
50    async fn stop(&self, await_termination: bool) -> Result<()>;
51}
52
53pub type SchedulerRef = Arc<dyn Scheduler>;
54
55/// A queued job with an optional completion future dispatched after the job exits.
56struct ScheduledJob {
57    job: Job,
58    completion: Option<Job>,
59}
60
61impl ScheduledJob {
62    /// Finishes the job before dispatching completion outside its runner.
63    async fn run(self) {
64        self.job.await;
65        if let Some(completion) = self.completion {
66            // The original job has exited. Notification must not occupy its runner.
67            common_runtime::spawn_global(completion);
68        }
69    }
70}
71
72/// Request scheduler based on local state.
73pub struct LocalScheduler {
74    /// Sends jobs to flume bounded channel
75    sender: RwLock<Option<async_channel::Sender<ScheduledJob>>>,
76    /// Task handles
77    handles: Mutex<Vec<JoinHandle<()>>>,
78    /// Token used to halt the scheduler
79    cancel_token: CancellationToken,
80    /// State of scheduler
81    state: Arc<AtomicU8>,
82}
83
84impl LocalScheduler {
85    /// Starts a new scheduler.
86    ///
87    /// concurrency: the number of bounded receiver
88    pub fn new(concurrency: usize) -> Self {
89        let (tx, rx) = async_channel::unbounded::<ScheduledJob>();
90        let token = CancellationToken::new();
91        let state = Arc::new(AtomicU8::new(STATE_RUNNING));
92
93        let mut handles = Vec::with_capacity(concurrency);
94
95        for _ in 0..concurrency {
96            let child = token.child_token();
97            let receiver = rx.clone();
98            let state_clone = state.clone();
99            let handle = common_runtime::spawn_global(async move {
100                while state_clone.load(Ordering::Relaxed) == STATE_RUNNING {
101                    tokio::select! {
102                        _ = child.cancelled() => {
103                            break;
104                        }
105                        req_opt = receiver.recv() =>{
106                            if let Ok(job) = req_opt {
107                                job.run().await;
108                            }
109                        }
110                    }
111                }
112                // When task scheduler is cancelled, we will wait all task finished
113                if state_clone.load(Ordering::Relaxed) == STATE_AWAIT_TERMINATION {
114                    // recv_async waits until all sender's been dropped.
115                    while let Ok(job) = receiver.recv().await {
116                        job.run().await;
117                    }
118                    state_clone.store(STATE_STOP, Ordering::Relaxed);
119                }
120            });
121            handles.push(handle);
122        }
123
124        Self {
125            sender: RwLock::new(Some(tx)),
126            cancel_token: token,
127            handles: Mutex::new(handles),
128            state,
129        }
130    }
131
132    #[inline]
133    fn is_running(&self) -> bool {
134        self.state.load(Ordering::Relaxed) == STATE_RUNNING
135    }
136
137    /// Enqueues work only while the scheduler and its sender remain available.
138    fn submit(&self, job: ScheduledJob) -> Result<()> {
139        ensure!(self.is_running(), InvalidSchedulerStateSnafu);
140        self.sender
141            .read()
142            .unwrap()
143            .as_ref()
144            .context(InvalidSchedulerStateSnafu)?
145            .try_send(job)
146            .map_err(|_| InvalidSenderSnafu {}.build())
147    }
148}
149
150#[async_trait::async_trait]
151impl Scheduler for LocalScheduler {
152    fn schedule(&self, job: Job) -> Result<()> {
153        self.submit(ScheduledJob {
154            job,
155            completion: None,
156        })
157    }
158
159    /// Queues the job and its completion together so completion cannot precede task exit.
160    fn schedule_with_completion(&self, job: Job, completion: Job) -> Result<()> {
161        self.submit(ScheduledJob {
162            job,
163            completion: Some(completion),
164        })
165    }
166
167    /// if await_termination is true, scheduler will wait all tasks finished before stopping
168    async fn stop(&self, await_termination: bool) -> Result<()> {
169        ensure!(self.is_running(), InvalidSchedulerStateSnafu);
170        let state = if await_termination {
171            STATE_AWAIT_TERMINATION
172        } else {
173            STATE_STOP
174        };
175        self.sender.write().unwrap().take();
176        self.state.store(state, Ordering::Relaxed);
177        self.cancel_token.cancel();
178
179        futures::future::join_all(self.handles.lock().await.drain(..))
180            .await
181            .into_iter()
182            .collect::<std::result::Result<Vec<_>, _>>()
183            .context(StopSchedulerSnafu)?;
184
185        Ok(())
186    }
187}
188
189impl Drop for LocalScheduler {
190    fn drop(&mut self) {
191        if self.state.load(Ordering::Relaxed) != STATE_STOP {
192            warn!(
193                "scheduler should be stopped before dropping, which means the state of scheduler must be STATE_STOP"
194            );
195
196            // We didn't call `stop()` so we cancel all background workers here.
197            self.sender.write().unwrap().take();
198            self.cancel_token.cancel();
199        }
200    }
201}
202
203#[cfg(test)]
204mod tests {
205    use std::sync::Arc;
206    use std::sync::atomic::AtomicI32;
207
208    use tokio::sync::Barrier;
209    use tokio::time::Duration;
210
211    use super::*;
212
213    #[tokio::test]
214    async fn test_sum_cap() {
215        let task_size = 1000;
216        let sum = Arc::new(AtomicI32::new(0));
217        let local = LocalScheduler::new(task_size);
218
219        for _ in 0..task_size {
220            let sum_clone = sum.clone();
221            local
222                .schedule(Box::pin(async move {
223                    sum_clone.fetch_add(1, Ordering::Relaxed);
224                }))
225                .unwrap();
226        }
227        local.stop(true).await.unwrap();
228        assert_eq!(sum.load(Ordering::Relaxed), 1000);
229    }
230
231    #[tokio::test]
232    async fn test_sum_consumer_num() {
233        let task_size = 1000;
234        let sum = Arc::new(AtomicI32::new(0));
235        let local = LocalScheduler::new(3);
236        let mut target = 0;
237        for _ in 0..task_size {
238            let sum_clone = sum.clone();
239            let ok = local
240                .schedule(Box::pin(async move {
241                    sum_clone.fetch_add(1, Ordering::Relaxed);
242                }))
243                .is_ok();
244            if ok {
245                target += 1;
246            }
247        }
248        local.stop(true).await.unwrap();
249        assert_eq!(sum.load(Ordering::Relaxed), target);
250    }
251
252    #[tokio::test]
253    async fn test_scheduler_many() {
254        let task_size = 1000;
255
256        let barrier = Arc::new(Barrier::new(task_size + 1));
257        let local: LocalScheduler = LocalScheduler::new(task_size);
258
259        for _ in 0..task_size {
260            let barrier_clone = barrier.clone();
261            local
262                .schedule(Box::pin(async move {
263                    barrier_clone.wait().await;
264                }))
265                .unwrap();
266        }
267        barrier.wait().await;
268        local.stop(true).await.unwrap();
269    }
270
271    #[tokio::test]
272    async fn test_scheduler_continuous_stop() {
273        let sum = Arc::new(AtomicI32::new(0));
274        let local = Arc::new(LocalScheduler::new(1000));
275
276        let barrier = Arc::new(Barrier::new(2));
277        let barrier_clone = barrier.clone();
278        let local_stop = local.clone();
279        tokio::spawn(async move {
280            tokio::time::sleep(Duration::from_millis(5)).await;
281            local_stop.stop(true).await.unwrap();
282            barrier_clone.wait().await;
283        });
284
285        let target = Arc::new(AtomicI32::new(0));
286        let local_task = local.clone();
287        let target_clone = target.clone();
288        let sum_clone = sum.clone();
289        tokio::spawn(async move {
290            loop {
291                let sum_c = sum_clone.clone();
292                let ok = local_task
293                    .schedule(Box::pin(async move {
294                        sum_c.fetch_add(1, Ordering::Relaxed);
295                    }))
296                    .is_ok();
297                if ok {
298                    target_clone.fetch_add(1, Ordering::Relaxed);
299                } else {
300                    break;
301                }
302                tokio::task::yield_now().await;
303            }
304        });
305        barrier.wait().await;
306        assert_eq!(sum.load(Ordering::Relaxed), target.load(Ordering::Relaxed));
307    }
308}
mito2/schedule/scheduler.rs

mito2/schedule/
scheduler.rs
