LCOV - code coverage report
Current view: top level - proxy/src/proxy - wake_compute.rs (source / functions) Coverage Total Hit
Test: f08493f498dc56383410c7d22f77751c183f1590.info Lines: 76.5 % 51 39
Test Date: 2024-02-22 21:37:45 Functions: 28.6 % 14 4

            Line data    Source code
       1              : use crate::console::{errors::WakeComputeError, provider::CachedNodeInfo};
       2              : use crate::context::RequestMonitoring;
       3              : use crate::metrics::{bool_to_str, NUM_WAKEUP_FAILURES};
       4              : use crate::proxy::retry::retry_after;
       5              : use hyper::StatusCode;
       6              : use std::ops::ControlFlow;
       7              : use tracing::{error, warn};
       8              : 
       9              : use super::connect_compute::ComputeConnectBackend;
      10              : use super::retry::ShouldRetry;
      11              : 
      12           22 : pub async fn wake_compute<B: ComputeConnectBackend>(
      13           22 :     num_retries: &mut u32,
      14           22 :     ctx: &mut RequestMonitoring,
      15           22 :     api: &B,
      16           22 : ) -> Result<CachedNodeInfo, WakeComputeError> {
      17              :     loop {
      18           26 :         let wake_res = api.wake_compute(ctx).await;
      19           26 :         match handle_try_wake(wake_res, *num_retries) {
      20            2 :             Err(e) => {
      21            2 :                 error!(error = ?e, num_retries, retriable = false, "couldn't wake compute node");
      22            2 :                 report_error(&e, false);
      23            2 :                 return Err(e);
      24              :             }
      25            4 :             Ok(ControlFlow::Continue(e)) => {
      26            0 :                 warn!(error = ?e, num_retries, retriable = true, "couldn't wake compute node");
      27            4 :                 report_error(&e, true);
      28              :             }
      29           20 :             Ok(ControlFlow::Break(n)) => return Ok(n),
      30              :         }
      31              : 
      32            4 :         let wait_duration = retry_after(*num_retries);
      33            4 :         *num_retries += 1;
      34            4 :         tokio::time::sleep(wait_duration).await;
      35              :     }
      36           22 : }
      37              : 
      38              : /// Attempts to wake up the compute node.
      39              : /// * Returns Ok(Continue(e)) if there was an error waking but retries are acceptable
      40              : /// * Returns Ok(Break(node)) if the wakeup succeeded
      41              : /// * Returns Err(e) if there was an error
      42           26 : pub fn handle_try_wake(
      43           26 :     result: Result<CachedNodeInfo, WakeComputeError>,
      44           26 :     num_retries: u32,
      45           26 : ) -> Result<ControlFlow<CachedNodeInfo, WakeComputeError>, WakeComputeError> {
      46           26 :     match result {
      47            6 :         Err(err) => match &err {
      48            6 :             WakeComputeError::ApiError(api) if api.should_retry(num_retries) => {
      49            4 :                 Ok(ControlFlow::Continue(err))
      50              :             }
      51            2 :             _ => Err(err),
      52              :         },
      53              :         // Ready to try again.
      54           20 :         Ok(new) => Ok(ControlFlow::Break(new)),
      55              :     }
      56           26 : }
      57              : 
      58            6 : fn report_error(e: &WakeComputeError, retry: bool) {
      59            6 :     use crate::console::errors::ApiError;
      60            6 :     let retry = bool_to_str(retry);
      61            6 :     let kind = match e {
      62            0 :         WakeComputeError::BadComputeAddress(_) => "bad_compute_address",
      63            0 :         WakeComputeError::ApiError(ApiError::Transport(_)) => "api_transport_error",
      64              :         WakeComputeError::ApiError(ApiError::Console {
      65              :             status: StatusCode::LOCKED,
      66            0 :             ref text,
      67            0 :         }) if text.contains("written data quota exceeded")
      68            0 :             || text.contains("the limit for current plan reached") =>
      69            0 :         {
      70            0 :             "quota_exceeded"
      71              :         }
      72              :         WakeComputeError::ApiError(ApiError::Console {
      73              :             status: StatusCode::LOCKED,
      74              :             ..
      75            0 :         }) => "api_console_locked",
      76              :         WakeComputeError::ApiError(ApiError::Console {
      77              :             status: StatusCode::BAD_REQUEST,
      78              :             ..
      79            4 :         }) => "api_console_bad_request",
      80            2 :         WakeComputeError::ApiError(ApiError::Console { status, .. })
      81            2 :             if status.is_server_error() =>
      82            0 :         {
      83            0 :             "api_console_other_server_error"
      84              :         }
      85            2 :         WakeComputeError::ApiError(ApiError::Console { .. }) => "api_console_other_error",
      86            0 :         WakeComputeError::TimeoutError => "timeout_error",
      87              :     };
      88            6 :     NUM_WAKEUP_FAILURES.with_label_values(&[retry, kind]).inc();
      89            6 : }
        

Generated by: LCOV version 2.1-beta