Skip to main content

ursula_control/
view.rs

1use std::collections::BTreeMap;
2use std::collections::BTreeSet;
3
4use serde::Deserialize;
5use serde::Serialize;
6use ursula_shard::RaftGroupId;
7
8use crate::model::NodeId;
9use crate::model::NodeState;
10
11#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
12pub struct PlacementNode {
13    pub node_id: NodeId,
14    pub client_url: String,
15    pub cluster_url: String,
16    pub state: NodeState,
17}
18
19#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
20pub struct GroupPlacementView {
21    pub raft_group_id: RaftGroupId,
22    pub voters: BTreeSet<NodeId>,
23    pub learners: BTreeSet<NodeId>,
24    pub draining: BTreeSet<NodeId>,
25    pub epoch: u64,
26    pub nodes: BTreeMap<NodeId, PlacementNode>,
27}
28
29impl GroupPlacementView {
30    pub fn hosts(&self, node_id: NodeId) -> bool {
31        self.voters.contains(&node_id) || self.learners.contains(&node_id)
32    }
33
34    pub fn serves_client_traffic(&self, node_id: NodeId) -> bool {
35        self.voters.contains(&node_id)
36            && !self.draining.contains(&node_id)
37            && self
38                .nodes
39                .get(&node_id)
40                .is_some_and(|node| node.state == NodeState::Active)
41    }
42
43    pub fn cluster_endpoints(&self) -> BTreeMap<NodeId, String> {
44        self.nodes
45            .iter()
46            .filter(|(node_id, _)| self.hosts(**node_id))
47            .map(|(node_id, node)| (*node_id, node.cluster_url.clone()))
48            .collect()
49    }
50
51    pub fn active_voter_client_url(&self, exclude: Option<NodeId>) -> Option<(NodeId, String)> {
52        let mut fallback: Option<(NodeId, String)> = None;
53        for node_id in &self.voters {
54            if !self.serves_client_traffic(*node_id) {
55                continue;
56            }
57            let Some(node) = self.nodes.get(node_id) else {
58                continue;
59            };
60            let candidate = (*node_id, node.client_url.clone());
61            if Some(*node_id) == exclude {
62                fallback.get_or_insert(candidate);
63            } else {
64                return Some(candidate);
65            }
66        }
67        fallback
68    }
69}