codex/codex-rs/code-mode-host/tests/grpc/process_metrics_tests.rs

//! Process observations are independent of session state and HTTP/2 keepalive.

use std::sync::Mutex;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;

use codex_code_mode::CodeModeProcessMetrics;
use codex_code_mode::CodeModeProcessMetricsSample;
use futures::StreamExt;
use pretty_assertions::assert_eq;
use tokio_stream::wrappers::TcpListenerStream;
use tonic::transport::Endpoint;

use super::*;

#[tokio::test]
async fn metrics_failures_and_deadlines_preserve_sessions_and_keepalive() -> Result<()> {
    let snapshot = CodeModeProcessMetrics {
        process_id: 42,
        process_start_time_seconds: 1_700_000_000.0,
        cpu_seconds_total: 1.5,
        resident_memory_bytes: 4096,
    };
    let expected = snapshot.clone();
    let calls = Arc::new(AtomicUsize::new(0));
    let collector_calls = calls.clone();
    let (release, blocked) = std::sync::mpsc::channel::<()>();
    let blocked = Mutex::new(blocked);
    let host = codex_code_mode_host::GrpcCodeModeHost::with_process_metrics(Arc::new(move || {
        match collector_calls.fetch_add(1, Ordering::SeqCst) {
            0 => Ok(snapshot.clone()),
            1 => Err("metrics unavailable".to_string()),
            2 => panic!("collector panic"),
            _ => {
                blocked
                    .lock()
                    .unwrap()
                    .recv()
                    .map_err(|error| error.to_string())?;
                Ok(snapshot.clone())
            }
        }
    }));
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
    let address = listener.local_addr()?;
    let connections = Arc::new(AtomicUsize::new(0));
    let accepted = connections.clone();
    let incoming = TcpListenerStream::new(listener).inspect(move |connection| {
        if connection.is_ok() {
            accepted.fetch_add(1, Ordering::SeqCst);
        }
    });
    let server = tokio_util::task::AbortOnDropHandle::new(tokio::spawn(async move {
        tonic::transport::Server::builder()
            .add_service(grpc::code_mode_host_server::CodeModeHostServer::new(host))
            .serve_with_incoming(incoming)
            .await
    }));
    // Multiple keepalive intervals elapse while the collector is blocked, even
    // on this single-threaded runtime. A stalled HTTP/2 driver would lose state.
    let channel = Endpoint::from_shared(format!("http://{address}"))?
        .http2_keep_alive_interval(Duration::from_millis(/*millis*/ 50))
        .keep_alive_timeout(Duration::from_millis(/*millis*/ 100))
        .keep_alive_while_idle(/*enabled*/ true)
        .connect()
        .await?;
    let provider = GrpcCodeModeSessionProvider::with_channel(channel);
    let session = provider
        .create_session()
        .await
        .map_err(anyhow::Error::msg)?;
    execute(
        &session,
        request("store('survives', 42);"),
        Arc::new(NoopCodeModeSessionDelegate),
    )
    .await?;
    assert_eq!(
        provider
            .sample_process_metrics()
            .await
            .map_err(anyhow::Error::msg)?,
        CodeModeProcessMetricsSample::Snapshot(expected)
    );
    for cell in 2..=4 {
        assert!(provider.sample_process_metrics().await.is_err());
        let response = execute(
            &session,
            request("text(load('survives'));"),
            Arc::new(NoopCodeModeSessionDelegate),
        )
        .await?;
        assert_eq!(
            response,
            text_response(&cell.to_string(), "42", response.code_mode_host_duration())
        );
    }
    // A canceled blocked RPC retains its collector quota: no unbounded worker
    // threads are added, and neither the sample nor overload opens another channel.
    assert!(provider.sample_process_metrics().await.is_err());
    assert_eq!(calls.load(Ordering::SeqCst), 4);
    assert_eq!(connections.load(Ordering::SeqCst), 1);
    release.send(())?;
    session.shutdown().await.map_err(anyhow::Error::msg)?;
    drop(server);
    Ok(())
}

/// A running server's collection precondition failure is not local host absence.
#[tokio::test]
async fn remote_metrics_precondition_is_a_failure() -> Result<()> {
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
    let address = listener.local_addr()?;
    let server = tokio_util::task::AbortOnDropHandle::new(tokio::spawn(async move {
        tonic::transport::Server::builder()
            .add_service(
                grpc::code_mode_host_server::CodeModeHostServer::with_interceptor(
                    codex_code_mode_host::GrpcCodeModeHost::new(),
                    |_request| {
                        Err(tonic::Status::failed_precondition(
                            "collector is misconfigured",
                        ))
                    },
                ),
            )
            .serve_with_incoming(TcpListenerStream::new(listener))
            .await
    }));
    let channel = Endpoint::from_shared(format!("http://{address}"))?
        .connect()
        .await?;
    let provider = GrpcCodeModeSessionProvider::with_channel(channel);
    let error = provider.sample_process_metrics().await.unwrap_err();
    assert!(error.contains("collector is misconfigured"), "{error}");
    drop(server);
    Ok(())
}

/// A permanently stuck collector is bounded and cannot hang normal runtime drop.
#[test]
fn stuck_collector_does_not_block_runtime_shutdown() -> Result<()> {
    const CHILD_ENV: &str = "CODE_MODE_METRICS_SHUTDOWN_TEST_CHILD";
    if let Some(shutdown_marker) = std::env::var_os(CHILD_ENV) {
        let runtime = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()?;
        runtime.block_on(async {
            let calls = Arc::new(AtomicUsize::new(0));
            let collector_calls = calls.clone();
            let host =
                codex_code_mode_host::GrpcCodeModeHost::with_process_metrics(Arc::new(move || {
                    collector_calls.fetch_add(1, Ordering::SeqCst);
                    loop {
                        std::thread::park();
                    }
                }));
            let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
            let address = listener.local_addr()?;
            let server = tokio_util::task::AbortOnDropHandle::new(tokio::spawn(async move {
                tonic::transport::Server::builder()
                    .add_service(grpc::code_mode_host_server::CodeModeHostServer::new(host))
                    .serve_with_incoming(TcpListenerStream::new(listener))
                    .await
            }));
            let channel = Endpoint::from_shared(format!("http://{address}"))?
                .connect()
                .await?;
            let provider = GrpcCodeModeSessionProvider::with_channel(channel);
            assert!(provider.sample_process_metrics().await.is_err());
            let busy = provider.sample_process_metrics().await.unwrap_err();
            assert!(busy.contains("collector is busy"), "{busy}");
            assert_eq!(calls.load(Ordering::SeqCst), 1);
            drop(server);
            Ok::<_, anyhow::Error>(())
        })?;
        drop(runtime);
        std::fs::write(shutdown_marker, "runtime dropped")?;
        return Ok(());
    }

    let directory = tempfile::tempdir()?;
    let shutdown_marker = directory.path().join("shutdown-completed");
    let mut child = std::process::Command::new(std::env::current_exe()?)
        .args([
            "--exact",
            "process_metrics_tests::stuck_collector_does_not_block_runtime_shutdown",
            "--nocapture",
        ])
        .env(CHILD_ENV, &shutdown_marker)
        .spawn()?;
    let deadline = std::time::Instant::now() + Duration::from_secs(/*secs*/ 10);
    loop {
        if let Some(status) = child.try_wait()? {
            assert!(status.success(), "metrics subprocess failed: {status}");
            assert_eq!(std::fs::read_to_string(shutdown_marker)?, "runtime dropped");
            return Ok(());
        }
        if std::time::Instant::now() >= deadline {
            child.kill()?;
            child.wait()?;
            anyhow::bail!("a stuck metrics collector blocked runtime shutdown");
        }
        std::thread::sleep(Duration::from_millis(/*millis*/ 10));
    }
}