azure_data_cosmos 0.32.0

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

//! Defines fault injection rules that combine conditions and results.

use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Instant;

use azure_core::http::StatusCode;

use super::condition::FaultInjectionCondition;
use super::result::FaultInjectionResult;

/// A fault injection rule that defines when and how to inject faults.
#[derive(Debug)]
pub struct FaultInjectionRule {
    /// The condition under which to inject the fault.
    pub condition: FaultInjectionCondition,
    /// The result to inject when the condition is met.
    pub result: FaultInjectionResult,
    /// The absolute time at which the rule becomes active.
    pub start_time: Instant,
    /// The absolute time at which the rule expires, if set.
    pub end_time: Option<Instant>,
    /// The total hit limit of the rule.
    pub hit_limit: Option<u32>,
    /// Unique identifier for the fault injection scenario.
    pub id: String,
    /// Whether the rule is currently enabled.
    enabled: Arc<AtomicBool>,
    /// Number of times the rule has been matched (including matches where no fault was injected).
    hit_count: Arc<AtomicU32>,
    /// HTTP status codes of responses for matched requests that passed through without fault injection.
    passthrough_statuses: Mutex<Vec<StatusCode>>,
}

/// Cloning snapshots the current `hit_count` and `enabled` state rather than
/// resetting them, so a clone of a rule that has been hit 5 times starts at 5.
impl Clone for FaultInjectionRule {
    fn clone(&self) -> Self {
        Self {
            condition: self.condition.clone(),
            result: self.result.clone(),
            start_time: self.start_time,
            end_time: self.end_time,
            hit_limit: self.hit_limit,
            id: self.id.clone(),
            enabled: Arc::new(AtomicBool::new(self.enabled.load(Ordering::SeqCst))),
            hit_count: Arc::new(AtomicU32::new(self.hit_count.load(Ordering::SeqCst))),
            passthrough_statuses: Mutex::new(self.passthrough_statuses.lock().unwrap().clone()),
        }
    }
}

impl FaultInjectionRule {
    /// Returns whether the rule is currently enabled.
    pub fn is_enabled(&self) -> bool {
        self.enabled.load(Ordering::SeqCst)
    }

    /// Enables the rule.
    pub fn enable(&self) {
        self.enabled.store(true, Ordering::SeqCst);
    }

    /// Disables the rule.
    pub fn disable(&self) {
        self.enabled.store(false, Ordering::SeqCst);
    }

    /// Returns the number of times this rule has been matched.
    ///
    /// The hit count is incremented each time the rule's condition matches a
    /// request, regardless of whether the fault was actually applied (e.g.,
    /// probability-based skipping still increments the count).
    pub fn hit_count(&self) -> u32 {
        self.hit_count.load(Ordering::SeqCst)
    }

    /// Increments the hit count by one.
    pub(super) fn increment_hit_count(&self) {
        self.hit_count.fetch_add(1, Ordering::SeqCst);
    }

    /// Resets the hit count to zero.
    pub fn reset_hit_count(&self) {
        self.hit_count.store(0, Ordering::SeqCst);
    }

    /// Returns a shared reference to the enabled flag for cross-path state sharing.
    pub(crate) fn shared_enabled(&self) -> Arc<AtomicBool> {
        Arc::clone(&self.enabled)
    }

    /// Returns a shared reference to the hit count for cross-path state sharing.
    pub(crate) fn shared_hit_count(&self) -> Arc<AtomicU32> {
        Arc::clone(&self.hit_count)
    }

    /// Records the HTTP status code of a response for a matched request that
    /// passed through without fault injection (spy/passthrough mode).
    pub(super) fn record_passthrough_status(&self, status: StatusCode) {
        self.passthrough_statuses.lock().unwrap().push(status);
    }

    /// Returns the HTTP status codes of responses for matched requests that
    /// passed through without fault injection.
    ///
    /// When a rule matches a request but does not inject a fault (e.g., no
    /// `error_type` or `custom_response` is set), the real service response
    /// status is recorded here. This enables "spy" rules that observe requests
    /// without modifying them.
    ///
    /// The history grows unbounded for the lifetime of the rule. This is
    /// designed for test scenarios with a bounded number of requests.
    pub fn passthrough_statuses(&self) -> Vec<StatusCode> {
        self.passthrough_statuses.lock().unwrap().clone()
    }
}

/// Builder for creating a fault injection rule.
pub struct FaultInjectionRuleBuilder {
    /// The condition under which to inject the fault.
    condition: FaultInjectionCondition,
    /// The result to inject when the condition is met.
    result: FaultInjectionResult,
    /// The absolute time at which the rule becomes active.
    start_time: Instant,
    /// The absolute time at which the rule expires.
    end_time: Option<Instant>,
    /// The total hit limit of the rule.
    hit_limit: Option<u32>,
    /// Unique identifier for the fault injection scenario.
    id: String,
}

impl FaultInjectionRuleBuilder {
    /// Creates a new FaultInjectionRuleBuilder with default values.
    ///
    /// By default the rule starts immediately and never expires.
    pub fn new(id: impl Into<String>, result: FaultInjectionResult) -> Self {
        Self {
            condition: FaultInjectionCondition::default(),
            result,
            start_time: Instant::now(),
            end_time: None,
            hit_limit: None,
            id: id.into(),
        }
    }

    /// Sets the condition for when to inject the fault.
    pub fn with_condition(mut self, condition: FaultInjectionCondition) -> Self {
        self.condition = condition;
        self
    }

    /// Sets the result to inject when the condition is met.
    pub fn with_result(mut self, result: FaultInjectionResult) -> Self {
        self.result = result;
        self
    }

    /// Sets the absolute time at which the rule becomes active.
    pub fn with_start_time(mut self, start_time: Instant) -> Self {
        self.start_time = start_time;
        self
    }

    /// Sets the absolute time at which the rule expires.
    pub fn with_end_time(mut self, end_time: Instant) -> Self {
        self.end_time = Some(end_time);
        self
    }

    /// Sets the total hit limit of the rule.
    pub fn with_hit_limit(mut self, hit_limit: u32) -> Self {
        self.hit_limit = Some(hit_limit);
        self
    }

    /// Builds the FaultInjectionRule.
    pub fn build(self) -> FaultInjectionRule {
        FaultInjectionRule {
            condition: self.condition,
            result: self.result,
            start_time: self.start_time,
            end_time: self.end_time,
            hit_limit: self.hit_limit,
            id: self.id,
            enabled: Arc::new(AtomicBool::new(true)),
            hit_count: Arc::new(AtomicU32::new(0)),
            passthrough_statuses: Mutex::new(Vec::new()),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::FaultInjectionRuleBuilder;
    use crate::fault_injection::{FaultInjectionErrorType, FaultInjectionResultBuilder};
    use std::time::Instant;

    fn create_test_error() -> crate::fault_injection::FaultInjectionResult {
        FaultInjectionResultBuilder::new()
            .with_error(FaultInjectionErrorType::Timeout)
            .build()
    }

    #[test]
    fn builder_default_values() {
        let before = Instant::now();
        let rule = FaultInjectionRuleBuilder::new("test-rule", create_test_error()).build();

        assert_eq!(rule.id, "test-rule");
        assert!(rule.start_time >= before);
        assert!(rule.start_time <= Instant::now());
        assert!(rule.end_time.is_none());
        assert!(rule.hit_limit.is_none());
        assert!(rule.condition.operation_type.is_none());
        assert!(rule.is_enabled());
        assert_eq!(rule.hit_count(), 0);
    }

    #[test]
    fn hit_count_increments() {
        let rule = FaultInjectionRuleBuilder::new("hit-test", create_test_error()).build();

        assert_eq!(rule.hit_count(), 0);
        rule.increment_hit_count();
        assert_eq!(rule.hit_count(), 1);
        rule.increment_hit_count();
        rule.increment_hit_count();
        assert_eq!(rule.hit_count(), 3);
    }

    #[test]
    fn reset_hit_count_clears_counter() {
        let rule = FaultInjectionRuleBuilder::new("reset-test", create_test_error()).build();

        rule.increment_hit_count();
        rule.increment_hit_count();
        assert_eq!(rule.hit_count(), 2);

        rule.reset_hit_count();
        assert_eq!(rule.hit_count(), 0);
    }
}