use std::fmt;
use serde::{Deserialize, Serialize};
use crate::cluster::{ClusterEpoch, ClusterNodeId, PartitionId};
use crate::grid::elasticity::{
validate_move_preserves_zone_quorum, MovePhase, NodeTopology, PartitionMove, RegionId,
ReshardPlan, ReshardPlanError, UpgradeGuard, UpgradeGuardError, UpgradeStep,
ZoneAwareReplicaSet,
};
use crate::grid::ReplicationConfig;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ScaleRecommendation {
Hold,
ScaleOut {
suggested: usize,
},
ScaleIn {
drain: Vec<ClusterNodeId>,
},
Rebalance,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct CapacitySample {
pub region: RegionId,
pub memory_pressure: f32,
pub replication_lag: u64,
pub hot_partition_skew: f32,
pub repair_debt: u64,
pub seconds_since_last_scale: u64,
pub scale_in_candidates: Vec<ClusterNodeId>,
}
impl CapacitySample {
pub fn new(region: impl Into<RegionId>) -> Self {
Self {
region: region.into(),
memory_pressure: 0.0,
replication_lag: 0,
hot_partition_skew: 0.0,
repair_debt: 0,
seconds_since_last_scale: u64::MAX,
scale_in_candidates: Vec::new(),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct CapacityThresholds {
pub scale_out_memory_pressure: f32,
pub scale_in_memory_pressure: f32,
pub replication_lag_limit: u64,
pub hot_partition_skew_limit: f32,
pub repair_debt_limit: u64,
pub minimum_dwell_secs: u64,
pub scale_out_suggested: usize,
}
impl Default for CapacityThresholds {
fn default() -> Self {
Self {
scale_out_memory_pressure: 0.85,
scale_in_memory_pressure: 0.25,
replication_lag_limit: 1_000,
hot_partition_skew_limit: 2.0,
repair_debt_limit: 100,
minimum_dwell_secs: 300,
scale_out_suggested: 1,
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct CapacitySignal {
pub region: RegionId,
pub memory_pressure: f32,
pub replication_lag: u64,
pub hot_partition_skew: f32,
pub repair_debt: u64,
pub recommendation: ScaleRecommendation,
}
pub fn evaluate_capacity(sample: CapacitySample, thresholds: CapacityThresholds) -> CapacitySignal {
let in_dwell_window = sample.seconds_since_last_scale < thresholds.minimum_dwell_secs;
let recommendation = if in_dwell_window {
ScaleRecommendation::Hold
} else if sample.memory_pressure >= thresholds.scale_out_memory_pressure
|| sample.replication_lag > thresholds.replication_lag_limit
{
ScaleRecommendation::ScaleOut {
suggested: thresholds.scale_out_suggested.max(1),
}
} else if sample.hot_partition_skew >= thresholds.hot_partition_skew_limit
|| sample.repair_debt > thresholds.repair_debt_limit
{
ScaleRecommendation::Rebalance
} else if sample.memory_pressure <= thresholds.scale_in_memory_pressure
&& !sample.scale_in_candidates.is_empty()
{
ScaleRecommendation::ScaleIn {
drain: sample.scale_in_candidates.clone(),
}
} else {
ScaleRecommendation::Hold
};
CapacitySignal {
region: sample.region,
memory_pressure: sample.memory_pressure,
replication_lag: sample.replication_lag,
hot_partition_skew: sample.hot_partition_skew,
repair_debt: sample.repair_debt,
recommendation,
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AutoscalerIntent {
ScaleOut {
node: ClusterNodeId,
topology: NodeTopology,
candidate: ZoneAwareReplicaSet,
backfill_sources: Vec<(PartitionId, ClusterNodeId, u64)>,
compat: UpgradeStep,
},
ScaleIn {
drain: ClusterNodeId,
remaining_voters: usize,
drain_targets: Vec<(PartitionId, ClusterNodeId, u64)>,
compat: UpgradeStep,
},
Rebalance {
plan: ReshardPlan,
compat: UpgradeStep,
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct AutoscalerAdmissionPolicy {
pub epoch: ClusterEpoch,
pub replication: ReplicationConfig,
pub upgrade_guard: UpgradeGuard,
pub max_concurrent_moves: usize,
}
impl AutoscalerAdmissionPolicy {
pub fn new(
epoch: ClusterEpoch,
replication: ReplicationConfig,
upgrade_guard: UpgradeGuard,
max_concurrent_moves: usize,
) -> Self {
Self {
epoch,
replication,
upgrade_guard,
max_concurrent_moves: max_concurrent_moves.max(1),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ScaleAction {
ScaleOut,
ScaleIn,
Rebalance,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AutoscalerAdmission {
pub action: ScaleAction,
pub plan: ReshardPlan,
pub quorum_eligible: bool,
pub removal_allowed: bool,
}
pub fn admit_autoscaler_intent(
intent: AutoscalerIntent,
policy: AutoscalerAdmissionPolicy,
) -> Result<AutoscalerAdmission, AutoscalerIntentError> {
policy
.replication
.validate()
.map_err(|error| AutoscalerIntentError::new(error.to_string()))?;
match intent {
AutoscalerIntent::ScaleOut {
node,
topology: _,
candidate,
backfill_sources,
compat,
} => {
policy
.upgrade_guard
.check(compat)
.map_err(AutoscalerIntentError::from)?;
validate_move_preserves_zone_quorum(&candidate, policy.replication.write_quorum)
.map_err(AutoscalerIntentError::from)?;
let moves = backfill_sources
.into_iter()
.map(|(partition, from, bytes)| {
PartitionMove::new(partition, from, node.clone(), bytes)
})
.collect();
Ok(AutoscalerAdmission {
action: ScaleAction::ScaleOut,
plan: ReshardPlan::new(policy.epoch, moves, policy.max_concurrent_moves),
quorum_eligible: false,
removal_allowed: false,
})
}
AutoscalerIntent::ScaleIn {
drain,
remaining_voters,
drain_targets,
compat,
} => {
policy
.upgrade_guard
.check(compat)
.map_err(AutoscalerIntentError::from)?;
if remaining_voters < policy.replication.write_quorum {
return Err(AutoscalerIntentError::new(
"autoscaler intent would break write quorum",
));
}
Ok(AutoscalerAdmission {
action: ScaleAction::ScaleIn,
plan: ReshardPlan::drain_node(
policy.epoch,
drain,
drain_targets,
policy.max_concurrent_moves,
),
quorum_eligible: true,
removal_allowed: false,
})
}
AutoscalerIntent::Rebalance { plan, compat } => {
policy
.upgrade_guard
.check(compat)
.map_err(AutoscalerIntentError::from)?;
Ok(AutoscalerAdmission {
action: ScaleAction::Rebalance,
plan,
quorum_eligible: true,
removal_allowed: true,
})
}
}
}
pub fn scale_out_counts_toward_quorum(plan: &ReshardPlan) -> bool {
!plan.moves.is_empty()
&& plan
.moves
.iter()
.all(|movement| movement.phase >= MovePhase::Commit)
}
pub fn scale_in_removal_allowed(plan: &ReshardPlan) -> bool {
!plan.moves.is_empty()
&& plan
.moves
.iter()
.all(|movement| movement.phase == MovePhase::Cleanup)
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AutoscalerIntentError {
message: String,
}
impl AutoscalerIntentError {
fn new(message: impl Into<String>) -> Self {
Self {
message: message.into(),
}
}
}
impl From<UpgradeGuardError> for AutoscalerIntentError {
fn from(error: UpgradeGuardError) -> Self {
Self::new(error.to_string())
}
}
impl From<ReshardPlanError> for AutoscalerIntentError {
fn from(error: ReshardPlanError) -> Self {
Self::new(error.to_string())
}
}
impl fmt::Display for AutoscalerIntentError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.message)
}
}
impl std::error::Error for AutoscalerIntentError {}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CapacityAutoscalerMetrics {
pub capacity_recommendation: ScaleRecommendation,
pub scale_actions_total: u64,
}