ahri_tre_lake/
dataset_executor.rs

1use std::{
2    collections::{HashMap, HashSet},
3    fs::File,
4    io::{Read, Seek, SeekFrom, Write},
5    os::unix::fs::{MetadataExt, PermissionsExt},
6    path::{Path, PathBuf},
7    sync::{Arc, Mutex, OnceLock, Weak},
8};
9
10use ahri_tre_core::{
11    CoreError, DatasetAttemptExecutor, DatasetExecutorIdentity, DatasetExecutorLease,
12};
13use rustix::fs::{FlockOperation, Mode, OFlags, flock, mkdirat, openat};
14use uuid::Uuid;
15
16use crate::{ScratchError, TrustedScratch};
17
18/// One process generation, held by Sessions and every outstanding Lake attempt.
19/// The stable lock file lives in durable coordination storage outside source scratch.
20pub struct DatasetExecutor {
21    identity: DatasetExecutorIdentity,
22    _lock: File,
23    active: Mutex<ExecutorActivity>,
24}
25
26impl std::fmt::Debug for DatasetExecutor {
27    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
28        f.debug_struct("DatasetExecutor")
29            .field("identity", &self.identity)
30            .finish_non_exhaustive()
31    }
32}
33
34#[derive(Default)]
35struct ExecutorActivity {
36    attempts: HashSet<Uuid>,
37    maintenance: bool,
38}
39
40type Executors = HashMap<(PathBuf, Uuid), Weak<DatasetExecutor>>;
41static EXECUTORS: OnceLock<Mutex<Executors>> = OnceLock::new();
42
43impl DatasetExecutor {
44    /// Opens the configured durable coordination root. Source scratch is separate.
45    pub fn acquire(root: &Path, datastore_id: Uuid) -> Result<Arc<Self>, ScratchError> {
46        TrustedScratch::open(root, [])?.acquire_dataset_executor(datastore_id)
47    }
48}
49
50impl TrustedScratch {
51    /// Acquires the configured, local-only executor lock for this Datastore.
52    /// A different coordinator root cannot replace a recorded database generation.
53    pub(crate) fn acquire_dataset_executor(
54        &self,
55        datastore_id: Uuid,
56    ) -> Result<Arc<DatasetExecutor>, ScratchError> {
57        let mut executors = EXECUTORS
58            .get_or_init(Default::default)
59            .lock()
60            .map_err(|_| ScratchError::Unavailable)?;
61        let key = (self.root.clone(), datastore_id);
62        if let Some(executor) = executors.get(&key).and_then(Weak::upgrade) {
63            return Ok(executor);
64        }
65        let directory = Path::new("locks");
66        match mkdirat(
67            &*self.root_directory,
68            directory,
69            Mode::from_bits_retain(0o700),
70        ) {
71            Ok(()) | Err(rustix::io::Errno::EXIST) => {}
72            Err(_) => return Err(ScratchError::Unavailable),
73        }
74        let directory = File::from(
75            openat(
76                &*self.root_directory,
77                directory,
78                OFlags::RDONLY | OFlags::DIRECTORY | OFlags::CLOEXEC | OFlags::NOFOLLOW,
79                Mode::empty(),
80            )
81            .map_err(|_| ScratchError::Unavailable)?,
82        );
83        let metadata = directory
84            .metadata()
85            .map_err(|_| ScratchError::Unavailable)?;
86        if metadata.uid() != rustix::process::geteuid().as_raw()
87            || metadata.permissions().mode() & 0o777 != 0o700
88        {
89            return Err(ScratchError::UnsafeRoot);
90        }
91        let filename = format!("{}.lock", datastore_id.simple());
92        let mut lock = File::from(
93            openat(
94                &directory,
95                Path::new(&filename),
96                OFlags::RDWR | OFlags::CREATE | OFlags::CLOEXEC | OFlags::NOFOLLOW,
97                Mode::from_bits_retain(0o600),
98            )
99            .map_err(|_| ScratchError::Unavailable)?,
100        );
101        let metadata = lock.metadata().map_err(|_| ScratchError::Unavailable)?;
102        if !metadata.is_file()
103            || metadata.nlink() != 1
104            || metadata.uid() != rustix::process::geteuid().as_raw()
105            || metadata.permissions().mode() & 0o777 != 0o600
106            || metadata.len() > 36
107        {
108            return Err(ScratchError::UnsafeAttempt);
109        }
110        flock(&lock, FlockOperation::NonBlockingLockExclusive)
111            .map_err(|_| ScratchError::Unavailable)?;
112        let mut stored = String::new();
113        lock.read_to_string(&mut stored)
114            .map_err(|_| ScratchError::Unavailable)?;
115        let coordinator_id = if stored.is_empty() {
116            let id = Uuid::new_v4();
117            lock.seek(SeekFrom::Start(0))
118                .map_err(|_| ScratchError::Unavailable)?;
119            lock.write_all(id.to_string().as_bytes())
120                .and_then(|()| lock.sync_all())
121                .map_err(|_| ScratchError::Unavailable)?;
122            id
123        } else {
124            Uuid::parse_str(&stored).map_err(|_| ScratchError::UnsafeAttempt)?
125        };
126        let executor = Arc::new(DatasetExecutor {
127            identity: DatasetExecutorIdentity {
128                coordinator_id,
129                generation_id: Uuid::new_v4(),
130            },
131            _lock: lock,
132            active: Mutex::new(ExecutorActivity::default()),
133        });
134        executors.insert(key, Arc::downgrade(&executor));
135        Ok(executor)
136    }
137}
138
139impl DatasetExecutorLease for DatasetExecutor {
140    fn identity(&self) -> DatasetExecutorIdentity {
141        self.identity
142    }
143
144    fn begin_attempt(
145        self: Arc<Self>,
146        attempt_id: Uuid,
147    ) -> Result<Box<dyn DatasetAttemptExecutor>, CoreError> {
148        {
149            let mut active = self
150                .active
151                .lock()
152                .map_err(|_| CoreError::Infrastructure("Dataset executor is unavailable".into()))?;
153            if active.maintenance || !active.attempts.insert(attempt_id) {
154                return Err(CoreError::Conflict(
155                    "Dataset executor is already owned".into(),
156                ));
157            }
158        }
159        Ok(Box::new(AttemptExecutor {
160            executor: self,
161            attempt_id,
162        }))
163    }
164
165    fn begin_maintenance(self: Arc<Self>) -> Result<Box<dyn DatasetAttemptExecutor>, CoreError> {
166        {
167            let mut active = self
168                .active
169                .lock()
170                .map_err(|_| CoreError::Infrastructure("Dataset executor is unavailable".into()))?;
171            if active.maintenance || !active.attempts.is_empty() {
172                return Err(CoreError::Conflict(
173                    "Dataset executor is already owned".into(),
174                ));
175            }
176            active.maintenance = true;
177        }
178        Ok(Box::new(MaintenanceExecutor(self)))
179    }
180
181    fn can_recover(&self, owner: DatasetExecutorIdentity, attempt_id: Uuid) -> bool {
182        owner.coordinator_id == self.identity.coordinator_id
183            && (owner.generation_id != self.identity.generation_id
184                || self.active.lock().is_ok_and(|active| {
185                    !active.maintenance && !active.attempts.contains(&attempt_id)
186                }))
187    }
188}
189
190struct AttemptExecutor {
191    executor: Arc<DatasetExecutor>,
192    attempt_id: Uuid,
193}
194impl DatasetAttemptExecutor for AttemptExecutor {}
195impl Drop for AttemptExecutor {
196    fn drop(&mut self) {
197        if let Ok(mut active) = self.executor.active.lock() {
198            active.attempts.remove(&self.attempt_id);
199        }
200    }
201}
202
203struct MaintenanceExecutor(Arc<DatasetExecutor>);
204impl DatasetAttemptExecutor for MaintenanceExecutor {}
205impl Drop for MaintenanceExecutor {
206    fn drop(&mut self) {
207        if let Ok(mut active) = self.0.active.lock() {
208            active.maintenance = false;
209        }
210    }
211}
212
213#[cfg(test)]
214mod tests {
215    use super::*;
216    use std::{process::Command, sync::Barrier, thread};
217
218    #[test]
219    fn executor_probe() {
220        let Ok(root) = std::env::var("AHRI_TEST_EXECUTOR_ROOT") else {
221            return;
222        };
223        let datastore =
224            Uuid::parse_str(&std::env::var("AHRI_TEST_EXECUTOR_DATASTORE").unwrap()).unwrap();
225        let scratch = TrustedScratch::open(Path::new(&root), []).unwrap();
226        let available = scratch.acquire_dataset_executor(datastore).is_ok();
227        assert_eq!(
228            available,
229            std::env::var("AHRI_TEST_EXECUTOR_AVAILABLE").unwrap() == "true"
230        );
231    }
232
233    #[test]
234    fn late_executor_holds_exclusion_after_session_owner_is_dropped() {
235        let root =
236            std::env::temp_dir().join(format!("dataset-executor-test-{}", Uuid::new_v4().simple()));
237        std::fs::create_dir(&root).unwrap();
238        std::fs::set_permissions(&root, std::fs::Permissions::from_mode(0o700)).unwrap();
239        let scratch = TrustedScratch::open(&root, []).unwrap();
240        let datastore = Uuid::new_v4();
241        let executor = scratch.acquire_dataset_executor(datastore).unwrap();
242        let identity = executor.identity();
243        let attempt_id = Uuid::new_v4();
244        let call = executor.clone().begin_attempt(attempt_id).unwrap();
245        assert!(!executor.can_recover(identity, attempt_id));
246        let barrier = Arc::new(Barrier::new(2));
247        let finished = barrier.clone();
248        let worker = thread::spawn(move || {
249            finished.wait();
250            drop(call);
251        });
252        drop(executor); // Session close cannot end an in-flight call's ownership.
253        let probe = |available: bool| {
254            let result = Command::new(std::env::current_exe().unwrap())
255                .args(["--exact", "dataset_executor::tests::executor_probe"])
256                .env("AHRI_TEST_EXECUTOR_ROOT", &root)
257                .env("AHRI_TEST_EXECUTOR_DATASTORE", datastore.to_string())
258                .env("AHRI_TEST_EXECUTOR_AVAILABLE", available.to_string())
259                .output()
260                .unwrap();
261            assert!(
262                result.status.success(),
263                "executor probe failed: {}",
264                String::from_utf8_lossy(&result.stderr)
265            );
266        };
267        probe(false);
268        barrier.wait();
269        worker.join().unwrap();
270        probe(true);
271        let replacement = scratch.acquire_dataset_executor(datastore).unwrap();
272        assert_eq!(
273            replacement.identity().coordinator_id,
274            identity.coordinator_id
275        );
276        assert_ne!(replacement.identity().generation_id, identity.generation_id);
277        assert!(replacement.can_recover(identity, attempt_id));
278        drop(replacement);
279        std::fs::remove_dir_all(root).unwrap();
280    }
281}
ahri_tre_lake/dataset_executor.rs

ahri_tre_lake/
dataset_executor.rs
