use std::sync::Arc;
use tokio::sync::Mutex;
use tokio::task::JoinHandle;

#[derive(Debug, Default)]
pub struct WorkerMetrics {
    pub total_processed: u64,
    pub failed_tasks: u64,
}

pub struct AsyncMetricsCollector {
    metrics: Arc<Mutex<WorkerMetrics>>,
}

impl AsyncMetricsCollector {
    pub fn new() -> Self {
        Self {
            metrics: Arc::new(Mutex::new(WorkerMetrics::default())),
        }
    }

    pub fn snapshot(&self) -> Arc<Mutex<WorkerMetrics>> {
        Arc::clone(&self.metrics)
    }

    /// Increments metrics safely inside a Tokio worker thread.
    pub async fn record_success(&self) {
        let metrics_ref = Arc::clone(&self.metrics);
        tokio::spawn(async move {
            tokio::time::sleep(std::time::Duration::from_millis(5)).await;
            let mut guard = metrics_ref.lock().await;
            guard.total_processed += 1;
        })
        .await
        .expect("Task panicked");
    }

    /// Demonstrates strict scope management ensuring non-Send data is dropped prior to `.await`.
    pub async fn process_with_scoped_non_send(&self, data: String) -> Result<usize, &'static str> {
        let len = {
            // std::rc::Rc is !Send and cannot cross thread boundaries or yield points
            let temp_buffer = std::rc::Rc::new(data);
            temp_buffer.len()
            // temp_buffer is explicitly dropped at block close
        };

        // .await point: compiler verifies no non-Send types remain live in state machine frame
        tokio::time::sleep(std::time::Duration::from_millis(1)).await;

        let mut guard = self.metrics.lock().await;
        guard.total_processed += 1;

        Ok(len)
    }
}

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

    #[tokio::test]
    async fn test_metrics_record_success() {
        let collector = AsyncMetricsCollector::new();
        collector.record_success().await;
        collector.record_success().await;

        let snap = collector.snapshot();
        let guard = snap.lock().await;
        assert_eq!(guard.total_processed, 2);
        assert_eq!(guard.failed_tasks, 0);
    }

    #[tokio::test]
    async fn test_scoped_non_send_across_await() {
        let collector = AsyncMetricsCollector::new();
        let res = collector.process_with_scoped_non_send("hello_tokio".into()).await;
        assert_eq!(res, Ok(11));

        let snap = collector.snapshot();
        let guard = snap.lock().await;
        assert_eq!(guard.total_processed, 1);
    }

    #[tokio::test]
    async fn test_concurrent_spawns() {
        let collector = Arc::new(AsyncMetricsCollector::new());
        let mut handles: Vec<JoinHandle<()>> = Vec::new();

        for _ in 0..5 {
            let col_clone = Arc::clone(&collector);
            handles.push(tokio::spawn(async move {
                col_clone.record_success().await;
            }));
        }

        for handle in handles {
            handle.await.unwrap();
        }

        let snap = collector.snapshot();
        let guard = snap.lock().await;
        assert_eq!(guard.total_processed, 5);
    }
}

Step-by-Step Explanation:

1. Compiler State Allocation: When the Rust compiler analyzes an async fn, it layout-analyzes all variables live across every .await point. Variables live across yield points become member fields in the generated future struct.
2. Send Trait Propagation: If any field in the generated future struct is !Send (such as Rc<T>, RefCell<T>, or raw pointers), the overall Future auto-trait implementation becomes !Send. tokio::spawn requires F: Send + 'static, causing a compile error if non-Send types persist across yield points.
3. Block Scope Isolation: Wrapping non-Send operations inside an inner scope block (let len = { let temp = Rc::new(...); temp.len() };) guarantees that temp is dropped before the execution reaches tokio::time::sleep(...).await.
4. Thread-Safe Synchronization: Replacing std::sync::Mutex with tokio::sync::Mutex ensures that holding a mutex guard across yield points suspends the async task without blocking worker OS threads.

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::oneshot;

#[derive(Debug, PartialEq, Eq, Clone)]
pub struct RpcResponse {
    pub body: String,
    pub attempts: usize,
}

#[derive(Debug, PartialEq, Eq)]
pub enum RpcError {
    NetworkFailure,
    MaxRetriesExceeded,
    Cancelled,
}

pub struct MockRpcClient {
    fail_count: Arc<AtomicUsize>,
}

impl MockRpcClient {
    pub fn new(failures_before_success: usize) -> Self {
        Self {
            fail_count: Arc::new(AtomicUsize::new(failures_before_success)),
        }
    }

    pub async fn call(&self) -> Result<String, RpcError> {
        let remaining = self.fail_count.load(Ordering::SeqCst);
        if remaining > 0 {
            self.fail_count.fetch_sub(1, Ordering::SeqCst);
            Err(RpcError::NetworkFailure)
        } else {
            Ok("SUCCESS_PAYLOAD".into())
        }
    }
}

/// Executes an RPC call with exponential retry backoff.
pub async fn rpc_with_retry(
    client: &MockRpcClient,
    max_retries: usize,
    initial_backoff: Duration,
) -> Result<RpcResponse, RpcError> {
    let mut attempt = 0;
    let mut backoff = initial_backoff;

    loop {
        attempt += 1;
        match client.call().await {
            Ok(body) => {
                return Ok(RpcResponse { body, attempts: attempt });
            }
            Err(_) => {
                if attempt > max_retries {
                    return Err(RpcError::MaxRetriesExceeded);
                }
                tokio::time::sleep(backoff).await;
                backoff *= 2;
            }
        }
    }
}

/// Executes RPC retry loop while racing against a cancellation signal safely.
pub async fn rpc_with_cancellation(
    client: &MockRpcClient,
    max_retries: usize,
    initial_backoff: Duration,
    mut shutdown_rx: oneshot::Receiver<()>,
) -> Result<RpcResponse, RpcError> {
    tokio::select! {
        res = rpc_with_retry(client, max_retries, initial_backoff) => res,
        _ = &mut shutdown_rx => Err(RpcError::Cancelled),
    }
}

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

    #[tokio::test]
    async fn test_rpc_retry_success_after_failures() {
        let client = MockRpcClient::new(2);
        let res = rpc_with_retry(&client, 3, Duration::from_millis(5)).await;
        assert!(res.is_ok());
        let resp = res.unwrap();
        assert_eq!(resp.body, "SUCCESS_PAYLOAD");
        assert_eq!(resp.attempts, 3);
    }

    #[tokio::test]
    async fn test_rpc_retry_exceeded() {
        let client = MockRpcClient::new(5);
        let res = rpc_with_retry(&client, 2, Duration::from_millis(5)).await;
        assert_eq!(res, Err(RpcError::MaxRetriesExceeded));
    }

    #[tokio::test]
    async fn test_rpc_cancellation_triggered() {
        let client = MockRpcClient::new(10);
        let (tx, rx) = oneshot::channel();

        // Trigger cancellation signal concurrently after 5ms
        tokio::spawn(async move {
            tokio::time::sleep(Duration::from_millis(5)).await;
            let _ = tx.send(());
        });

        let res = rpc_with_cancellation(&client, 5, Duration::from_millis(50), rx).await;
        assert_eq!(res, Err(RpcError::Cancelled));
    }
}

Step-by-Step Explanation:

1. Stateful Async Loops: async fn supports standard imperative control structures like loop, while, and match. At each .await point inside the loop (such as client.call().await or tokio::time::sleep(backoff).await), the compiler saves attempt and backoff inside the generated future struct frame.
2. Exponential Backoff: Multiplying duration backoff *= 2 prevents network retry storms by progressively increasing interval waits between retries.
3. Cancellation Safety in tokio::select!: tokio::select! concurrently polls multiple branch futures. When one branch completes (e.g. shutdown_rx receives a signal), all incomplete futures in other branches are dropped. In Rust, dropping an uncompleted Future immediately invokes its Drop implementation, cancelling pending tasks and freeing associated resources safely without leakages.