diff --git a/Makefile.toml b/Makefile.toml index cc95794396..bf6ce76672 100644 --- a/Makefile.toml +++ b/Makefile.toml @@ -668,6 +668,7 @@ cargo-test-r run --package golem-service-base --test '*' -- --nocapture --report cargo-test-r run --package golem-registry-service --test '*' -- --nocapture --report-time $JUNIT_OPTS cargo-test-r run --package golem-worker-service --test '*' -- --nocapture --report-time $JUNIT_OPTS RUST_LOG=debug cargo-test-r run --package golem-debugging-service --test '*' -- --report-time $JUNIT_OPTS +cargo-test-r run --package golem-shard-manager --test integration -- --nocapture --report-time $JUNIT_OPTS ''' [tasks.integration-tests-group6] diff --git a/golem-shard-manager/config/shard-manager.sample.env b/golem-shard-manager/config/shard-manager.sample.env index 383004f127..dc918c3f8d 100644 --- a/golem-shard-manager/config/shard-manager.sample.env +++ b/golem-shard-manager/config/shard-manager.sample.env @@ -4,6 +4,7 @@ GOLEM__HTTP_PORT=8081 GOLEM__NUMBER_OF_SHARDS=1024 GOLEM__REBALANCE_THRESHOLD=0.1 GOLEM__RUNTIME_METRICS_SAMPLING_INTERVAL="5s" +GOLEM__SHARD_LEASE_DURATION="1m" GOLEM__DB__TYPE="Sqlite" GOLEM__DB__CONFIG__DATABASE="golem_shard_manager.db" GOLEM__DB__CONFIG__FOREIGN_KEYS=false @@ -98,6 +99,7 @@ GOLEM__HTTP_PORT=8081 GOLEM__NUMBER_OF_SHARDS=1024 GOLEM__REBALANCE_THRESHOLD=0.1 GOLEM__RUNTIME_METRICS_SAMPLING_INTERVAL="5s" +GOLEM__SHARD_LEASE_DURATION="1m" GOLEM__DB__TYPE="Sqlite" GOLEM__DB__CONFIG__DATABASE="golem_shard_manager.db" GOLEM__DB__CONFIG__FOREIGN_KEYS=false diff --git a/golem-shard-manager/config/shard-manager.toml b/golem-shard-manager/config/shard-manager.toml index c08176c617..1efbb99748 100644 --- a/golem-shard-manager/config/shard-manager.toml +++ b/golem-shard-manager/config/shard-manager.toml @@ -3,6 +3,7 @@ http_port = 8081 number_of_shards = 1024 rebalance_threshold = 0.1 runtime_metrics_sampling_interval = "5s" +shard_lease_duration = "1m" [db] type = "Sqlite" @@ -144,6 +145,7 @@ type = "Disabled" # number_of_shards = 1024 # rebalance_threshold = 0.1 # runtime_metrics_sampling_interval = "5s" +# shard_lease_duration = "1m" # # [db] # type = "Sqlite" diff --git a/golem-shard-manager/db/migration/postgres/003_shard_lease_state.sql b/golem-shard-manager/db/migration/postgres/003_shard_lease_state.sql new file mode 100644 index 0000000000..fd2ff62a8a --- /dev/null +++ b/golem-shard-manager/db/migration/postgres/003_shard_lease_state.sql @@ -0,0 +1,3 @@ +-- The persisted shard manager state changed shape +-- Any state written by an earlier version is dropped; it is rebuilt as executors register. +DELETE FROM shard_manager_state; diff --git a/golem-shard-manager/db/migration/sqlite/003_shard_lease_state.sql b/golem-shard-manager/db/migration/sqlite/003_shard_lease_state.sql new file mode 100644 index 0000000000..464b8123ad --- /dev/null +++ b/golem-shard-manager/db/migration/sqlite/003_shard_lease_state.sql @@ -0,0 +1,3 @@ +-- The persisted shard manager state changed shape +-- Any state written by an earlier version is dropped; it is rebuilt as executors register. +DELETE FROM shard_manager_state; \ No newline at end of file diff --git a/golem-shard-manager/src/config.rs b/golem-shard-manager/src/config.rs index 4ab3ba8b16..623e783ce6 100644 --- a/golem-shard-manager/src/config.rs +++ b/golem-shard-manager/src/config.rs @@ -40,6 +40,8 @@ pub struct ShardManagerConfig { pub grpc: GrpcApiConfig, pub number_of_shards: usize, pub rebalance_threshold: f64, + #[serde(with = "humantime_serde")] + pub shard_lease_duration: Duration, pub registry_service: GrpcRegistryServiceConfig, pub resource_definition_fetcher: ResourceDefinitionFetcherConfig, pub quota: QuotaServiceConfig, @@ -80,6 +82,11 @@ impl SafeDisplay for ShardManagerConfig { "rebalance threshold: {}", self.rebalance_threshold ); + let _ = writeln!( + &mut result, + "shard lease duration: {:?}", + self.shard_lease_duration + ); let _ = writeln!(&mut result, "registry service:"); let _ = writeln!( &mut result, @@ -113,6 +120,7 @@ impl Default for ShardManagerConfig { grpc: GrpcApiConfig::default(), number_of_shards: 1024, rebalance_threshold: 0.1, + shard_lease_duration: Duration::from_secs(60), registry_service: GrpcRegistryServiceConfig::default(), resource_definition_fetcher: ResourceDefinitionFetcherConfig::default(), quota: QuotaServiceConfig::default(), diff --git a/golem-shard-manager/src/error.rs b/golem-shard-manager/src/error.rs index 19b1c2d478..f40ab0304c 100644 --- a/golem-shard-manager/src/error.rs +++ b/golem-shard-manager/src/error.rs @@ -85,6 +85,11 @@ impl From for golem::shardmanager::v1::ShardManagerError { err.to_string(), api::error_code::INTERNAL_FILESYSTEM_ERROR, ), + ShardManagerError::Internal(details) => error( + shard_manager_error::Error::Unknown, + details, + api::error_code::INTERNAL_UNKNOWN, + ), } } } diff --git a/golem-shard-manager/src/grpc.rs b/golem-shard-manager/src/grpc.rs index 9deac1aa1a..b5c8367cda 100644 --- a/golem-shard-manager/src/grpc.rs +++ b/golem-shard-manager/src/grpc.rs @@ -12,9 +12,10 @@ // See the License for the specific language governing permissions and // limitations under the License. -use crate::RoutingTable; +use crate::ShardLeaseState; use crate::error::ShardManagerTraceErrorKind; use crate::quota::QuotaService; +use crate::sharding::ExecutorAddr; use crate::sharding::error::ShardManagerError; use crate::sharding::shard_management::ShardManagement; use golem_api_grpc::proto::golem; @@ -40,10 +41,10 @@ impl ShardManagerServiceImpl { } } - async fn get_routing_table_internal(&self) -> RoutingTable { - let routing_table = self.shard_management.current_snapshot().await; - debug!("Providing routing table: {}", routing_table); - routing_table + async fn get_routing_table_internal(&self) -> ShardLeaseState { + let shard_state = self.shard_management.current_snapshot().await; + debug!("Providing routing table: {}", shard_state); + shard_state } async fn register_internal( @@ -51,8 +52,12 @@ impl ShardManagerServiceImpl { pod: Pod, pod_name: Option, ) -> Result<(), ShardManagerError> { - debug!("Received request to register pod: {}", pod); - self.shard_management.register_pod(pod, pod_name).await; + debug!("Received request to register executor at: {}", pod); + let executor_id = self + .shard_management + .register_executor(ExecutorAddr::from(pod), pod_name) + .await; + debug!(executor_id = %executor_id, addr = %pod, "Registered executor"); Ok(()) } } diff --git a/golem-shard-manager/src/lib.rs b/golem-shard-manager/src/lib.rs index a3d1d8dc82..d4c4313abe 100644 --- a/golem-shard-manager/src/lib.rs +++ b/golem-shard-manager/src/lib.rs @@ -40,7 +40,10 @@ pub use sharding::healthcheck::HealthCheck; pub use sharding::persistence::{DbRoutingTablePersistence, RoutingTablePersistence}; pub use sharding::shard_management::ShardManagement; pub use sharding::worker_executor::WorkerExecutorService; -pub use sharding::{PodState, RoutingTable, RoutingTableEntry}; +pub use sharding::{ + ExecutorAddr, ExecutorAddrs, ExecutorId, ExecutorLease, ExecutorShards, ShardAssignmentEntry, + ShardEpoch, ShardLeaseRevision, ShardLeaseState, +}; use std::net::{Ipv4Addr, SocketAddrV4}; use std::sync::Arc; use tokio::net::TcpListener; @@ -70,6 +73,16 @@ pub async fn run( ) -> anyhow::Result { debug!("Initializing shard manager"); + anyhow::ensure!( + !shard_manager_config.shard_lease_duration.is_zero(), + "shard_lease_duration must be greater than zero" + ); + anyhow::ensure!( + chrono::Duration::from_std(shard_manager_config.shard_lease_duration).is_ok(), + "shard_lease_duration {:?} is out of range", + shard_manager_config.shard_lease_duration + ); + let (health_reporter, health_service) = tonic_health::server::health_reporter(); health_reporter .set_serving::>() @@ -190,6 +203,7 @@ pub async fn run( worker_executors.clone(), health_check.clone(), shard_manager_config.rebalance_threshold, + shard_manager_config.shard_lease_duration, join_set, ) .await?, diff --git a/golem-shard-manager/src/sharding/error.rs b/golem-shard-manager/src/sharding/error.rs index a089fb9b6e..289b441b9e 100644 --- a/golem-shard-manager/src/sharding/error.rs +++ b/golem-shard-manager/src/sharding/error.rs @@ -39,6 +39,8 @@ pub enum ShardManagerError { MigrationError(#[from] anyhow::Error), #[error("IO error {0}")] IoError(#[from] std::io::Error), + #[error("Internal error: {0}")] + Internal(String), } impl IsRetriableError for ShardManagerError { @@ -54,6 +56,7 @@ impl IsRetriableError for ShardManagerError { ShardManagerError::RepoError(_) => false, ShardManagerError::MigrationError(_) => false, ShardManagerError::IoError(_) => false, + ShardManagerError::Internal(_) => false, } } diff --git a/golem-shard-manager/src/sharding/healthcheck.rs b/golem-shard-manager/src/sharding/healthcheck.rs index 52899b7b98..0489ac7279 100644 --- a/golem-shard-manager/src/sharding/healthcheck.rs +++ b/golem-shard-manager/src/sharding/healthcheck.rs @@ -13,6 +13,7 @@ // limitations under the License. use super::error::HealthCheckError; +use super::model::{ExecutorAddr, ExecutorId}; use super::worker_executor::WorkerExecutorService; use async_trait::async_trait; use golem_common::model::{Pod, RetryConfig}; @@ -28,19 +29,22 @@ pub trait HealthCheck: Send + Sync { async fn health_check(&self, pod: Pod, pod_name: Option) -> bool; } -/// Executes healthcheck on all the given worker executors, and returns a set of unhealthy ones -pub async fn get_unhealthy_pods( +/// Executes healthcheck on all the given worker executors, and returns the set of unhealthy ones +pub async fn get_unhealthy_executors( health_check: &Arc, - pods: &[(Pod, Option)], -) -> HashSet { - let futures: Vec<_> = pods + executors: &[(ExecutorId, ExecutorAddr, Option)], +) -> HashSet { + let futures: Vec<_> = executors .iter() - .map(|(pod, pod_name)| { + .map(|(executor_id, addr, pod_name)| { let health_check = health_check.clone(); Box::pin(async move { - match health_check.health_check(*pod, pod_name.clone()).await { + match health_check + .health_check(Pod::from(*addr), pod_name.clone()) + .await + { true => None, - false => Some(*pod), + false => Some(*executor_id), } }) }) diff --git a/golem-shard-manager/src/sharding/healthcheck_loop.rs b/golem-shard-manager/src/sharding/healthcheck_loop.rs index b2c71332f6..a3107beaae 100644 --- a/golem-shard-manager/src/sharding/healthcheck_loop.rs +++ b/golem-shard-manager/src/sharding/healthcheck_loop.rs @@ -15,7 +15,7 @@ use super::healthcheck::HealthCheck; use super::shard_management::ShardManagement; use crate::config::HealthCheckConfig; -use crate::sharding::healthcheck::get_unhealthy_pods; +use crate::sharding::healthcheck::get_unhealthy_executors; use std::sync::Arc; use tokio::task::JoinSet; use tracing::{Instrument, debug, warn}; @@ -44,20 +44,26 @@ async fn run_health_check( health_check: &Arc, ) { debug!("Scheduled to conduct health check"); - let routing_table = shard_management.current_snapshot().await; - debug!("Checking health of registered pods..."); - let failed_pods = get_unhealthy_pods(health_check, &routing_table.get_pods_with_names()).await; - if failed_pods.is_empty() { - debug!("All registered pods are healthy") + let shard_state = shard_management.current_snapshot().await; + debug!("Checking health of registered executors..."); + let executors = shard_state.get_executors_with_addrs(); + let failed_executors = get_unhealthy_executors(health_check, &executors).await; + if failed_executors.is_empty() { + debug!("All registered executors are healthy") } else { - warn!( - "The following pods were found to be unhealthy: {:?}", - failed_pods - ); - for failed_pod in failed_pods { - shard_management.unregister_pod(failed_pod).await; + for (executor_id, addr, pod_name) in executors + .into_iter() + .filter(|(id, _, _)| failed_executors.contains(id)) + { + warn!( + executor_id = %executor_id, + addr = %addr, + pod_name = pod_name.as_deref().unwrap_or(""), + "Executor was found to be unhealthy; unregistering" + ); + shard_management.unregister_executor(executor_id).await; } } - debug!("Finished checking health of registered pods"); + debug!("Finished checking health of registered executors"); } diff --git a/golem-shard-manager/src/sharding/mod.rs b/golem-shard-manager/src/sharding/mod.rs index d8c9cd9aa6..5bce731724 100644 --- a/golem-shard-manager/src/sharding/mod.rs +++ b/golem-shard-manager/src/sharding/mod.rs @@ -21,4 +21,7 @@ pub mod rebalancing; pub mod shard_management; pub mod worker_executor; -pub use model::{PodState, RoutingTable, RoutingTableEntry}; +pub use model::{ + ExecutorAddr, ExecutorAddrs, ExecutorId, ExecutorLease, ExecutorShards, ShardAssignmentEntry, + ShardEpoch, ShardLeaseRevision, ShardLeaseState, +}; diff --git a/golem-shard-manager/src/sharding/model.rs b/golem-shard-manager/src/sharding/model.rs index 55b7d623e8..10dab0ef5b 100644 --- a/golem-shard-manager/src/sharding/model.rs +++ b/golem-shard-manager/src/sharding/model.rs @@ -12,214 +12,518 @@ // See the License for the specific language governing permissions and // limitations under the License. +use super::error::ShardManagerError; use super::rebalancing::Rebalance; -use core::cmp::Ordering; +use chrono::{DateTime, Utc}; use desert_rust::BinaryCodec; use golem_api_grpc::proto::golem; use golem_common::model::{Pod, ShardId}; -use std::collections::{BTreeMap, BTreeSet, HashSet}; +use std::collections::{BTreeMap, BTreeSet}; use std::fmt; use std::fmt::{Debug, Display, Formatter}; +use std::net::IpAddr; +use std::time::Duration; +use tracing::warn; +use uuid::Uuid; -#[derive(Clone, Debug, Eq, PartialEq, BinaryCodec)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, BinaryCodec)] +#[desert(transparent)] +pub struct ShardEpoch(pub u64); + +impl ShardEpoch { + pub fn initial() -> Self { + Self(0) + } + + pub fn next(self) -> Self { + Self(self.0.checked_add(1).expect("ShardEpoch overflow")) + } +} + +impl Display for ShardEpoch { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.0) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, BinaryCodec)] +#[desert(transparent)] +pub struct ExecutorId(pub Uuid); + +impl ExecutorId { + pub fn generate() -> Self { + Self(Uuid::now_v7()) + } +} + +impl Display for ExecutorId { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.0) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, BinaryCodec)] #[desert(evolution())] -pub struct PodState { +pub struct ExecutorAddr { + pub ip: IpAddr, + pub port: u16, +} + +impl From for ExecutorAddr { + fn from(pod: Pod) -> Self { + Self { + ip: pod.ip, + port: pod.port, + } + } +} + +impl From for Pod { + fn from(addr: ExecutorAddr) -> Self { + Pod { + ip: addr.ip, + port: addr.port, + } + } +} + +impl Display for ExecutorAddr { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "{}:{}", self.ip, self.port) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, BinaryCodec)] +#[desert(transparent)] +pub struct ShardLeaseRevision(pub u64); + +impl ShardLeaseRevision { + pub const INITIAL: Self = Self(0); + + pub fn next(self) -> Result { + self.0 + .checked_add(1) + .map(Self) + .ok_or_else(|| ShardManagerError::Internal("ShardLeaseRevision overflow".to_string())) + } +} + +impl Display for ShardLeaseRevision { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.0) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, BinaryCodec)] +#[desert(evolution())] +pub struct ShardAssignmentEntry { + pub executor_id: ExecutorId, + pub epoch: ShardEpoch, +} + +#[derive(Clone, Debug, PartialEq, Eq, BinaryCodec)] +#[desert(evolution())] +pub struct ExecutorLease { + pub addr: ExecutorAddr, + pub granted_at: DateTime, + pub expires_at: DateTime, pub pod_name: Option, - pub assigned_shards: BTreeSet, } -#[derive(Clone, Debug, Eq, PartialEq, BinaryCodec)] +#[derive(Clone, Debug, PartialEq, Eq, BinaryCodec)] #[desert(evolution())] -pub struct RoutingTable { +pub struct ShardLeaseState { pub number_of_shards: usize, - pub pod_states: BTreeMap, + pub revision: ShardLeaseRevision, + pub shard_assignments: BTreeMap, + pub shard_epochs: BTreeMap, + pub executor_leases: BTreeMap, + pub pending_rebalance: BTreeSet, +} + +#[derive(Clone, Debug)] +pub struct ExecutorShards { + pub executor_id: ExecutorId, + pub shard_ids: BTreeSet, } -impl RoutingTable { +pub type ExecutorAddrs = BTreeMap; + +impl ShardLeaseState { pub fn new(number_of_shards: usize) -> Self { Self { number_of_shards, - pod_states: BTreeMap::new(), + revision: ShardLeaseRevision::INITIAL, + shard_assignments: BTreeMap::new(), + shard_epochs: BTreeMap::new(), + executor_leases: BTreeMap::new(), + pending_rebalance: BTreeSet::new(), } } - pub fn get_entries(&self) -> BTreeSet { - self.pod_states - .clone() - .into_iter() - .map(|(pod, state)| RoutingTableEntry::new(pod, state.assigned_shards)) - .collect() + pub fn get_executors(&self) -> impl Iterator { + self.executor_leases.keys() } - pub fn get_entries_vec(&self) -> Vec { - self.pod_states - .clone() - .into_iter() - .map(|(pod, state)| RoutingTableEntry::new(pod, state.assigned_shards)) + pub fn has_executor(&self, executor_id: ExecutorId) -> bool { + self.executor_leases.contains_key(&executor_id) + } + + pub fn executor_count(&self) -> usize { + self.executor_leases.len() + } + + pub fn get_executors_with_addrs(&self) -> Vec<(ExecutorId, ExecutorAddr, Option)> { + self.executor_leases + .iter() + .map(|(id, lease)| (*id, lease.addr, lease.pod_name.clone())) .collect() } - pub fn get_pods(&self) -> HashSet { - self.pod_states.clone().into_keys().collect() + pub fn addr_for(&self, executor_id: ExecutorId) -> Option { + self.executor_leases + .get(&executor_id) + .map(|lease| lease.addr) } - pub fn get_pods_with_names(&self) -> Vec<(Pod, Option)> { - self.pod_states + pub fn executor_addrs(&self) -> ExecutorAddrs { + self.executor_leases .iter() - .map(|(pod, pod_state)| (*pod, pod_state.pod_name.clone())) + .map(|(id, lease)| (*id, lease.addr)) .collect() } - pub fn rebalance(&mut self, rebalance: Rebalance) { - for (pod, shard_ids) in &rebalance.get_assignments().assignments { - if let Some(pod_state) = self.pod_states.get_mut(pod) { - pod_state.assigned_shards.extend(shard_ids); + pub fn executor_for_addr(&self, addr: ExecutorAddr) -> Option { + self.executor_leases + .iter() + .find(|(_, lease)| lease.addr == addr) + .map(|(id, _)| *id) + } + + pub fn shards_for_executor(&self, executor_id: ExecutorId) -> Option> { + if !self.has_executor(executor_id) { + return None; + } + Some( + self.shard_assignments + .iter() + .filter(|(_, entry)| entry.executor_id == executor_id) + .map(|(shard_id, _)| *shard_id) + .collect(), + ) + } + + pub fn executor_shard_sets(&self) -> Vec { + let mut by_executor: BTreeMap> = self + .executor_leases + .keys() + .map(|id| (*id, BTreeSet::new())) + .collect(); + for (shard_id, entry) in &self.shard_assignments { + if let Some(shard_ids) = by_executor.get_mut(&entry.executor_id) { + shard_ids.insert(*shard_id); } } - for (pod, shard_ids) in &rebalance.get_unassignments().unassignments { - if let Some(pod_state) = self.pod_states.get_mut(pod) { - pod_state - .assigned_shards - .retain(|shard_id| !shard_ids.contains(shard_id)); + by_executor + .into_iter() + .map(|(executor_id, shard_ids)| ExecutorShards { + executor_id, + shard_ids, + }) + .collect() + } + + // If new added; returns None, if replaced; returns the replaced ExecutorId. + pub fn add_executor( + &mut self, + executor_id: ExecutorId, + addr: ExecutorAddr, + pod_name: Option, + now: DateTime, + lease_ttl: Duration, + ) -> Option { + let replaced = match self.executor_for_addr(addr) { + Some(previous) if previous != executor_id => { + self.executor_leases.remove(&previous); + Some(previous) + } + _ => None, + }; + + self.executor_leases.insert( + executor_id, + ExecutorLease { + addr, + granted_at: now, + expires_at: now + lease_ttl, + pod_name, + }, + ); + + if let Some(previous) = replaced { + // The predecessor's shards change owner, so each of them advances its epoch. + let transferred: Vec = self + .shard_assignments + .iter() + .filter(|(_, entry)| entry.executor_id == previous) + .map(|(shard_id, _)| *shard_id) + .collect(); + for shard_id in transferred { + self.assign_shard(executor_id, shard_id); } } + + debug_assert!(self.check_invariants().is_ok()); + replaced } - pub fn get_unassigned_shards(&self) -> BTreeSet { - let mut unassigned_shards: BTreeSet = (0..self.number_of_shards) - .map(|shard_id| ShardId::new(shard_id as i64)) + pub fn remove_executor(&mut self, executor_id: ExecutorId) -> BTreeSet { + if self.executor_leases.remove(&executor_id).is_none() { + return BTreeSet::new(); + } + let orphaned: BTreeSet = self + .shard_assignments + .iter() + .filter(|(_, entry)| entry.executor_id == executor_id) + .map(|(shard_id, _)| *shard_id) .collect(); - - for pod_state in self.pod_states.values() { - unassigned_shards.retain(|shard_id| !pod_state.assigned_shards.contains(shard_id)); + for shard_id in &orphaned { + self.shard_assignments.remove(shard_id); } + self.pending_rebalance.extend(orphaned.iter().copied()); - unassigned_shards + debug_assert!(self.check_invariants().is_ok()); + orphaned } - pub fn get_shards(&self, pod: Pod) -> Option> { - self.pod_states - .get(&pod) - .map(|ps| ps.assigned_shards.clone()) + pub fn contains_shard(&self, shard_id: ShardId) -> bool { + shard_id.value() >= 0 && (shard_id.value() as usize) < self.number_of_shards } - pub fn get_pod_count(&self) -> usize { - self.pod_states.len() + pub fn get_unassigned_shards(&self) -> BTreeSet { + (0..self.number_of_shards) + .map(|shard_id| ShardId::new(shard_id as i64)) + .filter(|shard_id| !self.shard_assignments.contains_key(shard_id)) + .collect() } - pub fn add_pod(&mut self, pod: Pod, pod_name: Option) { - self.pod_states.insert( - pod, - PodState { - pod_name, - assigned_shards: BTreeSet::new(), - }, + pub fn assign_shard(&mut self, executor_id: ExecutorId, shard_id: ShardId) -> ShardEpoch { + debug_assert!( + self.has_executor(executor_id), + "assigning shard {shard_id} to executor {executor_id} without a lease" ); + debug_assert!( + self.contains_shard(shard_id), + "assigning shard {shard_id} outside 0..{}", + self.number_of_shards + ); + let epoch = match self.shard_assignments.get(&shard_id) { + Some(entry) if entry.executor_id == executor_id => entry.epoch, + _ => match self.shard_epochs.get(&shard_id) { + Some(last) => last.next(), + None => ShardEpoch::initial(), + }, + }; + self.shard_assignments + .insert(shard_id, ShardAssignmentEntry { executor_id, epoch }); + self.shard_epochs.insert(shard_id, epoch); + self.pending_rebalance.remove(&shard_id); + epoch } - pub fn remove_pod(&mut self, pod: Pod) { - self.pod_states.remove(&pod); + pub fn unassign_shard(&mut self, owner: ExecutorId, shard_id: ShardId) -> bool { + match self.shard_assignments.get(&shard_id) { + Some(entry) if entry.executor_id == owner => { + self.shard_assignments.remove(&shard_id); + true + } + _ => false, + } } - pub fn has_pod(&self, pod: Pod) -> bool { - self.pod_states.contains_key(&pod) + pub fn epoch_for_shard(&self, shard_id: ShardId) -> Option { + self.shard_assignments + .get(&shard_id) + .map(|entry| entry.epoch) } -} -impl From for golem::shardmanager::RoutingTable { - fn from(routing_table: RoutingTable) -> golem::shardmanager::RoutingTable { - golem::shardmanager::RoutingTable { - number_of_shards: routing_table.number_of_shards as u32, - shard_assignments: routing_table - .pod_states - .into_iter() - .flat_map(|(pod, pod_state)| { - pod_state - .assigned_shards - .into_iter() - .map(move |shard_id| (pod, shard_id)) - }) - .map(|(pod, shard_id)| golem::shardmanager::RoutingTableEntry { - pod: Some(pod.into()), - shard_id: Some(shard_id.into()), - }) - .collect(), + pub fn apply_rebalance(&mut self, rebalance: &Rebalance) { + for (executor_id, shard_ids) in &rebalance.get_assignments().assignments { + if !self.has_executor(*executor_id) { + warn!( + executor_id = %executor_id, + shards = shard_ids.len(), + "Skipping planned shard assignments: executor no longer holds a lease" + ); + continue; + } + for shard_id in shard_ids { + self.assign_shard(*executor_id, *shard_id); + } + } + for (executor_id, shard_ids) in &rebalance.get_unassignments().unassignments { + for shard_id in shard_ids { + self.unassign_shard(*executor_id, *shard_id); + } } + debug_assert!(self.check_invariants().is_ok()); } -} -impl Display for RoutingTable { - fn fmt(&self, f: &mut Formatter) -> fmt::Result { - write!( - f, - "{{ number_of_shards: {}, shard_assignments: [{}] }}", - self.number_of_shards, - pod_states_map_to_string(&self.pod_states) - ) + pub fn take_pending_rebalance(&mut self) -> BTreeSet { + std::mem::take(&mut self.pending_rebalance) } -} -pub struct RoutingTableEntry { - pub pod: Pod, - pub shard_ids: BTreeSet, -} - -impl RoutingTableEntry { - pub fn new(pod: Pod, shard_ids: BTreeSet) -> Self { - Self { pod, shard_ids } - } - pub fn get_shard_count(&self) -> usize { - self.shard_ids.len() + pub fn housekeep(&mut self, now: DateTime) -> Vec<(ExecutorId, BTreeSet)> { + let expired: Vec = self + .executor_leases + .iter() + .filter(|(_, lease)| now >= lease.expires_at) + .map(|(id, _)| *id) + .collect(); + expired + .into_iter() + .map(|executor_id| (executor_id, self.remove_executor(executor_id))) + .collect() } -} -impl PartialEq for RoutingTableEntry { - fn eq(&self, other: &Self) -> bool { - self.shard_ids.len() == other.shard_ids.len() && self.pod == other.pod + pub fn bump_revision(&mut self) -> Result { + self.revision = self.revision.next()?; + Ok(self.revision) } -} - -impl Eq for RoutingTableEntry {} -impl PartialOrd for RoutingTableEntry { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) + pub fn check_invariants(&self) -> Result<(), String> { + let mut addrs = BTreeSet::new(); + for (executor_id, lease) in &self.executor_leases { + if !addrs.insert(lease.addr) { + return Err(format!( + "address {} is leased by more than one executor (including {executor_id})", + lease.addr + )); + } + } + for (shard_id, entry) in &self.shard_assignments { + if !self.has_executor(entry.executor_id) { + return Err(format!( + "shard {shard_id} is assigned to executor {} which holds no lease", + entry.executor_id + )); + } + if !self.contains_shard(*shard_id) { + return Err(format!( + "shard {shard_id} is outside 0..{}", + self.number_of_shards + )); + } + if self.pending_rebalance.contains(shard_id) { + return Err(format!( + "shard {shard_id} is both assigned and pending rebalance" + )); + } + match self.shard_epochs.get(shard_id) { + Some(last) if *last == entry.epoch => {} + Some(last) => { + return Err(format!( + "shard {shard_id} is assigned with epoch {} but its recorded highest epoch is {last}", + entry.epoch + )); + } + None => { + return Err(format!( + "shard {shard_id} is assigned with epoch {} but has no recorded epoch", + entry.epoch + )); + } + } + } + for shard_id in self.shard_epochs.keys() { + if !self.contains_shard(*shard_id) { + return Err(format!( + "shard {shard_id} with a recorded epoch is outside 0..{}", + self.number_of_shards + )); + } + } + for shard_id in &self.pending_rebalance { + if !self.contains_shard(*shard_id) { + return Err(format!( + "pending shard {shard_id} is outside 0..{}", + self.number_of_shards + )); + } + } + Ok(()) } } -impl Ord for RoutingTableEntry { - fn cmp(&self, other: &Self) -> Ordering { - match self.shard_ids.len().cmp(&other.shard_ids.len()) { - Ordering::Equal => self.pod.cmp(&other.pod), - other => other, +impl From for golem::shardmanager::RoutingTable { + fn from(shard_state: ShardLeaseState) -> golem::shardmanager::RoutingTable { + golem::shardmanager::RoutingTable { + number_of_shards: shard_state.number_of_shards as u32, + shard_assignments: shard_state + .shard_assignments + .iter() + .filter_map(|(shard_id, entry)| { + shard_state + .executor_leases + .get(&entry.executor_id) + .map(|lease| (*shard_id, lease.addr)) + }) + .map(|(shard_id, addr)| golem::shardmanager::RoutingTableEntry { + pod: Some(Pod::from(addr).into()), + shard_id: Some(shard_id.into()), + }) + .collect(), } } } -impl Display for RoutingTableEntry { +impl Display for ShardLeaseState { fn fmt(&self, f: &mut Formatter) -> fmt::Result { - let shard_ids: Vec = self.shard_ids.iter().map(|elem| elem.to_string()).collect(); + let by_executor = self.executor_shard_sets(); + let executors: Vec = by_executor + .iter() + .map(|entry| { + let lease = &self.executor_leases[&entry.executor_id]; + shard_assignments_to_string( + &format!("{} {}", entry.executor_id, lease.addr), + lease.pod_name.as_deref(), + entry.shard_ids.iter(), + ) + }) + .collect(); + let pending: Vec = shard_ids_to_ranges(self.pending_rebalance.iter()) + .iter() + .map(|rng| rng.to_string()) + .collect(); write!( f, - "{{ pod: {}, shard_ids: [{}] }}", - self.pod, - shard_ids.join(", ") + "{{ number_of_shards: {}, revision: {}, executors: [{}], pending_rebalance: [{}] }}", + self.number_of_shards, + self.revision, + executors.join(", "), + pending.join(", ") ) } } #[derive(Clone, Debug)] pub struct Assignments { - pub assignments: BTreeMap>, + pub assignments: BTreeMap>, } impl Assignments { - pub fn assign(&mut self, pod: Pod, shard_id: ShardId) { - self.assignments.entry(pod).or_default().insert(shard_id); + pub fn assign(&mut self, executor_id: ExecutorId, shard_id: ShardId) { + self.assignments + .entry(executor_id) + .or_default() + .insert(shard_id); } - pub fn unassign(&mut self, pod: Pod, shard_id: ShardId) { - self.assignments.entry(pod).or_default().remove(&shard_id); + pub fn unassign(&mut self, executor_id: ExecutorId, shard_id: ShardId) { + self.assignments + .entry(executor_id) + .or_default() + .remove(&shard_id); } pub fn new() -> Self { @@ -241,18 +545,25 @@ impl Default for Assignments { impl Display for Assignments { fn fmt(&self, f: &mut Formatter) -> fmt::Result { - write!(f, "[{}]", pod_shard_ids_map_to_string(&self.assignments)) + write!( + f, + "[{}]", + executor_shard_ids_map_to_string(&self.assignments) + ) } } #[derive(Clone, Debug)] pub struct Unassignments { - pub unassignments: BTreeMap>, + pub unassignments: BTreeMap>, } impl Unassignments { - pub fn unassign(&mut self, pod: Pod, shard_id: ShardId) { - self.unassignments.entry(pod).or_default().insert(shard_id); + pub fn unassign(&mut self, executor_id: ExecutorId, shard_id: ShardId) { + self.unassignments + .entry(executor_id) + .or_default() + .insert(shard_id); } pub fn new() -> Self { @@ -274,35 +585,29 @@ impl Default for Unassignments { impl Display for Unassignments { fn fmt(&self, f: &mut Formatter) -> fmt::Result { - write!(f, "[{}]", pod_shard_ids_map_to_string(&self.unassignments)) + write!( + f, + "[{}]", + executor_shard_ids_map_to_string(&self.unassignments) + ) } } -fn pod_states_map_to_string(pod_states: &BTreeMap) -> String { - let elements: Vec = pod_states +fn executor_shard_ids_map_to_string( + shards_by_executor: &BTreeMap>, +) -> String { + let elements: Vec = shards_by_executor .iter() - .map(|(pod, pod_state)| { - pod_shard_assignments_to_string( - pod, - pod_state.pod_name.clone(), - pod_state.assigned_shards.iter(), - ) + .map(|(executor_id, shard_ids)| { + shard_assignments_to_string(executor_id, None, shard_ids.iter()) }) .collect(); elements.join(", ") } -fn pod_shard_ids_map_to_string(pod_states: &BTreeMap>) -> String { - let elements: Vec = pod_states - .iter() - .map(|(pod, shard_ids)| pod_shard_assignments_to_string(pod, None, shard_ids.iter())) - .collect(); - elements.join(", ") -} - -pub fn pod_shard_assignments_to_string<'a, T: Iterator>( - pod: &Pod, - pod_name: Option, +pub fn shard_assignments_to_string<'a, T: Iterator>( + label: &dyn Display, + pod_name: Option<&str>, shard_ids: T, ) -> String { let ranges: Vec = shard_ids_to_ranges(shard_ids); @@ -311,7 +616,7 @@ pub fn pod_shard_assignments_to_string<'a, T: Iterator>( .map(|rng| format!("{rng}").to_string()) .collect(); format!( - "{pod} {}: [{}]", + "{label} {}: [{}]", pod_name.unwrap_or_default(), strings.join(", ") ) @@ -372,3 +677,514 @@ fn shard_ids_to_ranges<'a, T: Iterator>(ids: T) -> Vec DateTime { + DateTime::from_timestamp(1_700_000_000, 0).unwrap() + } + + fn executor(idx: u128) -> ExecutorId { + ExecutorId(Uuid::from_u128(idx)) + } + + fn addr(idx: u8) -> ExecutorAddr { + ExecutorAddr { + ip: IpAddr::V4(Ipv4Addr::new(10, 0, 0, idx)), + port: 9000 + idx as u16, + } + } + + fn shard(id: i64) -> ShardId { + ShardId::new(id) + } + + fn shards(ids: &[i64]) -> BTreeSet { + ids.iter().copied().map(ShardId::new).collect() + } + + fn shard_state_with( + number_of_shards: usize, + executors: &[(u128, u8, &[i64])], + ) -> ShardLeaseState { + let mut shard_state = ShardLeaseState::new(number_of_shards); + for (idx, addr_idx, shard_ids) in executors { + shard_state.add_executor(executor(*idx), addr(*addr_idx), None, t0(), TTL); + for shard_id in *shard_ids { + shard_state.assign_shard(executor(*idx), shard(*shard_id)); + } + } + shard_state + } + + #[test] + fn new_state_is_empty_with_initial_revision() { + let shard_state = ShardLeaseState::new(8); + assert_eq!(shard_state.number_of_shards, 8); + assert_eq!(shard_state.revision, ShardLeaseRevision::INITIAL); + assert!(shard_state.shard_assignments.is_empty()); + assert!(shard_state.shard_epochs.is_empty()); + assert!(shard_state.executor_leases.is_empty()); + assert!(shard_state.pending_rebalance.is_empty()); + assert_eq!(shard_state.executor_count(), 0); + assert_eq!( + shard_state.get_unassigned_shards(), + shards(&[0, 1, 2, 3, 4, 5, 6, 7]) + ); + assert!(shard_state.check_invariants().is_ok()); + } + + #[test] + fn generated_executor_ids_are_unique_and_time_ordered() { + let ids: Vec = (0..64).map(|_| ExecutorId::generate()).collect(); + let unique: BTreeSet = ids.iter().copied().collect(); + assert_eq!(unique.len(), ids.len()); + assert!(ids.windows(2).all(|pair| pair[0] < pair[1])); + } + + #[test] + fn add_executor_grants_lease_with_expected_window() { + let mut shard_state = ShardLeaseState::new(4); + let replaced = + shard_state.add_executor(executor(1), addr(1), Some("pod-1".to_string()), t0(), TTL); + assert_eq!(replaced, None); + + let lease = &shard_state.executor_leases[&executor(1)]; + assert_eq!(lease.addr, addr(1)); + assert_eq!(lease.granted_at, t0()); + assert_eq!(lease.expires_at, t0() + chrono::Duration::seconds(60)); + assert_eq!(lease.pod_name.as_deref(), Some("pod-1")); + + assert!(shard_state.has_executor(executor(1))); + assert_eq!(shard_state.addr_for(executor(1)), Some(addr(1))); + assert_eq!(shard_state.executor_for_addr(addr(1)), Some(executor(1))); + assert_eq!( + shard_state.shards_for_executor(executor(1)), + Some(BTreeSet::new()) + ); + assert_eq!(shard_state.shards_for_executor(executor(2)), None); + assert_eq!( + shard_state.get_executors_with_addrs(), + vec![(executor(1), addr(1), Some("pod-1".to_string()))] + ); + } + + #[test] + fn add_executor_at_reused_address_replaces_lease_and_transfers_shards() { + let mut shard_state = shard_state_with(4, &[(1, 1, &[0, 1]), (2, 2, &[2])]); + let mut later = t0() + chrono::Duration::seconds(10); + + // bump shard 1 once so we can tell "advanced" from "reset": moving it to executor 2 and + // back advances its epoch to 1 + shard_state.assign_shard(executor(2), shard(1)); + shard_state.assign_shard(executor(1), shard(1)); + assert_eq!(shard_state.epoch_for_shard(shard(1)), Some(ShardEpoch(2))); + + let replaced = shard_state.add_executor(executor(3), addr(1), None, later, TTL); + assert_eq!(replaced, Some(executor(1))); + + assert!(!shard_state.has_executor(executor(1))); + assert!(shard_state.has_executor(executor(3))); + assert_eq!(shard_state.executor_for_addr(addr(1)), Some(executor(3))); + assert_eq!(shard_state.executor_count(), 2); + + assert_eq!( + shard_state.shards_for_executor(executor(3)), + Some(shards(&[0, 1])) + ); + assert_eq!(shard_state.epoch_for_shard(shard(0)), Some(ShardEpoch(1))); + assert_eq!(shard_state.epoch_for_shard(shard(1)), Some(ShardEpoch(3))); + assert_eq!( + shard_state.shard_epochs.get(&shard(1)), + Some(&ShardEpoch(3)) + ); + assert_eq!( + shard_state.shards_for_executor(executor(2)), + Some(shards(&[2])) + ); + assert_eq!(shard_state.epoch_for_shard(shard(2)), Some(ShardEpoch(0))); + assert!(shard_state.pending_rebalance.is_empty()); + assert_eq!(shard_state.executor_leases[&executor(3)].granted_at, later); + + // a second registration of the same id at the same address only refreshes the lease + later += chrono::Duration::seconds(10); + let replaced = shard_state.add_executor(executor(3), addr(1), None, later, TTL); + assert_eq!(replaced, None); + assert_eq!( + shard_state.shards_for_executor(executor(3)), + Some(shards(&[0, 1])) + ); + assert_eq!(shard_state.epoch_for_shard(shard(0)), Some(ShardEpoch(1))); + assert_eq!(shard_state.executor_leases[&executor(3)].granted_at, later); + assert!(shard_state.check_invariants().is_ok()); + } + + #[test] + fn remove_executor_orphans_shards_into_pending_rebalance() { + let mut shard_state = shard_state_with(4, &[(1, 1, &[0, 1]), (2, 2, &[2])]); + + let orphaned = shard_state.remove_executor(executor(1)); + assert_eq!(orphaned, shards(&[0, 1])); + assert!(!shard_state.has_executor(executor(1))); + assert_eq!(shard_state.pending_rebalance, shards(&[0, 1])); + assert_eq!(shard_state.get_unassigned_shards(), shards(&[0, 1, 3])); + assert_eq!(shard_state.shards_for_executor(executor(1)), None); + assert_eq!( + shard_state.shards_for_executor(executor(2)), + Some(shards(&[2])) + ); + assert!(shard_state.check_invariants().is_ok()); + + // unknown ids are a no-op + assert_eq!(shard_state.remove_executor(executor(42)), BTreeSet::new()); + assert_eq!(shard_state.executor_count(), 1); + } + + #[test] + fn assign_shard_epoch_semantics() { + let mut shard_state = shard_state_with(4, &[(1, 1, &[]), (2, 2, &[])]); + + // first assignment starts at the initial epoch + assert_eq!( + shard_state.assign_shard(executor(1), shard(0)), + ShardEpoch::initial() + ); + assert_eq!(shard_state.epoch_for_shard(shard(0)), Some(ShardEpoch(0))); + + // re-assigning to the same owner is idempotent + assert_eq!( + shard_state.assign_shard(executor(1), shard(0)), + ShardEpoch(0) + ); + + // moving to another owner advances the epoch; stored == returned + assert_eq!( + shard_state.assign_shard(executor(2), shard(0)), + ShardEpoch(1) + ); + assert_eq!(shard_state.epoch_for_shard(shard(0)), Some(ShardEpoch(1))); + assert_eq!( + shard_state.shards_for_executor(executor(1)), + Some(BTreeSet::new()) + ); + assert_eq!( + shard_state.shards_for_executor(executor(2)), + Some(shards(&[0])) + ); + + // after a full unassignment the next assignment continues past the highest epoch ever + // issued for the shard - an epoch is never reused in a shard's history + assert!(shard_state.unassign_shard(executor(2), shard(0))); + assert_eq!(shard_state.epoch_for_shard(shard(0)), None); + assert_eq!( + shard_state.shard_epochs.get(&shard(0)), + Some(&ShardEpoch(1)) + ); + assert_eq!( + shard_state.assign_shard(executor(1), shard(0)), + ShardEpoch(2) + ); + assert_eq!(shard_state.epoch_for_shard(shard(0)), Some(ShardEpoch(2))); + } + + #[test] + fn epochs_stay_unique_across_eviction_and_reassignment() { + let mut shard_state = shard_state_with(4, &[(1, 1, &[0, 1]), (2, 2, &[]), (3, 3, &[])]); + assert_eq!(shard_state.epoch_for_shard(shard(0)), Some(ShardEpoch(0))); + + // executor 1 is evicted (health check / lease expiry): its shards are orphaned ... + shard_state.remove_executor(executor(1)); + assert_eq!(shard_state.epoch_for_shard(shard(0)), None); + assert_eq!(shard_state.pending_rebalance, shards(&[0, 1])); + + // ... and later handed to other executors with epochs the evicted owner never held + assert_eq!( + shard_state.assign_shard(executor(2), shard(0)), + ShardEpoch(1) + ); + assert_eq!( + shard_state.assign_shard(executor(3), shard(1)), + ShardEpoch(1) + ); + + // a second eviction keeps advancing + shard_state.remove_executor(executor(2)); + assert_eq!( + shard_state.assign_shard(executor(3), shard(0)), + ShardEpoch(2) + ); + + // housekeep-driven eviction behaves the same way + let expired = t0() + chrono::Duration::seconds(3600); + assert_eq!( + shard_state.housekeep(expired), + vec![(executor(3), shards(&[0, 1]))] + ); + shard_state.add_executor(executor(4), addr(4), None, expired, TTL); + assert_eq!( + shard_state.assign_shard(executor(4), shard(0)), + ShardEpoch(3) + ); + assert_eq!( + shard_state.assign_shard(executor(4), shard(1)), + ShardEpoch(2) + ); + assert!(shard_state.check_invariants().is_ok()); + } + + #[test] + fn unassign_shard_is_guarded_by_owner() { + let mut shard_state = shard_state_with(4, &[(1, 1, &[0]), (2, 2, &[])]); + + assert!(!shard_state.unassign_shard(executor(2), shard(0))); + assert_eq!( + shard_state.shards_for_executor(executor(1)), + Some(shards(&[0])) + ); + + assert!(shard_state.unassign_shard(executor(1), shard(0))); + assert_eq!( + shard_state.shards_for_executor(executor(1)), + Some(BTreeSet::new()) + ); + assert!(!shard_state.unassign_shard(executor(1), shard(0))); + } + + #[test] + fn assign_shard_clears_pending_rebalance() { + let mut shard_state = shard_state_with(4, &[(1, 1, &[0, 1]), (2, 2, &[])]); + shard_state.remove_executor(executor(1)); + assert_eq!(shard_state.pending_rebalance, shards(&[0, 1])); + + shard_state.assign_shard(executor(2), shard(0)); + assert_eq!(shard_state.pending_rebalance, shards(&[1])); + assert!(shard_state.check_invariants().is_ok()); + } + + #[test] + fn apply_rebalance_moves_bump_epoch_once_and_unassignments_are_owner_guarded() { + let mut shard_state = shard_state_with(4, &[(1, 1, &[0, 1, 2, 3]), (2, 2, &[])]); + + // move shards 0 and 1 from executor 1 to executor 2, as a Rebalance would express it + let mut assignments = Assignments::new(); + assignments.assign(executor(2), shard(0)); + assignments.assign(executor(2), shard(1)); + let mut unassignments = Unassignments::new(); + unassignments.unassign(executor(1), shard(0)); + unassignments.unassign(executor(1), shard(1)); + let rebalance = Rebalance::new(assignments, unassignments); + + shard_state.apply_rebalance(&rebalance); + + assert_eq!( + shard_state.shards_for_executor(executor(1)), + Some(shards(&[2, 3])) + ); + assert_eq!( + shard_state.shards_for_executor(executor(2)), + Some(shards(&[0, 1])) + ); + assert_eq!(shard_state.epoch_for_shard(shard(0)), Some(ShardEpoch(1))); + assert_eq!(shard_state.epoch_for_shard(shard(1)), Some(ShardEpoch(1))); + assert_eq!(shard_state.epoch_for_shard(shard(2)), Some(ShardEpoch(0))); + assert!(shard_state.check_invariants().is_ok()); + + // applying the same plan again is a no-op (idempotent assignments, guarded unassignments) + shard_state.apply_rebalance(&rebalance); + assert_eq!( + shard_state.shards_for_executor(executor(2)), + Some(shards(&[0, 1])) + ); + assert_eq!(shard_state.epoch_for_shard(shard(0)), Some(ShardEpoch(1))); + } + + #[test] + fn executor_shard_sets_covers_every_lease_in_id_order() { + let shard_state = shard_state_with(6, &[(3, 3, &[4]), (1, 1, &[0, 2]), (2, 2, &[])]); + let sets = shard_state.executor_shard_sets(); + let ids: Vec = sets.iter().map(|s| s.executor_id).collect(); + assert_eq!(ids, vec![executor(1), executor(2), executor(3)]); + assert_eq!(sets[0].shard_ids, shards(&[0, 2])); + assert_eq!(sets[1].shard_ids, BTreeSet::new()); + assert_eq!(sets[2].shard_ids, shards(&[4])); + assert_eq!(shard_state.get_unassigned_shards(), shards(&[1, 3, 5])); + } + + #[test] + fn housekeep_evicts_only_expired_leases() { + let mut shard_state = ShardLeaseState::new(4); + shard_state.add_executor(executor(1), addr(1), None, t0(), Duration::from_secs(10)); + shard_state.add_executor(executor(2), addr(2), None, t0(), Duration::from_secs(60)); + shard_state.assign_shard(executor(1), shard(0)); + shard_state.assign_shard(executor(2), shard(1)); + + assert!( + shard_state + .housekeep(t0() + chrono::Duration::seconds(9)) + .is_empty() + ); + assert_eq!(shard_state.executor_count(), 2); + + // expiry is inclusive + let evicted = shard_state.housekeep(t0() + chrono::Duration::seconds(10)); + assert_eq!(evicted, vec![(executor(1), shards(&[0]))]); + assert!(!shard_state.has_executor(executor(1))); + assert!(shard_state.has_executor(executor(2))); + assert_eq!(shard_state.pending_rebalance, shards(&[0])); + assert_eq!( + shard_state.shards_for_executor(executor(2)), + Some(shards(&[1])) + ); + + assert_eq!(shard_state.take_pending_rebalance(), shards(&[0])); + assert!(shard_state.pending_rebalance.is_empty()); + assert!(shard_state.check_invariants().is_ok()); + } + + #[test] + fn bump_revision_increments_and_reports_overflow() { + let mut shard_state = ShardLeaseState::new(1); + assert_eq!(shard_state.bump_revision().unwrap(), ShardLeaseRevision(1)); + assert_eq!(shard_state.bump_revision().unwrap(), ShardLeaseRevision(2)); + assert_eq!(shard_state.revision, ShardLeaseRevision(2)); + + shard_state.revision = ShardLeaseRevision(u64::MAX); + match shard_state.bump_revision() { + Err(ShardManagerError::Internal(msg)) => assert!(msg.contains("overflow")), + other => panic!("expected Internal error, got {other:?}"), + } + assert_eq!(shard_state.revision, ShardLeaseRevision(u64::MAX)); + } + + #[test] + fn binary_codec_roundtrip_is_exact() { + let mut shard_state = shard_state_with(8, &[(1, 1, &[0, 1, 2]), (2, 2, &[5])]); + shard_state + .executor_leases + .get_mut(&executor(2)) + .unwrap() + .pod_name = Some("worker-executor-1".to_string()); + shard_state + .executor_leases + .get_mut(&executor(2)) + .unwrap() + .granted_at = DateTime::from_timestamp(1_700_000_000, 123_456_789).unwrap(); + shard_state.assign_shard(executor(2), shard(0)); + shard_state.remove_executor(executor(1)); + shard_state.bump_revision().unwrap(); + + let bytes = golem_common::serialization::serialize(&shard_state).unwrap(); + let decoded: ShardLeaseState = golem_common::serialization::deserialize(&bytes).unwrap(); + assert_eq!(decoded, shard_state); + } + + #[test] + fn proto_conversion_emits_one_entry_per_routable_shard() { + let mut shard_state = shard_state_with(8, &[(1, 1, &[0, 1]), (2, 2, &[5])]); + // violate invariant 1 on purpose: an assignment whose executor holds no lease + shard_state.shard_assignments.insert( + shard(7), + ShardAssignmentEntry { + executor_id: executor(99), + epoch: ShardEpoch(3), + }, + ); + assert!(shard_state.check_invariants().is_err()); + + let proto: golem::shardmanager::RoutingTable = shard_state.into(); + assert_eq!(proto.number_of_shards, 8); + assert_eq!(proto.shard_assignments.len(), 3); + for entry in &proto.shard_assignments { + assert!(entry.pod.is_some()); + assert!(entry.shard_id.is_some()); + } + let mut routed: Vec = proto + .shard_assignments + .iter() + .map(|entry| entry.shard_id.unwrap().value) + .collect(); + routed.sort(); + assert_eq!(routed, vec![0, 1, 5]); + } + + #[test] + fn check_invariants_detects_each_violation() { + let good = shard_state_with(4, &[(1, 1, &[0]), (2, 2, &[1])]); + assert!(good.check_invariants().is_ok()); + + let mut lease_less_owner = good.clone(); + lease_less_owner.shard_assignments.insert( + shard(2), + ShardAssignmentEntry { + executor_id: executor(99), + epoch: ShardEpoch(0), + }, + ); + assert!(lease_less_owner.check_invariants().is_err()); + + let mut duplicate_addr = good.clone(); + duplicate_addr + .executor_leases + .get_mut(&executor(2)) + .unwrap() + .addr = addr(1); + assert!(duplicate_addr.check_invariants().is_err()); + + let mut out_of_range = good.clone(); + out_of_range.shard_assignments.insert( + shard(4), + ShardAssignmentEntry { + executor_id: executor(1), + epoch: ShardEpoch(0), + }, + ); + assert!(out_of_range.check_invariants().is_err()); + + let mut pending_and_assigned = good.clone(); + pending_and_assigned.pending_rebalance.insert(shard(0)); + assert!(pending_and_assigned.check_invariants().is_err()); + + let mut pending_out_of_range = good.clone(); + pending_out_of_range.pending_rebalance.insert(shard(9)); + assert!(pending_out_of_range.check_invariants().is_err()); + + let mut epoch_mismatch = good.clone(); + epoch_mismatch.shard_epochs.insert(shard(0), ShardEpoch(7)); + assert!(epoch_mismatch.check_invariants().is_err()); + + let mut epoch_missing = good.clone(); + epoch_missing.shard_epochs.remove(&shard(0)); + assert!(epoch_missing.check_invariants().is_err()); + + let mut epoch_out_of_range = good; + epoch_out_of_range + .shard_epochs + .insert(shard(9), ShardEpoch(0)); + assert!(epoch_out_of_range.check_invariants().is_err()); + } + + #[test] + fn display_is_readable() { + let mut shard_state = shard_state_with(8, &[(1, 1, &[0, 1, 2, 5])]); + shard_state + .executor_leases + .get_mut(&executor(1)) + .unwrap() + .pod_name = Some("worker-executor-0".to_string()); + let rendered = shard_state.to_string(); + assert!(rendered.contains("number_of_shards: 8")); + assert!( + rendered.contains("10.0.0.1:9001 worker-executor-0: [<0>..<2>, <5>]"), + "{rendered}" + ); + assert!(rendered.contains("pending_rebalance: []")); + } +} diff --git a/golem-shard-manager/src/sharding/persistence.rs b/golem-shard-manager/src/sharding/persistence.rs index 171ac651d1..7d66bb81cf 100644 --- a/golem-shard-manager/src/sharding/persistence.rs +++ b/golem-shard-manager/src/sharding/persistence.rs @@ -13,11 +13,11 @@ // limitations under the License. use super::error::ShardManagerError; -use super::model::RoutingTable; +use super::model::ShardLeaseState; use anyhow::anyhow; use async_trait::async_trait; use conditional_trait_gen::trait_gen; -use golem_common::serialization::{deserialize, serialize}; +use golem_common::serialization::{serialize, try_deserialize}; use golem_service_base::db::postgres::PostgresPool; use golem_service_base::db::sqlite::SqlitePool; use golem_service_base::db::{Pool, PoolApi}; @@ -28,8 +28,8 @@ const PERSISTENCE_SVC: &str = "persistence"; #[async_trait] pub trait RoutingTablePersistence: Send + Sync { - async fn write(&self, routing_table: &RoutingTable) -> Result<(), ShardManagerError>; - async fn read(&self) -> Result; + async fn write(&self, shard_state: &ShardLeaseState) -> Result<(), ShardManagerError>; + async fn read(&self) -> Result; } pub struct DbRoutingTablePersistence { @@ -49,8 +49,8 @@ impl DbRoutingTablePersistence { #[trait_gen(PostgresPool -> PostgresPool, SqlitePool)] #[async_trait] impl RoutingTablePersistence for DbRoutingTablePersistence { - async fn write(&self, routing_table: &RoutingTable) -> Result<(), ShardManagerError> { - let encoded = serialize(&routing_table).map_err(ShardManagerError::SerializationError)?; + async fn write(&self, shard_state: &ShardLeaseState) -> Result<(), ShardManagerError> { + let encoded = serialize(shard_state).map_err(ShardManagerError::SerializationError)?; self.pool .with_rw(PERSISTENCE_SVC, "write") @@ -67,7 +67,7 @@ impl RoutingTablePersistence for DbRoutingTablePersistence { Ok(()) } - async fn read(&self) -> Result { + async fn read(&self) -> Result { let row = self .pool .with_ro(PERSISTENCE_SVC, "read") @@ -81,11 +81,109 @@ impl RoutingTablePersistence for DbRoutingTablePersistence { let bytes: Vec = row .try_get("state") .map_err(|err| RepoError::InternalError(anyhow!(err)))?; - let routing_table: RoutingTable = - deserialize(&bytes).map_err(ShardManagerError::SerializationError)?; - Ok(routing_table) + decode_shard_state(&bytes) } else { - Ok(RoutingTable::new(self.number_of_shards)) + Ok(ShardLeaseState::new(self.number_of_shards)) + } + } +} + +/// Decodes a persisted state blob and refuses to load one that violates the state invariants. +pub(crate) fn decode_shard_state(bytes: &[u8]) -> Result { + let shard_state: ShardLeaseState = try_deserialize(bytes) + .map_err(ShardManagerError::SerializationError)? + .ok_or_else(|| { + ShardManagerError::SerializationError( + "persisted shard lease state is empty or has an unknown serialization version" + .to_string(), + ) + })?; + shard_state.check_invariants().map_err(|violation| { + ShardManagerError::SerializationError(format!( + "persisted shard lease state violates invariants: {violation}" + )) + })?; + Ok(shard_state) +} + +#[cfg(test)] +mod tests { + use test_r::test; + + use super::*; + use crate::sharding::model::{ExecutorAddr, ExecutorId, ShardAssignmentEntry, ShardEpoch}; + use chrono::{DateTime, Utc}; + use golem_common::model::{Pod, ShardId}; + use std::net::{IpAddr, Ipv4Addr}; + use std::time::Duration; + use uuid::Uuid; + + const TTL: Duration = Duration::from_secs(60); + + fn t0() -> DateTime { + DateTime::from_timestamp(1_700_000_000, 0).unwrap() + } + + fn pod(last_octet: u8, port: u16) -> Pod { + Pod { + ip: IpAddr::V4(Ipv4Addr::new(10, 0, 0, last_octet)), + port, + } + } + + #[test] + fn roundtrips() { + let mut shard_state = ShardLeaseState::new(16); + shard_state.add_executor( + ExecutorId(Uuid::from_u128(1)), + ExecutorAddr::from(pod(1, 9010)), + Some("worker-executor-0".to_string()), + t0(), + TTL, + ); + shard_state.assign_shard(ExecutorId(Uuid::from_u128(1)), ShardId::new(3)); + shard_state.bump_revision().unwrap(); + + let bytes = serialize(&shard_state).unwrap(); + let decoded = decode_shard_state(&bytes).unwrap(); + assert_eq!(decoded, shard_state); + } + + #[test] + fn state_violating_invariants_is_rejected() { + let mut shard_state = ShardLeaseState::new(16); + shard_state.shard_assignments.insert( + ShardId::new(0), + ShardAssignmentEntry { + executor_id: ExecutorId(Uuid::from_u128(7)), + epoch: ShardEpoch::initial(), + }, + ); + let bytes = serialize(&shard_state).unwrap(); + match decode_shard_state(&bytes) { + Err(ShardManagerError::SerializationError(msg)) => { + assert!(msg.contains("violates invariants"), "{msg}"); + } + other => panic!("expected SerializationError, got {other:?}"), + } + } + + #[test] + fn empty_blob_is_rejected() { + match decode_shard_state(&[]) { + Err(ShardManagerError::SerializationError(msg)) => { + assert!(msg.contains("empty"), "{msg}"); + } + other => panic!("expected SerializationError, got {other:?}"), + } + } + + #[test] + fn truncated_blob_is_rejected() { + let bytes = [3u8, 0u8]; + match decode_shard_state(&bytes) { + Err(ShardManagerError::SerializationError(_)) => {} + other => panic!("expected SerializationError, got {other:?}"), } } } diff --git a/golem-shard-manager/src/sharding/rebalancing.rs b/golem-shard-manager/src/sharding/rebalancing.rs index a57550e95b..cf4a0a1e09 100644 --- a/golem-shard-manager/src/sharding/rebalancing.rs +++ b/golem-shard-manager/src/sharding/rebalancing.rs @@ -12,9 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. -use super::model::{Assignments, RoutingTable, Unassignments}; -use golem_common::model::{Pod, ShardId}; -use std::collections::{BTreeSet, HashSet}; +use super::model::{Assignments, ExecutorShards, ShardLeaseState, Unassignments}; +use golem_common::model::ShardId; +use std::collections::HashSet; use std::fmt; use std::fmt::{Display, Formatter}; use tracing::trace; @@ -33,80 +33,82 @@ impl Rebalance { } } - /// Constructs a rebalance plan from the current state of the routing table. + /// Constructs a rebalance plan from the current shard lease state. /// /// The `threshold` parameter is used to reduce the number of shard reassignments by - /// allowing a given number of shards to be over or under the optimal count per pod. + /// allowing a given number of shards to be over or under the optimal count per executor. /// - /// The optimal count (balanced state) is number_of_shards/pod_count. - /// Threshold is a percentage of the optimal count, so for 10 pods with 1000 shards, - /// and a threshold of 10%, pods with shard count between 90 and 110 will be considered + /// The optimal count (balanced state) is number_of_shards/executor_count. + /// Threshold is a percentage of the optimal count, so for 10 executors with 1000 shards, + /// and a threshold of 10%, executors with shard count between 90 and 110 will be considered /// balanced. - pub fn from_routing_table(routing_table: &RoutingTable, threshold: f64) -> Self { + /// + /// Executors are visited in `ExecutorId` order (see [`ShardLeaseState::executor_shard_sets`]). + pub fn from_shard_state(shard_state: &ShardLeaseState, threshold: f64) -> Self { let mut assignments = Assignments::new(); let mut unassignments = Unassignments::new(); - let pod_count = routing_table.get_pod_count(); - if pod_count == 0 { + let executor_count = shard_state.executor_count(); + if executor_count == 0 { return Rebalance { assignments, unassignments, }; } - let mut routing_table_entries = routing_table.get_entries_vec(); - let initial_target_pods: Vec = routing_table_entries + let mut executors: Vec = shard_state.executor_shard_sets(); + let initial_target_executors: Vec = executors .iter() .enumerate() .filter(|&(_idx, entry)| entry.shard_ids.is_empty()) .map(|(idx, _entry)| idx) .collect(); - let optimal_count = routing_table.number_of_shards / pod_count; + let optimal_count = shard_state.number_of_shards / executor_count; let upper_threshold = (optimal_count as f64 * (1.0 + threshold)).ceil() as usize; let lower_threshold = (optimal_count as f64 * (1.0 - threshold)).floor() as usize; // Distributing unassigned shards evenly - let unassigned_shards = routing_table.get_unassigned_shards(); + let unassigned_shards = shard_state.get_unassigned_shards(); let mut unassigned_shards_iter = unassigned_shards.into_iter(); - // First assign to and distribute among empty pods, until all of them reach the optimal count - if !initial_target_pods.is_empty() { - let pod_count = initial_target_pods.len(); - let last_pod_idx = pod_count - 1; + // First assign to and distribute among empty executors, until all of them reach the optimal count + if !initial_target_executors.is_empty() { + let executor_count = initial_target_executors.len(); + let last_executor_idx = executor_count - 1; let mut idx = 0; for shard in unassigned_shards_iter.by_ref() { - let target_idx = initial_target_pods[idx]; - let routing_table_entry = &mut routing_table_entries[target_idx]; + let target_idx = initial_target_executors[idx]; + let executor = &mut executors[target_idx]; trace!( - "Assigning shard to originally empty pod: {} to {}", + "Assigning shard to originally empty executor: {} to {}", shard, target_idx ); - assignments.assign(routing_table_entry.pod, shard); - routing_table_entry.shard_ids.insert(shard); + assignments.assign(executor.executor_id, shard); + executor.shard_ids.insert(shard); - // If the last pod is at optimal count, then all pods are at optimal count - if idx == last_pod_idx && routing_table_entry.shard_ids.len() == optimal_count { + // If the last executor is at optimal count, then all executors are at optimal count + if idx == last_executor_idx && executor.shard_ids.len() == optimal_count { break; } - idx = (idx + 1) % pod_count; + idx = (idx + 1) % executor_count; } } - // Now assign to and distribute among all pods + // Now assign to and distribute among all executors { let mut idx = 0; for shard in unassigned_shards_iter { trace!("Assigning shard: {} to {}", shard, idx); - let routing_table_entry = &mut routing_table_entries[idx]; - assignments.assign(routing_table_entry.pod, shard); - routing_table_entry.shard_ids.insert(shard); - idx = (idx + 1) % pod_count; + let executor = &mut executors[idx]; + assignments.assign(executor.executor_id, shard); + executor.shard_ids.insert(shard); + idx = (idx + 1) % executor_count; } } - if pod_count == 1 { + if executor_count == 1 { return Rebalance { assignments, unassignments, @@ -116,25 +118,25 @@ impl Rebalance { // We redistribute shards from each entry having more than the optimal count // to the last one until it becomes balanced, and repeat if we have more than one unbalanced entry. // We also apply a threshold to the optimal count, to reduce the number of shard reassignments. - for target_idx in 0..routing_table_entries.len() { - for (idx, entry) in routing_table_entries.iter().enumerate() { + for target_idx in 0..executors.len() { + for (idx, entry) in executors.iter().enumerate() { trace!( - "Pod {} has {} shards: {:?}", - idx, - entry.shard_ids.len(), - entry.shard_ids + executor = idx, + shard_count = entry.shard_ids.len(), + shards = ?entry.shard_ids, + "Executor shard count before rebalancing step" ); } - if routing_table_entries[target_idx].shard_ids.len() < lower_threshold { - trace!("Found a pod with too few shards: {}", target_idx); + if executors[target_idx].shard_ids.len() < lower_threshold { + trace!("Found an executor with too few shards: {}", target_idx); loop { trace!("Target count: {}..{}", lower_threshold, upper_threshold); - let current_target_len = routing_table_entries[target_idx].shard_ids.len(); + let current_target_len = executors[target_idx].shard_ids.len(); if current_target_len < lower_threshold { - // Finding a source pod which has more than enough shards - if let Some((source_idx, _)) = routing_table_entries + // Finding a source executor which has more than enough shards + if let Some((source_idx, _)) = executors .iter() .enumerate() .filter(|(idx, entry)| { @@ -143,24 +145,18 @@ impl Rebalance { }) .max_by(|(_, a), (_, b)| a.shard_ids.len().cmp(&b.shard_ids.len())) { - let shard_id = *routing_table_entries[source_idx] - .shard_ids - .iter() - .next() - .unwrap(); + let shard_id = *executors[source_idx].shard_ids.iter().next().unwrap(); // this is guaranteed by check (**) trace!( "Moving first shard from {} to {}: {}", source_idx, target_idx, shard_id ); - routing_table_entries[source_idx] - .shard_ids - .remove(&shard_id); - - routing_table_entries[target_idx].shard_ids.insert(shard_id); - assignments.assign(routing_table_entries[target_idx].pod, shard_id); - unassignments.unassign(routing_table_entries[source_idx].pod, shard_id); - assignments.unassign(routing_table_entries[source_idx].pod, shard_id); + executors[source_idx].shard_ids.remove(&shard_id); + + executors[target_idx].shard_ids.insert(shard_id); + assignments.assign(executors[target_idx].executor_id, shard_id); + unassignments.unassign(executors[source_idx].executor_id, shard_id); + assignments.unassign(executors[source_idx].executor_id, shard_id); } else { trace!("Target reached a balanced state"); // target reached a balanced state @@ -188,37 +184,10 @@ impl Rebalance { &self.unassignments } - pub fn get_pods(&self) -> HashSet { - let mut pods = HashSet::new(); - for pod in self.assignments.assignments.keys() { - pods.insert(*pod); - } - for pod in self.unassignments.unassignments.keys() { - pods.insert(*pod); - } - pods - } - pub fn is_empty(&self) -> bool { self.assignments.assignments.is_empty() && self.unassignments.unassignments.is_empty() } - pub fn empty() -> Self { - Rebalance { - assignments: Assignments::new(), - unassignments: Unassignments::new(), - } - } - - pub fn remove_pods(&mut self, pods: &HashSet) { - self.assignments - .assignments - .retain(|pod, _| !pods.contains(pod)); - self.unassignments - .unassignments - .retain(|pod, _| !pods.contains(pod)); - } - pub fn remove_shards(&mut self, shard_ids: &HashSet) { for assigned_shard_ids in self.assignments.assignments.values_mut() { assigned_shard_ids.retain(|shard_id| !shard_ids.contains(shard_id)); @@ -242,17 +211,6 @@ impl Rebalance { .assignments .retain(|_, shards| !shards.is_empty()); } - - pub fn add_assignments(&mut self, pod: &Pod, mut shard_ids: BTreeSet) { - let empty = BTreeSet::new(); - let unassignments = self.unassignments.unassignments.get(pod).unwrap_or(&empty); - shard_ids.retain(|shard_id| !unassignments.contains(shard_id)); - self.assignments - .assignments - .entry(*pod) - .or_default() - .append(&mut shard_ids); - } } impl Display for Rebalance { @@ -271,85 +229,105 @@ mod tests { use tracing_test::traced_test; - use golem_common::model::{Pod, ShardId}; + use golem_common::model::ShardId; use super::Rebalance; - use crate::sharding::model::RoutingTable; + use crate::sharding::model::{ExecutorAddr, ExecutorId, ShardLeaseState}; + use chrono::{DateTime, Utc}; use std::net::{IpAddr, Ipv4Addr}; + use std::time::Duration; + use uuid::Uuid; struct TestConfig { number_of_shards: usize, - number_of_pods: usize, + number_of_executors: usize, initial_assignments: Vec<(usize, Vec)>, } - fn pod(idx: usize) -> Pod { - Pod { + /// Executor ids are derived from the index so that `BTreeMap` iteration + /// order matches the index order (`Uuid::from_u128` is big-endian) - the rebalancing + /// algorithm is order sensitive and these expectations were written for that order. + fn executor(idx: usize) -> ExecutorId { + ExecutorId(Uuid::from_u128(idx as u128)) + } + + fn addr(idx: usize) -> ExecutorAddr { + ExecutorAddr { ip: IpAddr::V4(Ipv4Addr::new(192, 168, 0, idx as u8)), port: (9000 + idx) as u16, } } + fn now() -> DateTime { + DateTime::from_timestamp(1_700_000_000, 0).unwrap() + } + fn shard_ids(ids: Vec) -> Vec { ids.into_iter().map(ShardId::new).collect() } - fn new_routing_table(config: TestConfig) -> RoutingTable { - let mut routing_table = RoutingTable::new(config.number_of_shards); - for i in 0..config.number_of_pods { - routing_table.add_pod(pod(i), None); + fn new_shard_state(config: TestConfig) -> ShardLeaseState { + let mut shard_state = ShardLeaseState::new(config.number_of_shards); + for i in 0..config.number_of_executors { + shard_state.add_executor(executor(i), addr(i), None, now(), Duration::from_secs(60)); } - for (pod_idx, shards) in config.initial_assignments { - assign_shards(&mut routing_table, pod(pod_idx), shards); + for (executor_idx, shards) in config.initial_assignments { + assign_shards(&mut shard_state, executor(executor_idx), shards); } - routing_table + shard_state } - fn assert_assignments_for_pod(rebalance: &Rebalance, pod: Pod, shards: Vec) { + fn assert_assignments_for_executor( + rebalance: &Rebalance, + executor: ExecutorId, + shards: Vec, + ) { assert_eq!( - get_assigned_ids(rebalance, pod), + get_assigned_ids(rebalance, executor), shard_ids(shards), - "assert_assignments_for_pod: {pod}\n{rebalance:#?}\n", + "assert_assignments_for_executor: {executor}\n{rebalance:#?}\n", ); } fn assert_assignments(rebalance: &Rebalance, assignments: Vec<(usize, Vec)>) { - for (pod_idx, shards) in assignments { - assert_assignments_for_pod(rebalance, pod(pod_idx), shards) + for (executor_idx, shards) in assignments { + assert_assignments_for_executor(rebalance, executor(executor_idx), shards) } } - fn assert_unassignments_for_pod(rebalance: &Rebalance, pod: Pod, shards: Vec) { + fn assert_unassignments_for_executor( + rebalance: &Rebalance, + executor: ExecutorId, + shards: Vec, + ) { assert_eq!( - get_unassigned_ids(rebalance, pod), + get_unassigned_ids(rebalance, executor), shard_ids(shards), - "assert_unassignments_for_pod: {pod}\n{rebalance:#?}\n", + "assert_unassignments_for_executor: {executor}\n{rebalance:#?}\n", ); } fn assert_unassignments(rebalance: &Rebalance, unassignments: Vec<(usize, Vec)>) { - for (pod_idx, shards) in unassignments { - assert_unassignments_for_pod(rebalance, pod(pod_idx), shards) + for (executor_idx, shards) in unassignments { + assert_unassignments_for_executor(rebalance, executor(executor_idx), shards) } } - fn assign_shard(routing_table: &mut RoutingTable, pod: Pod, shard_id: i64) { - if let Some(pod_state) = routing_table.pod_states.get_mut(&pod) { - pod_state.assigned_shards.insert(ShardId::new(shard_id)); - } + fn assign_shard(shard_state: &mut ShardLeaseState, executor: ExecutorId, shard_id: i64) { + shard_state.assign_shard(executor, ShardId::new(shard_id)); } - fn assign_shards(routing_table: &mut RoutingTable, pod: Pod, shard_ids: Vec) { + fn assign_shards(shard_state: &mut ShardLeaseState, executor: ExecutorId, shard_ids: Vec) { for shar_id in shard_ids { - assign_shard(routing_table, pod, shar_id) + assign_shard(shard_state, executor, shar_id) } } - fn get_assigned_ids(rebalance: &Rebalance, pod: Pod) -> Vec { + fn get_assigned_ids(rebalance: &Rebalance, executor: ExecutorId) -> Vec { let mut assigned_ids = rebalance .get_assignments() .assignments - .get(&pod) + .get(&executor) .cloned() .unwrap_or_default() .iter() @@ -359,11 +337,11 @@ mod tests { assigned_ids } - fn get_unassigned_ids(rebalance: &Rebalance, pod: Pod) -> Vec { + fn get_unassigned_ids(rebalance: &Rebalance, executor: ExecutorId) -> Vec { let mut assigned_ids = rebalance .get_unassignments() .unassignments - .get(&pod) + .get(&executor) .cloned() .unwrap_or_default() .iter() @@ -376,39 +354,39 @@ mod tests { #[test] #[traced_test] fn rebalance_empty_table() { - let routing_table = new_routing_table(TestConfig { + let shard_state = new_shard_state(TestConfig { number_of_shards: 1000, - number_of_pods: 0, + number_of_executors: 0, initial_assignments: vec![], }); - let rebalance = Rebalance::from_routing_table(&routing_table, 0.0); + let rebalance = Rebalance::from_shard_state(&shard_state, 0.0); assert!(rebalance.is_empty()); } #[test] #[traced_test] fn rebalance_single_pod_no_unassigned() { - let routing_table = new_routing_table(TestConfig { + let shard_state = new_shard_state(TestConfig { number_of_shards: 4, - number_of_pods: 1, + number_of_executors: 1, initial_assignments: vec![(0, vec![0, 1, 2, 3])], }); - let rebalance = Rebalance::from_routing_table(&routing_table, 0.0); + let rebalance = Rebalance::from_shard_state(&shard_state, 0.0); assert!(rebalance.is_empty()); } #[test] #[traced_test] fn rebalance_single_pod_unassigned() { - let routing_table = new_routing_table(TestConfig { + let shard_state = new_shard_state(TestConfig { number_of_shards: 6, - number_of_pods: 1, + number_of_executors: 1, initial_assignments: vec![(0, vec![0, 3])], }); - let rebalance = Rebalance::from_routing_table(&routing_table, 0.0); + let rebalance = Rebalance::from_shard_state(&shard_state, 0.0); assert!(rebalance.get_unassignments().is_empty()); assert_assignments(&rebalance, vec![(0, vec![1, 2, 4, 5])]); @@ -417,9 +395,9 @@ mod tests { #[test] #[traced_test] fn rebalance_three_balanced_pods_no_unassigned() { - let routing_table = new_routing_table(TestConfig { + let shard_state = new_shard_state(TestConfig { number_of_shards: 9, - number_of_pods: 3, + number_of_executors: 3, initial_assignments: vec![ // (0, vec![0, 1, 2]), @@ -428,16 +406,16 @@ mod tests { ], }); - let rebalance = Rebalance::from_routing_table(&routing_table, 0.0); + let rebalance = Rebalance::from_shard_state(&shard_state, 0.0); assert!(rebalance.is_empty()); } #[test] #[traced_test] fn rebalance_three_balanced_pods_unassigned() { - let routing_table = new_routing_table(TestConfig { + let shard_state = new_shard_state(TestConfig { number_of_shards: 9, - number_of_pods: 3, + number_of_executors: 3, initial_assignments: vec![ // (0, vec![0, 1]), @@ -446,7 +424,7 @@ mod tests { ], }); - let rebalance = Rebalance::from_routing_table(&routing_table, 0.0); + let rebalance = Rebalance::from_shard_state(&shard_state, 0.0); assert!(rebalance.get_unassignments().is_empty()); assert_assignments( @@ -463,9 +441,9 @@ mod tests { #[test] #[traced_test] fn rebalance_one_new_pod() { - let routing_table = new_routing_table(TestConfig { + let shard_state = new_shard_state(TestConfig { number_of_shards: 9, - number_of_pods: 3, + number_of_executors: 3, initial_assignments: vec![ // (0, vec![0, 1, 2, 3]), @@ -473,7 +451,7 @@ mod tests { ], }); - let rebalance = Rebalance::from_routing_table(&routing_table, 0.0); + let rebalance = Rebalance::from_shard_state(&shard_state, 0.0); assert_assignments( &rebalance, @@ -499,9 +477,9 @@ mod tests { #[test] #[traced_test] fn rebalance_one_new_pod_with_threshold() { - let routing_table = new_routing_table(TestConfig { + let shard_state = new_shard_state(TestConfig { number_of_shards: 9, - number_of_pods: 3, + number_of_executors: 3, initial_assignments: vec![ // (0, vec![0, 1, 2, 3]), @@ -509,7 +487,7 @@ mod tests { ], }); - let rebalance = Rebalance::from_routing_table(&routing_table, 0.33); + let rebalance = Rebalance::from_shard_state(&shard_state, 0.33); assert_assignments( &rebalance, @@ -537,9 +515,9 @@ mod tests { fn rebalance_one_new_pod_after_removing_two() { // 3,4,5 and 9,10,11 are unassigned // pod3 is empty - let routing_table = new_routing_table(TestConfig { + let shard_state = new_shard_state(TestConfig { number_of_shards: 12, - number_of_pods: 3, + number_of_executors: 3, initial_assignments: vec![ // (0, vec![0, 1, 2]), @@ -547,7 +525,7 @@ mod tests { ], }); - let rebalance = Rebalance::from_routing_table(&routing_table, 0.0); + let rebalance = Rebalance::from_shard_state(&shard_state, 0.0); assert_assignments( &rebalance, @@ -573,13 +551,13 @@ mod tests { #[test] #[traced_test] fn rebalance_two_new_pods() { - let routing_table = new_routing_table(TestConfig { + let shard_state = new_shard_state(TestConfig { number_of_shards: 9, - number_of_pods: 3, + number_of_executors: 3, initial_assignments: vec![(0, vec![0, 1, 2, 3, 4, 5, 6, 7, 8])], }); - let rebalance = Rebalance::from_routing_table(&routing_table, 0.0); + let rebalance = Rebalance::from_shard_state(&shard_state, 0.0); assert_assignments( &rebalance, @@ -608,9 +586,9 @@ mod tests { // pod1 and pod2 has 4-4 shards because previously we had 3 pods for 12 shards // 4,5,6,11 are unassigned // pod3 and pod4 are empty - let routing_table = new_routing_table(TestConfig { + let shard_state = new_shard_state(TestConfig { number_of_shards: 12, - number_of_pods: 4, + number_of_executors: 4, initial_assignments: vec![ // (0, vec![0, 1, 2, 3]), @@ -618,7 +596,7 @@ mod tests { ], }); - let rebalance = Rebalance::from_routing_table(&routing_table, 0.0); + let rebalance = Rebalance::from_shard_state(&shard_state, 0.0); assert_assignments( &rebalance, @@ -648,13 +626,13 @@ mod tests { fn two_empty_pods_one_filled() { // pod2 is empty // pod3 is new and empty - let routing_table = new_routing_table(TestConfig { + let shard_state = new_shard_state(TestConfig { number_of_shards: 9, - number_of_pods: 3, + number_of_executors: 3, initial_assignments: vec![(0, vec![3, 4, 5])], }); - let rebalance = Rebalance::from_routing_table(&routing_table, 0.0); + let rebalance = Rebalance::from_shard_state(&shard_state, 0.0); assert_assignments( &rebalance, @@ -670,13 +648,13 @@ mod tests { #[test] #[traced_test] fn initial_assign_is_ordered_and_no_rebalance_needed() { - let routing_table = new_routing_table(TestConfig { + let shard_state = new_shard_state(TestConfig { number_of_shards: 8, - number_of_pods: 4, + number_of_executors: 4, initial_assignments: vec![], }); - let rebalance = Rebalance::from_routing_table(&routing_table, 0.0); + let rebalance = Rebalance::from_shard_state(&shard_state, 0.0); assert_assignments( &rebalance, @@ -695,13 +673,13 @@ mod tests { #[test] #[traced_test] fn initial_assign_is_ordered_and_no_rebalance_needed_with_less_then_opt() { - let routing_table = new_routing_table(TestConfig { + let shard_state = new_shard_state(TestConfig { number_of_shards: 14, - number_of_pods: 4, + number_of_executors: 4, initial_assignments: vec![], }); - let rebalance = Rebalance::from_routing_table(&routing_table, 0.0); + let rebalance = Rebalance::from_shard_state(&shard_state, 0.0); assert_assignments( &rebalance, @@ -720,13 +698,13 @@ mod tests { #[test] #[traced_test] fn initial_assign_is_ordered_and_no_rebalance_with_some_saturated_pod() { - let routing_table = new_routing_table(TestConfig { + let shard_state = new_shard_state(TestConfig { number_of_shards: 8, - number_of_pods: 4, + number_of_executors: 4, initial_assignments: vec![(0, vec![0, 1])], }); - let rebalance = Rebalance::from_routing_table(&routing_table, 0.0); + let rebalance = Rebalance::from_shard_state(&shard_state, 0.0); assert_assignments( &rebalance, diff --git a/golem-shard-manager/src/sharding/shard_management.rs b/golem-shard-manager/src/sharding/shard_management.rs index 7a8eaf9b8f..08dd9a905f 100644 --- a/golem-shard-manager/src/sharding/shard_management.rs +++ b/golem-shard-manager/src/sharding/shard_management.rs @@ -13,76 +13,81 @@ // limitations under the License. use super::error::ShardManagerError; -use super::healthcheck::{HealthCheck, get_unhealthy_pods}; -use super::model::{Assignments, RoutingTable}; +use super::healthcheck::{HealthCheck, get_unhealthy_executors}; +use super::model::{Assignments, ExecutorAddr, ExecutorAddrs, ExecutorId, ShardLeaseState}; use super::persistence::RoutingTablePersistence; use super::rebalancing::Rebalance; use super::worker_executor::{ WorkerExecutorService, assign_shards, revoke_shards, set_shard_assignments, }; use async_rwlock::RwLock; -use golem_common::model::{Pod, ShardId}; +use chrono::Utc; +use golem_common::model::ShardId; use itertools::Itertools; -use std::collections::{BTreeSet, HashMap, HashSet}; +use std::collections::{BTreeMap, BTreeSet, HashSet}; use std::sync::Arc; +use std::time::Duration; use tokio::sync::{Mutex, Notify}; use tokio::task::JoinSet; use tracing::{Instrument, debug, info, warn}; #[derive(Clone)] pub struct ShardManagement { - routing_table: Arc>, + shard_state: Arc>, change: Arc, updates: Arc>, } impl ShardManagement { - /// Initializes the shard management with an initial routing table and optionally - /// a pending rebalance, both read from the persistence service. + /// Initializes the shard management with the persisted shard lease state. + /// + /// Executors found in the persisted state are health checked once: unhealthy ones are + /// removed, healthy ones receive their authoritative full shard assignment (they might be + /// lagging after interleaved shard-manager and executor restarts). pub async fn new( persistence_service: Arc, worker_executors: Arc, health_check: Arc, threshold: f64, + lease_ttl: Duration, join_set: &mut JoinSet>, ) -> Result { - let routing_table = persistence_service.read().await?; + let shard_state = persistence_service.read().await?; info!("Initial healthcheck started"); - let pods = routing_table.get_pods_with_names(); - - let unhealthy_pods = get_unhealthy_pods(&health_check, &pods).await; - let healthy_pods = pods - .into_iter() - .filter(|(p, _)| !unhealthy_pods.contains(p)) + let executors = shard_state.get_executors_with_addrs(); + let unhealthy_executors = get_unhealthy_executors(&health_check, &executors).await; + let healthy_executors: HashSet = executors + .iter() + .map(|(id, _, _)| *id) + .filter(|id| !unhealthy_executors.contains(id)) .collect(); info!("Initial healthcheck finished"); let change = Arc::new(Notify::new()); - // NOTE: We consider all healthy pods as new pods to trigger full assigment, given they might be lagging: - // this can happen with interleaved shard-manager and worker restarts let updates = Arc::new(Mutex::new(ShardManagementChanges::new( - healthy_pods, - unhealthy_pods, + healthy_executors, + unhealthy_executors, ))); - let routing_table = Arc::new(RwLock::new(routing_table)); + let shard_state = Arc::new(RwLock::new(shard_state)); { let change = change.clone(); let updates = updates.clone(); - let routing_table = routing_table.clone(); + let shard_state = shard_state.clone(); join_set.spawn( async move { Self::worker( - routing_table, + shard_state, change, updates, persistence_service, worker_executors, threshold, + lease_ttl, ) .await; Ok(()) @@ -94,100 +99,110 @@ impl ShardManagement { change.notify_one(); Ok(ShardManagement { - routing_table, + shard_state, change, updates, }) } - /// Registers a new pod to be added - pub async fn register_pod(&self, pod: Pod, pod_name: Option) { - debug!(pod=%pod, "Registering pod"); - self.updates.lock().await.add_new_pod(pod, pod_name); + /// Registers a new executor instance listening at `addr`. + /// + /// Every registration is a new executor instance and receives a fresh [`ExecutorId`], + /// which is returned. If another instance is still registered at the same address it is + /// replaced on the next rebalance pass. + pub async fn register_executor( + &self, + addr: ExecutorAddr, + pod_name: Option, + ) -> ExecutorId { + let executor_id = ExecutorId::generate(); + debug!(executor_id = %executor_id, addr = %addr, "Registering executor"); + self.updates + .lock() + .await + .add_executor_registration(executor_id, addr, pod_name); self.change.notify_one(); + executor_id } - /// Marks a pod to be removed - pub async fn unregister_pod(&self, pod: Pod) { - debug!(pod=%pod, "Unregistering pod"); - self.updates.lock().await.remove_pod(pod); + /// Marks an executor to be removed + pub async fn unregister_executor(&self, executor_id: ExecutorId) { + debug!(executor_id = %executor_id, "Unregistering executor"); + self.updates.lock().await.remove_executor(executor_id); self.change.notify_one(); } - /// Gets the current snapshot of the routing table - pub async fn current_snapshot(&self) -> RoutingTable { - self.routing_table.read().await.clone() + /// Gets the current snapshot of the shard lease state + pub async fn current_snapshot(&self) -> ShardLeaseState { + self.shard_state.read().await.clone() } async fn worker( - routing_table: Arc>, + shard_state: Arc>, change: Arc, updates: Arc>, persistence_service: Arc, worker_executors: Arc, threshold: f64, + lease_ttl: Duration, ) { loop { debug!("Shard management loop awaiting changes"); change.notified().await; - let (new_pods, removed_pods, retry_full_assignment_pods) = updates.lock().await.reset(); + let (new_executors, removed_executors, full_assignment_requests) = + updates.lock().await.reset(); debug!( - new_pods = new_pods.keys().join(", "), - removed_pods = removed_pods.iter().join(", "), - retry_pods = retry_full_assignment_pods.iter().join(", "), + new_executors = new_executors + .values() + .map(|r| format!("{} ({})", r.executor_id, r.addr)) + .join(", "), + removed_executors = removed_executors.iter().join(", "), + full_assignment_requests = full_assignment_requests.iter().join(", "), "Shard management loop woken up", ); - // Getting a write lock while + // - registrations and removals are applied to the state and got persisted, // - the rebalance plan is calculated, - // - new and removed pods are added to the routing table and got persisted, - // but the rebalance plan is NOT applied yet. The lock is then release for apply. - let (mut rebalance, full_assignment_pods) = { - let mut current_routing_table = routing_table.write().await; - - for pod in removed_pods { - current_routing_table.remove_pod(pod); - info!(pod= %pod, "Pod removed"); - } - - let mut send_full_assignment = Vec::new(); - for (pod, pod_name) in new_pods { - if current_routing_table.has_pod(pod) { - // This pod has already an assignment - we have to send the full list of assigned shards to it - send_full_assignment.push(pod); - info!(pod= %pod, "Pod returned"); - } else { - // New pod, adding with empty assignment - current_routing_table.add_pod(pod, pod_name); - info!(pod= %pod, "Pod added"); - } + // but the rebalance plan is NOT applied yet. The lock is then released for apply. + let (mut rebalance, full_assignment_executors, addrs) = { + let mut current_shard_state = shard_state.write().await; + + // Shards orphaned by lease removals since the last pass. The rebalance plan + // below recomputes all unassigned shards from scratch, so this is only logged. + let pending = current_shard_state.take_pending_rebalance(); + if !pending.is_empty() { + debug!( + shards = pending.iter().join(", "), + "Redistributing shards orphaned since the last pass" + ); } - let rebalance = Rebalance::from_routing_table(¤t_routing_table, threshold); - let mut full_assignment_pods: HashSet = HashSet::new(); - - for pod in send_full_assignment { - full_assignment_pods.insert(pod); - } + let full_assignment_executors = apply_executor_changes( + &mut current_shard_state, + new_executors, + removed_executors, + full_assignment_requests, + lease_ttl, + ); - for pod in retry_full_assignment_pods { - if current_routing_table.has_pod(pod) { - full_assignment_pods.insert(pod); - } - } + let rebalance = Rebalance::from_shard_state(¤t_shard_state, threshold); + let addrs = current_shard_state.executor_addrs(); + current_shard_state + .bump_revision() + .expect("Failed to bump shard lease state revision"); persistence_service - .write(¤t_routing_table) + .write(¤t_shard_state) .await - .expect("Failed to persist routing table after pod changes"); + .expect("Failed to persist shard lease state after executor changes"); - (rebalance, full_assignment_pods) + (rebalance, full_assignment_executors, addrs) }; debug!(rebalance=%rebalance, "Applying rebalance plan"); let rebalance_failures = - Self::execute_rebalance(worker_executors.clone(), &mut rebalance).await; + Self::execute_rebalance(worker_executors.clone(), &mut rebalance, &addrs).await; let mut needs_retry = false; if !rebalance_failures.failed_assignments.is_empty() { @@ -205,9 +220,9 @@ impl ShardManagement { { let mut updates_guard = updates.lock().await; - for (pod, _) in &rebalance_failures.failed_assignments { - if full_assignment_pods.contains(pod) { - updates_guard.retry_full_assignment(*pod); + for (executor_id, _) in &rebalance_failures.failed_assignments { + if full_assignment_executors.contains(executor_id) { + updates_guard.retry_full_assignment(*executor_id); } } } @@ -216,30 +231,36 @@ impl ShardManagement { if !rebalance_failures.failed_unassignments.is_empty() { warn!( - failed_pods = rebalance_failures + failed_executors = rebalance_failures .failed_unassignments .iter() - .map(|(pod, _)| pod) + .map(|(executor_id, _)| executor_id) .join(", "), "Some shards could not be unassigned and rebalance will be retried" ); needs_retry = true; } - routing_table.write().await.rebalance(rebalance); - - let routing_table_snapshot = routing_table.read().await.clone(); + let shard_state_snapshot = { + let mut current_shard_state = shard_state.write().await; + current_shard_state.apply_rebalance(&rebalance); + current_shard_state + .bump_revision() + .expect("Failed to bump shard lease state revision"); + current_shard_state.clone() + }; persistence_service - .write(&routing_table_snapshot) + .write(&shard_state_snapshot) .await - .expect("Failed to persist routing table after rebalance"); + .expect("Failed to persist shard lease state after rebalance"); let mut full_assignments = Assignments::new(); - for pod in &full_assignment_pods { - if let Some(mut shard_ids) = routing_table_snapshot.get_shards(*pod) { + for executor_id in &full_assignment_executors { + if let Some(mut shard_ids) = shard_state_snapshot.shards_for_executor(*executor_id) + { full_assignments .assignments - .entry(*pod) + .entry(*executor_id) .or_default() .append(&mut shard_ids); } @@ -250,25 +271,26 @@ impl ShardManagement { } else { set_shard_assignments( worker_executors.clone(), - routing_table_snapshot.number_of_shards, + shard_state_snapshot.number_of_shards, &full_assignments, + &addrs, ) .await }; if !failed_full_assignments.is_empty() { warn!( - failed_pods = failed_full_assignments + failed_executors = failed_full_assignments .iter() - .map(|(pod, _)| pod) + .map(|(executor_id, _)| executor_id) .join(", "), - "Some pods could not receive authoritative shard assignment and will be retried" + "Some executors could not receive authoritative shard assignment and will be retried" ); { let mut updates_guard = updates.lock().await; - for (pod, _) in &failed_full_assignments { - updates_guard.retry_full_assignment(*pod); + for (executor_id, _) in &failed_full_assignments { + updates_guard.retry_full_assignment(*executor_id); } } needs_retry = true; @@ -283,6 +305,7 @@ impl ShardManagement { async fn execute_rebalance( worker_executors: Arc, rebalance: &mut Rebalance, + addrs: &ExecutorAddrs, ) -> RebalanceFailures { info!("Beginning rebalance..."); @@ -292,8 +315,12 @@ impl ShardManagement { "Executing shard unassignments", ); } - let failed_unassignments = - revoke_shards(worker_executors.clone(), rebalance.get_unassignments()).await; + let failed_unassignments = revoke_shards( + worker_executors.clone(), + rebalance.get_unassignments(), + addrs, + ) + .await; let failed_shards = failed_unassignments .iter() .flat_map(|(_, shard_ids)| shard_ids.clone()) @@ -314,7 +341,7 @@ impl ShardManagement { } let failed_assignments = - assign_shards(worker_executors.clone(), rebalance.get_assignments()).await; + assign_shards(worker_executors.clone(), rebalance.get_assignments(), addrs).await; RebalanceFailures { failed_assignments, @@ -323,53 +350,292 @@ impl ShardManagement { } } +/// Applies the executor changes queued since the last pass and returns the executors that must +/// receive their full shard assignment. +/// +/// New executors are added before removed ones are dropped, so that a restart re-registering at +/// an address whose previous instance was also reported unhealthy still takes over its shards; +/// the removal of the replaced instance is then a no-op. +fn apply_executor_changes( + shard_state: &mut ShardLeaseState, + new_executors: BTreeMap, + removed_executors: HashSet, + full_assignment_requests: HashSet, + lease_ttl: Duration, +) -> HashSet { + // Every lease granted in this pass starts at the same instant. + let now = Utc::now(); + let mut full_assignment_executors: HashSet = HashSet::new(); + + for registration in new_executors.into_values() { + let ExecutorRegistration { + executor_id, + addr, + pod_name, + } = registration; + match shard_state.add_executor(executor_id, addr, pod_name, now, lease_ttl) { + Some(replaced) => { + // A new instance at a known address: its predecessor's shards were transferred + // to it, and it has to receive the full list of its assigned shards. + full_assignment_executors.insert(executor_id); + info!( + executor_id = %executor_id, + replaced_executor_id = %replaced, + addr = %addr, + "Executor replaced at address" + ); + } + None => { + info!(executor_id = %executor_id, addr = %addr, "Executor added"); + } + } + } + + for executor_id in removed_executors { + if !shard_state.has_executor(executor_id) { + debug!( + executor_id = %executor_id, + "Executor to be removed is no longer registered" + ); + continue; + } + let released = shard_state.remove_executor(executor_id); + info!( + executor_id = %executor_id, + released_shards = released.len(), + "Executor removed" + ); + } + + for executor_id in full_assignment_requests { + if shard_state.has_executor(executor_id) { + full_assignment_executors.insert(executor_id); + } + } + + full_assignment_executors +} + #[derive(Debug)] struct RebalanceFailures { - failed_assignments: Vec<(Pod, BTreeSet)>, - failed_unassignments: Vec<(Pod, BTreeSet)>, + failed_assignments: Vec<(ExecutorId, BTreeSet)>, + failed_unassignments: Vec<(ExecutorId, BTreeSet)>, } +#[derive(Debug, Clone)] +struct ExecutorRegistration { + executor_id: ExecutorId, + addr: ExecutorAddr, + pod_name: Option, +} + +/// Changes accumulated between two passes of the shard management loop. #[derive(Debug)] struct ShardManagementChanges { - new_pods: HashMap>, - removed_pods: HashSet, - retry_full_assignment_pods: HashSet, + new_executors: BTreeMap, + removed_executors: HashSet, + full_assignment_requests: HashSet, } impl ShardManagementChanges { - pub fn new(new_pods: HashMap>, removed_pods: HashSet) -> Self { + pub fn new( + full_assignment_requests: HashSet, + removed_executors: HashSet, + ) -> Self { ShardManagementChanges { - new_pods, - removed_pods, - retry_full_assignment_pods: HashSet::new(), + new_executors: BTreeMap::new(), + removed_executors, + full_assignment_requests, + } + } + + /// Queues a registration. `executor_id` is always freshly minted, so it cannot already be + /// queued for removal or for a full assignment. + pub fn add_executor_registration( + &mut self, + executor_id: ExecutorId, + addr: ExecutorAddr, + pod_name: Option, + ) { + self.new_executors.insert( + addr, + ExecutorRegistration { + executor_id, + addr, + pod_name, + }, + ); + } + + pub fn remove_executor(&mut self, executor_id: ExecutorId) { + self.new_executors + .retain(|_, registration| registration.executor_id != executor_id); + self.full_assignment_requests.remove(&executor_id); + self.removed_executors.insert(executor_id); + } + + pub fn retry_full_assignment(&mut self, executor_id: ExecutorId) { + if !self.removed_executors.contains(&executor_id) { + self.full_assignment_requests.insert(executor_id); + } + } + + #[allow(clippy::type_complexity)] + pub fn reset( + &mut self, + ) -> ( + BTreeMap, + HashSet, + HashSet, + ) { + let new_executors = std::mem::take(&mut self.new_executors); + let removed = std::mem::take(&mut self.removed_executors); + let full = std::mem::take(&mut self.full_assignment_requests); + (new_executors, removed, full) + } +} + +#[cfg(test)] +mod tests { + use test_r::test; + + use super::*; + use crate::sharding::model::ShardEpoch; + use golem_common::model::ShardId; + use std::net::{IpAddr, Ipv4Addr}; + use uuid::Uuid; + + const TTL: Duration = Duration::from_secs(60); + + fn t0() -> chrono::DateTime { + chrono::DateTime::from_timestamp(1_700_000_000, 0).unwrap() + } + + fn executor(idx: u128) -> ExecutorId { + ExecutorId(Uuid::from_u128(idx)) + } + + fn addr(idx: u8) -> ExecutorAddr { + ExecutorAddr { + ip: IpAddr::V4(Ipv4Addr::new(10, 0, 0, idx)), + port: 9000 + idx as u16, } } - pub fn add_new_pod(&mut self, pod: Pod, pod_name: Option) { - self.removed_pods.remove(&pod); - self.retry_full_assignment_pods.remove(&pod); - self.new_pods.insert(pod, pod_name); + fn shards(ids: &[i64]) -> BTreeSet { + ids.iter().copied().map(ShardId::new).collect() } - pub fn remove_pod(&mut self, pod: Pod) { - self.new_pods.remove(&pod); - self.retry_full_assignment_pods.remove(&pod); - self.removed_pods.insert(pod); + fn executor_registration(executor_id: ExecutorId, addr: ExecutorAddr) -> ExecutorRegistration { + ExecutorRegistration { + executor_id, + addr, + pod_name: None, + } } - pub fn retry_full_assignment(&mut self, pod: Pod) { - if !self.removed_pods.contains(&pod) { - self.retry_full_assignment_pods.insert(pod); + #[test] + fn removal_and_reregistration_at_the_same_address_in_one_pass_transfer_the_shards() { + let mut shard_state = ShardLeaseState::new(4); + shard_state.add_executor(executor(1), addr(1), None, t0(), TTL); + shard_state.add_executor(executor(2), addr(2), None, t0(), TTL); + for shard_id in [0, 1] { + shard_state.assign_shard(executor(1), ShardId::new(shard_id)); + } + for shard_id in [2, 3] { + shard_state.assign_shard(executor(2), ShardId::new(shard_id)); } + + // the health check reported executor 1 unhealthy and the restarted process at the same + // address registered (as executor 3) before the loop woke up + let new_executors = + BTreeMap::from([(addr(1), executor_registration(executor(3), addr(1)))]); + let removed = HashSet::from([executor(1)]); + + let full = apply_executor_changes( + &mut shard_state, + new_executors, + removed, + HashSet::new(), + TTL, + ); + + assert_eq!(full, HashSet::from([executor(3)])); + assert!(!shard_state.has_executor(executor(1))); + assert_eq!( + shard_state.shards_for_executor(executor(3)), + Some(shards(&[0, 1])) + ); + assert_eq!( + shard_state.shards_for_executor(executor(2)), + Some(shards(&[2, 3])) + ); + assert_eq!( + shard_state.epoch_for_shard(ShardId::new(0)), + Some(ShardEpoch(1)) + ); + assert!(shard_state.get_unassigned_shards().is_empty()); + assert!(shard_state.pending_rebalance.is_empty()); + assert!(shard_state.check_invariants().is_ok()); + } + + #[test] + fn removal_of_an_executor_releases_its_shards_and_full_requests_are_filtered() { + let mut shard_state = ShardLeaseState::new(4); + shard_state.add_executor(executor(1), addr(1), None, t0(), TTL); + shard_state.add_executor(executor(2), addr(2), None, t0(), TTL); + shard_state.assign_shard(executor(1), ShardId::new(0)); + shard_state.assign_shard(executor(2), ShardId::new(1)); + + let full = apply_executor_changes( + &mut shard_state, + BTreeMap::new(), + HashSet::from([executor(1)]), + HashSet::from([executor(1), executor(2), executor(9)]), + TTL, + ); + + assert_eq!(full, HashSet::from([executor(2)])); + assert!(!shard_state.has_executor(executor(1))); + assert_eq!(shard_state.pending_rebalance, shards(&[0])); + assert_eq!(shard_state.get_unassigned_shards(), shards(&[0, 2, 3])); + } + + #[test] + fn brand_new_executor_is_added_without_a_full_assignment() { + let mut shard_state = ShardLeaseState::new(4); + shard_state.add_executor(executor(1), addr(1), None, t0(), TTL); + + let full = apply_executor_changes( + &mut shard_state, + BTreeMap::from([(addr(2), executor_registration(executor(2), addr(2)))]), + HashSet::new(), + HashSet::new(), + TTL, + ); + + assert!(full.is_empty()); + assert!(shard_state.has_executor(executor(2))); + assert_eq!(shard_state.executor_count(), 2); } - pub fn reset(&mut self) -> (HashMap>, HashSet, HashSet) { - let new = self.new_pods.clone(); - let removed = self.removed_pods.clone(); - let retry = self.retry_full_assignment_pods.clone(); - self.new_pods.clear(); - self.removed_pods.clear(); - self.retry_full_assignment_pods.clear(); - (new, removed, retry) + #[test] + fn unregistering_drops_a_queued_registration_of_the_same_executor() { + let mut changes = ShardManagementChanges::new(HashSet::new(), HashSet::new()); + changes.add_executor_registration(executor(1), addr(1), None); + changes.add_executor_registration(executor(2), addr(1), None); // last registration per address wins + changes.add_executor_registration(executor(3), addr(3), None); + changes.remove_executor(executor(3)); + changes.retry_full_assignment(executor(3)); // ignored: queued for removal + changes.retry_full_assignment(executor(4)); + + let (new_executors, removed, full) = changes.reset(); + assert_eq!(new_executors.len(), 1); + assert_eq!(new_executors[&addr(1)].executor_id, executor(2)); + assert_eq!(removed, HashSet::from([executor(3)])); + assert_eq!(full, HashSet::from([executor(4)])); + + let (new_executors, removed, full) = changes.reset(); + assert!(new_executors.is_empty() && removed.is_empty() && full.is_empty()); } } diff --git a/golem-shard-manager/src/sharding/worker_executor.rs b/golem-shard-manager/src/sharding/worker_executor.rs index 313db8345d..1709cec21f 100644 --- a/golem-shard-manager/src/sharding/worker_executor.rs +++ b/golem-shard-manager/src/sharding/worker_executor.rs @@ -13,9 +13,12 @@ // limitations under the License. use super::error::{HealthCheckError, ShardManagerError}; -use super::model::{Assignments, Unassignments, pod_shard_assignments_to_string}; +use super::model::{ + Assignments, ExecutorAddrs, ExecutorId, Unassignments, shard_assignments_to_string, +}; use crate::config::WorkerExecutorServiceConfig; use async_trait::async_trait; +use futures::future::BoxFuture; use golem_api_grpc::proto::golem; use golem_api_grpc::proto::golem::workerexecutor::v1::worker_executor_client::WorkerExecutorClient; use golem_common::model::Pod; @@ -23,7 +26,7 @@ use golem_common::model::ShardId; use golem_common::retries::with_retriable_errors; use golem_service_base::error::worker_executor::WorkerExecutorError; use golem_service_base::grpc::client::MultiTargetGrpcClient; -use std::collections::BTreeSet; +use std::collections::{BTreeMap, BTreeSet}; use std::sync::Arc; use tokio::time::error::Elapsed; use tokio::time::timeout; @@ -34,7 +37,7 @@ use tonic_health::pb::health_check_response::ServingStatus; use tonic_health::pb::health_client::HealthClient; use tonic_health::pb::{HealthCheckRequest, HealthCheckResponse}; use tonic_tracing_opentelemetry::middleware::client::OtelGrpcService; -use tracing::info; +use tracing::{info, warn}; #[async_trait] pub trait WorkerExecutorService: Send + Sync { @@ -62,74 +65,94 @@ pub trait WorkerExecutorService: Send + Sync { /// Sends revoke requests to all worker executors based on an `Unassignments` plan pub async fn revoke_shards( - worker_executors: Arc, + worker_executors: Arc, unassignments: &Unassignments, -) -> Vec<(Pod, BTreeSet)> { - let futures: Vec<_> = unassignments - .unassignments - .iter() - .map(|(pod, shard_ids)| { + addrs: &ExecutorAddrs, +) -> Vec<(ExecutorId, BTreeSet)> { + fan_out( + &unassignments.unassignments, + addrs, + "revoke_shards", + |pod, shard_ids| { let worker_executors = worker_executors.clone(); - Box::pin(async move { - match worker_executors.revoke_shards(pod, shard_ids).await { - Ok(_) => None, - Err(_) => Some((*pod, shard_ids.clone())), - } - }) - }) - .collect(); - futures::future::join_all(futures) - .await - .into_iter() - .flatten() - .collect() + Box::pin(async move { worker_executors.revoke_shards(&pod, shard_ids).await }) + }, + ) + .await } -/// Sends assign requests to all worker executors based on an `Assignments` plan +/// Sends assign requests to all worker executors based on an `Assignments` plan. pub async fn assign_shards( worker_executors: Arc, assignments: &Assignments, -) -> Vec<(Pod, BTreeSet)> { - let futures: Vec<_> = assignments - .assignments - .iter() - .map(|(pod, shard_ids)| { + addrs: &ExecutorAddrs, +) -> Vec<(ExecutorId, BTreeSet)> { + fan_out( + &assignments.assignments, + addrs, + "assign_shards", + |pod, shard_ids| { let worker_executors = worker_executors.clone(); - Box::pin(async move { - match worker_executors.assign_shards(pod, shard_ids).await { - Ok(_) => None, - Err(_) => Some((*pod, shard_ids.clone())), - } - }) - }) - .collect(); - futures::future::join_all(futures) - .await - .into_iter() - .flatten() - .collect() + Box::pin(async move { worker_executors.assign_shards(&pod, shard_ids).await }) + }, + ) + .await } -/// Reconciles executors to the routing-table shard assignments. +/// Reconciles executors to the authoritative shard assignments. pub async fn set_shard_assignments( worker_executors: Arc, number_of_shards: usize, assignments: &Assignments, -) -> Vec<(Pod, BTreeSet)> { - let futures: Vec<_> = assignments - .assignments - .iter() - .map(|(pod, shard_ids)| { + addrs: &ExecutorAddrs, +) -> Vec<(ExecutorId, BTreeSet)> { + fan_out( + &assignments.assignments, + addrs, + "set_shard_assignment", + |pod, shard_ids| { let worker_executors = worker_executors.clone(); Box::pin(async move { - match worker_executors - .set_shard_assignment(pod, number_of_shards, shard_ids) + worker_executors + .set_shard_assignment(&pod, number_of_shards, shard_ids) .await - { - Ok(_) => None, - Err(_) => Some((*pod, shard_ids.clone())), - } }) + }, + ) + .await +} + +async fn fan_out<'a, F>( + plan: &'a BTreeMap>, + addrs: &ExecutorAddrs, + operation: &'static str, + call: F, +) -> Vec<(ExecutorId, BTreeSet)> +where + F: Fn(Pod, &'a BTreeSet) -> BoxFuture<'a, Result<(), ShardManagerError>>, +{ + let futures: Vec<_> = plan + .iter() + .map(|(executor_id, shard_ids)| { + let call = addrs + .get(executor_id) + .map(|addr| call(Pod::from(*addr), shard_ids)); + async move { + match call { + None => { + warn!( + executor_id = %executor_id, + operation, + "Executor has no known address; reporting the operation as failed" + ); + Some((*executor_id, shard_ids.clone())) + } + Some(call) => match call.await { + Ok(_) => None, + Err(_) => Some((*executor_id, shard_ids.clone())), + }, + } + } }) .collect(); futures::future::join_all(futures) @@ -152,7 +175,7 @@ impl WorkerExecutorService for WorkerExecutorServiceDefault { shard_ids: &BTreeSet, ) -> Result<(), ShardManagerError> { info!( - assigned_shards = pod_shard_assignments_to_string(pod, None, shard_ids.iter()), + assigned_shards = shard_assignments_to_string(pod, None, shard_ids.iter()), "Assigning shards", ); @@ -199,7 +222,7 @@ impl WorkerExecutorService for WorkerExecutorServiceDefault { shard_ids: &BTreeSet, ) -> Result<(), ShardManagerError> { info!( - revoked_shards = pod_shard_assignments_to_string(pod, None, shard_ids.iter()), + revoked_shards = shard_assignments_to_string(pod, None, shard_ids.iter()), "Revoking shards", ); @@ -221,7 +244,7 @@ impl WorkerExecutorService for WorkerExecutorServiceDefault { shard_ids: &BTreeSet, ) -> Result<(), ShardManagerError> { info!( - assigned_shards = pod_shard_assignments_to_string(pod, None, shard_ids.iter()), + assigned_shards = shard_assignments_to_string(pod, None, shard_ids.iter()), number_of_shards, "Setting authoritative shard assignment", ); diff --git a/golem-shard-manager/tests/persistence.rs b/golem-shard-manager/tests/persistence.rs index 83cee92876..ee673fcc94 100644 --- a/golem-shard-manager/tests/persistence.rs +++ b/golem-shard-manager/tests/persistence.rs @@ -13,21 +13,25 @@ // limitations under the License. use async_trait::async_trait; +use chrono::{DateTime, Utc}; use golem_common::config::{DbPostgresConfig, DbSqliteConfig}; -use golem_common::model::{Pod, ShardId}; +use golem_common::model::ShardId; use golem_service_base::migration::{IncludedMigrationsDir, Migrations}; use golem_shard_manager::{ - DbRoutingTablePersistence, PodState, RoutingTable, RoutingTablePersistence, + DbRoutingTablePersistence, ExecutorAddr, ExecutorId, RoutingTablePersistence, + ShardLeaseRevision, ShardLeaseState, }; use golem_test_framework::components::rdb::docker_postgres::DockerPostgresRdb; -use std::collections::{BTreeMap, BTreeSet}; use std::net::{IpAddr, Ipv4Addr}; use std::sync::Arc; +use std::time::Duration; use tempfile::TempDir; use test_r::{define_matrix_dimension, test, test_dep}; use url::Url; use uuid::Uuid; +const LEASE_TTL: Duration = Duration::from_secs(60); + #[async_trait] trait GetRoutingTablePersistence: std::fmt::Debug + Send + Sync { async fn get_persistence(&self) -> Arc; @@ -151,13 +155,16 @@ async fn read_returns_default_when_empty( #[dimension(persistence)] persistence: &Arc, ) { let persistence = persistence.get_persistence().await; - let routing_table = persistence + let shard_state = persistence .read() .await - .expect("Reading default routing table should succeed"); + .expect("Reading default shard lease state should succeed"); - assert_eq!(routing_table.number_of_shards, 16); - assert!(routing_table.pod_states.is_empty()); + assert_eq!(shard_state.number_of_shards, 16); + assert_eq!(shard_state.revision, ShardLeaseRevision::INITIAL); + assert!(shard_state.shard_assignments.is_empty()); + assert!(shard_state.executor_leases.is_empty()); + assert!(shard_state.pending_rebalance.is_empty()); } #[test] @@ -166,7 +173,7 @@ async fn write_then_read_roundtrip( #[dimension(persistence)] persistence: &Arc, ) { let persistence = persistence.get_persistence().await; - let expected = sample_routing_table(16); + let expected = sample_shard_state(16); persistence .write(&expected) @@ -187,8 +194,8 @@ async fn last_write_wins( #[dimension(persistence)] persistence: &Arc, ) { let persistence = persistence.get_persistence().await; - let first = sample_routing_table(16); - let second = replacement_routing_table(16); + let first = sample_shard_state(16); + let second = replacement_shard_state(16); persistence .write(&first) @@ -207,50 +214,68 @@ async fn last_write_wins( assert_eq!(actual, second); } -fn sample_routing_table(number_of_shards: usize) -> RoutingTable { - let mut pod_states = BTreeMap::new(); - pod_states.insert( - Pod { - ip: IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)), - port: 9010, - }, - PodState { - pod_name: None, - assigned_shards: BTreeSet::from([ShardId::new(0), ShardId::new(1), ShardId::new(2)]), - }, - ); - pod_states.insert( - Pod { - ip: IpAddr::V4(Ipv4Addr::new(10, 0, 0, 2)), - port: 9011, - }, - PodState { - pod_name: None, - assigned_shards: BTreeSet::from([ShardId::new(3), ShardId::new(4)]), - }, - ); +fn granted_at() -> DateTime { + DateTime::from_timestamp(1_700_000_000, 0).expect("valid timestamp") +} - RoutingTable { - number_of_shards, - pod_states, +fn executor(idx: u128) -> ExecutorId { + ExecutorId(Uuid::from_u128(idx)) +} + +fn addr(last_octet: u8, port: u16) -> ExecutorAddr { + ExecutorAddr { + ip: IpAddr::V4(Ipv4Addr::new(10, 0, 0, last_octet)), + port, } } -fn replacement_routing_table(number_of_shards: usize) -> RoutingTable { - let mut pod_states = BTreeMap::new(); - pod_states.insert( - Pod { - ip: IpAddr::V4(Ipv4Addr::new(10, 0, 0, 3)), - port: 9012, - }, - PodState { - pod_name: None, - assigned_shards: BTreeSet::from([ShardId::new(5), ShardId::new(6), ShardId::new(7)]), - }, +fn shard_state_with_executors( + number_of_shards: usize, + executors: &[(ExecutorId, ExecutorAddr, Option<&str>, &[i64])], +) -> ShardLeaseState { + let mut shard_state = ShardLeaseState::new(number_of_shards); + for (executor_id, addr, pod_name, shard_ids) in executors { + shard_state.add_executor( + *executor_id, + *addr, + pod_name.map(str::to_string), + granted_at(), + LEASE_TTL, + ); + for shard_id in *shard_ids { + shard_state.assign_shard(*executor_id, ShardId::new(*shard_id)); + } + } + shard_state + .bump_revision() + .expect("revision bump should succeed"); + shard_state +} + +fn sample_shard_state(number_of_shards: usize) -> ShardLeaseState { + let mut shard_state = shard_state_with_executors( + number_of_shards, + &[ + ( + executor(1), + addr(1, 9010), + Some("worker-executor-0"), + &[0, 1, 2], + ), + (executor(2), addr(2, 9011), None, &[3, 4]), + ], ); + // make sure a non-initial epoch and an orphaned shard survive the roundtrip too + shard_state.assign_shard(executor(2), ShardId::new(0)); + shard_state.add_executor(executor(3), addr(3, 9012), None, granted_at(), LEASE_TTL); + shard_state.assign_shard(executor(3), ShardId::new(9)); + shard_state.remove_executor(executor(3)); + shard_state +} - RoutingTable { +fn replacement_shard_state(number_of_shards: usize) -> ShardLeaseState { + shard_state_with_executors( number_of_shards, - pod_states, - } + &[(executor(3), addr(3, 9012), None, &[5, 6, 7])], + ) } diff --git a/golem-shard-manager/tests/shard_management.rs b/golem-shard-manager/tests/shard_management.rs index a9100d46fe..9d7cdd91de 100644 --- a/golem-shard-manager/tests/shard_management.rs +++ b/golem-shard-manager/tests/shard_management.rs @@ -13,12 +13,13 @@ // limitations under the License. use async_trait::async_trait; +use chrono::{DateTime, Utc}; use golem_common::model::{Pod, ShardId}; use golem_shard_manager::{ - HealthCheck, HealthCheckError, PodState, RoutingTable, RoutingTablePersistence, - ShardManagement, ShardManagerError, WorkerExecutorService, + ExecutorAddr, ExecutorId, HealthCheck, HealthCheckError, RoutingTablePersistence, ShardEpoch, + ShardLeaseState, ShardManagement, ShardManagerError, WorkerExecutorService, }; -use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::collections::{BTreeSet, HashMap}; use std::net::{IpAddr, Ipv4Addr}; use std::sync::Arc; use std::time::Duration; @@ -26,36 +27,39 @@ use test_r::test; use tokio::sync::Mutex; use tokio::task::JoinSet; use tokio::time::Instant; +use uuid::Uuid; + +const LEASE_TTL: Duration = Duration::from_secs(60); #[derive(Clone, Debug)] struct TestPersistence { - state: Arc>, - writes: Arc>>, + shard_state: Arc>, + writes: Arc>>, } impl TestPersistence { - fn new(initial: RoutingTable) -> Self { + fn new(initial: ShardLeaseState) -> Self { Self { - state: Arc::new(Mutex::new(initial)), + shard_state: Arc::new(Mutex::new(initial)), writes: Arc::new(Mutex::new(Vec::new())), } } - async fn latest(&self) -> RoutingTable { - self.state.lock().await.clone() + async fn latest(&self) -> ShardLeaseState { + self.shard_state.lock().await.clone() } } #[async_trait] impl RoutingTablePersistence for TestPersistence { - async fn write(&self, routing_table: &RoutingTable) -> Result<(), ShardManagerError> { - *self.state.lock().await = routing_table.clone(); - self.writes.lock().await.push(routing_table.clone()); + async fn write(&self, shard_state: &ShardLeaseState) -> Result<(), ShardManagerError> { + *self.shard_state.lock().await = shard_state.clone(); + self.writes.lock().await.push(shard_state.clone()); Ok(()) } - async fn read(&self) -> Result { - Ok(self.state.lock().await.clone()) + async fn read(&self) -> Result { + Ok(self.shard_state.lock().await.clone()) } } @@ -192,25 +196,42 @@ fn pod(last_octet: u8, port: u16) -> Pod { } } -fn routing_table_with_pods( +fn executor(idx: u128) -> ExecutorId { + ExecutorId(Uuid::from_u128(idx)) +} + +fn granted_at() -> DateTime { + DateTime::from_timestamp(1_700_000_000, 0).expect("valid timestamp") +} + +fn shard_state_with_executors( number_of_shards: usize, - pods: Vec<(Pod, &str, &[i64])>, -) -> RoutingTable { - let mut pod_states = BTreeMap::new(); - for (pod, pod_name, shard_ids) in pods { - pod_states.insert( - pod, - PodState { - pod_name: Some(pod_name.to_string()), - assigned_shards: shard_ids.iter().copied().map(ShardId::new).collect(), - }, + pods: Vec<(ExecutorId, Pod, &str, &[i64])>, +) -> ShardLeaseState { + let mut shard_state = ShardLeaseState::new(number_of_shards); + for (executor_id, pod, pod_name, shard_ids) in pods { + shard_state.add_executor( + executor_id, + ExecutorAddr::from(pod), + Some(pod_name.to_string()), + granted_at(), + LEASE_TTL, ); + for shard_id in shard_ids { + shard_state.assign_shard(executor_id, ShardId::new(*shard_id)); + } } + shard_state +} - RoutingTable { - number_of_shards, - pod_states, - } +/// Shards of the executor registered at `pod` in the persisted state (there is exactly one). +fn shards_at(shard_state: &ShardLeaseState, pod: Pod) -> BTreeSet { + let executor_id = shard_state + .executor_for_addr(ExecutorAddr::from(pod)) + .unwrap_or_else(|| panic!("no executor registered at {pod}")); + shard_state + .shards_for_executor(executor_id) + .expect("executor should hold a lease") } fn shard_ids(ids: &[i64]) -> BTreeSet { @@ -238,14 +259,14 @@ async fn wait_for_local_assignment( } async fn new_shard_management( - routing_table: RoutingTable, + shard_state: ShardLeaseState, worker_executors: Arc, ) -> ( ShardManagement, TestPersistence, JoinSet>, ) { - let persistence = TestPersistence::new(routing_table); + let persistence = TestPersistence::new(shard_state); let health_check = Arc::new(TestHealthCheck::all_healthy()); let mut join_set = JoinSet::new(); @@ -254,6 +275,7 @@ async fn new_shard_management( worker_executors, health_check, 0.0, + LEASE_TTL, &mut join_set, ) .await @@ -273,11 +295,11 @@ async fn shard_manager_restart_clears_stale_executor_shards() { worker_executors.set_local_assignment(stale_pod, &[0]).await; let (_shard_management, _persistence, mut join_set) = new_shard_management( - routing_table_with_pods( + shard_state_with_executors( 1, vec![ - (authoritative_pod, "worker-executor-0", &[0]), - (stale_pod, "worker-executor-1", &[]), + (executor(1), authoritative_pod, "worker-executor-0", &[0]), + (executor(2), stale_pod, "worker-executor-1", &[]), ], ), worker_executors.clone(), @@ -308,34 +330,23 @@ async fn shard_manager_restart_recovers_from_partially_applied_rebalance() { .await; let (_shard_management, persistence, mut join_set) = new_shard_management( - routing_table_with_pods( + shard_state_with_executors( 1, vec![ - (persisted_owner, "worker-executor-0", &[0]), - (stale_new_owner, "worker-executor-1", &[]), + (executor(1), persisted_owner, "worker-executor-0", &[0]), + (executor(2), stale_new_owner, "worker-executor-1", &[]), ], ), worker_executors.clone(), ) .await; - let routing_table = persistence.latest().await; + let shard_state = persistence.latest().await; assert_eq!( - routing_table - .pod_states - .get(&persisted_owner) - .expect("persisted owner missing") - .assigned_shards, + shards_at(&shard_state, persisted_owner), [0].into_iter().map(ShardId::new).collect() ); - assert!( - routing_table - .pod_states - .get(&stale_new_owner) - .expect("stale new owner missing") - .assigned_shards - .is_empty() - ); + assert!(shards_at(&shard_state, stale_new_owner).is_empty()); assert_eq!( worker_executors.local_assignment(persisted_owner).await, [0].into_iter().map(ShardId::new).collect() @@ -358,11 +369,11 @@ async fn reconnecting_pod_clears_stale_local_shards() { .await; let (shard_management, persistence, mut join_set) = new_shard_management( - routing_table_with_pods( + shard_state_with_executors( 1, vec![ - (existing_pod, "worker-executor-0", &[]), - (pod(2, 9001), "worker-executor-1", &[0]), + (executor(1), existing_pod, "worker-executor-0", &[]), + (executor(2), pod(2, 9001), "worker-executor-1", &[0]), ], ), worker_executors.clone(), @@ -378,8 +389,8 @@ async fn reconnecting_pod_clears_stale_local_shards() { [0].into_iter().map(ShardId::new).collect() ); - shard_management - .register_pod(existing_pod, Some("worker-executor-0".to_string())) + let new_executor_id = shard_management + .register_executor(existing_pod.into(), Some("worker-executor-0".to_string())) .await; tokio::time::sleep(Duration::from_millis(50)).await; @@ -388,15 +399,16 @@ async fn reconnecting_pod_clears_stale_local_shards() { BTreeSet::new() ); - let routing_table = persistence.latest().await; - assert!( - routing_table - .pod_states - .get(&existing_pod) - .expect("existing pod missing") - .assigned_shards - .is_empty() + let shard_state = persistence.latest().await; + assert!(shards_at(&shard_state, existing_pod).is_empty()); + // the re-registered instance replaced the previous lease at the same address + assert_ne!(new_executor_id, executor(1)); + assert!(!shard_state.has_executor(executor(1))); + assert_eq!( + shard_state.executor_for_addr(existing_pod.into()), + Some(new_executor_id) ); + assert_eq!(shard_state.executor_count(), 2); join_set.abort_all(); } @@ -410,34 +422,23 @@ async fn reconciliation_clears_duplicate_local_shard_owner() { worker_executors.set_local_assignment(stale_pod, &[0]).await; let (_shard_management, persistence, mut join_set) = new_shard_management( - routing_table_with_pods( + shard_state_with_executors( 1, vec![ - (authoritative_pod, "worker-executor-0", &[0]), - (stale_pod, "worker-executor-1", &[]), + (executor(1), authoritative_pod, "worker-executor-0", &[0]), + (executor(2), stale_pod, "worker-executor-1", &[]), ], ), worker_executors.clone(), ) .await; - let routing_table = persistence.latest().await; + let shard_state = persistence.latest().await; assert_eq!( - routing_table - .pod_states - .get(&authoritative_pod) - .expect("authoritative pod missing") - .assigned_shards, + shards_at(&shard_state, authoritative_pod), [0].into_iter().map(ShardId::new).collect() ); - assert!( - routing_table - .pod_states - .get(&stale_pod) - .expect("stale pod missing") - .assigned_shards - .is_empty() - ); + assert!(shards_at(&shard_state, stale_pod).is_empty()); assert_eq!( worker_executors.local_assignment(authoritative_pod).await, [0].into_iter().map(ShardId::new).collect() @@ -463,35 +464,24 @@ async fn failed_assignment_is_retried_from_unassigned_shards() { worker_executors.fail_next_assignments(new_pod, 1).await; let (shard_management, persistence, mut join_set) = new_shard_management( - routing_table_with_pods(4, vec![(old_pod, "worker-executor-0", &[0, 1, 2, 3])]), + shard_state_with_executors( + 4, + vec![(executor(1), old_pod, "worker-executor-0", &[0, 1, 2, 3])], + ), worker_executors.clone(), ) .await; shard_management - .register_pod(new_pod, Some("worker-executor-1".to_string())) + .register_executor(new_pod.into(), Some("worker-executor-1".to_string())) .await; wait_for_local_assignment(&worker_executors, old_pod, shard_ids(&[2, 3])).await; wait_for_local_assignment(&worker_executors, new_pod, shard_ids(&[0, 1])).await; - let routing_table = persistence.latest().await; - assert_eq!( - routing_table - .pod_states - .get(&old_pod) - .expect("old pod missing") - .assigned_shards, - shard_ids(&[2, 3]) - ); - assert_eq!( - routing_table - .pod_states - .get(&new_pod) - .expect("new pod missing") - .assigned_shards, - shard_ids(&[0, 1]) - ); + let shard_state = persistence.latest().await; + assert_eq!(shards_at(&shard_state, old_pod), shard_ids(&[2, 3])); + assert_eq!(shards_at(&shard_state, new_pod), shard_ids(&[0, 1])); join_set.abort_all(); } @@ -509,13 +499,16 @@ async fn failed_revoke_is_retried_without_assigning_to_new_executor_first() { worker_executors.fail_next_revocations(old_pod, 1).await; let (shard_management, persistence, mut join_set) = new_shard_management( - routing_table_with_pods(4, vec![(old_pod, "worker-executor-0", &[0, 1, 2, 3])]), + shard_state_with_executors( + 4, + vec![(executor(1), old_pod, "worker-executor-0", &[0, 1, 2, 3])], + ), worker_executors.clone(), ) .await; shard_management - .register_pod(new_pod, Some("worker-executor-1".to_string())) + .register_executor(new_pod.into(), Some("worker-executor-1".to_string())) .await; wait_for_local_assignment(&worker_executors, old_pod, shard_ids(&[2, 3])).await; @@ -530,23 +523,9 @@ async fn failed_revoke_is_retried_without_assigning_to_new_executor_first() { shard_ids(&[0, 1]) ); - let routing_table = persistence.latest().await; - assert_eq!( - routing_table - .pod_states - .get(&old_pod) - .expect("old pod missing") - .assigned_shards, - shard_ids(&[2, 3]) - ); - assert_eq!( - routing_table - .pod_states - .get(&new_pod) - .expect("new pod missing") - .assigned_shards, - shard_ids(&[0, 1]) - ); + let shard_state = persistence.latest().await; + assert_eq!(shards_at(&shard_state, old_pod), shard_ids(&[2, 3])); + assert_eq!(shards_at(&shard_state, new_pod), shard_ids(&[0, 1])); join_set.abort_all(); } @@ -558,11 +537,11 @@ async fn failed_reconnect_reconciliation_is_retried() { let worker_executors = Arc::new(TestWorkerExecutors::default()); let (shard_management, _persistence, mut join_set) = new_shard_management( - routing_table_with_pods( + shard_state_with_executors( 1, vec![ - (existing_pod, "worker-executor-0", &[]), - (pod(2, 9001), "worker-executor-1", &[0]), + (executor(1), existing_pod, "worker-executor-0", &[]), + (executor(2), pod(2, 9001), "worker-executor-1", &[0]), ], ), worker_executors.clone(), @@ -577,10 +556,88 @@ async fn failed_reconnect_reconciliation_is_retried() { .await; shard_management - .register_pod(existing_pod, Some("worker-executor-0".to_string())) + .register_executor(existing_pod.into(), Some("worker-executor-0".to_string())) .await; wait_for_local_assignment(&worker_executors, existing_pod, BTreeSet::new()).await; join_set.abort_all(); } + +#[test] +// A new executor instance registering at an address that already holds a lease replaces the +// previous lease, inherits its shards with advanced epochs and receives the authoritative +// assignment; shards never become unassigned and no other executor is disturbed. +async fn same_address_reregistration_transfers_shards_and_reconciles() { + let restarted_pod = pod(1, 9000); + let other_pod = pod(2, 9001); + let worker_executors = Arc::new(TestWorkerExecutors::default()); + worker_executors + .set_local_assignment(restarted_pod, &[0, 1]) + .await; + worker_executors + .set_local_assignment(other_pod, &[2, 3]) + .await; + + let (shard_management, persistence, mut join_set) = new_shard_management( + shard_state_with_executors( + 4, + vec![ + (executor(1), restarted_pod, "worker-executor-0", &[0, 1]), + (executor(2), other_pod, "worker-executor-1", &[2, 3]), + ], + ), + worker_executors.clone(), + ) + .await; + + // the restarted process comes up with an empty local assignment + worker_executors + .set_local_assignment(restarted_pod, &[]) + .await; + + let new_executor_id = shard_management + .register_executor(restarted_pod.into(), Some("worker-executor-0".to_string())) + .await; + + wait_for_local_assignment(&worker_executors, restarted_pod, shard_ids(&[0, 1])).await; + assert_eq!( + worker_executors.local_assignment(other_pod).await, + shard_ids(&[2, 3]) + ); + + let shard_state = persistence.latest().await; + assert!(!shard_state.has_executor(executor(1))); + assert_eq!( + shard_state.executor_for_addr(restarted_pod.into()), + Some(new_executor_id) + ); + assert_eq!( + shard_state.shards_for_executor(new_executor_id), + Some(shard_ids(&[0, 1])) + ); + assert_eq!( + shard_state.shards_for_executor(executor(2)), + Some(shard_ids(&[2, 3])) + ); + assert_eq!( + shard_state.epoch_for_shard(ShardId::new(0)), + Some(ShardEpoch(1)) + ); + assert_eq!( + shard_state.epoch_for_shard(ShardId::new(2)), + Some(ShardEpoch(0)) + ); + assert!(shard_state.pending_rebalance.is_empty()); + assert!(shard_state.get_unassigned_shards().is_empty()); + + // every persisted state along the way kept all four shards routable + for written in persistence.writes.lock().await.iter() { + assert!( + written.get_unassigned_shards().is_empty(), + "shards became unassigned during re-registration: {written}" + ); + } + + join_set.abort_all(); +}