Replace the shard manager's pod-keyed routing table with an executor-keyed ShardLeaseState

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]

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

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"

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;

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;

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(),

golem-shard-manager/src/error.rs

@@ -85,6 +85,11 @@ impl From<ShardManagerError> 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,
+            ),
         }
     }
 }

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,19 +41,23 @@ 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(
         &self,
         pod: Pod,
         pod_name: Option<String>,
     ) -> 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(())
     }
 }

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<RunDetails> {
     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::<ShardManagerServiceServer<ShardManagerServiceImpl>>()
@@ -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?,

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,
         }
     }
 

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<String>) -> 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<dyn HealthCheck>,
-    pods: &[(Pod, Option<String>)],
-) -> HashSet<Pod> {
-    let futures: Vec<_> = pods
+    executors: &[(ExecutorId, ExecutorAddr, Option<String>)],
+) -> HashSet<ExecutorId> {
+    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),
                 }
             })
         })

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<dyn HealthCheck>,
 ) {
     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");
 }

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,
+};