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

use super::{super::super::regions::IpAddrWithPort, ChooseOptions, Chooser, ChooserFeedback, ChosenResults};
use rand::{seq::SliceRandom, thread_rng};

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

/// 随机选择器
///
/// 基于一个选择器实例，但将其返回的选择结果打乱
#[derive(Debug, Clone)]
pub struct ShuffledChooser<C: ?Sized> {
    chooser: C,
}

impl<C> ShuffledChooser<C> {
    /// 创建随机选择器
    #[inline]
    pub fn new(chooser: C) -> Self {
        Self { chooser }
    }
}

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

impl<C: Chooser> Chooser for ShuffledChooser<C> {
    #[inline]
    fn choose(&self, ips: &[IpAddrWithPort], opts: ChooseOptions) -> ChosenResults {
        let mut ips = self.chooser.choose(ips, opts);
        ips.shuffle(&mut thread_rng());
        ips
    }

    #[inline]
    #[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 mut ips = self.chooser.async_choose(ips, opts).await;
            ips.shuffle(&mut thread_rng());
            ips
        })
    }

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

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

#[cfg(test)]
mod tests {
    use super::{
        super::{
            super::{ResponseError, ResponseErrorKind, RetriedStatsInfo},
            IpChooser,
        },
        *,
    };
    use qiniu_http::Extensions;
    use std::{
        collections::HashSet,
        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_shuffled_chooser() {
        env_logger::builder().is_test(true).try_init().ok();

        let ip_chooser: ShuffledChooser<IpChooser> = Default::default();
        assert_eq!(
            make_set(ip_chooser.choose(IPS_WITHOUT_PORT, Default::default())),
            make_set(IPS_WITHOUT_PORT)
        );
        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!(
            make_set(ip_chooser.choose(IPS_WITHOUT_PORT, Default::default())),
            make_set(&[IpAddrWithPort::new(IpAddr::V4(Ipv4Addr::new(192, 168, 1, 3)), None)]),
        );

        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(),
            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,
            None,
        ));
        assert_eq!(
            make_set(ip_chooser.choose(IPS_WITHOUT_PORT, Default::default())),
            make_set(&[
                IpAddrWithPort::new(IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1)), None),
                IpAddrWithPort::new(IpAddr::V4(Ipv4Addr::new(192, 168, 1, 2)), None),
            ])
        );
    }

    fn make_set(ips: impl AsRef<[IpAddrWithPort]>) -> HashSet<IpAddrWithPort> {
        HashSet::from_iter(ips.as_ref().iter().copied())
    }
}