feat(proxy): Return prefixed errors to testodrome (#11561)

Testodrome measures uptime based on the failed requests and errors. In
case of testodrome request we send back error based on the service. This
will help us distinguish error types in testodrome and rely on the
uptime SLI.

proxy/src/stream.rs

@@ -6,11 +6,13 @@ use bytes::BytesMut;
 use pq_proto::framed::{ConnectionError, Framed};
 use pq_proto::{BeMessage, FeMessage, FeStartupPacket, ProtocolError};
 use rustls::ServerConfig;
+use serde::{Deserialize, Serialize};
 use thiserror::Error;
 use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
 use tokio_rustls::server::TlsStream;
 use tracing::debug;
 
+use crate::control_plane::messages::ColdStartInfo;
 use crate::error::{ErrorKind, ReportableError, UserFacingError};
 use crate::metrics::Metrics;
 use crate::tls::TlsServerEndPoint;
@@ -100,6 +102,44 @@ impl ReportableError for ReportedError {
     }
 }
 
+#[derive(Serialize, Deserialize, Debug)]
+enum ErrorTag {
+    #[serde(rename = "proxy")]
+    Proxy,
+    #[serde(rename = "compute")]
+    Compute,
+    #[serde(rename = "client")]
+    Client,
+    #[serde(rename = "controlplane")]
+    ControlPlane,
+    #[serde(rename = "other")]
+    Other,
+}
+
+impl From<ErrorKind> for ErrorTag {
+    fn from(error_kind: ErrorKind) -> Self {
+        match error_kind {
+            ErrorKind::User => Self::Client,
+            ErrorKind::ClientDisconnect => Self::Client,
+            ErrorKind::RateLimit => Self::Proxy,
+            ErrorKind::ServiceRateLimit => Self::Proxy, // considering rate limit as proxy error for SLI
+            ErrorKind::Quota => Self::Proxy,
+            ErrorKind::Service => Self::Proxy,
+            ErrorKind::ControlPlane => Self::ControlPlane,
+            ErrorKind::Postgres => Self::Other,
+            ErrorKind::Compute => Self::Compute,
+        }
+    }
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+#[serde(rename_all = "snake_case")]
+struct ProbeErrorData {
+    tag: ErrorTag,
+    msg: String,
+    cold_start_info: Option<ColdStartInfo>,
+}
+
 impl<S: AsyncWrite + Unpin> PqStream<S> {
     /// Write the message into an internal buffer, but don't flush the underlying stream.
     pub(crate) fn write_message_noflush(
@@ -125,26 +165,54 @@ impl<S: AsyncWrite + Unpin> PqStream<S> {
         Ok(self)
     }
 
-    /// Write the error message using [`Self::write_message`], then re-throw it.
+    /// Writes message with the given error kind to the stream.
+    /// Used only for probe queries
+    async fn write_format_message(
+        &mut self,
+        msg: &str,
+        error_kind: ErrorKind,
+        ctx: Option<&crate::context::RequestContext>,
+    ) -> String {
+        let formatted_msg = match ctx {
+            Some(ctx) if ctx.get_testodrome_id().is_some() => {
+                serde_json::to_string(&ProbeErrorData {
+                    tag: ErrorTag::from(error_kind),
+                    msg: msg.to_string(),
+                    cold_start_info: Some(ctx.cold_start_info()),
+                })
+                .unwrap_or_default()
+            }
+            _ => msg.to_string(),
+        };
+
+        // already error case, ignore client IO error
+        self.write_message(&BeMessage::ErrorResponse(&formatted_msg, None))
+            .await
+            .inspect_err(|e| debug!("write_message failed: {e}"))
+            .ok();
+
+        formatted_msg
+    }
+
+    /// Write the error message using [`Self::write_format_message`], then re-throw it.
     /// Allowing string literals is safe under the assumption they might not contain any runtime info.
     /// This method exists due to `&str` not implementing `Into<anyhow::Error>`.
+    /// If `ctx` is provided and has testodrome_id set, error messages will be prefixed according to error kind.
     pub async fn throw_error_str<T>(
         &mut self,
         msg: &'static str,
         error_kind: ErrorKind,
+        ctx: Option<&crate::context::RequestContext>,
     ) -> Result<T, ReportedError> {
-        // TODO: only log this for actually interesting errors
-        tracing::info!(
-            kind = error_kind.to_metric_label(),
-            msg,
-            "forwarding error to user"
-        );
+        self.write_format_message(msg, error_kind, ctx).await;
 
-        // already error case, ignore client IO error
-        self.write_message(&BeMessage::ErrorResponse(msg, None))
-            .await
-            .inspect_err(|e| debug!("write_message failed: {e}"))
-            .ok();
+        if error_kind != ErrorKind::RateLimit && error_kind != ErrorKind::User {
+            tracing::info!(
+                kind = error_kind.to_metric_label(),
+                msg,
+                "forwarding error to user"
+            );
+        }
 
         Err(ReportedError {
             source: anyhow::anyhow!(msg),
@@ -152,26 +220,28 @@ impl<S: AsyncWrite + Unpin> PqStream<S> {
         })
     }
 
-    /// Write the error message using [`Self::write_message`], then re-throw it.
+    /// Write the error message using [`Self::write_format_message`], then re-throw it.
     /// Trait [`UserFacingError`] acts as an allowlist for error types.
-    pub(crate) async fn throw_error<T, E>(&mut self, error: E) -> Result<T, ReportedError>
+    /// If `ctx` is provided and has testodrome_id set, error messages will be prefixed according to error kind.
+    pub(crate) async fn throw_error<T, E>(
+        &mut self,
+        error: E,
+        ctx: Option<&crate::context::RequestContext>,
+    ) -> Result<T, ReportedError>
     where
         E: UserFacingError + Into<anyhow::Error>,
     {
         let error_kind = error.get_error_kind();
         let msg = error.to_string_client();
-        tracing::info!(
-            kind=error_kind.to_metric_label(),
-            error=%error,
-            msg,
-            "forwarding error to user"
-        );
-
-        // already error case, ignore client IO error
-        self.write_message(&BeMessage::ErrorResponse(&msg, None))
-            .await
-            .inspect_err(|e| debug!("write_message failed: {e}"))
-            .ok();
+        self.write_format_message(&msg, error_kind, ctx).await;
+        if error_kind != ErrorKind::RateLimit && error_kind != ErrorKind::User {
+            tracing::info!(
+                kind=error_kind.to_metric_label(),
+                error=%error,
+                msg,
+                "forwarding error to user",
+            );
+        }
 
         Err(ReportedError {
             source: anyhow::anyhow!(error),