Line data    Source code

       1              : use std::error::Error;
       2              : use std::io;
       3              : 
       4              : use tokio::time;
       5              : 
       6              : use crate::compute;
       7              : use crate::config::RetryConfig;
       8              : 
       9              : pub(crate) trait CouldRetry {
      10              :     /// Returns true if the error could be retried
      11              :     fn could_retry(&self) -> bool;
      12              : }
      13              : 
      14              : pub(crate) trait ShouldRetryWakeCompute {
      15              :     /// Returns true if we need to invalidate the cache for this node.
      16              :     /// If false, we can continue retrying with the current node cache.
      17              :     fn should_retry_wake_compute(&self) -> bool;
      18              : }
      19              : 
      20            9 : pub(crate) fn should_retry(err: &impl CouldRetry, num_retries: u32, config: RetryConfig) -> bool {
      21            9 :     num_retries < config.max_retries && err.could_retry()
      22            9 : }
      23              : 
      24              : impl CouldRetry for io::Error {
      25            0 :     fn could_retry(&self) -> bool {
      26              :         use std::io::ErrorKind;
      27            0 :         matches!(
      28            0 :             self.kind(),
      29              :             ErrorKind::ConnectionRefused | ErrorKind::AddrNotAvailable | ErrorKind::TimedOut
      30              :         )
      31            0 :     }
      32              : }
      33              : 
      34              : impl CouldRetry for postgres_client::error::DbError {
      35            0 :     fn could_retry(&self) -> bool {
      36              :         use postgres_client::error::SqlState;
      37            0 :         matches!(
      38            0 :             self.code(),
      39              :             &SqlState::CONNECTION_FAILURE
      40              :                 | &SqlState::CONNECTION_EXCEPTION
      41              :                 | &SqlState::CONNECTION_DOES_NOT_EXIST
      42              :                 | &SqlState::SQLCLIENT_UNABLE_TO_ESTABLISH_SQLCONNECTION,
      43              :         )
      44            0 :     }
      45              : }
      46              : impl ShouldRetryWakeCompute for postgres_client::error::DbError {
      47            0 :     fn should_retry_wake_compute(&self) -> bool {
      48              :         use postgres_client::error::SqlState;
      49              :         // Here are errors that happens after the user successfully authenticated to the database.
      50              :         // TODO: there are pgbouncer errors that should be retried, but they are not listed here.
      51            0 :         !matches!(
      52            0 :             self.code(),
      53              :             &SqlState::TOO_MANY_CONNECTIONS
      54              :                 | &SqlState::OUT_OF_MEMORY
      55              :                 | &SqlState::SYNTAX_ERROR
      56              :                 | &SqlState::T_R_SERIALIZATION_FAILURE
      57              :                 | &SqlState::INVALID_CATALOG_NAME
      58              :                 | &SqlState::INVALID_SCHEMA_NAME
      59              :                 | &SqlState::INVALID_PARAMETER_VALUE
      60              :         )
      61            0 :     }
      62              : }
      63              : 
      64              : impl CouldRetry for postgres_client::Error {
      65            0 :     fn could_retry(&self) -> bool {
      66            0 :         if let Some(io_err) = self.source().and_then(|x| x.downcast_ref()) {
      67            0 :             io::Error::could_retry(io_err)
      68            0 :         } else if let Some(db_err) = self.source().and_then(|x| x.downcast_ref()) {
      69            0 :             postgres_client::error::DbError::could_retry(db_err)
      70              :         } else {
      71            0 :             false
      72              :         }
      73            0 :     }
      74              : }
      75              : impl ShouldRetryWakeCompute for postgres_client::Error {
      76            0 :     fn should_retry_wake_compute(&self) -> bool {
      77            0 :         if let Some(db_err) = self.source().and_then(|x| x.downcast_ref()) {
      78            0 :             postgres_client::error::DbError::should_retry_wake_compute(db_err)
      79              :         } else {
      80              :             // likely an IO error. Possible the compute has shutdown and the
      81              :             // cache is stale.
      82            0 :             true
      83              :         }
      84            0 :     }
      85              : }
      86              : 
      87              : impl CouldRetry for compute::ConnectionError {
      88            0 :     fn could_retry(&self) -> bool {
      89            0 :         match self {
      90            0 :             compute::ConnectionError::Postgres(err) => err.could_retry(),
      91            0 :             compute::ConnectionError::CouldNotConnect(err) => err.could_retry(),
      92            0 :             compute::ConnectionError::WakeComputeError(err) => err.could_retry(),
      93            0 :             _ => false,
      94              :         }
      95            0 :     }
      96              : }
      97              : impl ShouldRetryWakeCompute for compute::ConnectionError {
      98            0 :     fn should_retry_wake_compute(&self) -> bool {
      99            0 :         match self {
     100            0 :             compute::ConnectionError::Postgres(err) => err.should_retry_wake_compute(),
     101              :             // the cache entry was not checked for validity
     102            0 :             compute::ConnectionError::TooManyConnectionAttempts(_) => false,
     103            0 :             _ => true,
     104              :         }
     105            0 :     }
     106              : }
     107              : 
     108           10 : pub(crate) fn retry_after(num_retries: u32, config: RetryConfig) -> time::Duration {
     109           10 :     config
     110           10 :         .base_delay
     111           10 :         .mul_f64(config.backoff_factor.powi((num_retries as i32) - 1))
     112           10 : }