sort_governor/service/
lease.rs

1//! An admitted sort and the global resources it holds.
2
3use std::path::{
4    Path,
5    PathBuf,
6};
7use std::sync::Arc;
8use std::sync::atomic::{
9    AtomicU32,
10    Ordering,
11};
12
13use serde::Serialize;
14use serde::de::DeserializeOwned;
15use tokio::sync::OwnedSemaphorePermit;
16
17use crate::engine::SortSession;
18use crate::plan::SortPlan;
19
20/// Decrements the active-sort counter when its lease is dropped.
21#[derive(Debug)]
22struct ActiveGuard {
23    active: Arc<AtomicU32>,
24}
25
26impl Drop for ActiveGuard {
27    fn drop(&mut self) {
28        self.active.fetch_sub(1, Ordering::Relaxed);
29    }
30}
31
32/// An admitted sort: the chosen [`SortPlan`], a private scratch directory,
33/// and the global resource permits held for the sort's lifetime.
34///
35/// Dropping the lease returns the file-descriptor permits and the
36/// concurrency slot. **Keep the lease alive until the value stream from
37/// [`SortSession::finish`] is fully consumed** — dropping it early returns
38/// the fd budget while the merge may still be reading run files.
39#[derive(Debug)]
40pub struct SortLease {
41    plan: SortPlan,
42    scratch_dir: PathBuf,
43    _fd_permit: Option<OwnedSemaphorePermit>,
44    _active: Option<ActiveGuard>,
45}
46
47impl SortLease {
48    /// An in-memory lease: no fd permits, no concurrency slot.
49    pub(crate) fn in_memory(plan: SortPlan, scratch_dir: PathBuf) -> Self {
50        Self {
51            plan,
52            scratch_dir,
53            _fd_permit: None,
54            _active: None,
55        }
56    }
57
58    /// An external lease holding `fd_permit` file-descriptor permits and a
59    /// slot in the active-sort count.
60    pub(crate) fn external(
61        plan: SortPlan,
62        scratch_dir: PathBuf,
63        fd_permit: OwnedSemaphorePermit,
64        active: Arc<AtomicU32>,
65    ) -> Self {
66        Self {
67            plan,
68            scratch_dir,
69            _fd_permit: Some(fd_permit),
70            _active: Some(ActiveGuard { active }),
71        }
72    }
73
74    /// The plan chosen for this sort.
75    #[must_use]
76    pub fn plan(&self) -> SortPlan {
77        self.plan
78    }
79
80    /// The private scratch directory this sort may spill into.
81    #[must_use]
82    pub fn scratch_dir(&self) -> &Path {
83        &self.scratch_dir
84    }
85
86    /// Open the sort session bound to this lease, consuming the lease into
87    /// the session so its fd permits live for as long as the sort's output
88    /// stream — there is no way to drop the lease early by accident.
89    #[must_use]
90    pub fn into_session<K, V>(self, dedup: bool) -> SortSession<K, V>
91    where
92        K: Ord + Clone + Serialize + DeserializeOwned + Send + 'static,
93        V: Serialize + DeserializeOwned + Send + 'static,
94    {
95        let plan = self.plan;
96        let dir = self.scratch_dir.clone();
97        SortSession::new(plan, dir, dedup).hold_resource(Box::new(self))
98    }
99}
sort_governor/service/lease.rs

sort_governor/service/
lease.rs
