1use crate::config::{AssistantConfig, RigConfig};
8use crate::node::Node;
9use fs4::fs_std::FileExt;
10use std::fs::File;
11use std::path::PathBuf;
12use std::time::{Duration, Instant};
13
14const DEFAULT_LOCK_TIMEOUT: Duration = Duration::from_secs(300);
23
24pub enum Acquire {
26 Skip(String),
28 Fail(anyhow::Error),
30}
31
32impl From<anyhow::Error> for Acquire {
33 fn from(e: anyhow::Error) -> Self {
34 Acquire::Fail(e)
35 }
36}
37
38pub struct Rig {
39 pub config: RigConfig,
40 pub base_dir: PathBuf,
42 _lock: Exclusive,
43}
44
45enum Exclusive {
49 Flock(#[allow(dead_code)] RigLock),
51 Lease(#[allow(dead_code)] crate::net::lease::LeaseClient),
52}
53
54impl Rig {
55 pub fn acquire() -> Result<Rig, Acquire> {
61 let Some(path) = RigConfig::locate()? else {
62 return Err(Acquire::Skip(format!(
63 "no rig: {} not found (set {} or create one to run on hardware)",
64 crate::config::CONFIG_FILE,
65 crate::config::ENV_VAR,
66 )));
67 };
68 let config = RigConfig::load(&path)?;
69 let base_dir = path.parent().unwrap_or(std::path::Path::new(".")).to_path_buf();
70
71 let timeout = std::env::var("BANC_LOCK_TIMEOUT_SECS")
72 .ok()
73 .and_then(|s| s.parse().ok())
74 .map(Duration::from_secs)
75 .unwrap_or(DEFAULT_LOCK_TIMEOUT);
76
77 let lock = if let Some(lease) = &config.rig.lease {
78 let token = read_token(&lease.token_file, &base_dir)?;
79 let holder = std::env::var("BANC_LEASE_HOLDER").unwrap_or_else(|_| {
80 let host = std::env::var("HOSTNAME").unwrap_or_else(|_| "?".into());
81 format!("{host}:{}", std::process::id())
82 });
83 Exclusive::Lease(
84 crate::net::lease::LeaseClient::acquire(&lease.addr, &token, &holder, timeout)
85 .map_err(|e| anyhow::anyhow!("acquiring rig lease: {e}"))?,
86 )
87 } else {
88 let lock_path = config
89 .rig
90 .lock_file
91 .clone()
92 .map(|p| if p.is_absolute() { p } else { base_dir.join(p) })
93 .unwrap_or_else(|| base_dir.join("target").join("banc.lock"));
94 Exclusive::Flock(RigLock::take(lock_path, timeout)?)
95 };
96
97 Ok(Rig { config, base_dir, _lock: lock })
98 }
99
100 pub async fn assistant(&self, name: &str) -> anyhow::Result<Node> {
102 let cfg: &AssistantConfig = self
103 .config
104 .assistant(name)
105 .ok_or_else(|| anyhow::anyhow!("no assistant '{name}' in rig config"))?;
106 let token = cfg
107 .token_file
108 .as_ref()
109 .map(|p| read_token(p, &self.base_dir))
110 .transpose()?;
111 Node::connect(cfg, token.as_deref()).await
112 }
113}
114
115fn read_token(path: &std::path::Path, base_dir: &std::path::Path) -> anyhow::Result<String> {
116 let path = if path.is_absolute() { path.to_path_buf() } else { base_dir.join(path) };
117 let token = std::fs::read_to_string(&path)
118 .map_err(|e| anyhow::anyhow!("reading token file {}: {e}", path.display()))?;
119 Ok(token.trim().to_owned())
120}
121
122struct RigLock {
126 file: File,
127}
128
129impl RigLock {
130 fn take(path: PathBuf, timeout: Duration) -> anyhow::Result<Self> {
131 if let Some(parent) = path.parent() {
132 std::fs::create_dir_all(parent)?;
133 }
134 let file = File::create(&path)
135 .map_err(|e| anyhow::anyhow!("creating rig lock {}: {e}", path.display()))?;
136 let deadline = Instant::now() + timeout;
137 loop {
138 if file.try_lock_exclusive()? {
139 return Ok(RigLock { file });
140 }
141 if Instant::now() >= deadline {
142 anyhow::bail!(
143 "rig lock {} held by another process for over {timeout:?}",
144 path.display()
145 );
146 }
147 std::thread::sleep(Duration::from_millis(200));
148 }
149 }
150}
151
152impl Drop for RigLock {
153 fn drop(&mut self) {
154 let _ = FileExt::unlock(&self.file);
155 }
156}
157
158#[cfg(test)]
159mod tests {
160 use super::*;
161
162 #[test]
163 fn lock_excludes_second_taker_until_dropped() {
164 let path = std::env::temp_dir()
165 .join(format!("banc-rig-lock-test-{}", std::process::id()));
166 let held = RigLock::take(path.clone(), Duration::ZERO).unwrap();
167 let contended = RigLock::take(path.clone(), Duration::ZERO);
168 assert!(contended.is_err(), "second take must fail while lock is held");
169 drop(held);
170 RigLock::take(path.clone(), Duration::ZERO).unwrap();
171 std::fs::remove_file(path).ok();
172 }
173}