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
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147

use super::{super::super::regions::IpAddrWithPort, ChooseOptions, Chooser, ChooserFeedback, ChosenResults};
use num_rational::Ratio;
use rand::{prelude::*, thread_rng};

#[cfg(feature = "async")]
use futures::future::BoxFuture;

const DEFAULT_RANDOM_CHOOSE_RATIO: Ratio<usize> = Ratio::new_raw(1, 2);

/// 永不空手的选择器
///
/// 确保 [`Chooser`] 实例不会因为所有可选择的 IP 地址都被屏蔽而导致 HTTP 客户端直接返回错误，
/// 在内置的 [`Chooser`] 没有返回结果时，将会随机返回一定比例的 IP 地址供 HTTP 客户端做一轮尝试。
#[derive(Debug, Clone)]
pub struct NeverEmptyHandedChooser<C: ?Sized> {
    random_choose_ratio: Ratio<usize>,
    inner_chooser: C,
}

impl<C> NeverEmptyHandedChooser<C> {
    /// 创建永不空手的选择器
    ///
    /// 需要提供在所有 IP 地址都被屏蔽的情况下，随机返回的 IP 地址的比例
    ///
    /// 需要注意，提供的随机比例的分母必须大于 0，且比值小于 1。
    #[inline]
    pub fn new(chooser: C, random_choose_ratio: Ratio<usize>) -> Self {
        assert!(random_choose_ratio.numer() <= random_choose_ratio.denom());
        assert!(*random_choose_ratio.denom() > 0);
        Self {
            inner_chooser: chooser,
            random_choose_ratio,
        }
    }
}

impl<C: Default> Default for NeverEmptyHandedChooser<C> {
    #[inline]
    fn default() -> Self {
        Self::new(Default::default(), DEFAULT_RANDOM_CHOOSE_RATIO)
    }
}

impl<C: Chooser> Chooser for NeverEmptyHandedChooser<C> {
    fn choose(&self, ips: &[IpAddrWithPort], opts: ChooseOptions) -> ChosenResults {
        let chosen = self.inner_chooser.choose(ips, opts);
        if chosen.is_empty() {
            self.random_choose(ips).into()
        } else {
            chosen
        }
    }

    #[inline]
    fn feedback(&self, feedback: ChooserFeedback) {
        self.inner_chooser.feedback(feedback)
    }

    #[cfg(feature = "async")]
    #[cfg_attr(feature = "docs", doc(cfg(feature = "async")))]
    fn async_choose<'a>(&'a self, ips: &'a [IpAddrWithPort], opts: ChooseOptions<'a>) -> BoxFuture<'a, ChosenResults> {
        Box::pin(async move {
            let chosen = self.inner_chooser.async_choose(ips, opts).await;
            if chosen.is_empty() {
                self.random_choose(ips).into()
            } else {
                chosen
            }
        })
    }

    #[inline]
    #[cfg(feature = "async")]
    #[cfg_attr(feature = "docs", doc(cfg(feature = "async")))]
    fn async_feedback<'a>(&'a self, feedback: ChooserFeedback<'a>) -> BoxFuture<'a, ()> {
        self.inner_chooser.async_feedback(feedback)
    }
}

impl<C> NeverEmptyHandedChooser<C> {
    fn random_choose(&self, ips: &[IpAddrWithPort]) -> Vec<IpAddrWithPort> {
        let chosen_len = (self.random_choose_ratio * ips.len()).ceil().to_integer();
        ips.choose_multiple(&mut thread_rng(), chosen_len).copied().collect()
    }
}

#[cfg(test)]
mod tests {
    use super::{
        super::{
            super::{ResponseError, ResponseErrorKind, RetriedStatsInfo},
            IpChooser,
        },
        *,
    };
    use qiniu_http::Extensions;
    use std::net::{IpAddr, Ipv4Addr};

    const IPS_WITHOUT_PORT: &[IpAddrWithPort] = &[
        IpAddrWithPort::new(IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1)), None),
        IpAddrWithPort::new(IpAddr::V4(Ipv4Addr::new(192, 168, 1, 2)), None),
        IpAddrWithPort::new(IpAddr::V4(Ipv4Addr::new(192, 168, 1, 3)), None),
    ];

    #[test]
    fn test_never_empty_handed_chooser() {
        env_logger::builder().is_test(true).try_init().ok();

        let ip_chooser: NeverEmptyHandedChooser<IpChooser> = Default::default();
        assert_eq!(
            ip_chooser.choose(IPS_WITHOUT_PORT, Default::default()).into_ip_addrs(),
            IPS_WITHOUT_PORT.to_vec()
        );
        ip_chooser.feedback(ChooserFeedback::new(
            &[
                IpAddrWithPort::new(IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1)), None),
                IpAddrWithPort::new(IpAddr::V4(Ipv4Addr::new(192, 168, 1, 2)), None),
            ],
            None,
            &RetriedStatsInfo::default(),
            &mut Extensions::default(),
            None,
            Some(&ResponseError::new(ResponseErrorKind::ParseResponseError, "Test Error")),
        ));
        assert_eq!(
            ip_chooser.choose(IPS_WITHOUT_PORT, Default::default()).into_ip_addrs(),
            [IpAddrWithPort::new(IpAddr::V4(Ipv4Addr::new(192, 168, 1, 3)), None)].to_vec(),
        );

        ip_chooser.feedback(ChooserFeedback::new(
            IPS_WITHOUT_PORT,
            None,
            &RetriedStatsInfo::default(),
            &mut Extensions::default(),
            None,
            Some(&ResponseError::new(ResponseErrorKind::ParseResponseError, "Test Error")),
        ));

        assert_eq!(
            ip_chooser
                .choose(IPS_WITHOUT_PORT, Default::default())
                .into_ip_addrs()
                .len(),
            2
        );
    }
}