mito2/schedule/
remote_job_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::fmt;
16use std::sync::{Arc, Mutex};
17use std::time::Instant;
18
19use common_telemetry::error;
20use common_time::TimeToLive;
21use serde::{Deserialize, Serialize};
22use snafu::{Location, ResultExt, Snafu};
23use store_api::storage::RegionId;
24use tokio::sync::mpsc::Sender;
25use uuid::Uuid;
26
27use crate::compaction::CompactionExecution;
28use crate::compaction::compactor::CompactionRegion;
29use crate::compaction::picker::PickerOutput;
30use crate::error::{CompactRegionSnafu, Error, ParseJobIdSnafu, Result};
31use crate::manifest::action::RegionEdit;
32use crate::metrics::{COMPACTION_FAILURE_COUNT, INFLIGHT_COMPACTION_COUNT};
33use crate::request::{
34    BackgroundNotify, CompactionFailed, CompactionFinished, OutputTx, WorkerRequest,
35    WorkerRequestWithTime,
36};
37
38pub type RemoteJobSchedulerRef = Arc<dyn RemoteJobScheduler>;
39
40#[cfg_attr(doc, aquamarine::aquamarine)]
41/// RemoteJobScheduler is a trait that defines the API to schedule remote jobs.
42/// For example, a compaction job can be scheduled remotely as the following workflow:
43/// ```mermaid
44///   participant User
45///   participant MitoEngine
46///   participant CompactionScheduler
47///   participant Plugins
48///   participant RemoteJobScheduler
49///
50///   User->>MitoEngine: Initiates compaction
51///   MitoEngine->>CompactionScheduler: schedule_compaction()
52///   CompactionScheduler->>Plugins: Handle plugins
53///   CompactionScheduler->>RemoteJobScheduler: schedule(CompactionJob)
54///   RemoteJobScheduler-->>CompactionScheduler: Returns Job UUID
55///   CompactionScheduler-->>MitoEngine: Task scheduled with Job UUID
56///   MitoEngine-->>User: Compaction task scheduled
57/// ```
58#[async_trait::async_trait]
59pub trait RemoteJobScheduler: Send + Sync + 'static {
60    /// Sends a job to the scheduler and returns a UUID for the job.
61    async fn schedule(
62        &self,
63        job: RemoteJob,
64        notifier: Box<dyn Notifier>,
65    ) -> Result<JobId, RemoteJobSchedulerError>;
66}
67
68#[derive(Snafu, Debug)]
69#[snafu(display("Internal error occurred in remote job scheduler: {}", reason))]
70pub struct RemoteJobSchedulerError {
71    #[snafu(implicit)]
72    pub location: Location,
73    pub reason: String,
74    // Keep the waiters in the error so that we can notify them when fallback to the local compaction.
75    pub waiters: Vec<OutputTx>,
76}
77
78/// Notifier is used to notify the mito engine when a remote job is completed.
79#[async_trait::async_trait]
80pub trait Notifier: Send + Sync + 'static {
81    /// Notify the mito engine that a remote job is completed.
82    async fn notify(&self, result: RemoteJobResult, waiters: Vec<OutputTx>);
83}
84
85/// Unique id for a remote job.
86#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
87pub struct JobId(Uuid);
88
89impl JobId {
90    /// Parses job id from string.
91    pub fn parse_str(input: &str) -> Result<JobId> {
92        Uuid::parse_str(input).map(JobId).context(ParseJobIdSnafu)
93    }
94}
95
96impl fmt::Display for JobId {
97    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
98        write!(f, "{}", self.0)
99    }
100}
101
102/// RemoteJob is a job that can be executed remotely. For example, a remote compaction job.
103#[allow(dead_code)]
104pub enum RemoteJob {
105    CompactionJob(CompactionJob),
106}
107
108/// CompactionJob is a remote job that compacts a set of files in a compaction service.
109#[allow(dead_code)]
110pub struct CompactionJob {
111    pub compaction_region: CompactionRegion,
112    pub picker_output: PickerOutput,
113    pub start_time: Instant,
114    pub ttl: TimeToLive,
115    /// Send the result of the compaction job to these waiters.
116    pub waiters: Vec<OutputTx>,
117}
118
119/// RemoteJobResult is the result of a remote job.
120#[allow(dead_code)]
121pub enum RemoteJobResult {
122    CompactionJobResult(CompactionJobResult),
123}
124
125/// CompactionJobResult is the result of a compaction job.
126#[allow(dead_code)]
127pub struct CompactionJobResult {
128    pub job_id: JobId,
129    pub region_id: RegionId,
130    pub start_time: Instant,
131    pub region_edit: Result<RegionEdit>,
132}
133
134/// DefaultNotifier is a default implementation of Notifier that sends WorkerRequest to the mito engine.
135pub(crate) struct DefaultNotifier {
136    /// The sender to send WorkerRequest to the mito engine. This is used to notify the mito engine when a remote job is completed.
137    pub(crate) request_sender: Sender<WorkerRequestWithTime>,
138    execution: Mutex<Option<CompactionExecution>>,
139}
140
141impl DefaultNotifier {
142    pub(crate) fn new(
143        request_sender: Sender<WorkerRequestWithTime>,
144        execution: CompactionExecution,
145    ) -> Self {
146        Self {
147            request_sender,
148            execution: Mutex::new(Some(execution)),
149        }
150    }
151
152    fn on_failure(&self, err: Arc<Error>, region_id: RegionId, mut waiters: Vec<OutputTx>) {
153        COMPACTION_FAILURE_COUNT.inc();
154        for waiter in waiters.drain(..) {
155            waiter.send(Err(err.clone()).context(CompactRegionSnafu { region_id }));
156        }
157    }
158}
159
160#[async_trait::async_trait]
161impl Notifier for DefaultNotifier {
162    async fn notify(&self, result: RemoteJobResult, waiters: Vec<OutputTx>) {
163        let Some(execution) = self
164            .execution
165            .lock()
166            .unwrap_or_else(|poisoned| poisoned.into_inner())
167            .take()
168        else {
169            error!("Remote compaction notifier invoked more than once");
170            return;
171        };
172        INFLIGHT_COMPACTION_COUNT.dec();
173        match result {
174            RemoteJobResult::CompactionJobResult(result) => {
175                let notify = {
176                    match result.region_edit {
177                        Ok(edit) => BackgroundNotify::CompactionFinished(CompactionFinished {
178                            region_id: result.region_id,
179                            execution,
180                            senders: waiters,
181                            start_time: result.start_time,
182                            edit,
183                        }),
184                        Err(err) => {
185                            error!(
186                                "Compaction failed for region {}: {:?}",
187                                result.region_id, err
188                            );
189                            let err = Arc::new(err);
190                            self.on_failure(err.clone(), result.region_id, waiters);
191                            BackgroundNotify::CompactionFailed(CompactionFailed {
192                                region_id: result.region_id,
193                                execution,
194                                err,
195                            })
196                        }
197                    }
198                };
199
200                if let Err(e) = self
201                    .request_sender
202                    .send(WorkerRequestWithTime::new(WorkerRequest::Background {
203                        region_id: result.region_id,
204                        notify,
205                    }))
206                    .await
207                {
208                    error!(
209                        "Failed to notify compaction job status for region {}, error: {:?}",
210                        result.region_id, e
211                    );
212                }
213            }
214        }
215    }
216}
217
218#[cfg(test)]
219mod tests {
220    use super::*;
221
222    #[test]
223    fn test_job_id() {
224        let id = Uuid::new_v4().to_string();
225        let job_id = JobId::parse_str(&id).unwrap();
226        assert_eq!(job_id.to_string(), id);
227    }
228}
mito2/schedule/remote_job_scheduler.rs

mito2/schedule/
remote_job_scheduler.rs
