bililive_core/retry/
context.rs

1use std::sync::atomic::AtomicUsize;
2use std::sync::atomic::Ordering::SeqCst;
3use std::sync::Arc;
4
5use crate::config::StreamConfig;
6
7/// Internal context for server picking during (re)connection.
8///
9/// Implements a round-robin policy for server selection.
10#[derive(Debug, Clone)]
11pub struct RetryContext {
12    config: StreamConfig,
13    cursor: Arc<AtomicUsize>,
14}
15
16impl RetryContext {
17    /// Get the stream config.
18    #[must_use]
19    pub const fn config(&self) -> &StreamConfig {
20        &self.config
21    }
22    /// Get the next server.
23    #[allow(clippy::missing_panics_doc)]
24    pub fn get(&mut self) -> &str {
25        let cursor: usize = self
26            .cursor
27            .fetch_update(SeqCst, SeqCst, |i| {
28                Some((i + 1) % self.config.servers().len())
29            })
30            .unwrap();
31        &*self.config.servers()[cursor]
32    }
33}
34
35impl From<StreamConfig> for RetryContext {
36    fn from(config: StreamConfig) -> Self {
37        Self {
38            config,
39            cursor: Arc::new(Default::default()),
40        }
41    }
42}