klukai_types/
members.rs

1use std::{collections::BTreeMap, net::SocketAddr, ops::Range, time::Duration};
2
3use circular_buffer::CircularBuffer;
4use serde::{Deserialize, Serialize};
5use tracing::{debug, trace};
6
7use crate::{
8    actor::{Actor, ActorId, ClusterId},
9    broadcast::Timestamp,
10};
11
12#[derive(Clone, Debug, Serialize, Deserialize)]
13pub struct MemberState {
14    pub addr: SocketAddr,
15    pub ts: Timestamp,
16    pub cluster_id: ClusterId,
17
18    pub ring: Option<u8>,
19    pub last_sync_ts: Option<Timestamp>,
20}
21
22impl MemberState {
23    pub fn new(addr: SocketAddr, ts: Timestamp, cluster_id: ClusterId) -> Self {
24        Self {
25            addr,
26            ts,
27            cluster_id,
28            ring: None,
29            last_sync_ts: None,
30        }
31    }
32
33    pub fn is_ring0(&self) -> bool {
34        self.ring == Some(0)
35    }
36}
37
38const RING_BUCKETS: [Range<u64>; 6] = [0..6, 6..15, 15..50, 50..100, 100..200, 200..300];
39
40#[derive(Debug, Default, Clone)]
41pub struct Rtt {
42    pub buf: CircularBuffer<20, u64>,
43}
44
45#[derive(Default)]
46pub struct Members {
47    pub states: BTreeMap<ActorId, MemberState>,
48    pub by_addr: BTreeMap<SocketAddr, ActorId>,
49    pub rtts: BTreeMap<SocketAddr, Rtt>,
50}
51
52#[derive(Debug, PartialEq)]
53pub enum MemberAddedResult {
54    NewMember,
55    Updated,
56    Ignored,
57}
58
59impl Members {
60    pub fn get(&self, id: &ActorId) -> Option<&MemberState> {
61        self.states.get(id)
62    }
63
64    pub fn update_sync_ts(&mut self, actor_id: &ActorId, ts: Timestamp) {
65        if let Some(state) = self.states.get_mut(actor_id) {
66            state.last_sync_ts = Some(ts);
67        }
68    }
69
70    // A result of `true` means that the effective list of
71    // cluster member addresses has changed
72    pub fn add_member(&mut self, actor: &Actor) -> MemberAddedResult {
73        let actor_id = actor.id();
74        let mut ret = MemberAddedResult::Ignored;
75
76        let member = self.states.entry(actor_id).or_insert_with(|| {
77            ret = MemberAddedResult::NewMember;
78            MemberState::new(actor.addr(), actor.ts(), actor.cluster_id())
79        });
80
81        trace!("member: {member:?}");
82
83        // The received timestamp is older than the previously known
84        // one.  If we just added the member this shouldn't ever
85        // trigger (because the timestamps would be the same).
86        if actor.ts().to_duration() < member.ts.to_duration() {
87            debug!("older timestamp, ignoring");
88            return MemberAddedResult::Ignored;
89        }
90
91        // If the new timestamp is newer than what we had on file we
92        // update the member, then set the return to "Update".
93        // Because a newly inserted member would always have the same
94        // timestamp this code doesn't run if we just inserted.
95        if actor.ts().to_duration() > member.ts.to_duration() {
96            member.addr = actor.addr();
97            member.ts = actor.ts();
98            member.cluster_id = actor.cluster_id();
99            ret = MemberAddedResult::Updated;
100        }
101
102        // If we just inserted, add the actor to the by_addr set and
103        // recalculate the RTT rings.
104        if ret == MemberAddedResult::NewMember {
105            self.by_addr.insert(actor.addr(), actor.id());
106            self.recalculate_rings(actor.addr());
107        }
108
109        ret
110    }
111
112    // A result of `true` means that the effective list of
113    // cluster member addresses has changed
114    pub fn remove_member(&mut self, actor: &Actor) -> bool {
115        let effectively_down = if let Some(member) = self.states.get(&actor.id()) {
116            member.ts == actor.ts()
117        } else {
118            // Shouldn't happen
119            false
120        };
121
122        if effectively_down {
123            self.by_addr.remove(&actor.addr());
124            self.states.remove(&actor.id());
125        }
126
127        effectively_down
128    }
129
130    pub fn add_rtt(&mut self, addr: SocketAddr, rtt: Duration) {
131        self.rtts
132            .entry(addr)
133            .or_default()
134            .buf
135            .push_front(rtt.subsec_millis() as u64 + (rtt.as_secs() * 1000));
136
137        self.recalculate_rings(addr)
138    }
139
140    /// For a given member, calculate the average RTT and update
141    /// `self.ring` with the index of the corresponding bucket in
142    /// `RING_BUCKETS`.
143    fn recalculate_rings(&mut self, addr: SocketAddr) {
144        if let Some(actor_id) = self.by_addr.get(&addr)
145            && let Some(avg) = self.rtts.get(&addr).and_then(|rtt| {
146                // If the ring buffer isn't empty
147                (!rtt.buf.is_empty()).then(|| {
148                    // We can only access the ring buffer via two
149                    // slices, so we sum both of them together
150                    (rtt.buf.as_slices().0.iter().sum::<u64>()
151                     + rtt.buf.as_slices().1.iter().sum::<u64>())
152                        // Then average over the full size of the ring
153                        // buffer for the average of recent RTTs
154                        / rtt.buf.len() as u64
155                })
156            })
157            && let Some(state) = self.states.get_mut(actor_id)
158        {
159            // We check which range-bucket the RTT is
160            // contained in, then update the stored index
161            for (ring, n) in RING_BUCKETS.iter().enumerate() {
162                if n.contains(&avg) {
163                    state.ring = Some(ring as u8);
164                    break;
165                }
166            }
167        }
168    }
169
170    /// Get member addresses where the ring index is `0` (meaning a
171    /// very small RTT)
172    pub fn ring0(&self, cluster_id: ClusterId) -> impl Iterator<Item = SocketAddr> + '_ {
173        self.states.values().filter_map(move |v| {
174            v.ring
175                .and_then(|ring| (v.cluster_id == cluster_id && ring == 0).then_some(v.addr))
176        })
177    }
178}
klukai_types/members.rs

klukai_types/
members.rs
