rs-zero 0.2.6

Rust-first microservice framework inspired by go-zero engineering practices
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

use std::{future::Future, time::Duration};

use tonic::Status;

use crate::rpc::{RpcRetryConfig, deadline::DeadlineBudget};

/// Returns whether the status should be retried by the given config.
pub fn is_retryable_status(status: &Status, config: &RpcRetryConfig) -> bool {
    config.enabled && config.retryable_codes.contains(&status.code())
}

/// Returns the exponential backoff for a zero-based retry attempt.
pub fn next_backoff(config: &RpcRetryConfig, retry_attempt: u32) -> Duration {
    let shift = retry_attempt.min(16);
    let multiplier = 2_u32.saturating_pow(shift);
    let backoff = config.initial_backoff.saturating_mul(multiplier);
    backoff.min(config.max_backoff)
}

/// Runs an async RPC operation with bounded retries.
pub async fn run_with_retry<T, F, Fut>(
    config: &RpcRetryConfig,
    mut operation: F,
) -> Result<T, Status>
where
    F: FnMut() -> Fut,
    Fut: Future<Output = Result<T, Status>>,
{
    let max_attempts = if config.enabled {
        config.max_attempts.max(1)
    } else {
        1
    };
    let mut attempt = 0;
    loop {
        let result = operation().await;
        match result {
            Ok(value) => return Ok(value),
            Err(status) if attempt + 1 < max_attempts && is_retryable_status(&status, config) => {
                tokio::time::sleep(next_backoff(config, attempt)).await;
                attempt += 1;
            }
            Err(status) => return Err(status),
        }
    }
}

/// Runs an async RPC operation with retries clipped to the remaining deadline.
pub async fn run_with_retry_budget<T, F, Fut>(
    config: &RpcRetryConfig,
    budget: &DeadlineBudget,
    mut operation: F,
) -> Result<T, Status>
where
    F: FnMut(std::time::Duration) -> Fut,
    Fut: Future<Output = Result<T, Status>>,
{
    if budget.expired() {
        return Err(Status::deadline_exceeded("rpc deadline exhausted"));
    }

    let max_attempts = if config.enabled {
        config.max_attempts.max(1)
    } else {
        1
    };
    let mut attempt = 0;
    loop {
        let remaining = budget.remaining();
        if remaining.is_zero() {
            return Err(Status::deadline_exceeded("rpc deadline exhausted"));
        }
        let result = operation(remaining).await;
        match result {
            Ok(value) => return Ok(value),
            Err(status) if attempt + 1 < max_attempts && is_retryable_status(&status, config) => {
                let sleep_for = next_backoff(config, attempt).min(budget.remaining());
                if sleep_for.is_zero() {
                    return Err(Status::deadline_exceeded("rpc deadline exhausted"));
                }
                tokio::time::sleep(sleep_for).await;
                attempt += 1;
            }
            Err(status) => return Err(status),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use tonic::{Code, Status};

    use super::{is_retryable_status, next_backoff};
    use crate::rpc::RpcRetryConfig;

    #[test]
    fn retry_policy_caps_backoff() {
        let config = RpcRetryConfig {
            initial_backoff: Duration::from_millis(10),
            max_backoff: Duration::from_millis(25),
            ..RpcRetryConfig::production_defaults()
        };

        assert_eq!(next_backoff(&config, 0), Duration::from_millis(10));
        assert_eq!(next_backoff(&config, 4), Duration::from_millis(25));
    }

    #[test]
    fn retry_policy_only_retries_configured_codes() {
        let config = RpcRetryConfig::production_defaults();

        assert!(is_retryable_status(
            &Status::new(Code::Unavailable, "down"),
            &config
        ));
        assert!(!is_retryable_status(
            &Status::new(Code::InvalidArgument, "bad"),
            &config
        ));
    }
}