ursula_control/
command.rs

1use std::collections::BTreeMap;
2use std::collections::BTreeSet;
3use std::fmt;
4
5use serde::Deserialize;
6use serde::Serialize;
7use ursula_shard::RaftGroupId;
8
9use crate::model::LearnerStatus;
10use crate::model::MigrationPhase;
11use crate::model::NodeId;
12use crate::model::NodeState;
13
14#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
15pub enum ControlCommand {
16    RegisterNode {
17        node_id: NodeId,
18        client_url: String,
19        cluster_url: String,
20        #[serde(default)]
21        labels: BTreeMap<String, String>,
22        now_ms: u64,
23    },
24    SetNodeState {
25        node_id: NodeId,
26        state: NodeState,
27        now_ms: u64,
28    },
29    SeedPlacement {
30        raft_group_id: RaftGroupId,
31        voters: BTreeSet<NodeId>,
32        now_ms: u64,
33    },
34    BeginMigration {
35        raft_group_id: RaftGroupId,
36        target_voters: BTreeSet<NodeId>,
37        retain_removed: bool,
38        now_ms: u64,
39    },
40    AdvanceMigration {
41        migration_id: u64,
42        phase: MigrationPhase,
43        now_ms: u64,
44    },
45    SetLearnerStatus {
46        migration_id: u64,
47        node_id: NodeId,
48        status: LearnerStatus,
49        now_ms: u64,
50    },
51    RecordMigrationError {
52        migration_id: u64,
53        error: String,
54        now_ms: u64,
55    },
56    CommitPlacement {
57        raft_group_id: RaftGroupId,
58        voters: BTreeSet<NodeId>,
59        learners: BTreeSet<NodeId>,
60        draining: BTreeSet<NodeId>,
61        now_ms: u64,
62    },
63    FinishMigration {
64        migration_id: u64,
65        success: bool,
66        now_ms: u64,
67    },
68    EvictLearner {
69        raft_group_id: RaftGroupId,
70        node_id: NodeId,
71        now_ms: u64,
72    },
73}
74
75impl ControlCommand {
76    pub fn now_ms(&self) -> u64 {
77        match self {
78            Self::RegisterNode { now_ms, .. }
79            | Self::SetNodeState { now_ms, .. }
80            | Self::SeedPlacement { now_ms, .. }
81            | Self::BeginMigration { now_ms, .. }
82            | Self::AdvanceMigration { now_ms, .. }
83            | Self::SetLearnerStatus { now_ms, .. }
84            | Self::RecordMigrationError { now_ms, .. }
85            | Self::CommitPlacement { now_ms, .. }
86            | Self::FinishMigration { now_ms, .. }
87            | Self::EvictLearner { now_ms, .. } => *now_ms,
88        }
89    }
90}
91
92impl fmt::Display for ControlCommand {
93    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
94        f.write_str(match self {
95            Self::RegisterNode { .. } => "register_node",
96            Self::SetNodeState { .. } => "set_node_state",
97            Self::SeedPlacement { .. } => "seed_placement",
98            Self::BeginMigration { .. } => "begin_migration",
99            Self::AdvanceMigration { .. } => "advance_migration",
100            Self::SetLearnerStatus { .. } => "set_learner_status",
101            Self::RecordMigrationError { .. } => "record_migration_error",
102            Self::CommitPlacement { .. } => "commit_placement",
103            Self::FinishMigration { .. } => "finish_migration",
104            Self::EvictLearner { .. } => "evict_learner",
105        })
106    }
107}
108
109#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
110pub enum ControlResponse {
111    Ok,
112    MigrationStarted { migration_id: u64 },
113    Rejected { reason: String },
114}
115
116impl ControlResponse {
117    pub fn is_rejected(&self) -> bool {
118        matches!(self, Self::Rejected { .. })
119    }
120}
121
122impl fmt::Display for ControlResponse {
123    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
124        f.write_str(match self {
125            Self::Ok => "ok",
126            Self::MigrationStarted { .. } => "migration_started",
127            Self::Rejected { .. } => "rejected",
128        })
129    }
130}
ursula_control/command.rs

ursula_control/
command.rs
