1use std::net::SocketAddr;
52use std::sync::Arc;
53use std::time::Duration;
54
55use parking_lot::Mutex;
56
57use crate::cluster::peer::{Peer, PeerEndpoint, PeerState};
58use crate::cluster::pool::ServerPool;
59use crate::hashkit::DynToken;
60use crate::net::auto_eject::AutoEject;
61
62#[derive(Clone, Debug, Eq, PartialEq)]
67pub struct PeerSnapshot {
68 pub idx: u32,
71 pub dc: String,
73 pub rack: String,
75 pub host: String,
77 pub port: u16,
79 pub tokens: Vec<u32>,
82 pub state: PeerState,
84 pub is_local: bool,
87}
88
89#[derive(Clone, Debug, Eq, PartialEq)]
95pub struct PeerSpec {
96 pub host: String,
98 pub port: u16,
100 pub dc: String,
102 pub rack: String,
104 pub tokens: Vec<u32>,
106 pub is_secure: bool,
108}
109
110#[derive(Copy, Clone, Debug, Eq, PartialEq)]
112pub enum ClusterChangeKind {
113 Add,
115 Remove,
117}
118
119#[derive(Clone, Debug, Eq, PartialEq)]
124pub struct ClusterChange {
125 pub kind: ClusterChangeKind,
127 pub peer_idx: Option<u32>,
129 pub peer: Option<PeerSpec>,
131}
132
133#[derive(Clone, Debug, Eq, PartialEq)]
136pub struct JoinPlan {
137 pub change: ClusterChange,
139}
140
141#[derive(Debug, thiserror::Error)]
143#[non_exhaustive]
144pub enum ClusterError {
145 #[error("peer not found: idx={idx}")]
147 PeerNotFound {
148 idx: u32,
150 },
151 #[error("cannot remove the local peer")]
154 CannotRemoveLocal,
155 #[error("peer with endpoint {addr} already exists")]
157 PeerAlreadyExists {
158 addr: String,
160 },
161 #[error("invalid request: {0}")]
164 Invalid(String),
165}
166
167pub trait ClusterAdmin: Send + Sync + std::fmt::Debug {
174 fn list_peers(&self) -> Vec<PeerSnapshot>;
176 fn cluster_join(&self, target: SocketAddr) -> Result<JoinPlan, ClusterError>;
183 fn cluster_leave(&self, peer_idx: u32) -> Result<JoinPlan, ClusterError>;
192 fn cluster_plan_pending(&self) -> Vec<ClusterChange>;
194 fn cluster_commit(&self) -> Result<(), ClusterError>;
201}
202
203#[derive(Debug, Default)]
207pub struct NoopClusterAdmin;
208
209impl ClusterAdmin for NoopClusterAdmin {
210 fn list_peers(&self) -> Vec<PeerSnapshot> {
211 Vec::new()
212 }
213 fn cluster_join(&self, _target: SocketAddr) -> Result<JoinPlan, ClusterError> {
214 Err(ClusterError::Invalid(
215 "cluster admin RPC not configured on this node".into(),
216 ))
217 }
218 fn cluster_leave(&self, _peer_idx: u32) -> Result<JoinPlan, ClusterError> {
219 Err(ClusterError::Invalid(
220 "cluster admin RPC not configured on this node".into(),
221 ))
222 }
223 fn cluster_plan_pending(&self) -> Vec<ClusterChange> {
224 Vec::new()
225 }
226 fn cluster_commit(&self) -> Result<(), ClusterError> {
227 Ok(())
228 }
229}
230
231#[derive(Debug)]
238pub struct PoolClusterAdmin {
239 pool: Arc<ServerPool>,
240 staged: Mutex<Vec<ClusterChange>>,
241}
242
243impl PoolClusterAdmin {
244 #[must_use]
246 pub fn new(pool: Arc<ServerPool>) -> Self {
247 Self {
248 pool,
249 staged: Mutex::new(Vec::new()),
250 }
251 }
252
253 #[must_use]
257 pub fn pool(&self) -> &Arc<ServerPool> {
258 &self.pool
259 }
260}
261
262impl ClusterAdmin for PoolClusterAdmin {
263 fn list_peers(&self) -> Vec<PeerSnapshot> {
264 let peers = self.pool.peers().read();
265 peers
266 .iter()
267 .map(|p| PeerSnapshot {
268 idx: p.idx(),
269 dc: p.dc().to_string(),
270 rack: p.rack().to_string(),
271 host: p.endpoint().host().to_string(),
272 port: p.endpoint().port(),
273 tokens: p.tokens().iter().map(DynToken::get_int).collect(),
274 state: p.state(),
275 is_local: p.is_local(),
276 })
277 .collect()
278 }
279
280 fn cluster_join(&self, target: SocketAddr) -> Result<JoinPlan, ClusterError> {
281 let host = target.ip().to_string();
282 let port = target.port();
283 let peers = self.pool.peers().read();
284 if peers
285 .iter()
286 .any(|p| p.endpoint().host() == host && p.endpoint().port() == port)
287 {
288 return Err(ClusterError::PeerAlreadyExists {
289 addr: target.to_string(),
290 });
291 }
292 let staged = self.staged.lock();
293 if staged
294 .iter()
295 .any(|c| matches!(&c.peer, Some(s) if s.host == host && s.port == port))
296 {
297 return Err(ClusterError::PeerAlreadyExists {
298 addr: target.to_string(),
299 });
300 }
301 let (dc, rack) = peers.iter().find(|p| p.is_local()).map_or_else(
303 || {
304 (
305 self.pool.config().dc.clone(),
306 self.pool.config().rack.clone(),
307 )
308 },
309 |p| (p.dc().to_string(), p.rack().to_string()),
310 );
311 drop(staged);
312 drop(peers);
313 let token_val = derive_token(&host, port);
314 let spec = PeerSpec {
315 host,
316 port,
317 dc,
318 rack,
319 tokens: vec![token_val],
320 is_secure: false,
321 };
322 let change = ClusterChange {
323 kind: ClusterChangeKind::Add,
324 peer_idx: None,
325 peer: Some(spec),
326 };
327 let plan = JoinPlan {
328 change: change.clone(),
329 };
330 self.staged.lock().push(change);
331 Ok(plan)
332 }
333
334 fn cluster_leave(&self, peer_idx: u32) -> Result<JoinPlan, ClusterError> {
335 let peers = self.pool.peers().read();
336 let target = peers
337 .iter()
338 .find(|p| p.idx() == peer_idx)
339 .ok_or(ClusterError::PeerNotFound { idx: peer_idx })?;
340 if target.is_local() {
341 return Err(ClusterError::CannotRemoveLocal);
342 }
343 drop(peers);
344 let change = ClusterChange {
345 kind: ClusterChangeKind::Remove,
346 peer_idx: Some(peer_idx),
347 peer: None,
348 };
349 let plan = JoinPlan {
350 change: change.clone(),
351 };
352 self.staged.lock().push(change);
353 Ok(plan)
354 }
355
356 fn cluster_plan_pending(&self) -> Vec<ClusterChange> {
357 self.staged.lock().clone()
358 }
359
360 fn cluster_commit(&self) -> Result<(), ClusterError> {
361 let mut staged = self.staged.lock();
362 if staged.is_empty() {
363 return Ok(());
364 }
365 let mut peers = self.pool.peers().write();
366 let mut auto_ejects = self.pool.auto_eject().write();
367 let local_dc = self.pool.config().dc.clone();
368 for change in staged.iter() {
369 match change.kind {
370 ClusterChangeKind::Add => {
371 let spec = change
372 .peer
373 .as_ref()
374 .ok_or_else(|| ClusterError::Invalid("Add change missing peer".into()))?;
375 if peers.iter().any(|p| {
376 p.endpoint().host() == spec.host && p.endpoint().port() == spec.port
377 }) {
378 return Err(ClusterError::PeerAlreadyExists {
379 addr: format!("{}:{}", spec.host, spec.port),
380 });
381 }
382 let new_idx = u32::try_from(peers.len()).unwrap_or(u32::MAX);
383 let is_same_dc = spec.dc == local_dc;
384 let new_peer = Peer::new(
385 new_idx,
386 PeerEndpoint::tcp(spec.host.clone(), spec.port),
387 spec.rack.clone(),
388 spec.dc.clone(),
389 spec.tokens
390 .iter()
391 .copied()
392 .map(DynToken::from_u32)
393 .collect(),
394 false,
395 is_same_dc,
396 spec.is_secure,
397 );
398 peers.push(new_peer);
399 let template = AutoEject::new(
400 self.pool.config().auto_eject_hosts,
401 self.pool.config().server_failure_limit,
402 Duration::from_millis(self.pool.config().server_retry_timeout_ms),
403 );
404 auto_ejects.push(template);
405 }
406 ClusterChangeKind::Remove => {
407 let idx = change
408 .peer_idx
409 .ok_or_else(|| ClusterError::Invalid("Remove change missing idx".into()))?;
410 let pos = peers
411 .iter()
412 .position(|p| p.idx() == idx)
413 .ok_or(ClusterError::PeerNotFound { idx })?;
414 if peers[pos].is_local() {
415 return Err(ClusterError::CannotRemoveLocal);
416 }
417 peers.remove(pos);
418 if pos < auto_ejects.len() {
419 auto_ejects.remove(pos);
420 }
421 }
422 }
423 }
424 staged.clear();
425 drop(peers);
426 drop(auto_ejects);
427 self.pool.rebuild_ring();
430 Ok(())
431 }
432}
433
434fn derive_token(host: &str, port: u16) -> u32 {
442 let mut hash: u64 = 0xcbf2_9ce4_8422_2325;
443 for &b in host.as_bytes() {
444 hash ^= u64::from(b);
445 hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
446 }
447 for byte in port.to_be_bytes() {
448 hash ^= u64::from(byte);
449 hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
450 }
451 u32::try_from(hash & 0xffff_ffff).unwrap_or(0)
452}
453
454#[cfg(test)]
455mod tests {
456 use super::*;
457 use crate::cluster::peer::PeerEndpoint;
458 use crate::cluster::pool::PoolConfig;
459
460 fn small_pool() -> Arc<ServerPool> {
461 let cfg = PoolConfig {
462 dc: "dc1".into(),
463 rack: "r1".into(),
464 ..PoolConfig::default()
465 };
466 let local = Peer::new(
467 0,
468 PeerEndpoint::tcp("127.0.0.1".into(), 8101),
469 "r1".into(),
470 "dc1".into(),
471 vec![DynToken::from_u32(0)],
472 true,
473 true,
474 false,
475 );
476 let remote = Peer::new(
477 1,
478 PeerEndpoint::tcp("127.0.0.1".into(), 8102),
479 "r1".into(),
480 "dc1".into(),
481 vec![DynToken::from_u32(2_147_483_648)],
482 false,
483 true,
484 false,
485 );
486 Arc::new(ServerPool::new(cfg, vec![local, remote]))
487 }
488
489 #[test]
490 fn list_peers_reports_every_peer() {
491 let admin = PoolClusterAdmin::new(small_pool());
492 let snaps = admin.list_peers();
493 assert_eq!(snaps.len(), 2);
494 let local = snaps.iter().find(|s| s.is_local).expect("local snapshot");
495 assert_eq!(local.idx, 0);
496 assert_eq!(local.port, 8101);
497 let remote = snaps.iter().find(|s| !s.is_local).expect("remote snapshot");
498 assert_eq!(remote.idx, 1);
499 assert_eq!(remote.tokens, vec![2_147_483_648]);
500 }
501
502 #[test]
503 fn join_stages_and_commit_appends_peer() {
504 let admin = PoolClusterAdmin::new(small_pool());
505 let target: SocketAddr = "127.0.0.1:8103".parse().unwrap();
506 let plan = admin.cluster_join(target).expect("plan");
507 assert_eq!(plan.change.kind, ClusterChangeKind::Add);
508 assert_eq!(admin.cluster_plan_pending().len(), 1);
509 assert_eq!(admin.list_peers().len(), 2);
511 admin.cluster_commit().expect("commit");
512 assert_eq!(admin.cluster_plan_pending().len(), 0);
513 let snaps = admin.list_peers();
514 assert_eq!(snaps.len(), 3);
515 assert!(snaps.iter().any(|s| s.port == 8103));
516 }
517
518 #[test]
519 fn join_rejects_duplicate_endpoint() {
520 let admin = PoolClusterAdmin::new(small_pool());
521 let target: SocketAddr = "127.0.0.1:8102".parse().unwrap();
522 let err = admin.cluster_join(target).expect_err("duplicate");
523 assert!(matches!(err, ClusterError::PeerAlreadyExists { .. }));
524 }
525
526 #[test]
527 fn join_rejects_duplicate_in_staging() {
528 let admin = PoolClusterAdmin::new(small_pool());
529 let target: SocketAddr = "127.0.0.1:8200".parse().unwrap();
530 admin.cluster_join(target).expect("first");
531 let err = admin.cluster_join(target).expect_err("staged dup");
532 assert!(matches!(err, ClusterError::PeerAlreadyExists { .. }));
533 }
534
535 #[test]
536 fn leave_stages_and_commit_removes_peer() {
537 let admin = PoolClusterAdmin::new(small_pool());
538 let plan = admin.cluster_leave(1).expect("plan");
539 assert_eq!(plan.change.kind, ClusterChangeKind::Remove);
540 assert_eq!(plan.change.peer_idx, Some(1));
541 assert_eq!(admin.list_peers().len(), 2);
543 admin.cluster_commit().expect("commit");
544 let snaps = admin.list_peers();
545 assert_eq!(snaps.len(), 1);
546 assert_eq!(snaps[0].idx, 0);
547 }
548
549 #[test]
550 fn leave_rejects_unknown_idx() {
551 let admin = PoolClusterAdmin::new(small_pool());
552 let err = admin.cluster_leave(99).expect_err("unknown");
553 assert!(matches!(err, ClusterError::PeerNotFound { idx: 99 }));
554 }
555
556 #[test]
557 fn leave_rejects_local_peer() {
558 let admin = PoolClusterAdmin::new(small_pool());
559 let err = admin.cluster_leave(0).expect_err("local");
560 assert!(matches!(err, ClusterError::CannotRemoveLocal));
561 }
562
563 #[test]
564 fn commit_with_empty_staging_is_ok() {
565 let admin = PoolClusterAdmin::new(small_pool());
566 admin.cluster_commit().expect("noop commit");
567 assert_eq!(admin.list_peers().len(), 2);
568 }
569
570 #[test]
571 fn commit_applies_mixed_batch_in_order() {
572 let admin = PoolClusterAdmin::new(small_pool());
573 admin.cluster_leave(1).expect("stage leave");
574 let target: SocketAddr = "10.0.0.1:8101".parse().unwrap();
575 admin.cluster_join(target).expect("stage join");
576 assert_eq!(admin.cluster_plan_pending().len(), 2);
577 admin.cluster_commit().expect("commit");
578 let snaps = admin.list_peers();
579 assert_eq!(snaps.len(), 2);
580 let new = snaps.iter().find(|s| s.host == "10.0.0.1").expect("added");
584 assert_eq!(new.port, 8101);
585 assert!(!snaps
587 .iter()
588 .any(|s| s.host == "127.0.0.1" && s.port == 8102));
589 }
590
591 #[test]
592 fn noop_admin_returns_empty_and_errors_on_mutations() {
593 let admin = NoopClusterAdmin;
594 assert!(admin.list_peers().is_empty());
595 assert!(admin.cluster_plan_pending().is_empty());
596 admin.cluster_commit().expect("noop commit");
597 let target: SocketAddr = "127.0.0.1:1".parse().unwrap();
598 assert!(matches!(
599 admin.cluster_join(target),
600 Err(ClusterError::Invalid(_))
601 ));
602 assert!(matches!(
603 admin.cluster_leave(0),
604 Err(ClusterError::Invalid(_))
605 ));
606 }
607
608 #[test]
609 fn join_inherits_dc_rack_from_config_when_no_local_peer() {
610 let cfg = PoolConfig {
614 dc: "dcX".into(),
615 rack: "rX".into(),
616 ..PoolConfig::default()
617 };
618 let remote = Peer::new(
619 0,
620 PeerEndpoint::tcp("127.0.0.1".into(), 9001),
621 "rX".into(),
622 "dcX".into(),
623 vec![DynToken::from_u32(0)],
624 false,
625 true,
626 false,
627 );
628 let pool = Arc::new(ServerPool::new(cfg, vec![remote]));
629 let admin = PoolClusterAdmin::new(pool);
630 let target: SocketAddr = "127.0.0.1:9002".parse().unwrap();
631 let plan = admin.cluster_join(target).expect("plan");
632 let spec = plan.change.peer.expect("spec");
633 assert_eq!(spec.dc, "dcX");
634 assert_eq!(spec.rack, "rX");
635 }
636
637 #[test]
638 fn commit_rejects_add_for_endpoint_already_present() {
639 let admin = PoolClusterAdmin::new(small_pool());
644 let dup = PeerSpec {
645 host: "127.0.0.1".into(),
646 port: 8102,
647 dc: "dc1".into(),
648 rack: "r1".into(),
649 tokens: vec![1],
650 is_secure: false,
651 };
652 admin.staged.lock().push(ClusterChange {
653 kind: ClusterChangeKind::Add,
654 peer_idx: None,
655 peer: Some(dup),
656 });
657 let err = admin.cluster_commit().expect_err("commit dup add");
658 assert!(matches!(err, ClusterError::PeerAlreadyExists { .. }));
659 }
660
661 #[test]
662 fn commit_rejects_remove_of_local_peer() {
663 let admin = PoolClusterAdmin::new(small_pool());
666 admin.staged.lock().push(ClusterChange {
667 kind: ClusterChangeKind::Remove,
668 peer_idx: Some(0),
669 peer: None,
670 });
671 let err = admin.cluster_commit().expect_err("commit remove local");
672 assert!(matches!(err, ClusterError::CannotRemoveLocal));
673 }
674
675 #[test]
676 fn derive_token_is_stable_per_endpoint() {
677 let a = derive_token("10.0.0.1", 8101);
678 let b = derive_token("10.0.0.1", 8101);
679 let c = derive_token("10.0.0.2", 8101);
680 assert_eq!(a, b);
681 assert_ne!(a, c);
682 }
683
684 #[test]
685 fn join_rebuilds_ring() {
686 let admin = PoolClusterAdmin::new(small_pool());
689 let target: SocketAddr = "10.0.0.5:8101".parse().unwrap();
690 admin.cluster_join(target).expect("plan");
691 admin.cluster_commit().expect("commit");
692 let pool = admin.pool();
693 let topology = pool.datacenters().read();
694 let dc1 = topology.iter().find(|d| d.name() == "dc1").expect("dc1");
695 let r1 = dc1.racks().iter().find(|r| r.name() == "r1").expect("r1");
696 let entries = r1.continuums();
700 let mut idxs: Vec<u32> = entries.iter().map(|e| e.peer_idx).collect();
701 idxs.sort_unstable();
702 idxs.dedup();
703 assert_eq!(idxs.len(), 3);
704 }
705}