kcode-intelligence-cancellation 0.1.0

//! UUID operation registration and hierarchical cancellation.

#![forbid(unsafe_code)]

use std::{
    collections::HashMap,
    sync::{Arc, Mutex},
};

use tokio::sync::watch;
use uuid::Uuid;

/// A cloneable registry of active UUID-named operations.
#[derive(Clone, Default)]
pub struct OperationRegistry {
    senders: Arc<Mutex<HashMap<Uuid, watch::Sender<bool>>>>,
}

/// A live operation registration removed from its registry when dropped.
pub struct ActiveOperation {
    id: Uuid,
    registry: OperationRegistry,
    cancellation: watch::Receiver<bool>,
    parent_cancellation: Option<watch::Receiver<bool>>,
}

/// A failure to inspect or register an operation.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum RegistryError {
    RegistryUnavailable,
    SameAsParent,
    OperationInProgress,
    ParentNotRunning,
}

impl OperationRegistry {
    /// Registers a request and optionally binds it to a currently running parent.
    pub fn register_request(
        &self,
        request_id: Uuid,
        parent_operation_id: Option<Uuid>,
    ) -> Result<ActiveOperation, RegistryError> {
        if let Some(parent_id) = parent_operation_id {
            if request_id == parent_id {
                return Err(RegistryError::SameAsParent);
            }
            let parent_cancellation = self
                .senders
                .lock()
                .map_err(|_| RegistryError::RegistryUnavailable)?
                .get(&parent_id)
                .map(watch::Sender::subscribe)
                .ok_or(RegistryError::ParentNotRunning)?;
            let mut operation = self.register(request_id)?;
            operation.parent_cancellation = Some(parent_cancellation);
            Ok(operation)
        } else {
            self.register(request_id)
        }
    }

    /// Sends direct cancellation to a registered operation.
    pub fn cancel(&self, id: Uuid) -> Result<bool, RegistryError> {
        let sender = self
            .senders
            .lock()
            .map_err(|_| RegistryError::RegistryUnavailable)?
            .get(&id)
            .cloned();
        Ok(sender.is_some_and(|sender| sender.send(true).is_ok()))
    }

    fn register(&self, id: Uuid) -> Result<ActiveOperation, RegistryError> {
        let (sender, cancellation) = watch::channel(false);
        let mut senders = self
            .senders
            .lock()
            .map_err(|_| RegistryError::RegistryUnavailable)?;
        if senders.contains_key(&id) {
            return Err(RegistryError::OperationInProgress);
        }
        senders.insert(id, sender);
        Ok(ActiveOperation {
            id,
            registry: self.clone(),
            cancellation,
            parent_cancellation: None,
        })
    }

    fn remove(&self, id: Uuid) {
        if let Ok(mut senders) = self.senders.lock() {
            senders.remove(&id);
        }
    }
}

impl ActiveOperation {
    /// Waits until this operation or its registered parent is cancelled.
    pub async fn cancelled(&mut self) {
        if let Some(parent) = &mut self.parent_cancellation {
            tokio::select! {
                _ = cancellation_requested(&mut self.cancellation) => {}
                _ = cancellation_requested(parent) => {}
            }
        } else {
            cancellation_requested(&mut self.cancellation).await;
        }
    }

    /// Reports only direct cancellation without waiting.
    ///
    /// This deliberately excludes cancellation inherited from a registered
    /// parent and must not be used as a general cancellation-state query.
    pub fn direct_cancellation_requested(&self) -> bool {
        *self.cancellation.borrow()
    }
}

async fn cancellation_requested(cancellation: &mut watch::Receiver<bool>) {
    if *cancellation.borrow() {
        return;
    }
    while cancellation.changed().await.is_ok() {
        if *cancellation.borrow() {
            return;
        }
    }
}

impl Drop for ActiveOperation {
    fn drop(&mut self) {
        self.registry.remove(self.id);
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use super::*;

    #[tokio::test]
    async fn direct_cancellation_stops_the_registered_operation() {
        let registry = OperationRegistry::default();
        let id = Uuid::new_v4();
        let mut operation = registry.register_request(id, None).unwrap();

        assert!(registry.cancel(id).unwrap());
        tokio::time::timeout(Duration::from_millis(50), operation.cancelled())
            .await
            .unwrap();
    }

    #[tokio::test]
    async fn child_inherits_registered_parent_cancellation() {
        let registry = OperationRegistry::default();
        let parent_id = Uuid::new_v4();
        let child_id = Uuid::new_v4();
        let _parent = registry.register_request(parent_id, None).unwrap();
        let mut child = registry
            .register_request(child_id, Some(parent_id))
            .unwrap();

        assert!(registry.cancel(parent_id).unwrap());
        tokio::time::timeout(Duration::from_millis(50), child.cancelled())
            .await
            .unwrap();
    }

    #[test]
    fn duplicate_ids_are_rejected() {
        let registry = OperationRegistry::default();
        let id = Uuid::new_v4();
        let _operation = registry.register_request(id, None).unwrap();

        assert!(matches!(
            registry.register_request(id, None),
            Err(RegistryError::OperationInProgress)
        ));
    }

    #[test]
    fn request_cannot_be_its_own_parent() {
        let registry = OperationRegistry::default();
        let id = Uuid::new_v4();

        assert!(matches!(
            registry.register_request(id, Some(id)),
            Err(RegistryError::SameAsParent)
        ));
    }

    #[test]
    fn missing_parent_is_rejected() {
        let registry = OperationRegistry::default();

        assert!(matches!(
            registry.register_request(Uuid::new_v4(), Some(Uuid::new_v4())),
            Err(RegistryError::ParentNotRunning)
        ));
    }

    #[test]
    fn dropping_an_operation_cleans_up_and_allows_reuse() {
        let registry = OperationRegistry::default();
        let id = Uuid::new_v4();
        let operation = registry.register_request(id, None).unwrap();
        drop(operation);

        assert!(!registry.cancel(id).unwrap());
        assert!(registry.register_request(id, None).is_ok());
    }

    #[test]
    fn nonblocking_probe_reports_only_direct_cancellation() {
        let registry = OperationRegistry::default();
        let parent_id = Uuid::new_v4();
        let child_id = Uuid::new_v4();
        let _parent = registry.register_request(parent_id, None).unwrap();
        let child = registry
            .register_request(child_id, Some(parent_id))
            .unwrap();

        assert!(registry.cancel(parent_id).unwrap());
        assert!(!child.direct_cancellation_requested());
        assert!(registry.cancel(child_id).unwrap());
        assert!(child.direct_cancellation_requested());
    }
}