use super::*;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ClusterDiscoveryEvent {
CandidateSeen(ClusterCandidate),
MemberLive(ClusterNodeId),
MemberLeaving {
node_id: ClusterNodeId,
generation: ClusterGeneration,
role: ClusterRole,
},
MemberSuspect(ClusterNodeId),
MemberDead(ClusterNodeId),
}
#[async_trait::async_trait]
pub trait ClusterDiscovery: fmt::Debug + Send + Sync {
async fn announce(&self, candidate: ClusterCandidate) -> Result<()>;
async fn mark_live(&self, node_id: ClusterNodeId) -> Result<()>;
async fn mark_suspect(&self, node_id: ClusterNodeId) -> Result<()>;
async fn mark_dead(&self, node_id: ClusterNodeId) -> Result<()>;
fn candidates(&self) -> Vec<ClusterCandidate>;
fn events(&self) -> Vec<ClusterDiscoveryEvent>;
}
#[derive(Debug, Default)]
struct InMemoryClusterDiscoveryState {
candidates: BTreeMap<ClusterNodeId, ClusterCandidate>,
events: Vec<ClusterDiscoveryEvent>,
}
#[derive(Debug, Default)]
pub struct InMemoryClusterDiscovery {
state: Mutex<InMemoryClusterDiscoveryState>,
}
impl InMemoryClusterDiscovery {
pub fn new() -> Self {
Self::default()
}
pub fn announce(&self, candidate: ClusterCandidate) {
let mut state = self.state.lock().expect("cluster discovery poisoned");
state
.candidates
.insert(candidate.node_id.clone(), candidate.clone());
state
.events
.push(ClusterDiscoveryEvent::CandidateSeen(candidate));
}
pub fn mark_live(&self, node_id: impl Into<ClusterNodeId>) {
self.push_liveness(ClusterDiscoveryEvent::MemberLive(node_id.into()));
}
pub fn mark_suspect(&self, node_id: impl Into<ClusterNodeId>) {
self.push_liveness(ClusterDiscoveryEvent::MemberSuspect(node_id.into()));
}
pub fn mark_dead(&self, node_id: impl Into<ClusterNodeId>) {
self.push_liveness(ClusterDiscoveryEvent::MemberDead(node_id.into()));
}
fn push_liveness(&self, event: ClusterDiscoveryEvent) {
self.state
.lock()
.expect("cluster discovery poisoned")
.events
.push(event);
}
pub fn candidates(&self) -> Vec<ClusterCandidate> {
self.state
.lock()
.expect("cluster discovery poisoned")
.candidates
.values()
.cloned()
.collect()
}
pub fn events(&self) -> Vec<ClusterDiscoveryEvent> {
self.state
.lock()
.expect("cluster discovery poisoned")
.events
.clone()
}
}
#[async_trait::async_trait]
impl ClusterDiscovery for InMemoryClusterDiscovery {
async fn announce(&self, candidate: ClusterCandidate) -> Result<()> {
InMemoryClusterDiscovery::announce(self, candidate);
Ok(())
}
async fn mark_live(&self, node_id: ClusterNodeId) -> Result<()> {
InMemoryClusterDiscovery::mark_live(self, node_id);
Ok(())
}
async fn mark_suspect(&self, node_id: ClusterNodeId) -> Result<()> {
InMemoryClusterDiscovery::mark_suspect(self, node_id);
Ok(())
}
async fn mark_dead(&self, node_id: ClusterNodeId) -> Result<()> {
InMemoryClusterDiscovery::mark_dead(self, node_id);
Ok(())
}
fn candidates(&self) -> Vec<ClusterCandidate> {
InMemoryClusterDiscovery::candidates(self)
}
fn events(&self) -> Vec<ClusterDiscoveryEvent> {
InMemoryClusterDiscovery::events(self)
}
}
#[derive(Debug)]
pub struct ChitchatStyleDiscovery {
seeds: Vec<String>,
inner: InMemoryClusterDiscovery,
}
impl ChitchatStyleDiscovery {
pub fn new<I, S>(seeds: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
Self {
seeds: seeds.into_iter().map(Into::into).collect(),
inner: InMemoryClusterDiscovery::new(),
}
}
pub fn seeds(&self) -> &[String] {
&self.seeds
}
pub fn seed_count(&self) -> usize {
self.seeds.len()
}
pub fn has_seeds(&self) -> bool {
!self.seeds.is_empty()
}
pub fn adapter_name(&self) -> &'static str {
"chitchat-style"
}
pub fn announce(&self, mut candidate: ClusterCandidate) {
candidate
.metadata
.entry("discovery.adapter".to_owned())
.or_insert_with(|| self.adapter_name().to_owned());
if self.has_seeds() {
candidate
.metadata
.entry("discovery.seeds".to_owned())
.or_insert_with(|| self.seeds.join(","));
}
self.inner.announce(candidate);
}
pub fn mark_live(&self, node_id: impl Into<ClusterNodeId>) {
self.inner.mark_live(node_id);
}
pub fn mark_suspect(&self, node_id: impl Into<ClusterNodeId>) {
self.inner.mark_suspect(node_id);
}
pub fn mark_dead(&self, node_id: impl Into<ClusterNodeId>) {
self.inner.mark_dead(node_id);
}
pub fn candidates(&self) -> Vec<ClusterCandidate> {
self.inner.candidates()
}
pub fn events(&self) -> Vec<ClusterDiscoveryEvent> {
self.inner.events()
}
}
impl Default for ChitchatStyleDiscovery {
fn default() -> Self {
Self::new(std::iter::empty::<String>())
}
}
#[async_trait::async_trait]
impl ClusterDiscovery for ChitchatStyleDiscovery {
async fn announce(&self, candidate: ClusterCandidate) -> Result<()> {
ChitchatStyleDiscovery::announce(self, candidate);
Ok(())
}
async fn mark_live(&self, node_id: ClusterNodeId) -> Result<()> {
ChitchatStyleDiscovery::mark_live(self, node_id);
Ok(())
}
async fn mark_suspect(&self, node_id: ClusterNodeId) -> Result<()> {
ChitchatStyleDiscovery::mark_suspect(self, node_id);
Ok(())
}
async fn mark_dead(&self, node_id: ClusterNodeId) -> Result<()> {
ChitchatStyleDiscovery::mark_dead(self, node_id);
Ok(())
}
fn candidates(&self) -> Vec<ClusterCandidate> {
ChitchatStyleDiscovery::candidates(self)
}
fn events(&self) -> Vec<ClusterDiscoveryEvent> {
ChitchatStyleDiscovery::events(self)
}
}