1use std::collections::HashMap;
9use std::fs::{File, OpenOptions};
10use std::io;
11use std::path::{Path, PathBuf};
12use std::time::{Duration, Instant};
13
14use super::backend_registry::BackendKey;
15
16pub const DEFAULT_SPAWN_ATTEMPTS_PER_WINDOW: u32 = 3;
18
19pub const DEFAULT_SPAWN_BUDGET_WINDOW: Duration = Duration::from_secs(30);
21
22pub fn acquire_spawn_lock(path: impl AsRef<Path>) -> Result<SpawnLockGuard, SpawnLockError> {
35 acquire_spawn_lock_with_hook(path.as_ref(), |_, _| {})
36}
37
38fn acquire_spawn_lock_with_hook<F>(
39 path: &Path,
40 mut before_lock: F,
41) -> Result<SpawnLockGuard, SpawnLockError>
42where
43 F: FnMut(&Path, &File),
44{
45 let path_buf = path.to_path_buf();
46 let file = open_lock_file(path).map_err(|source| SpawnLockError::Open {
47 path: path_buf.clone(),
48 source,
49 })?;
50
51 before_lock(path, &file);
52
53 try_lock_file(&file).map_err(|source| {
54 if is_lock_conflict(&source) {
55 SpawnLockError::AlreadyLocked {
56 path: path_buf.clone(),
57 }
58 } else {
59 SpawnLockError::Lock {
60 path: path_buf.clone(),
61 source,
62 }
63 }
64 })?;
65
66 let opened_identity =
67 file_identity(&file).map_err(|source| lock_identity_error(&path_buf, &file, source))?;
68 let current_identity = match path_identity(path) {
69 Ok(identity) => identity,
70 Err(source) if source.kind() == io::ErrorKind::NotFound => {
71 let _ = try_unlock_file(&file);
72 return Err(SpawnLockError::DeletedOrRecreated {
73 path: path_buf,
74 opened_identity,
75 current_identity: None,
76 });
77 }
78 Err(source) => return Err(lock_identity_error(&path_buf, &file, source)),
79 };
80
81 if opened_identity != current_identity {
82 let _ = try_unlock_file(&file);
83 return Err(SpawnLockError::DeletedOrRecreated {
84 path: path_buf,
85 opened_identity,
86 current_identity,
87 });
88 }
89
90 Ok(SpawnLockGuard {
91 file,
92 path: path_buf,
93 identity: opened_identity,
94 })
95}
96
97fn lock_identity_error(path: &Path, file: &File, source: io::Error) -> SpawnLockError {
98 let _ = try_unlock_file(file);
99 SpawnLockError::Identity {
100 path: path.to_path_buf(),
101 source,
102 }
103}
104
105#[must_use = "dropping the guard releases the backend spawn lock immediately"]
107#[derive(Debug)]
108pub struct SpawnLockGuard {
109 file: File,
110 path: PathBuf,
111 identity: Option<SpawnLockFileIdentity>,
112}
113
114impl SpawnLockGuard {
115 pub fn path(&self) -> &Path {
117 &self.path
118 }
119
120 pub fn file_identity(&self) -> Option<SpawnLockFileIdentity> {
122 self.identity
123 }
124}
125
126impl Drop for SpawnLockGuard {
127 fn drop(&mut self) {
128 let _ = try_unlock_file(&self.file);
129 }
130}
131
132#[derive(Clone, Copy, Debug, PartialEq, Eq)]
134pub struct SpawnLockFileIdentity {
135 pub device: u64,
137 pub file: u64,
139}
140
141#[derive(Debug, thiserror::Error)]
143pub enum SpawnLockError {
144 #[error("failed to open backend spawn lock file {path}: {source}")]
146 Open {
147 path: PathBuf,
149 #[source]
151 source: io::Error,
152 },
153 #[error("backend spawn lock file {path} is already locked")]
155 AlreadyLocked {
156 path: PathBuf,
158 },
159 #[error("failed to lock backend spawn lock file {path}: {source}")]
161 Lock {
162 path: PathBuf,
164 #[source]
166 source: io::Error,
167 },
168 #[error("backend spawn lock file {path} was deleted or recreated during acquisition")]
170 DeletedOrRecreated {
171 path: PathBuf,
173 opened_identity: Option<SpawnLockFileIdentity>,
175 current_identity: Option<SpawnLockFileIdentity>,
177 },
178 #[error("failed to verify backend spawn lock file identity for {path}: {source}")]
180 Identity {
181 path: PathBuf,
183 #[source]
185 source: io::Error,
186 },
187}
188
189#[derive(Clone, Copy, Debug, PartialEq, Eq)]
191pub struct SpawnBudgetConfig {
192 pub max_attempts: u32,
194 pub window: Duration,
196}
197
198impl SpawnBudgetConfig {
199 pub fn new(max_attempts: u32, window: Duration) -> Self {
201 Self {
202 max_attempts: max_attempts.max(1),
203 window: if window.is_zero() {
204 Duration::from_millis(1)
205 } else {
206 window
207 },
208 }
209 }
210}
211
212impl Default for SpawnBudgetConfig {
213 fn default() -> Self {
214 Self {
215 max_attempts: DEFAULT_SPAWN_ATTEMPTS_PER_WINDOW,
216 window: DEFAULT_SPAWN_BUDGET_WINDOW,
217 }
218 }
219}
220
221#[derive(Debug)]
223pub struct SpawnCoordinator {
224 config: SpawnBudgetConfig,
225 states: HashMap<BackendKey, SpawnBudgetState>,
226}
227
228impl SpawnCoordinator {
229 pub fn new() -> Self {
231 Self::with_config(SpawnBudgetConfig::default())
232 }
233
234 pub fn with_config(config: SpawnBudgetConfig) -> Self {
236 Self {
237 config,
238 states: HashMap::new(),
239 }
240 }
241
242 pub fn try_begin(
248 &mut self,
249 key: BackendKey,
250 now: Instant,
251 ) -> Result<SpawnPermit, SpawnBeginError> {
252 let state = self
253 .states
254 .entry(key.clone())
255 .or_insert_with(|| SpawnBudgetState::new(now));
256 state.refresh(now, self.config.window);
257
258 if state.in_flight {
259 return Err(SpawnBeginError::AlreadyInProgress);
260 }
261
262 if state.attempts_used >= self.config.max_attempts {
263 let is_storm_trip = !state.storm_signaled;
264 state.storm_signaled = true;
265 return Err(SpawnBeginError::BudgetExhausted {
266 retry_after: retry_after(state.window_started_at, now, self.config.window),
267 remaining: 0,
268 is_storm_trip,
269 });
270 }
271
272 state.attempts_used += 1;
273 state.in_flight = true;
274 Ok(SpawnPermit {
275 key,
276 attempt_number: state.attempts_used,
277 remaining_after_begin: self.config.max_attempts - state.attempts_used,
278 })
279 }
280
281 pub fn finish(&mut self, key: &BackendKey, outcome: SpawnOutcome, now: Instant) {
283 let Some(state) = self.states.get_mut(key) else {
284 return;
285 };
286 state.refresh(now, self.config.window);
287 state.in_flight = false;
288 if outcome == SpawnOutcome::Success {
289 state.window_started_at = now;
290 state.attempts_used = 0;
291 state.storm_signaled = false;
292 }
293 }
294
295 pub fn snapshot(&mut self, key: BackendKey, now: Instant) -> SpawnBudgetSnapshot {
297 let state = self
298 .states
299 .entry(key.clone())
300 .or_insert_with(|| SpawnBudgetState::new(now));
301 state.refresh(now, self.config.window);
302 snapshot_for(key, state, self.config, now)
303 }
304}
305
306impl Default for SpawnCoordinator {
307 fn default() -> Self {
308 Self::new()
309 }
310}
311
312#[derive(Clone, Debug, PartialEq, Eq)]
314pub struct SpawnPermit {
315 pub key: BackendKey,
317 pub attempt_number: u32,
319 pub remaining_after_begin: u32,
321}
322
323#[derive(Clone, Copy, Debug, PartialEq, Eq)]
325pub enum SpawnOutcome {
326 Success,
328 Failed,
330}
331
332#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
334pub enum SpawnBeginError {
335 #[error("backend spawn already in progress")]
337 AlreadyInProgress,
338 #[error("backend spawn budget exhausted; retry after {retry_after:?}")]
340 BudgetExhausted {
341 retry_after: Duration,
343 remaining: u32,
345 is_storm_trip: bool,
352 },
353}
354
355#[derive(Clone, Debug, PartialEq, Eq)]
357pub struct SpawnBudgetSnapshot {
358 pub key: BackendKey,
360 pub attempts_used: u32,
362 pub remaining: u32,
364 pub in_flight: bool,
366 pub retry_after: Option<Duration>,
368}
369
370#[derive(Clone, Debug)]
371struct SpawnBudgetState {
372 window_started_at: Instant,
373 attempts_used: u32,
374 in_flight: bool,
375 storm_signaled: bool,
383}
384
385impl SpawnBudgetState {
386 fn new(now: Instant) -> Self {
387 Self {
388 window_started_at: now,
389 attempts_used: 0,
390 in_flight: false,
391 storm_signaled: false,
392 }
393 }
394
395 fn refresh(&mut self, now: Instant, window: Duration) {
396 if elapsed_since(self.window_started_at, now) >= window {
397 self.window_started_at = now;
398 self.attempts_used = 0;
399 self.in_flight = false;
400 self.storm_signaled = false;
401 }
402 }
403}
404
405fn snapshot_for(
406 key: BackendKey,
407 state: &SpawnBudgetState,
408 config: SpawnBudgetConfig,
409 now: Instant,
410) -> SpawnBudgetSnapshot {
411 let remaining = config.max_attempts.saturating_sub(state.attempts_used);
412 SpawnBudgetSnapshot {
413 key,
414 attempts_used: state.attempts_used,
415 remaining,
416 in_flight: state.in_flight,
417 retry_after: (remaining == 0)
418 .then(|| retry_after(state.window_started_at, now, config.window)),
419 }
420}
421
422fn retry_after(window_started_at: Instant, now: Instant, window: Duration) -> Duration {
423 window.saturating_sub(elapsed_since(window_started_at, now))
424}
425
426fn elapsed_since(started_at: Instant, now: Instant) -> Duration {
427 now.checked_duration_since(started_at)
428 .unwrap_or(Duration::ZERO)
429}
430
431fn open_lock_file(path: &Path) -> io::Result<File> {
432 let mut options = OpenOptions::new();
433 options.read(true).write(true).create(true);
434 configure_lock_file_options(&mut options);
435 options.open(path)
436}
437
438#[cfg(unix)]
439fn configure_lock_file_options(options: &mut OpenOptions) {
440 use std::os::unix::fs::OpenOptionsExt;
441
442 options.mode(0o600);
443}
444
445#[cfg(windows)]
446fn configure_lock_file_options(options: &mut OpenOptions) {
447 use std::os::windows::fs::OpenOptionsExt;
448 use winapi::um::winnt::{FILE_SHARE_DELETE, FILE_SHARE_READ, FILE_SHARE_WRITE};
449
450 options.share_mode(FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE);
451}
452
453#[cfg(not(any(unix, windows)))]
454fn configure_lock_file_options(_options: &mut OpenOptions) {}
455
456#[cfg(unix)]
457fn try_lock_file(file: &File) -> io::Result<()> {
458 use std::os::unix::io::AsRawFd;
459
460 let result = unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_EX | libc::LOCK_NB) };
461 if result == 0 {
462 Ok(())
463 } else {
464 Err(io::Error::last_os_error())
465 }
466}
467
468#[cfg(unix)]
469fn try_unlock_file(file: &File) -> io::Result<()> {
470 use std::os::unix::io::AsRawFd;
471
472 let result = unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_UN) };
473 if result == 0 {
474 Ok(())
475 } else {
476 Err(io::Error::last_os_error())
477 }
478}
479
480#[cfg(unix)]
481fn is_lock_conflict(error: &io::Error) -> bool {
482 error.raw_os_error() == Some(libc::EWOULDBLOCK) || error.raw_os_error() == Some(libc::EAGAIN)
483}
484
485#[cfg(unix)]
486fn file_identity(file: &File) -> io::Result<Option<SpawnLockFileIdentity>> {
487 use std::os::unix::fs::MetadataExt;
488
489 let metadata = file.metadata()?;
490 Ok(Some(SpawnLockFileIdentity {
491 device: metadata.dev(),
492 file: metadata.ino(),
493 }))
494}
495
496#[cfg(unix)]
497fn path_identity(path: &Path) -> io::Result<Option<SpawnLockFileIdentity>> {
498 use std::os::unix::fs::MetadataExt;
499
500 let metadata = path.metadata()?;
501 Ok(Some(SpawnLockFileIdentity {
502 device: metadata.dev(),
503 file: metadata.ino(),
504 }))
505}
506
507#[cfg(windows)]
508fn try_lock_file(file: &File) -> io::Result<()> {
509 use std::mem;
510 use std::os::windows::io::AsRawHandle;
511 use winapi::um::fileapi::LockFileEx;
512 use winapi::um::minwinbase::{LOCKFILE_EXCLUSIVE_LOCK, LOCKFILE_FAIL_IMMEDIATELY, OVERLAPPED};
513 use winapi::um::winnt::HANDLE;
514
515 let mut overlapped: OVERLAPPED = unsafe { mem::zeroed() };
516 let result = unsafe {
517 LockFileEx(
518 file.as_raw_handle() as HANDLE,
519 LOCKFILE_EXCLUSIVE_LOCK | LOCKFILE_FAIL_IMMEDIATELY,
520 0,
521 u32::MAX,
522 u32::MAX,
523 &mut overlapped,
524 )
525 };
526 if result == 0 {
527 Err(io::Error::last_os_error())
528 } else {
529 Ok(())
530 }
531}
532
533#[cfg(windows)]
534fn try_unlock_file(file: &File) -> io::Result<()> {
535 use std::mem;
536 use std::os::windows::io::AsRawHandle;
537 use winapi::um::fileapi::UnlockFileEx;
538 use winapi::um::minwinbase::OVERLAPPED;
539 use winapi::um::winnt::HANDLE;
540
541 let mut overlapped: OVERLAPPED = unsafe { mem::zeroed() };
542 let result = unsafe {
543 UnlockFileEx(
544 file.as_raw_handle() as HANDLE,
545 0,
546 u32::MAX,
547 u32::MAX,
548 &mut overlapped,
549 )
550 };
551 if result == 0 {
552 Err(io::Error::last_os_error())
553 } else {
554 Ok(())
555 }
556}
557
558#[cfg(windows)]
559fn is_lock_conflict(error: &io::Error) -> bool {
560 use winapi::shared::winerror::ERROR_LOCK_VIOLATION;
561
562 error.raw_os_error() == Some(ERROR_LOCK_VIOLATION as i32)
563}
564
565#[cfg(windows)]
566fn file_identity(file: &File) -> io::Result<Option<SpawnLockFileIdentity>> {
567 use std::mem::MaybeUninit;
568 use std::os::windows::io::AsRawHandle;
569 use winapi::um::fileapi::{GetFileInformationByHandle, BY_HANDLE_FILE_INFORMATION};
570 use winapi::um::winnt::HANDLE;
571
572 let mut info = MaybeUninit::<BY_HANDLE_FILE_INFORMATION>::uninit();
573 let result =
574 unsafe { GetFileInformationByHandle(file.as_raw_handle() as HANDLE, info.as_mut_ptr()) };
575 if result == 0 {
576 return Err(io::Error::last_os_error());
577 }
578
579 let info = unsafe { info.assume_init() };
580 Ok(Some(SpawnLockFileIdentity {
581 device: info.dwVolumeSerialNumber as u64,
582 file: ((info.nFileIndexHigh as u64) << 32) | info.nFileIndexLow as u64,
583 }))
584}
585
586#[cfg(windows)]
587fn path_identity(path: &Path) -> io::Result<Option<SpawnLockFileIdentity>> {
588 let mut options = OpenOptions::new();
589 options.read(true).write(true);
590 configure_lock_file_options(&mut options);
591 let file = options.open(path)?;
592 file_identity(&file)
593}
594
595#[cfg(not(any(unix, windows)))]
596fn try_lock_file(_file: &File) -> io::Result<()> {
597 Err(io::Error::new(
598 io::ErrorKind::Unsupported,
599 "backend spawn file locks are supported only on Unix and Windows",
600 ))
601}
602
603#[cfg(not(any(unix, windows)))]
604fn try_unlock_file(_file: &File) -> io::Result<()> {
605 Ok(())
606}
607
608#[cfg(not(any(unix, windows)))]
609fn is_lock_conflict(_error: &io::Error) -> bool {
610 false
611}
612
613#[cfg(not(any(unix, windows)))]
614fn file_identity(_file: &File) -> io::Result<Option<SpawnLockFileIdentity>> {
615 Ok(None)
616}
617
618#[cfg(not(any(unix, windows)))]
619fn path_identity(_path: &Path) -> io::Result<Option<SpawnLockFileIdentity>> {
620 Ok(None)
621}
622
623#[cfg(test)]
624mod tests {
625 use std::fs;
626
627 use super::*;
628
629 #[test]
630 #[cfg(any(unix, windows))]
631 fn acquire_spawn_lock_detects_lock_file_replacement_between_open_and_lock() {
632 let tmp = tempfile::tempdir().unwrap();
633 let lock_path = tmp.path().join("backend.spawn.lock");
634 let replaced_path = tmp.path().join("backend.spawn.lock.replaced");
635
636 let err = acquire_spawn_lock_with_hook(&lock_path, |path, _file| {
637 fs::rename(path, &replaced_path).unwrap();
638 fs::write(path, b"replacement lock file").unwrap();
639 })
640 .unwrap_err();
641
642 let SpawnLockError::DeletedOrRecreated {
643 path,
644 opened_identity: Some(opened_identity),
645 current_identity: Some(current_identity),
646 } = err
647 else {
648 panic!("expected deleted/recreated error, got {err:?}");
649 };
650
651 assert_eq!(path, lock_path);
652 assert_ne!(opened_identity, current_identity);
653
654 let _guard = acquire_spawn_lock(&lock_path).unwrap();
655 }
656}