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}