do_over/

context.rs

//! Policy context for passing metadata through execution.
//!
//! The context module provides types for passing request-scoped data through
//! policy execution, useful for logging, tracing, and correlation IDs.
//!
//! # Examples
//!
//! ```rust
//! use do_over::context::{Context, ContextKey};
//!
//! // Define a context key
//! static REQUEST_ID: ContextKey<String> = ContextKey::new("request_id");
//!
//! // Create a context with data
//! let mut ctx = Context::new();
//! ctx.insert(&REQUEST_ID, "req-123".to_string());
//!
//! // Retrieve data
//! if let Some(id) = ctx.get(&REQUEST_ID) {
//!     println!("Request ID: {}", id);
//! }
//! ```

use std::any::{Any, TypeId};
use std::collections::HashMap;
use std::sync::Arc;

/// A key for storing values in a Context.
///
/// Keys are identified by a name and the type they store.
///
/// # Examples
///
/// ```rust
/// use do_over::context::ContextKey;
///
/// static USER_ID: ContextKey<u64> = ContextKey::new("user_id");
/// static TRACE_ID: ContextKey<String> = ContextKey::new("trace_id");
/// ```
pub struct ContextKey<T> {
    name: &'static str,
    _marker: std::marker::PhantomData<T>,
}

impl<T> ContextKey<T> {
    /// Create a new context key.
    ///
    /// # Arguments
    ///
    /// * `name` - A descriptive name for debugging
    pub const fn new(name: &'static str) -> Self {
        Self {
            name,
            _marker: std::marker::PhantomData,
        }
    }

    /// Get the name of this key.
    pub fn name(&self) -> &'static str {
        self.name
    }
}

/// A type-safe container for request-scoped data.
///
/// Context allows you to pass metadata through policy execution without
/// modifying function signatures.
///
/// # Examples
///
/// ```rust
/// use do_over::context::{Context, ContextKey};
///
/// static CORRELATION_ID: ContextKey<String> = ContextKey::new("correlation_id");
/// static RETRY_COUNT: ContextKey<u32> = ContextKey::new("retry_count");
///
/// let mut ctx = Context::new();
/// ctx.insert(&CORRELATION_ID, "abc-123".to_string());
/// ctx.insert(&RETRY_COUNT, 0u32);
///
/// assert_eq!(ctx.get(&CORRELATION_ID), Some(&"abc-123".to_string()));
/// assert_eq!(ctx.get(&RETRY_COUNT), Some(&0u32));
/// ```
#[derive(Default)]
pub struct Context {
    values: HashMap<(TypeId, &'static str), Arc<dyn Any + Send + Sync>>,
}

impl Context {
    /// Create a new empty context.
    pub fn new() -> Self {
        Self::default()
    }

    /// Insert a value into the context.
    ///
    /// # Arguments
    ///
    /// * `key` - The context key
    /// * `value` - The value to store
    ///
    /// # Examples
    ///
    /// ```rust
    /// use do_over::context::{Context, ContextKey};
    ///
    /// static KEY: ContextKey<String> = ContextKey::new("key");
    ///
    /// let mut ctx = Context::new();
    /// ctx.insert(&KEY, "value".to_string());
    /// ```
    pub fn insert<T: Send + Sync + 'static>(&mut self, key: &ContextKey<T>, value: T) {
        let type_id = TypeId::of::<T>();
        self.values
            .insert((type_id, key.name), Arc::new(value));
    }

    /// Get a value from the context.
    ///
    /// # Arguments
    ///
    /// * `key` - The context key
    ///
    /// # Returns
    ///
    /// A reference to the value if it exists, or None.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use do_over::context::{Context, ContextKey};
    ///
    /// static KEY: ContextKey<String> = ContextKey::new("key");
    ///
    /// let mut ctx = Context::new();
    /// ctx.insert(&KEY, "value".to_string());
    ///
    /// assert_eq!(ctx.get(&KEY), Some(&"value".to_string()));
    /// ```
    pub fn get<T: Send + Sync + 'static>(&self, key: &ContextKey<T>) -> Option<&T> {
        let type_id = TypeId::of::<T>();
        self.values
            .get(&(type_id, key.name))
            .and_then(|v| v.downcast_ref::<T>())
    }

    /// Remove a value from the context.
    ///
    /// # Arguments
    ///
    /// * `key` - The context key
    ///
    /// # Returns
    ///
    /// The removed value if it existed (cloned before removal).
    pub fn remove<T: Send + Sync + Clone + 'static>(&mut self, key: &ContextKey<T>) -> Option<T> {
        let type_id = TypeId::of::<T>();
        // Clone the value before removing (required because Arc contains trait object)
        let value = self.values
            .get(&(type_id, key.name))
            .and_then(|v| v.downcast_ref::<T>())
            .cloned();
        self.values.remove(&(type_id, key.name));
        value
    }

    /// Check if the context contains a key.
    pub fn contains<T: Send + Sync + 'static>(&self, key: &ContextKey<T>) -> bool {
        let type_id = TypeId::of::<T>();
        self.values.contains_key(&(type_id, key.name))
    }

    /// Clear all values from the context.
    pub fn clear(&mut self) {
        self.values.clear();
    }
}

impl Clone for Context {
    fn clone(&self) -> Self {
        Self {
            values: self.values.clone(),
        }
    }
}

/// Common context keys for resilience operations.
pub mod keys {
    use super::ContextKey;

    /// Correlation ID for distributed tracing.
    pub static CORRELATION_ID: ContextKey<String> = ContextKey::new("correlation_id");

    /// Current retry attempt number.
    pub static RETRY_ATTEMPT: ContextKey<u32> = ContextKey::new("retry_attempt");

    /// Operation name or identifier.
    pub static OPERATION_NAME: ContextKey<String> = ContextKey::new("operation_name");

    /// Start time of the operation.
    pub static START_TIME: ContextKey<std::time::Instant> = ContextKey::new("start_time");
}

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

    static STRING_KEY: ContextKey<String> = ContextKey::new("string");
    static INT_KEY: ContextKey<i32> = ContextKey::new("int");

    #[test]
    fn test_insert_and_get() {
        let mut ctx = Context::new();
        ctx.insert(&STRING_KEY, "hello".to_string());
        ctx.insert(&INT_KEY, 42);

        assert_eq!(ctx.get(&STRING_KEY), Some(&"hello".to_string()));
        assert_eq!(ctx.get(&INT_KEY), Some(&42));
    }

    #[test]
    fn test_get_missing_key() {
        let ctx = Context::new();
        assert_eq!(ctx.get(&STRING_KEY), None);
    }

    #[test]
    fn test_contains() {
        let mut ctx = Context::new();
        assert!(!ctx.contains(&STRING_KEY));

        ctx.insert(&STRING_KEY, "value".to_string());
        assert!(ctx.contains(&STRING_KEY));
    }

    #[test]
    fn test_remove() {
        let mut ctx = Context::new();
        ctx.insert(&STRING_KEY, "value".to_string());

        let removed = ctx.remove(&STRING_KEY);
        assert_eq!(removed, Some("value".to_string()));
        assert!(!ctx.contains(&STRING_KEY));
    }

    #[test]
    fn test_clear() {
        let mut ctx = Context::new();
        ctx.insert(&STRING_KEY, "value".to_string());
        ctx.insert(&INT_KEY, 42);

        ctx.clear();
        assert!(!ctx.contains(&STRING_KEY));
        assert!(!ctx.contains(&INT_KEY));
    }

    #[test]
    fn test_clone() {
        let mut ctx = Context::new();
        ctx.insert(&STRING_KEY, "value".to_string());

        let ctx2 = ctx.clone();
        assert_eq!(ctx2.get(&STRING_KEY), Some(&"value".to_string()));
    }
}