okbot_agent/
slots.rs

1//! A machine-wide limit on running agent processes, shared by every bot on the machine.
2//!
3//! Each slot is a lock file in a shared directory; a turn holds an exclusive lock on one of them.
4//! Locks are released by the operating system when the process exits, so a crashed bot never
5//! leaks a slot.
6
7use std::fs::{File, OpenOptions};
8use std::path::PathBuf;
9use std::time::Duration;
10
11/// `max` slots in `dir`.
12#[derive(Debug, Clone)]
13pub struct MachineSlots {
14    dir: PathBuf,
15    max: usize,
16}
17
18/// A held slot; dropping it frees the slot.
19#[derive(Debug)]
20pub struct Slot(#[allow(dead_code)] File);
21
22impl MachineSlots {
23    /// Slots in `dir` (created if missing), at most `max` agents at once.
24    pub fn new(dir: impl Into<PathBuf>, max: usize) -> Self {
25        Self {
26            dir: dir.into(),
27            max: max.max(1),
28        }
29    }
30
31    fn try_acquire(&self) -> std::io::Result<Option<Slot>> {
32        std::fs::create_dir_all(&self.dir)?;
33        for i in 0..self.max {
34            let file = OpenOptions::new()
35                .create(true)
36                .truncate(false)
37                .write(true)
38                .open(self.dir.join(format!("slot-{i}.lock")))?;
39            match file.try_lock() {
40                Ok(()) => return Ok(Some(Slot(file))),
41                Err(std::fs::TryLockError::WouldBlock) => {}
42                Err(std::fs::TryLockError::Error(e)) => return Err(e),
43            }
44        }
45        Ok(None)
46    }
47
48    /// Waits for a free slot.
49    pub async fn acquire(&self) -> std::io::Result<Slot> {
50        loop {
51            let this = self.clone();
52            let got = tokio::task::spawn_blocking(move || this.try_acquire())
53                .await
54                .map_err(std::io::Error::other)??;
55            if let Some(slot) = got {
56                return Ok(slot);
57            }
58            tokio::time::sleep(Duration::from_millis(250)).await;
59        }
60    }
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66
67    #[tokio::test]
68    async fn limits_holders_across_handles() {
69        let dir = tempfile::tempdir().unwrap();
70        let a = MachineSlots::new(dir.path(), 2);
71        let b = MachineSlots::new(dir.path(), 2); // another bot, same machine
72        let s1 = a.acquire().await.unwrap();
73        let _s2 = b.acquire().await.unwrap();
74        assert!(a.try_acquire().unwrap().is_none(), "both slots taken");
75        drop(s1);
76        assert!(
77            b.try_acquire().unwrap().is_some(),
78            "a freed slot is reusable"
79        );
80    }
81}
okbot_agent/slots.rs

okbot_agent/
slots.rs
