diff --git a/crates/kild-tmux-shim/src/commands.rs b/crates/kild-tmux-shim/src/commands.rs index 02df0cad..7a0711e5 100644 --- a/crates/kild-tmux-shim/src/commands.rs +++ b/crates/kild-tmux-shim/src/commands.rs @@ -300,28 +300,41 @@ fn handle_split_window(args: SplitWindowArgs<'_>) -> Result { ); let sid = session_id()?; - let mut registry = state::load(&sid)?; + let mut locked = state::load_and_lock(&sid)?; // Determine which window the target pane belongs to let parent_pane_id = resolve_pane_id(args.target); - let window_id = registry + let window_id = locked + .registry() .panes .get(&parent_pane_id) .map(|p| p.window_id.clone()) .unwrap_or_else(|| "0".to_string()); - let pane_id = create_pty_pane(&mut registry, &window_id, &args.command)?; - state::save(&sid, ®istry)?; + let pane_id = create_pty_pane(locked.registry_mut(), &window_id, &args.command)?; - if args.print_info { - let fmt = args.format.unwrap_or("#{pane_id}"); - let session_name = ®istry.session_name; - let window_name = registry + // Capture format values before save consumes the guard (only when needed) + let print_values = if args.print_info { + let session_name = locked.registry().session_name.clone(); + let window_name = locked + .registry() .windows .get(&window_id) - .map(|w| w.name.as_str()) - .unwrap_or("main"); - let output = expand_format(fmt, &pane_id, session_name, &window_id, window_name, ""); + .map(|w| w.name.clone()) + .unwrap_or_else(|| "main".to_string()); + Some(( + args.format.unwrap_or("#{pane_id}"), + session_name, + window_name, + )) + } else { + None + }; + + locked.save()?; + + if let Some((fmt, session_name, window_name)) = print_values { + let output = expand_format(fmt, &pane_id, &session_name, &window_id, &window_name, ""); println!("{}", output); } @@ -451,10 +464,11 @@ fn handle_kill_pane(args: KillPaneArgs<'_>) -> Result { debug!(event = "shim.kill_pane_started", target = ?args.target); let sid = session_id()?; - let mut registry = state::load(&sid)?; + let mut locked = state::load_and_lock(&sid)?; let pane_id = resolve_pane_id(args.target); - let pane = registry + let pane = locked + .registry() .panes .get(&pane_id) .ok_or_else(|| ShimError::state(format!("pane {} not found in registry", pane_id)))?; @@ -492,8 +506,8 @@ fn handle_kill_pane(args: KillPaneArgs<'_>) -> Result { } } - registry.remove_pane(&pane_id); - state::save(&sid, ®istry)?; + locked.registry_mut().remove_pane(&pane_id); + locked.save()?; debug!(event = "shim.kill_pane_completed", pane_id = pane_id); Ok(0) @@ -574,17 +588,17 @@ fn handle_select_pane(args: SelectPaneArgs<'_>) -> Result { // Only need state if we have style or title to store if args.style.is_some() || args.title.is_some() { let sid = session_id()?; - let mut registry = state::load(&sid)?; + let mut locked = state::load_and_lock(&sid)?; let pane_id = resolve_pane_id(args.target); - if let Some(pane) = registry.panes.get_mut(&pane_id) { + if let Some(pane) = locked.registry_mut().panes.get_mut(&pane_id) { if let Some(style) = args.style { pane.border_style = style.to_string(); } if let Some(title) = args.title { pane.title = title.to_string(); } - state::save(&sid, ®istry)?; + locked.save()?; } } @@ -604,15 +618,15 @@ fn handle_set_option(args: SetOptionArgs<'_>) -> Result { // Store pane-scoped options in the pane entry if matches!(args.scope, OptionScope::Pane) { let sid = session_id()?; - let mut registry = state::load(&sid)?; + let mut locked = state::load_and_lock(&sid)?; let pane_id = resolve_pane_id(args.target); - if let Some(pane) = registry.panes.get_mut(&pane_id) { + if let Some(pane) = locked.registry_mut().panes.get_mut(&pane_id) { // Store known pane options if args.key == "pane-border-style" || args.key.ends_with("-style") { pane.border_style = args.value; } - state::save(&sid, ®istry)?; + locked.save()?; } } @@ -655,21 +669,21 @@ fn handle_new_session(args: NewSessionArgs<'_>) -> Result { debug!(event = "shim.new_session_started", name = ?args.session_name); let sid = session_id()?; - let mut registry = state::load(&sid)?; + let mut locked = state::load_and_lock(&sid)?; + // Read phase: derive names and IDs from current state let session_name = args .session_name .map(|s| s.to_string()) - .unwrap_or_else(|| format!("kild_{}", registry.sessions.len())); - + .unwrap_or_else(|| format!("kild_{}", locked.registry().sessions.len())); let window_name = args .window_name .map(|s| s.to_string()) .unwrap_or_else(|| "main".to_string()); + let window_id = format!("{}", locked.registry().windows.len()); - // Allocate a new window - let window_id = format!("{}", registry.windows.len()); - registry.windows.insert( + // Write phase: mutate registry + locked.registry_mut().windows.insert( window_id.clone(), WindowEntry { name: window_name, @@ -678,10 +692,10 @@ fn handle_new_session(args: NewSessionArgs<'_>) -> Result { ); // Create initial pane in the new window - let pane_id = create_pty_pane(&mut registry, &window_id, &[])?; + let pane_id = create_pty_pane(locked.registry_mut(), &window_id, &[])?; // Register session - registry.sessions.insert( + locked.registry_mut().sessions.insert( session_name.clone(), SessionEntry { name: session_name.clone(), @@ -689,7 +703,7 @@ fn handle_new_session(args: NewSessionArgs<'_>) -> Result { }, ); - state::save(&sid, ®istry)?; + locked.save()?; if args.print_info { let fmt = args.format.unwrap_or("#{pane_id}"); @@ -708,15 +722,23 @@ fn handle_new_window(args: NewWindowArgs<'_>) -> Result { debug!(event = "shim.new_window_started", name = ?args.name); let sid = session_id()?; - let mut registry = state::load(&sid)?; + let mut locked = state::load_and_lock(&sid)?; let window_name = args .name .map(|s| s.to_string()) .unwrap_or_else(|| "window".to_string()); - let window_id = format!("{}", registry.windows.len()); - registry.windows.insert( + // Read phase: derive IDs from current state + let window_id = format!("{}", locked.registry().windows.len()); + let session_key = args + .target + .and_then(|t| t.split(':').next()) + .unwrap_or(&locked.registry().session_name) + .to_string(); + + // Write phase: mutate registry + locked.registry_mut().windows.insert( window_id.clone(), WindowEntry { name: window_name.clone(), @@ -724,20 +746,13 @@ fn handle_new_window(args: NewWindowArgs<'_>) -> Result { }, ); - let pane_id = create_pty_pane(&mut registry, &window_id, &[])?; - - // Add window to the target session (or default session) - let session_key = args - .target - .and_then(|t| t.split(':').next()) - .unwrap_or(®istry.session_name) - .to_string(); + let pane_id = create_pty_pane(locked.registry_mut(), &window_id, &[])?; - if let Some(session) = registry.sessions.get_mut(&session_key) { + if let Some(session) = locked.registry_mut().sessions.get_mut(&session_key) { session.windows.push(window_id.clone()); } - state::save(&sid, ®istry)?; + locked.save()?; if args.print_info { let fmt = args.format.unwrap_or("#{pane_id}"); @@ -789,13 +804,13 @@ fn handle_break_pane(args: BreakPaneArgs<'_>) -> Result { debug!(event = "shim.break_pane_started", source = ?args.source); let sid = session_id()?; - let mut registry = state::load(&sid)?; + let mut locked = state::load_and_lock(&sid)?; let pane_id = resolve_pane_id(args.source); - if let Some(pane) = registry.panes.get_mut(&pane_id) { + if let Some(pane) = locked.registry_mut().panes.get_mut(&pane_id) { pane.hidden = true; } - state::save(&sid, ®istry)?; + locked.save()?; debug!(event = "shim.break_pane_completed", pane_id = pane_id); Ok(0) @@ -805,13 +820,13 @@ fn handle_join_pane(args: JoinPaneArgs<'_>) -> Result { debug!(event = "shim.join_pane_started", source = ?args.source); let sid = session_id()?; - let mut registry = state::load(&sid)?; + let mut locked = state::load_and_lock(&sid)?; let pane_id = resolve_pane_id(args.source); - if let Some(pane) = registry.panes.get_mut(&pane_id) { + if let Some(pane) = locked.registry_mut().panes.get_mut(&pane_id) { pane.hidden = false; } - state::save(&sid, ®istry)?; + locked.save()?; debug!(event = "shim.join_pane_completed", pane_id = pane_id); Ok(0) diff --git a/crates/kild-tmux-shim/src/state.rs b/crates/kild-tmux-shim/src/state.rs index 92a5881f..4a095431 100644 --- a/crates/kild-tmux-shim/src/state.rs +++ b/crates/kild-tmux-shim/src/state.rs @@ -4,8 +4,14 @@ //! //! - **Reads**: Use `load_shared()` to acquire a shared (read-only) lock. //! Multiple readers can hold shared locks concurrently. -//! - **Writes**: Use `load()` to acquire an exclusive lock, mutate the registry, -//! then call `save()` (which re-acquires the exclusive lock). +//! - **Read-modify-write**: Use `load_and_lock()` to acquire an exclusive lock +//! and return a `LockedRegistry` guard. The lock is held across the entire +//! load → mutate → save cycle, preventing pane ID races from concurrent +//! `split-window` calls. Call `locked.save()` to persist and release. +//! - **Init only**: `save()` acquires its own exclusive lock. Only used by +//! `init_registry()` where atomic load-modify-save is not needed. +//! - **Test only**: `load()` acquires an exclusive lock for a single read. +//! Gated behind `#[cfg(test)]`. //! //! Locks are automatically released when the `Flock` handle is dropped (RAII). @@ -174,6 +180,9 @@ pub fn load_shared(session_id: &str) -> Result { } /// Load the registry with an exclusive (write) lock. +/// Production callers should use `load_and_lock()` instead for atomic +/// load-modify-save. This remains available for tests. +#[cfg(test)] pub fn load(session_id: &str) -> Result { let data_path = panes_path(session_id)?; let _lock = acquire_lock(session_id, LockMode::Exclusive)?; @@ -192,6 +201,8 @@ pub fn load(session_id: &str) -> Result { Ok(registry) } +/// Standalone save with its own lock. For load-modify-save workflows, +/// prefer `LockedRegistry::save()` which holds the lock atomically. pub fn save(session_id: &str, registry: &PaneRegistry) -> Result<(), ShimError> { let data_path = panes_path(session_id)?; let _lock = acquire_lock(session_id, LockMode::Exclusive)?; @@ -209,6 +220,63 @@ pub fn save(session_id: &str, registry: &PaneRegistry) -> Result<(), ShimError> Ok(()) } +/// RAII guard holding both the exclusive flock and the deserialized registry. +/// The lock is held for the entire duration the guard is alive. +/// Drop releases the lock via `Flock`'s `Drop` impl. +pub struct LockedRegistry { + registry: PaneRegistry, + session_id: String, + _lock: Flock, +} + +impl LockedRegistry { + /// Access the registry for reading. + pub fn registry(&self) -> &PaneRegistry { + &self.registry + } + + /// Access the registry for mutation. + pub fn registry_mut(&mut self) -> &mut PaneRegistry { + &mut self.registry + } + + /// Write the (possibly mutated) registry to disk while still holding the lock. + /// Consumes self — lock is released when this returns. + pub fn save(self) -> Result<(), ShimError> { + let data_path = panes_path(&self.session_id)?; + let content = + serde_json::to_string_pretty(&self.registry).map_err(|e| ShimError::StateError { + message: format!("failed to serialize pane registry: {}", e), + })?; + let mut file = fs::File::create(&data_path).map_err(|e| ShimError::StateError { + message: format!("failed to write {}: {}", data_path.display(), e), + })?; + file.write_all(content.as_bytes())?; + file.flush()?; + Ok(()) + } +} + +/// Acquire an exclusive lock and read the registry atomically. +/// Returns a guard that keeps the lock alive until dropped or saved. +pub fn load_and_lock(session_id: &str) -> Result { + let data_path = panes_path(session_id)?; + let lock = acquire_lock(session_id, LockMode::Exclusive)?; + let content = fs::read_to_string(&data_path).map_err(|e| ShimError::StateError { + message: format!("failed to read {}: {}", data_path.display(), e), + })?; + let registry: PaneRegistry = + serde_json::from_str(&content).map_err(|e| ShimError::StateError { + message: format!("failed to parse pane registry: {}", e), + })?; + registry.validate()?; + Ok(LockedRegistry { + registry, + session_id: session_id.to_string(), + _lock: lock, + }) +} + pub fn allocate_pane_id(registry: &mut PaneRegistry) -> String { let id = format!("%{}", registry.next_pane_id); registry.next_pane_id += 1; @@ -689,4 +757,73 @@ mod tests { let dir = state_dir(&test_id).unwrap(); fs::remove_dir_all(&dir).ok(); } + + #[test] + fn test_load_and_lock_basic_roundtrip() { + let test_id = format!("test-{}", uuid::Uuid::new_v4()); + init_registry(&test_id, "daemon-abc-123").unwrap(); + + let mut locked = load_and_lock(&test_id).unwrap(); + assert_eq!(locked.registry().next_pane_id, 1); + + let pane_id = allocate_pane_id(locked.registry_mut()); + assert_eq!(pane_id, "%1"); + + locked.save().unwrap(); + + let reloaded = load(&test_id).unwrap(); + assert_eq!(reloaded.next_pane_id, 2); + + let dir = state_dir(&test_id).unwrap(); + fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn test_load_and_lock_concurrent_allocations_unique() { + use std::collections::HashSet; + use std::sync::Arc; + use std::thread; + + let test_id = format!("test-{}", uuid::Uuid::new_v4()); + init_registry(&test_id, "daemon-abc-123").unwrap(); + + let num_threads = 8; + let test_id = Arc::new(test_id); + let mut handles = vec![]; + + for _ in 0..num_threads { + let id = Arc::clone(&test_id); + handles.push(thread::spawn(move || { + let mut locked = load_and_lock(&id).unwrap(); + let pane_id = allocate_pane_id(locked.registry_mut()); + locked.save().unwrap(); + pane_id + })); + } + + let pane_ids: Vec = handles.into_iter().map(|h| h.join().unwrap()).collect(); + let unique: HashSet<&String> = pane_ids.iter().collect(); + + assert_eq!( + unique.len(), + num_threads, + "Expected {} unique pane IDs, got {}: {:?}", + num_threads, + unique.len(), + pane_ids + ); + + let final_registry = load(&test_id).unwrap(); + // Initial next_pane_id was 1, plus 8 allocations = 9 + assert_eq!( + final_registry.next_pane_id, + 1 + num_threads as u32, + "next_pane_id should be {} after {} allocations from initial 1", + 1 + num_threads as u32, + num_threads + ); + + let dir = state_dir(&test_id).unwrap(); + fs::remove_dir_all(&dir).ok(); + } }