use serde::Serialize;
use thiserror::Error;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum UpgradeStrategy {
InheritedSocket,
ReusePort,
}
impl UpgradeStrategy {
pub fn platform_default() -> Self {
if cfg!(windows) {
Self::ReusePort
} else {
Self::InheritedSocket
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum UpgradePhase {
Prepared,
NewReady,
OldDraining,
Complete,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpgradePlan {
generation: u64,
member_id: String,
strategy: UpgradeStrategy,
}
impl UpgradePlan {
pub fn new(generation: u64, member_id: impl Into<String>) -> Self {
Self {
generation,
member_id: member_id.into(),
strategy: UpgradeStrategy::platform_default(),
}
}
pub fn with_strategy(mut self, strategy: UpgradeStrategy) -> Self {
self.strategy = strategy;
self
}
pub fn prepare(self) -> Result<GracefulUpgrade, UpgradeError> {
if self.generation == 0 {
return Err(UpgradeError::InvalidGeneration);
}
if self.member_id.trim().is_empty() {
return Err(UpgradeError::MissingMemberId);
}
Ok(GracefulUpgrade {
plan: self,
phase: UpgradePhase::Prepared,
old_accepting: true,
new_ready: false,
in_flight: 0,
completed: 0,
})
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GracefulUpgrade {
plan: UpgradePlan,
phase: UpgradePhase,
old_accepting: bool,
new_ready: bool,
in_flight: usize,
completed: usize,
}
impl GracefulUpgrade {
pub fn mark_new_ready(&mut self) {
self.new_ready = true;
self.phase = UpgradePhase::NewReady;
}
pub fn start_draining_old(&mut self) -> Result<(), UpgradeError> {
if !self.new_ready {
return Err(UpgradeError::ReplacementNotReady);
}
self.old_accepting = false;
self.phase = UpgradePhase::OldDraining;
Ok(())
}
pub fn record_request(&mut self) -> bool {
if !self.old_accepting {
return false;
}
self.in_flight = self.in_flight.saturating_add(1);
true
}
pub fn finish_request(&mut self) {
if self.in_flight > 0 {
self.in_flight -= 1;
self.completed = self.completed.saturating_add(1);
}
}
pub fn complete(mut self) -> Result<UpgradeReport, UpgradeError> {
if !self.new_ready {
return Err(UpgradeError::ReplacementNotReady);
}
if self.in_flight > 0 {
return Err(UpgradeError::InFlightRequestsRemaining(self.in_flight));
}
self.phase = UpgradePhase::Complete;
Ok(UpgradeReport {
generation: self.plan.generation,
member_id: self.plan.member_id,
strategy: self.plan.strategy,
phase: self.phase,
completed_requests: self.completed,
dropped_requests: 0,
})
}
pub fn phase(&self) -> UpgradePhase {
self.phase
}
pub fn membership_stable(&self) -> bool {
!self.plan.member_id.trim().is_empty()
}
pub fn in_flight(&self) -> usize {
self.in_flight
}
pub fn old_accepting(&self) -> bool {
self.old_accepting
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct UpgradeReport {
pub generation: u64,
pub member_id: String,
pub strategy: UpgradeStrategy,
pub phase: UpgradePhase,
pub completed_requests: usize,
pub dropped_requests: usize,
}
#[derive(Debug, Error, PartialEq, Eq)]
pub enum UpgradeError {
#[error("upgrade generation must be greater than zero")]
InvalidGeneration,
#[error("upgrade requires a non-empty member id")]
MissingMemberId,
#[error("replacement process is not ready")]
ReplacementNotReady,
#[error("{0} in-flight request(s) remain during upgrade")]
InFlightRequestsRemaining(usize),
}