telemetry [12]: report gRPC client transport failures

Cargo.toml

@@ -40,6 +40,7 @@ envy = "0.4.2"
 futures = "0.3.32"
 glob = "0.3.1"
 gpt = "4.1.0"
+h2 = "0.4.20"
 hex = "0.4.0"
 hostname = "0.4.0"
 humantime = "2.3.0"

crates/trident/Cargo.toml

@@ -84,6 +84,7 @@ pytest_gen = { path = "../pytest_gen", optional = true }
 pytest = { path = "../pytest", optional = true }
 
 [dev-dependencies]
+h2 = { workspace = true }
 indoc = { workspace = true }
 maplit = { workspace = true }
 mockito = { workspace = true }

crates/trident/src/grpc_client/mod.rs

@@ -1,20 +1,23 @@
-use std::process::ExitCode;
+use std::{cell::Cell, process::ExitCode};
 
 use anyhow::{bail, Context, Error};
 use log::error;
 use tokio::fs;
 use tokio::runtime::Builder;
+use tonic::Code;
+
+use trident_api::error::{InternalError, TridentError};
 
 use crate::{
-    cli::{ClientArgs, ClientCommands, TridentExitCodes},
-    ExitKind, TRIDENT_VERSION,
+    cli::{self, ClientArgs, ClientCommands, TridentExitCodes},
+    logging::operation_context,
+    run_command_if, ExitKind, OperationSource, TRIDENT_VERSION,
 };
 
-use crate::cli;
-
 mod error;
 mod tridentclient;
 
+use error::TridentClientError;
 use tridentclient::{RebootHandling, TridentClient};
 
 pub fn client_main(args: &ClientArgs) -> ExitCode {
@@ -24,7 +27,53 @@ pub fn client_main(args: &ClientArgs) -> ExitCode {
         return TridentExitCodes::SetupFailed.into();
     };
 
-    match runtime.block_on(run_client(args)) {
+    // `setup_tracing()` (see `main.rs`) treats grpc-client the same as any
+    // other command -- it's a first-class telemetry participant, not just
+    // a transport-only escape hatch, so it fires `command_start` like
+    // every other command. `command_error` is more selective (see
+    // `is_transport_failure` below). `run_client`'s errors are plain
+    // `anyhow::Error` (not `TridentError`), so they're wrapped in a
+    // generic `InternalError::Internal` here purely to get them into
+    // `run_command_if`'s `Result<_, TridentError>` shape -- the original
+    // anyhow context chain is preserved as the error's source and still
+    // printed in full below.
+    // `CommandKind` gives the same stage/finalize-granular naming
+    // (`install_stage`, `update_finalize`, `rollback_stage`, etc.) the CLI
+    // and daemon use for servicing telemetry -- otherwise every
+    // grpc-client update/install/rollback would collapse to the generic
+    // `client_update`/`client_install`/`client_rollback` regardless of
+    // which operations were actually requested. `Rollback { check: true,
+    // .. }` is deliberately excluded from stage/finalize naming (mirrors
+    // `main.rs`'s own `Commands::Rollback { check: true, .. }`
+    // special-casing) and from `is_servicing` below -- a dry-run check
+    // never stages or finalizes anything.
+    let kind = args.command.kind();
+
+    // `run_client` (the actual RPC) runs *inside* this closure, not before
+    // it, so `command_start` (fired by `run_command_if` the moment this
+    // closure is entered) actually brackets the RPC instead of always
+    // following it -- otherwise every client-side event the RPC itself
+    // fires, and the timestamp of `command_start` itself, would be
+    // reported after the call had already finished. `is_transport_failure`
+    // is computed from the raw `anyhow::Error` chain here, inside the
+    // closure, and stashed via `transport_failure` for `run_command_if`'s
+    // `should_report` predicate below, which only ever sees the already-
+    // wrapped `TridentError` and has no way to inspect that chain itself.
+    let transport_failure = Cell::new(false);
+    let result = run_command_if(
+        &kind,
+        OperationSource::GrpcClient,
+        || {
+            let client_result = runtime.block_on(run_client(args));
+            transport_failure.set(is_transport_failure(&client_result));
+            client_result.map_err(|e| {
+                TridentError::with_source(InternalError::Internal("grpc-client command failed"), e)
+            })
+        },
+        |_error| transport_failure.get() && kind.is_servicing(),
+    );
+
+    match result {
         Err(e) => {
             error!("Client failed: {:?}", e);
             return TridentExitCodes::Failed.into();
@@ -41,6 +90,44 @@ pub fn client_main(args: &ClientArgs) -> ExitCode {
     TridentExitCodes::Success.into()
 }
 
+/// The daemon fires its own, correctly-classified `command_error` for any
+/// request it actually received and acted on -- including one it rejected
+/// outright (see e.g. `services::reject_invalid_argument`). A genuine
+/// transport-level failure -- the daemon never received or finished
+/// answering this request at all -- has no other reporter, and is what
+/// this checks for:
+/// - `ConnectionError`: the initial connection attempt itself failed
+///   (socket not found, connection refused).
+/// - `RequestError`/`ResponseError` whose wrapped `Status` retains a local
+///   error source (including HTTP/2 resets with non-`Unavailable` codes),
+///   or is `Code::Unavailable`, *except* for the daemon's own admission-control
+///   rejections -- connection-lock or servicing-lock contention (see
+///   `try_acquire_read_lock`/`try_acquire_write_lock`/`servicing_request`/
+///   `reading_request` in `server::tridentserver`), which deliberately
+///   also return `Code::Unavailable` since a busy daemon is retryable the
+///   same way a broken connection is. Those are recognized by their fixed
+///   message text (`CONNECTION_LOCK_BUSY_MESSAGE`/`SERVICING_LOCK_BUSY_MESSAGE`)
+///   and excluded here, since the daemon DID receive and answer this
+///   request, unlike tonic's own `Code::Unavailable` for a connection that
+///   broke mid-call (e.g. the daemon process died or the socket was closed
+///   while a request/response was in flight).
+fn is_transport_failure(client_result: &Result<ExitKind, Error>) -> bool {
+    client_result.as_ref().err().is_some_and(|e| {
+        e.chain()
+            .any(|cause| match cause.downcast_ref::<TridentClientError>() {
+                Some(TridentClientError::ConnectionError(..)) => true,
+                Some(TridentClientError::RequestError(_, status))
+                | Some(TridentClientError::ResponseError(_, status)) => {
+                    (std::error::Error::source(status).is_some()
+                        || status.code() == Code::Unavailable)
+                        && status.message() != operation_context::CONNECTION_LOCK_BUSY_MESSAGE
+                        && status.message() != operation_context::SERVICING_LOCK_BUSY_MESSAGE
+                }
+                _ => false,
+            })
+    })
+}
+
 async fn run_client(args: &ClientArgs) -> Result<ExitKind, Error> {
     let mut client = TridentClient::connect(&args.server)
         .await
@@ -138,3 +225,62 @@ async fn run_client(args: &ClientArgs) -> Result<ExitKind, Error> {
 
     Ok(ExitKind::Done)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_non_unavailable_transport_causes_are_reported() {
+        for code in [
+            Code::Cancelled,
+            Code::Internal,
+            Code::ResourceExhausted,
+            Code::Unknown,
+        ] {
+            let mut status = tonic::Status::new(code, "transport failure");
+            status.set_source(std::sync::Arc::new(std::io::Error::new(
+                std::io::ErrorKind::ConnectionReset,
+                "connection reset",
+            )));
+            for error in [
+                TridentClientError::RequestError("request".to_string(), status.clone()),
+                TridentClientError::ResponseError("stream".to_string(), status.clone()),
+            ] {
+                assert!(is_transport_failure(&Err(
+                    Error::new(error).context("RPC failed")
+                )));
+            }
+        }
+    }
+
+    #[test]
+    fn test_daemon_statuses_and_servicing_errors_are_not_transport_failures() {
+        for status in [
+            tonic::Status::cancelled("daemon cancelled request"),
+            tonic::Status::internal("daemon error"),
+            tonic::Status::invalid_argument("invalid request"),
+            tonic::Status::unavailable(operation_context::CONNECTION_LOCK_BUSY_MESSAGE),
+            tonic::Status::unavailable(operation_context::SERVICING_LOCK_BUSY_MESSAGE),
+        ] {
+            for error in [
+                TridentClientError::RequestError("request".to_string(), status.clone()),
+                TridentClientError::ResponseError("stream".to_string(), status.clone()),
+            ] {
+                assert!(!is_transport_failure(&Err(error.into())));
+            }
+        }
+        assert!(!is_transport_failure(&Err(
+            TridentClientError::ServicingError("servicing failed".to_string()).into()
+        )));
+        assert!(!is_transport_failure(&Ok(ExitKind::Done)));
+    }
+
+    #[test]
+    fn test_source_less_unavailable_transport_failure_is_still_reported() {
+        let status = tonic::Status::unavailable("connection closed");
+        assert!(is_transport_failure(&Err(
+            TridentClientError::ResponseError("stream".to_string(), status).into()
+        )));
+    }
+}

crates/trident/src/grpc_client/tridentclient.rs

@@ -285,7 +285,10 @@ async fn handle_servicing_response_stream(
         }
     }
 
-    Ok(ExitKind::Done)
+    Err(TridentClientError::ResponseError(
+        "servicing stream".to_string(),
+        tonic::Status::unavailable("Servicing stream ended before the final response"),
+    ))
 }
 
 async fn handle_servicing_response(
@@ -364,3 +367,119 @@ async fn handle_servicing_response(
 
     Ok(ControlFlow::Continue(()))
 }
+
+#[cfg(test)]
+mod tests {
+    use std::time::Duration;
+
+    use hyper::{body::Bytes, Response};
+    use tokio::{net::TcpListener, sync::oneshot, task::JoinSet, time::timeout};
+
+    use super::*;
+
+    async fn assert_mid_stream_reset_is_transport_failure(
+        reason: h2::Reason,
+        expected_code: tonic::Code,
+    ) {
+        let mut server = JoinSet::new();
+        timeout(Duration::from_secs(10), async {
+            let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+            let address = listener.local_addr().unwrap();
+            let (reset, reset_requested) = oneshot::channel();
+            server.spawn(async move {
+                let (socket, _) = listener.accept().await.unwrap();
+                let mut connection = h2::server::handshake(socket).await.unwrap();
+                let (_request, mut respond) = connection.accept().await.unwrap().unwrap();
+                let response = Response::builder()
+                    .status(200)
+                    .header("content-type", "application/grpc")
+                    .body(())
+                    .unwrap();
+                let mut stream = respond.send_response(response, false).unwrap();
+                // An empty protobuf message establishes a readable gRPC stream.
+                stream
+                    .send_data(Bytes::from_static(&[0, 0, 0, 0, 0]), false)
+                    .unwrap();
+                tokio::select! {
+                    requested = reset_requested => {
+                        requested.unwrap();
+                        stream.send_reset(reason);
+                    }
+                    _ = connection.accept() => panic!("connection closed before the reset"),
+                }
+                while let Some(request) = connection.accept().await {
+                    request.unwrap();
+                }
+            });
+
+            let channel = Endpoint::from_shared(format!("http://{address}"))
+                .unwrap()
+                .connect()
+                .await
+                .unwrap();
+            let mut client = StreamingServiceClient::new(channel);
+            let mut stream = client
+                .stream_disk(StreamDiskRequest::default())
+                .await
+                .unwrap()
+                .into_inner();
+            assert!(stream.message().await.unwrap().is_some());
+            reset.send(()).unwrap();
+
+            let error = handle_servicing_response_stream(stream).await.unwrap_err();
+            let TridentClientError::ResponseError(_, ref status) = error else {
+                panic!("expected a mid-stream response error: {error}");
+            };
+            assert_eq!(status.code(), expected_code);
+            assert_eq!(
+                std::error::Error::source(status).is_some(),
+                reason != h2::Reason::CANCEL,
+            );
+            assert!(super::super::is_transport_failure(&Err(error.into())));
+            drop(client);
+            server.join_next().await.unwrap().unwrap();
+        })
+        .await
+        .expect("the stream reset test must finish within its deadline");
+    }
+
+    #[tokio::test]
+    async fn test_mid_stream_http2_cancel_is_a_transport_failure() {
+        // Hyper 1.8.1 maps CANCEL to EOF instead of retaining a status source.
+        assert_mid_stream_reset_is_transport_failure(h2::Reason::CANCEL, tonic::Code::Unavailable)
+            .await;
+    }
+
+    #[tokio::test]
+    async fn test_mid_stream_http2_internal_reset_is_a_transport_failure() {
+        assert_mid_stream_reset_is_transport_failure(
+            h2::Reason::INTERNAL_ERROR,
+            tonic::Code::Internal,
+        )
+        .await;
+    }
+
+    #[tokio::test]
+    async fn test_completed_servicing_response_preserves_exit_kind() {
+        for reboot in [
+            RebootStatus::RebootNotRequired,
+            RebootStatus::RebootStarted,
+            RebootStatus::RebootRequired,
+        ] {
+            let response = ServicingResponse {
+                response: Some(ResponseBody::Completed(trident_proto::v1::Completed {
+                    status: StatusCode::Success.into(),
+                    reboot_status: reboot.into(),
+                    ..Default::default()
+                })),
+                ..Default::default()
+            };
+            let result = handle_servicing_response(response).await.unwrap();
+            if reboot == RebootStatus::RebootRequired {
+                assert!(matches!(result, ControlFlow::Break(ExitKind::NeedsReboot)));
+            } else {
+                assert!(matches!(result, ControlFlow::Break(ExitKind::Done)));
+            }
+        }
+    }
+}

crates/trident/src/logging/operation_context.rs

@@ -6,24 +6,27 @@
 //! `Trident::*`, etc.) needing to pass them explicitly.
 //!
 //! A thread-local (rather than e.g. a `tracing` span) is enough here
-//! because both places that actually install this context run the entire
+//! because all three places that install this context run the entire
 //! command synchronously on a single, dedicated thread for the command's
 //! whole duration:
 //! - CLI: `run_trident`'s command dispatch (synchronous, main thread),
 //!   tagged [`OperationSource::Cli`].
 //! - gRPC/daemon: `servicing_request`'s closure runs inside
 //!   `tokio::task::spawn_blocking`, which gives it its own OS thread for
 //!   as long as the closure runs, tagged [`OperationSource::Daemon`].
+//! - `grpc-client`: `client_main`'s dispatch runs the whole RPC via
+//!   `runtime.block_on(...)` on the calling (main) thread -- nothing in
+//!   the client path spawns a separate task -- tagged
+//!   [`OperationSource::GrpcClient`], wrapped in `run_command_if` just
+//!   like the other two entry points wrap their dispatch in
+//!   `run_command`.
 //!
-//! [`OperationSource::GrpcClient`] is defined for a third entry point --
-//! `grpc_client`'s command dispatch (the CLI acting as a client of a
-//! running daemon) -- but nothing currently wraps that dispatch with
-//! `run_with_operation`, so client-side telemetry is not enriched with
-//! `operation_id`/`command`/`source` today. Not considered worth closing:
-//! the daemon side of that same request already produces the interesting
-//! `source = daemon` telemetry, so a client-side `source = grpc-client`
-//! event would mostly just duplicate it from a less informative vantage
-//! point.
+//! All three entry points attach `operation_id`/`command`/`source`.
+//! `GrpcClient` does not attach persisted `installation_id`/`servicing_id`:
+//! `client_main` is never handed the `TraceStream` built by `main()`.
+//! Those fields remain absent, rather than substituting an unconfirmed
+//! installation identity. The daemon side of the request attaches the
+//! datastore-backed IDs to its own `source = daemon` telemetry.
 //!
 //! `operation_id` is a fresh, random ID generated once per command
 //! invocation (distinct from the persisted `installation_id`/