stasis/application/runtime/
queue_ownership_rebalance_job_handler.rs1use std::collections::BTreeSet;
2use std::sync::Arc;
3
4use async_trait::async_trait;
5use chrono::Utc;
6use serde::Deserialize;
7use serde_json::json;
8
9use crate::application::runtime::in_memory_runtime::{JobExecutionOutcome, JobHandler};
10use crate::domain::errors::Result;
11use crate::domain::runtime::cluster_node::ClusterNodeHeartbeat;
12use crate::domain::runtime::job::Job;
13use crate::ports::outbound::runtime::cluster_node_store::ClusterNodeStore;
14
15#[derive(Clone)]
16pub struct QueueOwnershipRebalanceJobHandler {
17 cluster_store: Arc<dyn ClusterNodeStore>,
18}
19
20#[derive(Deserialize)]
21struct QueueRebalancePayload {
22 queue: String,
23 desired_owners: Vec<String>,
24 strategy: Option<String>,
25 reason: Option<String>,
26}
27
28impl QueueOwnershipRebalanceJobHandler {
29 pub fn new(cluster_store: Arc<dyn ClusterNodeStore>) -> Self {
30 Self { cluster_store }
31 }
32
33 fn parse_payload(raw: &str) -> std::result::Result<QueueRebalancePayload, String> {
34 let payload: QueueRebalancePayload = serde_json::from_str(raw)
35 .map_err(|err| format!("invalid queue rebalance payload json: {err}"))?;
36
37 if payload.queue.trim().is_empty() {
38 return Err("queue must not be empty".to_string());
39 }
40 if payload.desired_owners.is_empty() {
41 return Err("desired_owners must not be empty".to_string());
42 }
43 if payload
44 .desired_owners
45 .iter()
46 .any(|owner| owner.trim().is_empty())
47 {
48 return Err("desired_owners must not contain empty values".to_string());
49 }
50
51 Ok(payload)
52 }
53
54 fn failure(message: impl Into<String>) -> JobExecutionOutcome {
55 JobExecutionOutcome::FatalFailure {
56 message: message.into(),
57 execution_id: None,
58 diagnostics: None,
59 }
60 }
61
62 fn retryable(message: impl Into<String>) -> JobExecutionOutcome {
63 JobExecutionOutcome::RetryableFailure {
64 message: message.into(),
65 execution_id: None,
66 diagnostics: None,
67 }
68 }
69
70 fn remaining_ttl_seconds(
71 lease_expires_at: chrono::DateTime<Utc>,
72 now: chrono::DateTime<Utc>,
73 ) -> i64 {
74 lease_expires_at
75 .signed_duration_since(now)
76 .num_seconds()
77 .max(1)
78 }
79}
80
81#[async_trait]
82impl JobHandler for QueueOwnershipRebalanceJobHandler {
83 fn job_type(&self) -> &'static str {
84 "workflow.stasis.cluster.queue_ownership_rebalance"
85 }
86
87 async fn execute(&self, job: &Job) -> Result<JobExecutionOutcome> {
88 let payload = match Self::parse_payload(&job.payload_ref) {
89 Ok(payload) => payload,
90 Err(message) => return Ok(Self::failure(message)),
91 };
92
93 let QueueRebalancePayload {
94 queue,
95 desired_owners,
96 strategy,
97 reason,
98 } = payload;
99
100 let now = Utc::now();
101 let nodes = match self.cluster_store.list().await {
102 Ok(nodes) => nodes,
103 Err(err) => {
104 return Ok(Self::retryable(format!(
105 "failed to load cluster nodes for rebalance: {err}"
106 )));
107 }
108 };
109
110 let active_nodes = nodes
111 .into_iter()
112 .filter(|node| node.lease_expires_at >= now)
113 .collect::<Vec<_>>();
114
115 let desired = desired_owners.into_iter().collect::<BTreeSet<_>>();
116
117 if let Some(missing) = desired
118 .iter()
119 .find(|owner| !active_nodes.iter().any(|node| &node.node_id == *owner))
120 {
121 return Ok(Self::failure(format!(
122 "desired owner is not active: {}",
123 missing
124 )));
125 }
126
127 let rebalance_metadata = json!({
128 "action": "queue_ownership_rebalance",
129 "queue": &queue,
130 "strategy": &strategy,
131 "reason": &reason,
132 })
133 .to_string();
134
135 let mut updated_nodes = 0usize;
136
137 for node in active_nodes {
138 let owns_queue = node.queue_ownership.iter().any(|q| q == &queue);
139 let should_own = desired.contains(&node.node_id);
140 if owns_queue == should_own {
141 continue;
142 }
143
144 let mut queues = node.queue_ownership.into_iter().collect::<BTreeSet<_>>();
145 if should_own {
146 queues.insert(queue.clone());
147 } else {
148 queues.remove(&queue);
149 }
150
151 let ttl = Self::remaining_ttl_seconds(node.lease_expires_at, now);
152 let result = self
153 .cluster_store
154 .heartbeat(ClusterNodeHeartbeat {
155 node_id: node.node_id,
156 heartbeat_at: now,
157 lease_ttl_seconds: ttl,
158 queue_ownership: Some(queues.into_iter().collect::<Vec<_>>()),
159 capability_tags: None,
160 metadata: Some(rebalance_metadata.clone()),
161 })
162 .await;
163
164 if result.is_err() {
165 return Ok(Self::retryable(
166 "failed to update node during queue rebalance",
167 ));
168 }
169
170 updated_nodes += 1;
171 }
172
173 let diagnostics = json!({
174 "status": "success",
175 "queue": queue,
176 "updated_nodes": updated_nodes,
177 "desired_owners": desired.into_iter().collect::<Vec<_>>(),
178 "strategy": strategy,
179 })
180 .to_string();
181
182 Ok(JobExecutionOutcome::Success {
183 sttp_output_node_id: "sttp:out:stasis:cluster:queue_ownership_rebalance".to_string(),
184 execution_id: None,
185 diagnostics: Some(diagnostics),
186 })
187 }
188}
189
190#[cfg(test)]
191mod tests {
192 use chrono::Utc;
193 use std::sync::Arc;
194
195 use crate::application::runtime::in_memory_runtime::JobHandler;
196 use crate::application::runtime::queue_ownership_rebalance_job_handler::QueueOwnershipRebalanceJobHandler;
197 use crate::domain::runtime::cluster_node::{ClusterNodeRole, NewClusterNode};
198 use crate::domain::runtime::job::{BackoffPolicy, Job, JobState};
199 use crate::infrastructure::runtime::in_memory_cluster_node_store::InMemoryClusterNodeStore;
200 use crate::ports::outbound::runtime::cluster_node_store::ClusterNodeStore;
201
202 fn sample_job(payload_ref: String) -> Job {
203 Job {
204 id: "job.cluster.rebalance.1".to_string(),
205 queue: "cluster-control".to_string(),
206 job_type: "workflow.stasis.cluster.queue_ownership_rebalance".to_string(),
207 payload_ref,
208 state: JobState::Enqueued,
209 priority: 100,
210 attempts: 0,
211 max_attempts: 3,
212 backoff_policy: BackoffPolicy::default(),
213 idempotency_key: "idem-rebalance".to_string(),
214 correlation_id: "corr-rebalance".to_string(),
215 causation_id: "cause-rebalance".to_string(),
216 trace_id: "trace-rebalance".to_string(),
217 sttp_input_node_id: "sttp:in:cluster:rebalance".to_string(),
218 sttp_output_node_id: None,
219 lease_owner: None,
220 lease_expires_at: None,
221 heartbeat_at: None,
222 scheduled_at: Utc::now(),
223 started_at: None,
224 finished_at: None,
225 last_error: None,
226 }
227 }
228
229 #[tokio::test]
230 async fn handler_rebalances_queue_ownership_across_nodes() {
231 let store = InMemoryClusterNodeStore::default();
232 let now = Utc::now();
233
234 store
235 .register(NewClusterNode {
236 node_id: "node.a".to_string(),
237 role: ClusterNodeRole::Worker,
238 region: "us-east".to_string(),
239 queue_ownership: vec!["priority".to_string()],
240 capability_tags: vec![],
241 heartbeat_at: now,
242 lease_ttl_seconds: 60,
243 metadata: None,
244 })
245 .await
246 .expect("register node a should succeed");
247
248 store
249 .register(NewClusterNode {
250 node_id: "node.b".to_string(),
251 role: ClusterNodeRole::Worker,
252 region: "us-east".to_string(),
253 queue_ownership: vec![],
254 capability_tags: vec![],
255 heartbeat_at: now,
256 lease_ttl_seconds: 60,
257 metadata: None,
258 })
259 .await
260 .expect("register node b should succeed");
261
262 let handler = QueueOwnershipRebalanceJobHandler::new(Arc::new(store.clone()));
263 let payload = serde_json::json!({
264 "queue": "priority",
265 "desired_owners": ["node.b"],
266 "strategy": "least-loaded",
267 "reason": "capacity",
268 })
269 .to_string();
270
271 let outcome = handler
272 .execute(&sample_job(payload))
273 .await
274 .expect("handler execution should succeed");
275
276 assert!(matches!(
277 outcome,
278 crate::application::runtime::in_memory_runtime::JobExecutionOutcome::Success { .. }
279 ));
280
281 let node_a = store
282 .get("node.a")
283 .await
284 .expect("node a get should succeed")
285 .expect("node a should exist");
286 let node_b = store
287 .get("node.b")
288 .await
289 .expect("node b get should succeed")
290 .expect("node b should exist");
291
292 assert!(!node_a.queue_ownership.iter().any(|q| q == "priority"));
293 assert!(node_b.queue_ownership.iter().any(|q| q == "priority"));
294 }
295}