systemd-resolved-rs 0.1.1

// SPDX-License-Identifier: LGPL-2.1-or-later
use crate::resolver::ResolveError;
use std::cell::RefCell;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;

thread_local! {
    static CURRENT: RefCell<Option<QueryCancellation>> = RefCell::new(None);
}

#[derive(Clone, Debug, Default)]
pub(crate) struct QueryCancellation(Arc<AtomicBool>);

impl QueryCancellation {
    pub(crate) fn cancel(&self) {
        self.0.store(true, Ordering::Release);
    }

    pub(crate) fn is_cancelled(&self) -> bool {
        self.0.load(Ordering::Acquire)
    }

    pub(crate) fn same_as(&self, other: &Self) -> bool {
        Arc::ptr_eq(&self.0, &other.0)
    }
}

#[derive(Debug)]
struct CurrentGuard(Option<QueryCancellation>);

impl Drop for CurrentGuard {
    fn drop(&mut self) {
        CURRENT.with(|current| {
            current.replace(self.0.take());
        });
    }
}

pub(crate) fn current() -> Option<QueryCancellation> {
    CURRENT.with(|current| current.borrow().clone())
}

pub(crate) fn with<T>(cancellation: QueryCancellation, operation: impl FnOnce() -> T) -> T {
    let previous = CURRENT.with(|current| current.replace(Some(cancellation)));
    let _guard = CurrentGuard(previous);
    operation()
}

pub(crate) fn with_optional<T>(
    cancellation: Option<QueryCancellation>,
    operation: impl FnOnce() -> T,
) -> T {
    match cancellation {
        Some(cancellation) => with(cancellation, operation),
        None => operation(),
    }
}

pub(crate) fn check() -> Result<(), ResolveError> {
    if current().is_some_and(|cancellation| cancellation.is_cancelled()) {
        Err(ResolveError::QueryAborted)
    } else {
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn nested_context_restores_the_outer_cancellation() {
        let outer = QueryCancellation::default();
        let inner = QueryCancellation::default();
        with(outer.clone(), || {
            assert!(!current().expect("outer context").is_cancelled());
            with(inner.clone(), || {
                inner.cancel();
                assert!(matches!(check(), Err(ResolveError::QueryAborted)));
            });
            assert!(!current().expect("restored outer context").is_cancelled());
        });
        assert!(current().is_none());
    }
}