mito2/compaction/scheduler/
coordinator.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
15#[cfg(test)]
16#[path = "coordinator_test.rs"]
17mod tests;
18
19use std::collections::{BTreeSet, HashMap, VecDeque};
20use std::future::Future;
21use std::pin::Pin;
22use std::sync::{Arc, Mutex, MutexGuard, Weak};
23use std::task::{Context, Poll};
24use std::time::Instant;
25
26use common_telemetry::{debug, error};
27use futures::FutureExt;
28use store_api::storage::RegionId;
29use tokio::sync::Notify;
30
31use crate::error::{InvalidSchedulerStateSnafu, Result};
32use crate::metrics::{
33    COMPACTION_ADMISSION_WAIT, COMPACTION_ADMITTED_UNITS, COMPACTION_READY_UNITS,
34};
35use crate::schedule::CancellableTaskState;
36use crate::schedule::scheduler::{Job, LocalScheduler, SchedulerRef};
37
38/// Prepared work consumed on submission to the coordinator, preventing repeated
39/// admission of the same unit. Its task owns the execution ID, cancellation state
40/// and input leases.
41pub(crate) struct ReadyUnit {
42    /// Time the unit became ready; feeds [`COMPACTION_ADMISSION_WAIT`].
43    queued_at: Instant,
44    /// Cancellation state shared with the unit's task and the region worker.
45    state: CancellableTaskState,
46    /// Immutable merge-work estimate; zero for metadata-only expiration.
47    input_bytes: u64,
48    /// Execution future that owns the reserved input leases.
49    job: Job,
50    /// Dispatched after `job` and its resources are dropped for accepted units,
51    /// including withdrawal, executor rejection, panic and shutdown draining.
52    /// Rejection by a closed coordinator drops this future without running it.
53    completion: Job,
54}
55
56impl ReadyUnit {
57    /// Creates ready work whose job owns the reserved inputs.
58    ///
59    /// Successful submission to [`CompactionCoordinator::schedule`] transfers
60    /// completion handling to the coordinator. If submission fails, neither
61    /// future is run, and the caller must finish its own bookkeeping.
62    pub(crate) fn new(
63        state: CancellableTaskState,
64        input_bytes: u64,
65        job: Job,
66        completion: Job,
67    ) -> Self {
68        Self {
69            queued_at: Instant::now(),
70            state,
71            input_bytes,
72            job,
73            completion,
74        }
75    }
76
77    fn release(self) {
78        drop(self.job);
79        common_runtime::spawn_global(self.completion);
80    }
81}
82
83/// Engine-wide compaction admission, in front of the generic FIFO executor.
84/// Only admitted jobs enter that executor, so its queue cannot hide contenders.
85pub(crate) struct CompactionCoordinator {
86    /// Executor that runs admitted units. [`Self::new_local`] derives its runner
87    /// count from the same value as `capacity`.
88    executor: SchedulerRef,
89    /// Maximum number of units holding capacity at once, including units the
90    /// executor has not polled yet. Clamped to at least one.
91    capacity: usize,
92    /// Admission accounting shared by every region worker. [`Self::lock`]
93    /// deliberately ignores poisoning because an unwind leaves the accounting
94    /// consistent.
95    state: Mutex<CoordinatorState>,
96    /// Signal that wakes [`Self::stop`] when `running` or the ready index changes.
97    drained: Notify,
98}
99
100/// Admission accounting shared by every region worker.
101#[derive(Default)]
102struct CoordinatorState {
103    /// Per-region service and ready units. Idle regions stay here until they are
104    /// unloaded, so batch boundaries cannot erase their service debt.
105    regions: HashMap<RegionId, RegionQueue>,
106    /// Iteration scales with current ready regions, not a hash table's peak capacity.
107    /// Idle regions retain service debt in `regions` without participating in admission.
108    ready_regions: BTreeSet<RegionId>,
109    /// Units holding capacity from admission through their released future,
110    /// including work that was authorized but never polled.
111    running: usize,
112    /// Set by [`CompactionCoordinator::stop`]; later submissions are rejected.
113    closed: bool,
114    /// Non-recursive dispatch guard: only the current drainer admits and submits.
115    dispatching: bool,
116    /// Monotonic floor of observed [`RegionQueue::service`], in the same
117    /// half-byte unit. It clamps a region that slept through a busy period.
118    virtual_time: u128,
119    /// Global admission tick. A larger value means a later admission; it breaks
120    /// service ties and is unrelated to the WAL or row `sequence`.
121    sequence: u64,
122}
123
124/// One region's ready units and accumulated service.
125struct RegionQueue {
126    /// Prepared units in fair-picker order. Admission pops the front.
127    ready: VecDeque<ReadyUnit>,
128    /// Units of this region currently holding capacity.
129    running: usize,
130    /// Weighted service consumed so far, in half-byte units: each admission adds
131    /// its charge, and only work refunded before its first poll is subtracted.
132    service: u128,
133    /// [`CoordinatorState::sequence`] of this region's last admission. Service
134    /// ties go to the region whose last admission is older.
135    last_served: u64,
136    /// Set by [`CompactionCoordinator::forget_region`]; the entry is pruned once
137    /// it has no ready or running work.
138    retired: bool,
139}
140
141/// The executor owns this future from admission until all task resources exit.
142/// Drop ordering also covers rejection, panic and shutdown, not only success.
143struct AdmittedUnit {
144    job: Option<Job>,
145    completion: Option<Job>,
146    state: CancellableTaskState,
147    coordinator: Weak<CompactionCoordinator>,
148    region_id: RegionId,
149    /// Immutable charge, including admission-time weight; refunded only before first poll.
150    charged_cost: u128,
151    started: bool,
152}
153
154impl std::fmt::Debug for CompactionCoordinator {
155    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
156        let state = self.lock();
157        f.debug_struct("CompactionCoordinator")
158            .field("capacity", &self.capacity)
159            .field("running", &state.running)
160            .field("regions", &state.regions.len())
161            .finish()
162    }
163}
164
165impl CompactionCoordinator {
166    /// Builds a coordinator over `executor` with an independent capacity.
167    ///
168    /// Production code should prefer [`Self::new_local`], which derives both
169    /// numbers from one value; the admission premise only holds while no more
170    /// units are admitted than the executor can run at once.
171    pub(crate) fn new(executor: SchedulerRef, capacity: usize) -> Self {
172        Self {
173            executor,
174            capacity: capacity.max(1),
175            state: Mutex::new(CoordinatorState::default()),
176            drained: Notify::new(),
177        }
178    }
179
180    /// Builds a coordinator whose executor runs at most `concurrency` units.
181    ///
182    /// Capacity and runner count come from the same value, so admitted work can
183    /// never wait in the executor's FIFO behind already-admitted contenders.
184    pub(crate) fn new_local(concurrency: usize) -> Self {
185        let concurrency = concurrency.max(1);
186        Self::new(Arc::new(LocalScheduler::new(concurrency)), concurrency)
187    }
188
189    fn lock(&self) -> MutexGuard<'_, CoordinatorState> {
190        self.state
191            .lock()
192            .unwrap_or_else(std::sync::PoisonError::into_inner)
193    }
194
195    /// Accepts a bounded batch's ready work, not an unbounded discovered backlog.
196    ///
197    /// # Errors
198    /// Returns an error if shutdown has started. The rejected unit's job and
199    /// completion are dropped without running; the caller handles the error.
200    pub(crate) fn schedule(self: &Arc<Self>, region_id: RegionId, unit: ReadyUnit) -> Result<()> {
201        {
202            let mut state = self.lock();
203            if state.closed {
204                return InvalidSchedulerStateSnafu.fail();
205            }
206            // A cancelled manual unit can still be in the worker's pending
207            // list when cancellation withdraws its coordinator-visible siblings.
208            if unit.state.is_cancelled() {
209                drop(state);
210                unit.release();
211                return Ok(());
212            }
213            let service = state.virtual_time;
214            let queue = state
215                .regions
216                .entry(region_id)
217                .or_insert_with(|| RegionQueue {
218                    ready: VecDeque::new(),
219                    running: 0,
220                    service,
221                    last_served: 0,
222                    retired: false,
223                });
224            queue.retired = false;
225            // Idle periods cannot bank unlimited credit against active regions.
226            if queue.ready.is_empty() {
227                queue.service = queue.service.max(service);
228            }
229            queue.ready.push_back(unit);
230            state.ready_regions.insert(region_id);
231            COMPACTION_READY_UNITS.inc();
232        }
233        self.dispatch();
234        Ok(())
235    }
236
237    /// Withdraws cancelled queued work without waiting for an execution slot.
238    /// Only the owning worker can cancel a unit; a reopened region's new units
239    /// have different cancellation state and are not affected by an old drop.
240    pub(crate) fn cancel_pending(self: &Arc<Self>, region_id: RegionId) {
241        let cancelled = {
242            let mut state = self.lock();
243            let Some(queue) = state.regions.get_mut(&region_id) else {
244                return;
245            };
246            let mut cancelled = Vec::new();
247            let mut ready = VecDeque::new();
248            for unit in queue.ready.drain(..) {
249                if unit.state.is_cancelled() {
250                    cancelled.push(unit);
251                } else {
252                    ready.push_back(unit);
253                }
254            }
255            queue.ready = ready;
256            if queue.ready.is_empty() {
257                state.ready_regions.remove(&region_id);
258            }
259            state.prune(region_id);
260            cancelled
261        };
262        for unit in cancelled {
263            Self::withdraw(unit);
264        }
265        self.drained.notify_one();
266    }
267
268    /// Keeps service debt between batches, but not after the region is unloaded.
269    /// In-flight old work can drain before its accounting entry is reclaimed.
270    pub(crate) fn forget_region(&self, region_id: RegionId) {
271        let mut state = self.lock();
272        if let Some(queue) = state.regions.get_mut(&region_id) {
273            queue.retired = true;
274        }
275        state.prune(region_id);
276    }
277
278    /// Drains the admission queue before stopping its executor. Workers remain
279    /// alive to acknowledge apply while these jobs drain.
280    pub(crate) async fn stop(&self) -> Result<()> {
281        self.lock().closed = true;
282        loop {
283            let notified = self.drained.notified();
284            let drained = {
285                let state = self.lock();
286                state.running == 0 && state.ready_regions.is_empty()
287            };
288            if drained {
289                break;
290            }
291            notified.await;
292        }
293        self.executor.stop(true).await
294    }
295
296    fn dispatch(self: &Arc<Self>) {
297        {
298            let mut state = self.lock();
299            if state.dispatching {
300                return;
301            }
302            state.dispatching = true;
303        }
304        loop {
305            let next = {
306                let mut state = self.lock();
307                let next = (state.running < self.capacity)
308                    .then(|| state.admit())
309                    .flatten();
310                if next.is_none() {
311                    state.dispatching = false;
312                }
313                next
314            };
315            let Some((region_id, unit, charged_cost)) = next else {
316                break;
317            };
318            let admitted = AdmittedUnit::new(self, region_id, unit, charged_cost);
319            // Rejection drops the future and returns its credit/resources. The
320            // existing Released path reports a missing terminal result as failure.
321            if let Err(err) = self.executor.schedule(Box::pin(admitted)) {
322                error!(err; "Failed to submit admitted compaction for region {}", region_id);
323            }
324        }
325    }
326
327    fn release(self: &Arc<Self>, region_id: RegionId, refund: u128) {
328        {
329            let mut state = self.lock();
330            state.running -= 1;
331            if let Some(queue) = state.regions.get_mut(&region_id) {
332                queue.running -= 1;
333                queue.service = queue.service.saturating_sub(refund);
334            }
335            state.prune(region_id);
336        }
337        self.drained.notify_one();
338        // dispatching serializes drainers without holding a lock across executor
339        // submission. In particular, synchronous rejection cannot recurse here.
340        self.dispatch();
341    }
342
343    fn withdraw(unit: ReadyUnit) {
344        COMPACTION_READY_UNITS.dec();
345        unit.release();
346    }
347}
348
349impl Drop for CompactionCoordinator {
350    fn drop(&mut self) {
351        let state = self
352            .state
353            .get_mut()
354            .unwrap_or_else(std::sync::PoisonError::into_inner);
355        for queue in state.regions.values_mut() {
356            for unit in queue.ready.drain(..) {
357                Self::withdraw(unit);
358            }
359        }
360    }
361}
362
363impl CoordinatorState {
364    /// Least-served selection charges bytes at admission, including jobs that
365    /// have not started. A backlog can boost weight by at most two; no region can
366    /// obtain unlimited weight. Positive byte cost prevents zero-sized work from
367    /// monopolizing ties. Queue order retains the region's fair picker priority.
368    fn admit(&mut self) -> Option<(RegionId, ReadyUnit, u128)> {
369        let region_id = self
370            .ready_regions
371            .iter()
372            .filter_map(|id| self.regions.get(id).map(|queue| (id, queue)))
373            .min_by_key(|(id, q)| (q.service, q.last_served, **id))
374            .map(|(id, _)| *id)?;
375        let queue = self.regions.get_mut(&region_id)?;
376        let weight = if queue.ready.len() > 1 { 2 } else { 1 };
377        let unit = queue.ready.pop_front()?;
378        if queue.ready.is_empty() {
379            self.ready_regions.remove(&region_id);
380        }
381        COMPACTION_READY_UNITS.dec();
382        COMPACTION_ADMITTED_UNITS.inc();
383        self.virtual_time = self.virtual_time.max(queue.service);
384        let charged_cost = charge(unit.input_bytes, weight);
385        queue.service = queue.service.saturating_add(charged_cost);
386        self.sequence = self.sequence.wrapping_add(1);
387        queue.last_served = self.sequence;
388        queue.running += 1;
389        self.running += 1;
390        debug!(
391            "Admit compaction unit: region={}, input_bytes={}, weight={}, service={}, ready={}, admitted={}",
392            region_id,
393            unit.input_bytes,
394            weight,
395            queue.service,
396            queue.ready.len(),
397            queue.running,
398        );
399        Some((region_id, unit, charged_cost))
400    }
401
402    fn prune(&mut self, region_id: RegionId) {
403        if self
404            .regions
405            .get(&region_id)
406            .is_some_and(|q| q.retired && q.running == 0 && q.ready.is_empty())
407        {
408            self.regions.remove(&region_id);
409        }
410    }
411}
412
413fn charge(bytes: u64, weight: u128) -> u128 {
414    // Half-byte units avoid rounding a small task's cost down to zero.
415    u128::from(bytes.max(1)) * 2 / weight
416}
417
418impl AdmittedUnit {
419    fn new(
420        coordinator: &Arc<CompactionCoordinator>,
421        region_id: RegionId,
422        unit: ReadyUnit,
423        charged_cost: u128,
424    ) -> Self {
425        COMPACTION_ADMISSION_WAIT.observe(unit.queued_at.elapsed().as_secs_f64());
426        let job = Box::pin(async move {
427            if let Err(panic) = std::panic::AssertUnwindSafe(unit.job).catch_unwind().await {
428                error!(
429                    "Compaction executor panicked: {}",
430                    crate::compaction::format_panic_reason(panic.as_ref())
431                );
432            }
433        });
434        Self {
435            job: Some(job),
436            completion: Some(unit.completion),
437            state: unit.state,
438            coordinator: Arc::downgrade(coordinator),
439            region_id,
440            charged_cost,
441            started: false,
442        }
443    }
444}
445
446impl Future for AdmittedUnit {
447    type Output = ();
448
449    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
450        if !self.started && self.state.is_cancelled() {
451            return Poll::Ready(());
452        }
453        self.started = true;
454        let Some(job) = self.job.as_mut() else {
455            return Poll::Ready(());
456        };
457        let poll = job.as_mut().poll(cx);
458        if poll.is_ready() {
459            // A completed future stays completed if the executor polls it again.
460            drop(self.job.take());
461        }
462        poll
463    }
464}
465
466impl Drop for AdmittedUnit {
467    fn drop(&mut self) {
468        // This ordering is essential even when the executor rejects the future
469        // before its first poll: Released must never precede input lease release.
470        drop(self.job.take());
471        COMPACTION_ADMITTED_UNITS.dec();
472        if let Some(coordinator) = self.coordinator.upgrade() {
473            let refund = if self.started { 0 } else { self.charged_cost };
474            coordinator.release(self.region_id, refund);
475        }
476        if let Some(completion) = self.completion.take() {
477            common_runtime::spawn_global(completion);
478        }
479    }
480}
mito2/compaction/scheduler/coordinator.rs

mito2/compaction/scheduler/
coordinator.rs
