trace2e_core/transport/

loopback.rs

//! # Loopback Transport Implementation
//!
//! This module provides a loopback transport implementation for testing and
//! development scenarios where multiple middleware instances need to communicate
//! within a single process or test environment. It simulates network communication
//! by routing M2M requests directly to local middleware instances.
//!
//! ## Features
//!
//! - **In-Process Communication**: Routes calls directly to registered middleware instances
//! - **Network Simulation**: Supports configurable delays and jitter to simulate network latency
//! - **Test Orchestration**: Provides utilities for spawning multiple middleware instances
//! - **Resource Pre-enrollment**: Supports pre-populating middleware with test resources
//!
//! ## Use Cases
//!
//! - Unit and integration testing of distributed traceability scenarios
//! - Development and debugging of multi-node workflows
//! - Performance testing with controlled network conditions
//! - Simulation of distributed systems in a single process
//!
//! ## Network Simulation
//!
//! The loopback transport can simulate network characteristics by introducing
//! configurable delays and jitter to M2M calls, allowing testing of timeout
//! handling and performance under various network conditions.

use std::{
    collections::VecDeque,
    future::Future,
    pin::Pin,
    sync::Arc,
    task::Poll,
    time::{Duration, Instant},
};

use dashmap::DashMap;
use tower::Service;

use crate::{
    traceability::{
        M2mApiDefaultStack, O2mApiDefaultStack, P2mApiDefaultStack,
        api::{M2mRequest, M2mResponse},
        error::TraceabilityError,
        init_middleware_with_enrolled_resources,
    },
    transport::eval_remote_ip,
};

/// Spawns multiple loopback middleware instances with no pre-enrolled resources.
///
/// Creates a set of middleware instances that can communicate with each other
/// through the loopback transport. Each middleware is identified by an IP address
/// for routing purposes.
///
/// # Arguments
///
/// * `ips` - Vector of IP addresses to assign to the middleware instances
///
/// # Returns
///
/// A queue of (P2M service, O2M service) tuples for each middleware instance.
pub async fn spawn_loopback_middlewares(
    ips: Vec<String>,
) -> VecDeque<(P2mApiDefaultStack<M2mLoopback>, O2mApiDefaultStack<M2mLoopback>)> {
    spawn_loopback_middlewares_with_enrolled_resources(ips, 0, 0, 0).await
}

/// Spawns multiple loopback middleware instances with pre-enrolled test resources.
///
/// Creates middleware instances and pre-populates them with the specified number
/// of processes, files, and network streams for testing purposes.
///
/// # Arguments
///
/// * `ips` - Vector of IP addresses for the middleware instances
/// * `process_count` - Number of processes to enroll per middleware
/// * `per_process_file_count` - Number of files to create per process
/// * `per_process_stream_count` - Number of network streams to create per process
pub async fn spawn_loopback_middlewares_with_enrolled_resources(
    ips: Vec<String>,
    process_count: u32,
    per_process_file_count: u32,
    per_process_stream_count: u32,
) -> VecDeque<(P2mApiDefaultStack<M2mLoopback>, O2mApiDefaultStack<M2mLoopback>)> {
    spawn_loopback_middlewares_with_entropy(
        ips,
        0,
        0,
        process_count,
        per_process_file_count,
        per_process_stream_count,
    )
    .await
}

/// Spawns loopback middleware instances with network simulation and pre-enrolled resources.
///
/// Creates middleware instances with configurable network delay simulation and
/// pre-populated test resources. This is the most comprehensive setup function
/// for testing complex distributed scenarios.
///
/// # Arguments
///
/// * `ips` - Vector of IP addresses for the middleware instances
/// * `base_delay_ms` - Base network delay in milliseconds
/// * `jitter_max_ms` - Maximum additional random delay in milliseconds
/// * `process_count` - Number of processes to enroll per middleware
/// * `per_process_file_count` - Number of files to create per process
/// * `per_process_stream_count` - Number of network streams to create per process
pub async fn spawn_loopback_middlewares_with_entropy(
    ips: Vec<String>,
    base_delay_ms: u64,
    jitter_max_ms: u64,
    process_count: u32,
    per_process_file_count: u32,
    per_process_stream_count: u32,
) -> VecDeque<(P2mApiDefaultStack<M2mLoopback>, O2mApiDefaultStack<M2mLoopback>)> {
    let m2m_loopback = M2mLoopback::new(base_delay_ms, jitter_max_ms);
    let mut middlewares = VecDeque::new();
    for ip in ips {
        let (m2m, p2m, o2m) = init_middleware_with_enrolled_resources(
            ip.clone(),
            None,
            0,
            m2m_loopback.clone(),
            false, // Disable resource validation for loopback tests
            process_count,
            per_process_file_count,
            per_process_stream_count,
        );
        m2m_loopback.register_middleware(ip.clone(), m2m).await;
        middlewares.push_back((p2m, o2m));
    }
    middlewares
}

/// Loopback transport service for in-process M2M communication.
///
/// `M2mLoopback` provides a transport implementation that routes M2M requests
/// to local middleware instances within the same process. It maintains a
/// registry of middleware instances indexed by IP address and supports
/// configurable network delay simulation.
///
/// ## Network Simulation
///
/// The service can simulate network latency by introducing delays before
/// processing requests. This includes both a base delay and random jitter
/// to simulate real network conditions.
///
/// ## Thread Safety
///
/// The service is thread-safe and can be safely cloned and used across
/// multiple concurrent tasks. All internal state is protected by appropriate
/// synchronization primitives.
#[derive(Clone)]
pub struct M2mLoopback {
    /// Registry of middleware instances indexed by IP address.
    middlewares: Arc<DashMap<String, M2mApiDefaultStack>>,
    /// Base network delay in milliseconds.
    base_delay_ms: u64,
    /// Maximum additional random delay in milliseconds.
    jitter_max_ms: u64,
    /// Timestamp of the last call for delay calculation.
    last_call_time: Arc<std::sync::Mutex<Option<Instant>>>,
}

impl Default for M2mLoopback {
    fn default() -> Self {
        Self::new(0, 0)
    }
}

impl M2mLoopback {
    /// Creates a new loopback transport with the specified delay characteristics.
    ///
    /// # Arguments
    ///
    /// * `base_delay_ms` - Base delay to add to all requests in milliseconds
    /// * `jitter_max_ms` - Maximum random additional delay in milliseconds
    pub fn new(base_delay_ms: u64, jitter_max_ms: u64) -> Self {
        Self {
            middlewares: Arc::new(DashMap::new()),
            base_delay_ms,
            jitter_max_ms,
            last_call_time: Arc::new(std::sync::Mutex::new(None)),
        }
    }

    /// Registers a middleware instance with the specified IP address.
    ///
    /// This allows the loopback transport to route requests to the appropriate
    /// middleware instance based on the target IP address extracted from requests.
    ///
    /// # Arguments
    ///
    /// * `ip` - IP address identifier for the middleware
    /// * `middleware` - The middleware service instance to register
    pub async fn register_middleware(&self, ip: String, middleware: M2mApiDefaultStack) {
        self.middlewares.insert(ip, middleware);
    }

    /// Retrieves a middleware instance for the specified IP address.
    ///
    /// # Arguments
    ///
    /// * `ip` - IP address of the target middleware
    ///
    /// # Returns
    ///
    /// The middleware service instance, or an error if not found.
    ///
    /// # Errors
    ///
    /// Returns `TransportFailedToContactRemote` if no middleware is registered
    /// for the specified IP address.
    pub fn get_middleware(&self, ip: String) -> Result<M2mApiDefaultStack, TraceabilityError> {
        self.middlewares
            .get(&ip)
            .map(|c| c.to_owned())
            .ok_or(TraceabilityError::TransportFailedToContactRemote(ip))
    }

    /// Calculates the delay to apply based on the configured delay parameters.
    ///
    /// Combines the base delay with a random jitter component to simulate
    /// realistic network latency characteristics.
    ///
    /// # Returns
    ///
    /// The total delay duration to apply to the current request.
    fn calculate_delay(&self) -> Duration {
        if self.base_delay_ms == 0 && self.jitter_max_ms == 0 {
            return Duration::from_millis(0);
        }

        let base = Duration::from_millis(self.base_delay_ms);
        if self.jitter_max_ms == 0 {
            return base;
        }

        // Simple entropy based on current time
        let seed = Instant::now().elapsed().as_nanos() as u64;
        let jitter_ms = seed % (self.jitter_max_ms + 1);
        base + Duration::from_millis(jitter_ms)
    }
}

impl Service<M2mRequest> for M2mLoopback {
    type Response = M2mResponse;
    type Error = TraceabilityError;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

    fn poll_ready(&mut self, cx: &mut std::task::Context<'_>) -> Poll<Result<(), Self::Error>> {
        let now = Instant::now();
        let delay = self.calculate_delay();

        if delay.is_zero() {
            return Poll::Ready(Ok(()));
        }

        // Check if we need to wait based on last call time
        if let Ok(mut last_time) = self.last_call_time.lock() {
            match *last_time {
                Some(last) if now.duration_since(last) < delay => {
                    // Still need to wait
                    let waker = cx.waker().clone();
                    let remaining = delay - now.duration_since(last);
                    tokio::spawn(async move {
                        tokio::time::sleep(remaining).await;
                        waker.wake();
                    });
                    return Poll::Pending;
                }
                _ => {
                    // Update last call time
                    *last_time = Some(now);
                }
            }
        }

        Poll::Ready(Ok(()))
    }

    fn call(&mut self, request: M2mRequest) -> Self::Future {
        let this = self.clone();
        let request_clone = request.clone();
        match request_clone {
            M2mRequest::BroadcastDeletion(_) => {
                Box::pin(async move {
                    // Spawn all middleware calls concurrently
                    let mut handles = Vec::new();
                    for entry in this.middlewares.iter() {
                        let mut middleware = entry.value().clone();
                        let request = request_clone.clone();
                        let handle = tokio::spawn(async move { middleware.call(request).await });
                        handles.push(handle);
                    }

                    // Collect all results and check for any errors
                    for handle in handles {
                        match handle.await {
                            Ok(result) => {
                                result?; // Check for service errors, ignore the response for the moment
                            }
                            Err(_) => {
                                return Err(TraceabilityError::TransportFailedToContactRemote(
                                    "BroadcastDeletion join failed".to_string(),
                                ));
                            }
                        }
                    }

                    // All middleware calls succeeded, return Ack
                    Ok(M2mResponse::Ack)
                })
            }
            M2mRequest::CheckSourceCompliance { sources, destination } => Box::pin(async move {
                // Partition sources by node_id
                let mut partitioned: std::collections::HashMap<String, Vec<_>> =
                    std::collections::HashMap::new();
                for source in sources {
                    partitioned
                        .entry(source.node_id().clone())
                        .or_insert_with(Vec::new)
                        .push(source);
                }

                // Spawn tasks for each remote node
                let mut handles = Vec::new();
                for (node_id, sources) in partitioned {
                    let mut middleware = this.get_middleware(node_id)?;
                    let dest = destination.clone();
                    let request = M2mRequest::CheckSourceCompliance {
                        sources: sources.into_iter().collect(),
                        destination: dest,
                    };
                    handles.push(tokio::spawn(async move { middleware.call(request).await }));
                }

                // Collect all results
                for handle in handles {
                    match handle.await.map_err(|_| TraceabilityError::InternalTrace2eError)?? {
                        M2mResponse::Ack => continue,
                        _ => return Err(TraceabilityError::InternalTrace2eError),
                    }
                }

                Ok(M2mResponse::Ack)
            }),
            _ => Box::pin(async move {
                this.get_middleware(eval_remote_ip(request.clone())?)?.call(request).await
            }),
        }
    }
}