From 9a52fc6db4dd22075e70b209c4d710faaa28020d Mon Sep 17 00:00:00 2001 From: Aditya Salunkhe Date: Sat, 29 Aug 2026 16:54:28 +0530 Subject: [PATCH 1/6] Make shard manager persistence compare-and-swap based across local and distributed backends --- .github/actions/setup-rust/action.yml | 7 + Cargo.lock | 31 +- Cargo.toml | 1 + cli/golem/src/launch.rs | 2 +- .../published-postgres/compose.yaml | 16 +- golem-shard-manager/Cargo.toml | 1 + .../config/shard-manager.sample.env | 111 ++- golem-shard-manager/config/shard-manager.toml | 167 +++- .../postgres/004_shard_state_revision.sql | 59 ++ .../sqlite/004_shard_state_revision.sql | 59 ++ golem-shard-manager/src/config.rs | 135 ++- golem-shard-manager/src/error.rs | 10 + golem-shard-manager/src/lib.rs | 60 +- golem-shard-manager/src/quota/mod.rs | 2 +- golem-shard-manager/src/quota/quota_repo.rs | 61 ++ .../src/quota/quota_service_tests.rs | 52 +- golem-shard-manager/src/sharding/error.rs | 20 +- .../src/sharding/persistence.rs | 189 ----- .../src/sharding/persistence/db.rs | 409 +++++++++ .../src/sharding/persistence/etcd.rs | 150 ++++ .../src/sharding/persistence/mod.rs | 259 ++++++ .../src/sharding/shard_management.rs | 206 +++-- golem-shard-manager/tests/persistence.rs | 800 +++++++++++++++++- golem-shard-manager/tests/shard_management.rs | 240 +++++- .../src/components/etcd/docker_etcd.rs | 138 +++ .../src/components/etcd/mod.rs | 15 + golem-test-framework/src/components/mod.rs | 1 + .../src/components/shard_manager/mod.rs | 13 +- integration-tests/tests/sharding.rs | 8 + local-run/start.sh | 4 +- 30 files changed, 2766 insertions(+), 460 deletions(-) create mode 100644 golem-shard-manager/db/migration/postgres/004_shard_state_revision.sql create mode 100644 golem-shard-manager/db/migration/sqlite/004_shard_state_revision.sql delete mode 100644 golem-shard-manager/src/sharding/persistence.rs create mode 100644 golem-shard-manager/src/sharding/persistence/db.rs create mode 100644 golem-shard-manager/src/sharding/persistence/etcd.rs create mode 100644 golem-shard-manager/src/sharding/persistence/mod.rs create mode 100644 golem-test-framework/src/components/etcd/docker_etcd.rs create mode 100644 golem-test-framework/src/components/etcd/mod.rs diff --git a/.github/actions/setup-rust/action.yml b/.github/actions/setup-rust/action.yml index 35c041c716..8aa02aa07e 100644 --- a/.github/actions/setup-rust/action.yml +++ b/.github/actions/setup-rust/action.yml @@ -39,6 +39,13 @@ runs: for target in ${{ inputs.rust-targets }}; do rustup target add "$target" done + - name: Install protoc + # `etcd-client`'s build script compiles its protos with tonic-prost-build, which shells out + # to `protoc`. The workspace's own protos are compiled with protox (pure Rust) and need no + # binary, so this is the only thing that requires it. + uses: arduino/setup-protoc@v3 + with: + repo-token: ${{ github.token }} - name: Setup Rust cache if: inputs.use-cache == 'true' uses: Swatinem/rust-cache@v2 diff --git a/Cargo.lock b/Cargo.lock index 2bfda99e40..d28e3f7935 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1716,7 +1716,7 @@ version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] @@ -3009,6 +3009,24 @@ version = "0.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5692dd7b5a1978a5aeb0ce83b7655c58ca8efdcb79d21036ea249da95afec2c6" +[[package]] +name = "etcd-client" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ef5da6e9a6ae89f4a91f80ba1caae45a5a924397a19947e18f5121a43285e9bc" +dependencies = [ + "http 1.5.0", + "prost 0.14.4", + "tokio", + "tokio-stream", + "tonic 0.14.6", + "tonic-build", + "tonic-prost", + "tonic-prost-build", + "tower 0.5.3", + "tower-service", +] + [[package]] name = "etcetera" version = "0.8.0" @@ -4203,6 +4221,7 @@ dependencies = [ "chrono", "conditional-trait-gen", "desert_rust", + "etcd-client", "futures", "golem-api-grpc", "golem-common", @@ -4991,7 +5010,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.5.10", + "socket2 0.6.5", "system-configuration", "tokio", "tower-service", @@ -7914,7 +7933,7 @@ dependencies = [ "quinn-udp", "rustc-hash", "rustls 0.23.43", - "socket2 0.5.10", + "socket2 0.6.5", "thiserror 2.0.20", "tokio", "tracing", @@ -7953,7 +7972,7 @@ dependencies = [ "cfg_aliases 0.2.2", "libc", "once_cell", - "socket2 0.5.10", + "socket2 0.6.5", "tracing", "windows-sys 0.61.2", ] @@ -10065,7 +10084,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.3.4", + "getrandom 0.4.3", "once_cell", "rustix 1.1.4", "windows-sys 0.61.2", @@ -12145,7 +12164,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 1067bf1f6c..3c0461e0b6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -106,6 +106,7 @@ dir-diff = "0.3.3" dirs = "6.0.0" dotenvy = "0.15.7" drop-stream = "0.3.2" +etcd-client = "0.19.0" evicting_cache_map = "0.4.0" expectrl = "0.8.0" fancy-regex = "0.14.0" diff --git a/cli/golem/src/launch.rs b/cli/golem/src/launch.rs index 64bd43a35a..8db7aa9272 100644 --- a/cli/golem/src/launch.rs +++ b/cli/golem/src/launch.rs @@ -313,7 +313,7 @@ fn shard_manager_config( port: 0, ..Default::default() }, - db: DbConfig::Sqlite(DbSqliteConfig { + persistence: golem_shard_manager::config::PersistenceConfig::Sqlite(DbSqliteConfig { database: args .data_dir .join("shard_manager.db") diff --git a/docker-examples/published-postgres/compose.yaml b/docker-examples/published-postgres/compose.yaml index b5f784a182..d6e97f1ddf 100644 --- a/docker-examples/published-postgres/compose.yaml +++ b/docker-examples/published-postgres/compose.yaml @@ -92,14 +92,14 @@ services: GOLEM__REGISTRY_SERVICE__HOST: golem-registry-service GOLEM__REGISTRY_SERVICE__PORT: ${REGISTRY_SERVICE_GRPC_PORT} - GOLEM__DB__TYPE: Postgres - GOLEM__DB__CONFIG__DATABASE: golem_db - GOLEM__DB__CONFIG__SCHEMA: golem_shard_manager - GOLEM__DB__CONFIG__MAX_CONNECTIONS: 10 - GOLEM__DB__CONFIG__HOST: postgres - GOLEM__DB__CONFIG__PORT: 5432 - GOLEM__DB__CONFIG__USERNAME: golem_user - GOLEM__DB__CONFIG__PASSWORD: golem_password + GOLEM__PERSISTENCE__TYPE: Postgres + GOLEM__PERSISTENCE__CONFIG__DATABASE: golem_db + GOLEM__PERSISTENCE__CONFIG__SCHEMA: golem_shard_manager + GOLEM__PERSISTENCE__CONFIG__MAX_CONNECTIONS: 10 + GOLEM__PERSISTENCE__CONFIG__HOST: postgres + GOLEM__PERSISTENCE__CONFIG__PORT: 5432 + GOLEM__PERSISTENCE__CONFIG__USERNAME: golem_user + GOLEM__PERSISTENCE__CONFIG__PASSWORD: golem_password depends_on: postgres: condition: service_healthy diff --git a/golem-shard-manager/Cargo.toml b/golem-shard-manager/Cargo.toml index f9f03b4940..52740100b9 100644 --- a/golem-shard-manager/Cargo.toml +++ b/golem-shard-manager/Cargo.toml @@ -32,6 +32,7 @@ chrono = { workspace = true } async-trait = { workspace = true } conditional-trait-gen = { workspace = true } desert_rust = { workspace = true } +etcd-client = { workspace = true } futures = { workspace = true } http = { workspace = true } humantime-serde = { workspace = true } diff --git a/golem-shard-manager/config/shard-manager.sample.env b/golem-shard-manager/config/shard-manager.sample.env index dc918c3f8d..cd31054f01 100644 --- a/golem-shard-manager/config/shard-manager.sample.env +++ b/golem-shard-manager/config/shard-manager.sample.env @@ -5,15 +5,15 @@ 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__DB__CONFIG__MAX_CONNECTIONS=10 GOLEM__GRPC__PORT=9092 GOLEM__GRPC__TLS__TYPE="Disabled" GOLEM__HEALTH_CHECK__DELAY="10s" GOLEM__HEALTH_CHECK__SILENT=false GOLEM__HEALTH_CHECK__MODE__TYPE="Grpc" +GOLEM__PERSISTENCE__TYPE="Sqlite" +GOLEM__PERSISTENCE__CONFIG__DATABASE="golem_shard_manager.db" +GOLEM__PERSISTENCE__CONFIG__FOREIGN_KEYS=false +GOLEM__PERSISTENCE__CONFIG__MAX_CONNECTIONS=10 GOLEM__QUOTA__DEFINITION_STALENESS_TTL="5m" GOLEM__QUOTA__LEASE_DURATION="1m" GOLEM__QUOTA__MIN_EXECUTORS=2 @@ -100,16 +100,111 @@ 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__DB__CONFIG__MAX_CONNECTIONS=10 GOLEM__GRPC__PORT=9092 GOLEM__GRPC__TLS__TYPE="Disabled" GOLEM__HEALTH_CHECK__DELAY="1s" GOLEM__HEALTH_CHECK__SILENT=false GOLEM__HEALTH_CHECK__MODE__TYPE="K8s" GOLEM__HEALTH_CHECK__MODE__CONFIG__NAMESPACE="namespace" +GOLEM__PERSISTENCE__TYPE="Sqlite" +GOLEM__PERSISTENCE__CONFIG__DATABASE="golem_shard_manager.db" +GOLEM__PERSISTENCE__CONFIG__FOREIGN_KEYS=false +GOLEM__PERSISTENCE__CONFIG__MAX_CONNECTIONS=10 +GOLEM__QUOTA__DEFINITION_STALENESS_TTL="5m" +GOLEM__QUOTA__LEASE_DURATION="1m" +GOLEM__QUOTA__MIN_EXECUTORS=2 +GOLEM__REGISTRY_SERVICE__CONNECT_TIMEOUT="10s" +GOLEM__REGISTRY_SERVICE__HOST="localhost" +GOLEM__REGISTRY_SERVICE__MAX_MESSAGE_SIZE=52428800 +GOLEM__REGISTRY_SERVICE__PORT=8080 +#GOLEM__REGISTRY_SERVICE__REQUEST_TIMEOUT= +GOLEM__REGISTRY_SERVICE__INVALIDATION_EVENT_SUBSCRIBER__INITIAL_BACKOFF="100ms" +GOLEM__REGISTRY_SERVICE__INVALIDATION_EVENT_SUBSCRIBER__MAX_BACKOFF="30s" +GOLEM__REGISTRY_SERVICE__RETRIES_ON_UNAVAILABLE__MAX_ATTEMPTS=5 +GOLEM__REGISTRY_SERVICE__RETRIES_ON_UNAVAILABLE__MAX_DELAY="2s" +GOLEM__REGISTRY_SERVICE__RETRIES_ON_UNAVAILABLE__MAX_JITTER_FACTOR=0.15 +GOLEM__REGISTRY_SERVICE__RETRIES_ON_UNAVAILABLE__MIN_DELAY="100ms" +GOLEM__REGISTRY_SERVICE__RETRIES_ON_UNAVAILABLE__MULTIPLIER=2.0 +GOLEM__REGISTRY_SERVICE__TLS__TYPE="Disabled" +GOLEM__RESOURCE_DEFINITION_FETCHER__CACHE_EVICTION_PERIOD="1m" +GOLEM__RESOURCE_DEFINITION_FETCHER__CACHE_MAX_CAPACITY=1024 +GOLEM__RESOURCE_DEFINITION_FETCHER__CACHE_TTL="5m" +GOLEM__TRACING__CONSOLE=false +GOLEM__TRACING__DTOR_FRIENDLY=false +#GOLEM__TRACING__FILE_DIR= +GOLEM__TRACING__FILE_NAME="shard-manager.log" +GOLEM__TRACING__FILE_TRUNCATE=true +GOLEM__TRACING__FILE__ANSI=false +GOLEM__TRACING__FILE__COMPACT=false +GOLEM__TRACING__FILE__ENABLED=false +GOLEM__TRACING__FILE__JSON=true +GOLEM__TRACING__FILE__JSON_FLATTEN=true +GOLEM__TRACING__FILE__JSON_FLATTEN_SPAN=true +GOLEM__TRACING__FILE__JSON_SOURCE_LOCATION=false +GOLEM__TRACING__FILE__PRETTY=false +GOLEM__TRACING__FILE__SPAN_EVENTS_ACTIVE=false +GOLEM__TRACING__FILE__SPAN_EVENTS_FULL=false +GOLEM__TRACING__FILE__WITHOUT_TIME=false +GOLEM__TRACING__OTLP__ENABLED=false +GOLEM__TRACING__OTLP__HOST="localhost" +GOLEM__TRACING__OTLP__PORT=4318 +GOLEM__TRACING__OTLP__SERVICE_NAME="golem" +GOLEM__TRACING__STDERR__ANSI=false +GOLEM__TRACING__STDERR__COMPACT=false +GOLEM__TRACING__STDERR__ENABLED=false +GOLEM__TRACING__STDERR__JSON=false +GOLEM__TRACING__STDERR__JSON_FLATTEN=false +GOLEM__TRACING__STDERR__JSON_FLATTEN_SPAN=false +GOLEM__TRACING__STDERR__JSON_SOURCE_LOCATION=false +GOLEM__TRACING__STDERR__PRETTY=false +GOLEM__TRACING__STDERR__SPAN_EVENTS_ACTIVE=false +GOLEM__TRACING__STDERR__SPAN_EVENTS_FULL=false +GOLEM__TRACING__STDERR__WITHOUT_TIME=false +GOLEM__TRACING__STDOUT__ANSI=true +GOLEM__TRACING__STDOUT__COMPACT=false +GOLEM__TRACING__STDOUT__ENABLED=true +GOLEM__TRACING__STDOUT__JSON=false +GOLEM__TRACING__STDOUT__JSON_FLATTEN=true +GOLEM__TRACING__STDOUT__JSON_FLATTEN_SPAN=true +GOLEM__TRACING__STDOUT__JSON_SOURCE_LOCATION=false +GOLEM__TRACING__STDOUT__PRETTY=false +GOLEM__TRACING__STDOUT__SPAN_EVENTS_ACTIVE=false +GOLEM__TRACING__STDOUT__SPAN_EVENTS_FULL=false +GOLEM__TRACING__STDOUT__WITHOUT_TIME=false +GOLEM__WORKER_EXECUTORS__ASSIGN_SHARDS_TIMEOUT="5s" +GOLEM__WORKER_EXECUTORS__CONNECT_TIMEOUT="10s" +GOLEM__WORKER_EXECUTORS__HEALTH_CHECK_TIMEOUT="2s" +GOLEM__WORKER_EXECUTORS__MAX_MESSAGE_SIZE=33554432 +#GOLEM__WORKER_EXECUTORS__REQUEST_TIMEOUT= +GOLEM__WORKER_EXECUTORS__REVOKE_SHARDS_TIMEOUT="5s" +GOLEM__WORKER_EXECUTORS__RETRIES__MAX_ATTEMPTS=5 +GOLEM__WORKER_EXECUTORS__RETRIES__MAX_DELAY="2s" +GOLEM__WORKER_EXECUTORS__RETRIES__MAX_JITTER_FACTOR=0.15 +GOLEM__WORKER_EXECUTORS__RETRIES__MIN_DELAY="100ms" +GOLEM__WORKER_EXECUTORS__RETRIES__MULTIPLIER=2.0 +GOLEM__WORKER_EXECUTORS__RETRIES_ON_UNAVAILABLE__MAX_ATTEMPTS=5 +GOLEM__WORKER_EXECUTORS__RETRIES_ON_UNAVAILABLE__MAX_DELAY="2s" +GOLEM__WORKER_EXECUTORS__RETRIES_ON_UNAVAILABLE__MAX_JITTER_FACTOR=0.15 +GOLEM__WORKER_EXECUTORS__RETRIES_ON_UNAVAILABLE__MIN_DELAY="100ms" +GOLEM__WORKER_EXECUTORS__RETRIES_ON_UNAVAILABLE__MULTIPLIER=2.0 +GOLEM__WORKER_EXECUTORS__TLS__TYPE="Disabled" + +### Generated from example config: with etcd persistence (distributed mode) + +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__GRPC__PORT=9092 +GOLEM__GRPC__TLS__TYPE="Disabled" +GOLEM__HEALTH_CHECK__DELAY="10s" +GOLEM__HEALTH_CHECK__SILENT=false +GOLEM__HEALTH_CHECK__MODE__TYPE="Grpc" +GOLEM__PERSISTENCE__TYPE="Etcd" +GOLEM__PERSISTENCE__CONFIG__CONNECT_TIMEOUT="10s" +GOLEM__PERSISTENCE__CONFIG__ENDPOINTS=["http://localhost:2379"] +GOLEM__PERSISTENCE__CONFIG__REQUEST_TIMEOUT="5s" GOLEM__QUOTA__DEFINITION_STALENESS_TTL="5m" GOLEM__QUOTA__LEASE_DURATION="1m" GOLEM__QUOTA__MIN_EXECUTORS=2 diff --git a/golem-shard-manager/config/shard-manager.toml b/golem-shard-manager/config/shard-manager.toml index 1efbb99748..c330cbf334 100644 --- a/golem-shard-manager/config/shard-manager.toml +++ b/golem-shard-manager/config/shard-manager.toml @@ -5,14 +5,6 @@ rebalance_threshold = 0.1 runtime_metrics_sampling_interval = "5s" shard_lease_duration = "1m" -[db] -type = "Sqlite" - -[db.config] -database = "golem_shard_manager.db" -foreign_keys = false -max_connections = 10 - [grpc] port = 9092 @@ -30,6 +22,14 @@ type = "Grpc" [health_check.mode.config] +[persistence] +type = "Sqlite" + +[persistence.config] +database = "golem_shard_manager.db" +foreign_keys = false +max_connections = 10 + [quota] definition_staleness_ttl = "5m" lease_duration = "1m" @@ -147,14 +147,148 @@ type = "Disabled" # runtime_metrics_sampling_interval = "5s" # shard_lease_duration = "1m" # -# [db] +# [grpc] +# port = 9092 +# +# [grpc.tls] +# type = "Disabled" +# +# [grpc.tls.config] +# +# [health_check] +# delay = "1s" +# silent = false +# +# [health_check.mode] +# type = "K8s" +# +# [health_check.mode.config] +# namespace = "namespace" +# +# [persistence] # type = "Sqlite" # -# [db.config] +# [persistence.config] # database = "golem_shard_manager.db" # foreign_keys = false # max_connections = 10 # +# [quota] +# definition_staleness_ttl = "5m" +# lease_duration = "1m" +# min_executors = 2 +# +# [registry_service] +# connect_timeout = "10s" +# host = "localhost" +# max_message_size = 52428800 +# port = 8080 +# +# [registry_service.invalidation_event_subscriber] +# initial_backoff = "100ms" +# max_backoff = "30s" +# +# [registry_service.retries_on_unavailable] +# max_attempts = 5 +# max_delay = "2s" +# max_jitter_factor = 0.15 +# min_delay = "100ms" +# multiplier = 2.0 +# +# [registry_service.tls] +# type = "Disabled" +# +# [registry_service.tls.config] +# +# [resource_definition_fetcher] +# cache_eviction_period = "1m" +# cache_max_capacity = 1024 +# cache_ttl = "5m" +# +# [tracing] +# console = false +# dtor_friendly = false +# file_name = "shard-manager.log" +# file_truncate = true +# +# [tracing.file] +# ansi = false +# compact = false +# enabled = false +# json = true +# json_flatten = true +# json_flatten_span = true +# json_source_location = false +# pretty = false +# span_events_active = false +# span_events_full = false +# without_time = false +# +# [tracing.otlp] +# enabled = false +# host = "localhost" +# port = 4318 +# service_name = "golem" +# +# [tracing.stderr] +# ansi = false +# compact = false +# enabled = false +# json = false +# json_flatten = false +# json_flatten_span = false +# json_source_location = false +# pretty = false +# span_events_active = false +# span_events_full = false +# without_time = false +# +# [tracing.stdout] +# ansi = true +# compact = false +# enabled = true +# json = false +# json_flatten = true +# json_flatten_span = true +# json_source_location = false +# pretty = false +# span_events_active = false +# span_events_full = false +# without_time = false +# +# [worker_executors] +# assign_shards_timeout = "5s" +# connect_timeout = "10s" +# health_check_timeout = "2s" +# max_message_size = 33554432 +# revoke_shards_timeout = "5s" +# +# [worker_executors.retries] +# max_attempts = 5 +# max_delay = "2s" +# max_jitter_factor = 0.15 +# min_delay = "100ms" +# multiplier = 2.0 +# +# [worker_executors.retries_on_unavailable] +# max_attempts = 5 +# max_delay = "2s" +# max_jitter_factor = 0.15 +# min_delay = "100ms" +# multiplier = 2.0 +# +# [worker_executors.tls] +# type = "Disabled" +# +# [worker_executors.tls.config] + +## Generated from example config: with etcd persistence (distributed mode) +# http_port = 8081 +# number_of_shards = 1024 +# rebalance_threshold = 0.1 +# runtime_metrics_sampling_interval = "5s" +# shard_lease_duration = "1m" +# # [grpc] # port = 9092 # @@ -164,14 +298,21 @@ type = "Disabled" # [grpc.tls.config] # # [health_check] -# delay = "1s" +# delay = "10s" # silent = false # # [health_check.mode] -# type = "K8s" +# type = "Grpc" # # [health_check.mode.config] -# namespace = "namespace" +# +# [persistence] +# type = "Etcd" +# +# [persistence.config] +# connect_timeout = "10s" +# endpoints = ["http://localhost:2379"] +# request_timeout = "5s" # # [quota] # definition_staleness_ttl = "5m" diff --git a/golem-shard-manager/db/migration/postgres/004_shard_state_revision.sql b/golem-shard-manager/db/migration/postgres/004_shard_state_revision.sql new file mode 100644 index 0000000000..c9223acc56 --- /dev/null +++ b/golem-shard-manager/db/migration/postgres/004_shard_state_revision.sql @@ -0,0 +1,59 @@ +-- The persisted shard lease state is now written with a compare-and-swap guard, and mirrored into +-- queryable tables. +-- +-- `revision` is a storage-level fencing token. It is incremented by exactly one on every +-- successful write and is deliberately UNRELATED to the domain-level ShardLeaseRevision inside +-- `state`: that one starts at 0 and is only bumped when the routing table meaningfully changes. +-- Never derive this column from it. +-- +-- Revision 0 is reserved to mean "no state stored", so a stored row always carries >= 1. The +-- column therefore has NO DEFAULT: a write that forgets to bind it must fail loudly. +-- +-- The table is recreated rather than altered: any row written before this migration has no +-- meaningful revision, and the state rebuilds itself from executor registrations (same reasoning +-- as 003). Recreating also avoids SQLite's requirement that ADD COLUMN supply a non-null default. +DROP TABLE shard_manager_state; + +CREATE TABLE shard_manager_state +( + id INTEGER PRIMARY KEY, + state BYTEA NOT NULL, + revision BIGINT NOT NULL +); + +-- Local-mode mirror of the state blob, for inspection with plain SQL (`executor_leases` is the +-- counterpart of the quota system's `quota_leases`). The blob is the source of truth: both tables +-- are rewritten wholesale in the same transaction as every write the shard manager makes, and it +-- never reads them back. Not mirrored: `pending_rebalance`, `shard_epochs` (the per-shard epoch +-- high-water marks, which outlive leases) and the in-blob `ShardLeaseRevision`. +-- +-- Anyone clearing the state by hand (`DELETE FROM shard_manager_state`) must clear these two +-- tables as well, or they keep describing leases that no longer exist until the next write. +-- +-- The shard manager refuses to persist a state whose assignments reference an executor without a +-- lease, so the foreign key below is belt-and-braces. Postgres always enforces it; SQLite only +-- with `foreign_keys = true` (the default is `false`). +CREATE TABLE executor_leases +( + executor_id UUID NOT NULL, + ip BYTEA NOT NULL, + port INTEGER NOT NULL, + granted_at TIMESTAMP NOT NULL, + expires_at TIMESTAMP NOT NULL, + pod_name TEXT, + + CONSTRAINT executor_leases_pk + PRIMARY KEY (executor_id) +); + +CREATE TABLE shard_assignments +( + shard_id INTEGER NOT NULL, + executor_id UUID NOT NULL, + epoch BIGINT NOT NULL, + + CONSTRAINT shard_assignments_pk + PRIMARY KEY (shard_id), + CONSTRAINT shard_assignments_executor_fk + FOREIGN KEY (executor_id) REFERENCES executor_leases +); diff --git a/golem-shard-manager/db/migration/sqlite/004_shard_state_revision.sql b/golem-shard-manager/db/migration/sqlite/004_shard_state_revision.sql new file mode 100644 index 0000000000..55f1468e68 --- /dev/null +++ b/golem-shard-manager/db/migration/sqlite/004_shard_state_revision.sql @@ -0,0 +1,59 @@ +-- The persisted shard lease state is now written with a compare-and-swap guard, and mirrored into +-- queryable tables. +-- +-- `revision` is a storage-level fencing token. It is incremented by exactly one on every +-- successful write and is deliberately UNRELATED to the domain-level ShardLeaseRevision inside +-- `state`: that one starts at 0 and is only bumped when the routing table meaningfully changes. +-- Never derive this column from it. +-- +-- Revision 0 is reserved to mean "no state stored", so a stored row always carries >= 1. The +-- column therefore has NO DEFAULT: a write that forgets to bind it must fail loudly. +-- +-- The table is recreated rather than altered: any row written before this migration has no +-- meaningful revision, and the state rebuilds itself from executor registrations (same reasoning +-- as 003). Recreating also avoids SQLite's requirement that ADD COLUMN supply a non-null default. +DROP TABLE shard_manager_state; + +CREATE TABLE shard_manager_state +( + id INTEGER PRIMARY KEY, + state BLOB NOT NULL, + revision BIGINT NOT NULL +); + +-- Local-mode mirror of the state blob, for inspection with plain SQL (`executor_leases` is the +-- counterpart of the quota system's `quota_leases`). The blob is the source of truth: both tables +-- are rewritten wholesale in the same transaction as every write the shard manager makes, and it +-- never reads them back. Not mirrored: `pending_rebalance`, `shard_epochs` (the per-shard epoch +-- high-water marks, which outlive leases) and the in-blob `ShardLeaseRevision`. +-- +-- Anyone clearing the state by hand (`DELETE FROM shard_manager_state`) must clear these two +-- tables as well, or they keep describing leases that no longer exist until the next write. +-- +-- The shard manager refuses to persist a state whose assignments reference an executor without a +-- lease, so the foreign key below is belt-and-braces. Postgres always enforces it; SQLite only +-- with `foreign_keys = true` (the default is `false`). +CREATE TABLE executor_leases +( + executor_id UUID NOT NULL, + ip BLOB NOT NULL, + port INTEGER NOT NULL, + granted_at TIMESTAMP NOT NULL, + expires_at TIMESTAMP NOT NULL, + pod_name TEXT, + + CONSTRAINT executor_leases_pk + PRIMARY KEY (executor_id) +); + +CREATE TABLE shard_assignments +( + shard_id INTEGER NOT NULL, + executor_id UUID NOT NULL, + epoch BIGINT NOT NULL, + + CONSTRAINT shard_assignments_pk + PRIMARY KEY (shard_id), + CONSTRAINT shard_assignments_executor_fk + FOREIGN KEY (executor_id) REFERENCES executor_leases +); diff --git a/golem-shard-manager/src/config.rs b/golem-shard-manager/src/config.rs index 623e783ce6..76baba4557 100644 --- a/golem-shard-manager/src/config.rs +++ b/golem-shard-manager/src/config.rs @@ -16,7 +16,7 @@ use crate::config::HealthCheckMode::K8s; use golem_common::SafeDisplay; use golem_common::config::{ - ConfigExample, ConfigLoader, DbConfig, DbSqliteConfig, HasConfigExamples, + ConfigExample, ConfigLoader, DbPostgresConfig, DbSqliteConfig, HasConfigExamples, }; use golem_common::model::{Empty, RetryConfig}; use golem_common::tracing::TracingConfig; @@ -31,7 +31,7 @@ use std::time::Duration; #[derive(Clone, Debug, Serialize, Deserialize)] pub struct ShardManagerConfig { pub tracing: TracingConfig, - pub db: DbConfig, + pub persistence: PersistenceConfig, pub worker_executors: WorkerExecutorServiceConfig, pub health_check: HealthCheckConfig, pub http_port: u16, @@ -52,8 +52,12 @@ impl SafeDisplay for ShardManagerConfig { let mut result = String::new(); let _ = writeln!(&mut result, "tracing:"); let _ = writeln!(&mut result, "{}", self.tracing.to_safe_string_indented()); - let _ = writeln!(&mut result, "db:"); - let _ = writeln!(&mut result, "{}", self.db.to_safe_string_indented()); + let _ = writeln!(&mut result, "persistence:"); + let _ = writeln!( + &mut result, + "{}", + self.persistence.to_safe_string_indented() + ); let _ = writeln!(&mut result, "worker executors:"); let _ = writeln!( &mut result, @@ -109,10 +113,7 @@ impl Default for ShardManagerConfig { fn default() -> Self { Self { tracing: TracingConfig::local_dev("shard-manager"), - db: DbConfig::Sqlite(DbSqliteConfig { - database: "golem_shard_manager.db".to_string(), - ..Default::default() - }), + persistence: PersistenceConfig::default(), worker_executors: WorkerExecutorServiceConfig::default(), health_check: HealthCheckConfig::default(), http_port: 8081, @@ -130,30 +131,122 @@ impl Default for ShardManagerConfig { impl HasConfigExamples for ShardManagerConfig { fn examples() -> Vec> { + let etcd_example: ConfigExample = ( + "with etcd persistence (distributed mode)", + Self { + persistence: PersistenceConfig::Etcd(EtcdConfig::default()), + ..Self::default() + }, + ); + #[cfg(feature = "kubernetes")] { - vec![( - "with k8s healthcheck", - Self { - health_check: HealthCheckConfig { - delay: Duration::from_secs(1), - mode: K8s(HealthCheckK8sConfig { - namespace: "namespace".to_string(), - }), - silent: false, + vec![ + ( + "with k8s healthcheck", + Self { + health_check: HealthCheckConfig { + delay: Duration::from_secs(1), + mode: K8s(HealthCheckK8sConfig { + namespace: "namespace".to_string(), + }), + silent: false, + }, + ..Self::default() }, - ..Self::default() - }, - )] + ), + etcd_example, + ] } #[cfg(not(feature = "kubernetes"))] { - Vec::new() + vec![etcd_example] + } + } +} + +/// Where the shard manager persists its state, which also selects the deployment mode. +/// +/// * `Postgres` / `Sqlite` - **local mode**: a single shard manager instance, with the shard +/// lease state and the quota state in one SQL database. +/// * `Etcd` - **distributed mode**: the shard lease state lives in etcd behind a +/// compare-and-swap on the key's `mod_revision`. Quota state is not durable in this mode. +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(tag = "type", content = "config")] +pub enum PersistenceConfig { + Postgres(DbPostgresConfig), + Sqlite(DbSqliteConfig), + Etcd(EtcdConfig), +} + +impl Default for PersistenceConfig { + fn default() -> Self { + Self::Sqlite(DbSqliteConfig { + database: "golem_shard_manager.db".to_string(), + ..Default::default() + }) + } +} + +impl SafeDisplay for PersistenceConfig { + fn to_safe_string(&self) -> String { + let mut result = String::new(); + match self { + PersistenceConfig::Postgres(postgres) => { + let _ = writeln!(&mut result, "postgres:"); + let _ = writeln!(&mut result, "{}", postgres.to_safe_string_indented()); + } + PersistenceConfig::Sqlite(sqlite) => { + let _ = writeln!(&mut result, "sqlite:"); + let _ = writeln!(&mut result, "{}", sqlite.to_safe_string_indented()); + } + PersistenceConfig::Etcd(etcd) => { + let _ = writeln!(&mut result, "etcd:"); + let _ = writeln!(&mut result, "{}", etcd.to_safe_string_indented()); + } + } + result + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct EtcdConfig { + /// Client URLs of the etcd cluster. + /// + /// TLS is not configurable, so only `http://` endpoints are accepted; the shard manager + /// refuses to start otherwise. + /// + /// An environment override must be bracketed, otherwise it is read as a single string and + /// fails to deserialize: + /// `GOLEM__PERSISTENCE__CONFIG__ENDPOINTS=["http://a:2379","http://b:2379"]` + pub endpoints: Vec, + #[serde(with = "humantime_serde")] + pub connect_timeout: Duration, + #[serde(with = "humantime_serde")] + pub request_timeout: Duration, +} + +impl Default for EtcdConfig { + fn default() -> Self { + Self { + endpoints: vec!["http://localhost:2379".to_string()], + connect_timeout: Duration::from_secs(10), + request_timeout: Duration::from_secs(5), } } } +impl SafeDisplay for EtcdConfig { + fn to_safe_string(&self) -> String { + let mut result = String::new(); + let _ = writeln!(&mut result, "endpoints: {}", self.endpoints.join(", ")); + let _ = writeln!(&mut result, "connect timeout: {:?}", self.connect_timeout); + let _ = writeln!(&mut result, "request timeout: {:?}", self.request_timeout); + result + } +} + #[derive(Clone, Debug, Serialize, Deserialize)] pub struct ResourceDefinitionFetcherConfig { pub cache_max_capacity: usize, diff --git a/golem-shard-manager/src/error.rs b/golem-shard-manager/src/error.rs index f40ab0304c..ceb8a86d13 100644 --- a/golem-shard-manager/src/error.rs +++ b/golem-shard-manager/src/error.rs @@ -70,11 +70,21 @@ impl From for golem::shardmanager::v1::ShardManagerError { details, api::error_code::INTERNAL_UNKNOWN, ), + ShardManagerError::ConcurrentModification => error( + shard_manager_error::Error::Unknown, + "Concurrent modification of the persisted shard state".to_string(), + api::error_code::CONCURRENT_UPDATE, + ), ShardManagerError::RepoError(err) => error( shard_manager_error::Error::Unknown, err.to_string(), api::error_code::INTERNAL_DEPENDENCY_FAILURE, ), + ShardManagerError::EtcdError(err) => error( + shard_manager_error::Error::Unknown, + err.to_string(), + api::error_code::INTERNAL_DEPENDENCY_FAILURE, + ), ShardManagerError::MigrationError(err) => error( shard_manager_error::Error::Unknown, err.to_string(), diff --git a/golem-shard-manager/src/lib.rs b/golem-shard-manager/src/lib.rs index d4c4313abe..0ba71cd1fc 100644 --- a/golem-shard-manager/src/lib.rs +++ b/golem-shard-manager/src/lib.rs @@ -22,8 +22,8 @@ pub(crate) mod sharding; use self::grpc::ShardManagerServiceImpl; #[cfg(feature = "kubernetes")] use crate::config::HealthCheckK8sConfig; -use crate::config::HealthCheckMode; -use crate::quota::{DbQuotaRepo, GrpcResourceDefinitionFetcher, QuotaService}; +use crate::config::{HealthCheckMode, PersistenceConfig}; +use crate::quota::{DbQuotaRepo, GrpcResourceDefinitionFetcher, InMemoryQuotaRepo, QuotaService}; use crate::registry_event_subscriber::ShardManagerRegistryInvalidationHandler; use crate::sharding::healthcheck::GrpcHealthCheck; use crate::sharding::worker_executor::WorkerExecutorServiceDefault; @@ -37,7 +37,10 @@ use include_dir::include_dir; use prometheus::Registry; pub use sharding::error::{HealthCheckError, ShardManagerError}; pub use sharding::healthcheck::HealthCheck; -pub use sharding::persistence::{DbRoutingTablePersistence, RoutingTablePersistence}; +pub use sharding::persistence::{ + DbRoutingTablePersistence, EtcdRoutingTablePersistence, ExternalRevision, NO_REVISION, + RoutingTablePersistence, STATE_KEY, +}; pub use sharding::shard_management::ShardManagement; pub use sharding::worker_executor::WorkerExecutorService; pub use sharding::{ @@ -107,48 +110,59 @@ pub async fn run( Arc, Arc, ) = { - use golem_common::config::DbConfig; use golem_service_base::db; use golem_service_base::migration::{IncludedMigrationsDir, Migrations}; - use include_dir::include_dir; - static DB_MIGRATIONS: include_dir::Dir = include_dir!("$CARGO_MANIFEST_DIR/db/migration"); let migrations = IncludedMigrationsDir::new(&DB_MIGRATIONS); - match &shard_manager_config.db { - DbConfig::Postgres(postgres) => { + match &shard_manager_config.persistence { + PersistenceConfig::Postgres(postgres) => { db::postgres::migrate(postgres, migrations.postgres_migrations()).await?; - let pool = - golem_service_base::db::postgres::PostgresPool::configured(postgres).await?; + let pool = db::postgres::PostgresPool::configured(postgres).await?; let pool_for_metrics = pool.clone(); join_set .spawn(async move { pool_for_metrics.run_metrics_loop("shard_manager").await }); - let persistence = Arc::new( - crate::sharding::persistence::DbRoutingTablePersistence::new( + ( + Arc::new(DbRoutingTablePersistence::new( pool.clone(), shard_manager_config.number_of_shards, - ), - ); - let quota_repo = Arc::new(DbQuotaRepo::logged(pool)); - (persistence, quota_repo) + )), + Arc::new(DbQuotaRepo::logged(pool)), + ) } - DbConfig::Sqlite(sqlite) => { + PersistenceConfig::Sqlite(sqlite) => { db::sqlite::migrate(sqlite, migrations.sqlite_migrations()).await?; - let pool = golem_service_base::db::sqlite::SqlitePool::configured(sqlite).await?; + let pool = db::sqlite::SqlitePool::configured(sqlite).await?; - let persistence = Arc::new( - crate::sharding::persistence::DbRoutingTablePersistence::new( + ( + Arc::new(DbRoutingTablePersistence::new( pool.clone(), shard_manager_config.number_of_shards, + )), + Arc::new(DbQuotaRepo::logged(pool)), + ) + } + PersistenceConfig::Etcd(etcd) => { + // Distributed mode. Quota state has no durable store in this mode: it lives in + // the quota service's memory and is rebuilt as executors re-acquire their leases + // after a restart. + // InMemoryQuotaRepo is a placeholder for now + ( + Arc::new( + EtcdRoutingTablePersistence::new( + etcd, + shard_manager_config.number_of_shards, + ) + .await?, ), - ); - let quota_repo = Arc::new(DbQuotaRepo::logged(pool)); - (persistence, quota_repo) + Arc::new(InMemoryQuotaRepo), + ) } } }; + let worker_executors = Arc::new(WorkerExecutorServiceDefault::new( shard_manager_config.worker_executors.clone(), )); diff --git a/golem-shard-manager/src/quota/mod.rs b/golem-shard-manager/src/quota/mod.rs index 71777561d7..4e19754157 100644 --- a/golem-shard-manager/src/quota/mod.rs +++ b/golem-shard-manager/src/quota/mod.rs @@ -20,6 +20,6 @@ mod quota_service_tests; mod quota_state; pub mod resource_definition_fetcher; -pub use quota_repo::{DbQuotaRepo, QuotaRepo}; +pub use quota_repo::{DbQuotaRepo, InMemoryQuotaRepo, QuotaRepo}; pub use quota_service::{QuotaError, QuotaService}; pub use resource_definition_fetcher::{GrpcResourceDefinitionFetcher, ResourceDefinitionFetcher}; diff --git a/golem-shard-manager/src/quota/quota_repo.rs b/golem-shard-manager/src/quota/quota_repo.rs index f421b72676..2f99240ec2 100644 --- a/golem-shard-manager/src/quota/quota_repo.rs +++ b/golem-shard-manager/src/quota/quota_repo.rs @@ -100,6 +100,67 @@ pub trait QuotaRepo: Send + Sync { ) -> Result<(), QuotaRepoError>; } +/// A [`QuotaRepo`] that persists nothing. +/// +/// Quota state then lives only in the quota service's memory: it starts empty and is rebuilt as +/// executors acquire their leases. Used in distributed (etcd) mode, which has no durable quota +/// repository, and by the quota service unit tests. +#[derive(Debug, Default)] +pub struct InMemoryQuotaRepo; + +#[async_trait] +impl QuotaRepo for InMemoryQuotaRepo { + async fn save_lease_change( + &self, + _resource: &QuotaResourceRecord, + _previous_resource_revision: i64, + _lease: &QuotaLeaseRecord, + _expired_pods: &[(Blob, i32)], + ) -> Result<(), QuotaRepoError> { + Ok(()) + } + + async fn save_lease_release( + &self, + _resource: &QuotaResourceRecord, + _previous_resource_revision: i64, + _pod_ip: Blob, + _pod_port: i32, + ) -> Result<(), QuotaRepoError> { + Ok(()) + } + + async fn save_resource( + &self, + _record: &QuotaResourceRecord, + _previous_revision: i64, + ) -> Result<(), QuotaRepoError> { + Ok(()) + } + + async fn delete_resource_and_leases( + &self, + _resource_definition_id: ResourceDefinitionId, + ) -> Result<(), QuotaRepoError> { + Ok(()) + } + + async fn get_all_resources(&self) -> Result, QuotaRepoError> { + Ok(Vec::new()) + } + + async fn get_all_leases(&self) -> Result, QuotaRepoError> { + Ok(Vec::new()) + } + + async fn delete_leases_for_resource( + &self, + _resource_definition_id: ResourceDefinitionId, + ) -> Result<(), QuotaRepoError> { + Ok(()) + } +} + static SPAN_NAME: &str = "quota repository"; pub struct LoggedQuotaRepo { diff --git a/golem-shard-manager/src/quota/quota_service_tests.rs b/golem-shard-manager/src/quota/quota_service_tests.rs index 33b3c9ec73..197b044fee 100644 --- a/golem-shard-manager/src/quota/quota_service_tests.rs +++ b/golem-shard-manager/src/quota/quota_service_tests.rs @@ -13,7 +13,7 @@ // limitations under the License. use super::quota_lease::QuotaLease; -use super::quota_repo::{QuotaLeaseRecord, QuotaRepo, QuotaRepoError, QuotaResourceRecord}; +use super::quota_repo::{InMemoryQuotaRepo, QuotaRepo}; use super::quota_service::{QuotaError, QuotaService}; use super::resource_definition_fetcher::{FetchError, ResourceDefinitionFetcher}; use crate::config::QuotaServiceConfig; @@ -26,7 +26,6 @@ use golem_common::model::quota::{ ResourceDefinitionId, ResourceDefinitionRevision, ResourceLimit, ResourceName, ResourceRateLimit, TimePeriod, }; -use golem_service_base::repo::Blob; use std::collections::HashMap; use std::net::{IpAddr, Ipv4Addr}; use std::sync::Arc; @@ -34,55 +33,6 @@ use std::time::Duration; use test_r::test; use tokio::sync::RwLock; -struct InMemoryQuotaRepo; - -#[async_trait] -impl QuotaRepo for InMemoryQuotaRepo { - async fn save_lease_change( - &self, - _resource: &QuotaResourceRecord, - _previous_resource_revision: i64, - _lease: &QuotaLeaseRecord, - _expired_pods: &[(Blob, i32)], - ) -> Result<(), QuotaRepoError> { - Ok(()) - } - async fn save_lease_release( - &self, - _resource: &QuotaResourceRecord, - _previous_resource_revision: i64, - _pod_ip: Blob, - _pod_port: i32, - ) -> Result<(), QuotaRepoError> { - Ok(()) - } - async fn save_resource( - &self, - _record: &QuotaResourceRecord, - _previous_revision: i64, - ) -> Result<(), QuotaRepoError> { - Ok(()) - } - async fn delete_resource_and_leases( - &self, - _id: ResourceDefinitionId, - ) -> Result<(), QuotaRepoError> { - Ok(()) - } - async fn get_all_resources(&self) -> Result, QuotaRepoError> { - Ok(Vec::new()) - } - async fn get_all_leases(&self) -> Result, QuotaRepoError> { - Ok(Vec::new()) - } - async fn delete_leases_for_resource( - &self, - _id: ResourceDefinitionId, - ) -> Result<(), QuotaRepoError> { - Ok(()) - } -} - fn test_repo() -> Arc { Arc::new(InMemoryQuotaRepo) } diff --git a/golem-shard-manager/src/sharding/error.rs b/golem-shard-manager/src/sharding/error.rs index 289b441b9e..15dfc61c38 100644 --- a/golem-shard-manager/src/sharding/error.rs +++ b/golem-shard-manager/src/sharding/error.rs @@ -33,8 +33,12 @@ pub enum ShardManagerError { WorkerExecutionError(WorkerExecutorError), #[error("Persistence serialization error {0}")] SerializationError(String), - #[error("Postgres error {0}")] + #[error("Concurrent modification: the persisted shard state was changed by another writer")] + ConcurrentModification, + #[error("DB error {0}")] RepoError(#[from] RepoError), + #[error("etcd error {0}")] + EtcdError(#[from] etcd_client::Error), #[error("Migration error {0}")] MigrationError(#[from] anyhow::Error), #[error("IO error {0}")] @@ -53,7 +57,21 @@ impl IsRetriableError for ShardManagerError { ShardManagerError::NoResult => true, ShardManagerError::WorkerExecutionError(_) => true, // TODO: can we define which ones are retryable? ShardManagerError::SerializationError(_) => false, + // Retrying a compare-and-swap with the same, now stale, previous revision can never + // succeed: recovery is a re-read followed by re-deriving the change, which is a + // different operation. Reporting this as retriable would turn a conflict into a spin. + ShardManagerError::ConcurrentModification => false, ShardManagerError::RepoError(_) => false, + ShardManagerError::EtcdError(err) => match err { + etcd_client::Error::GRpcStatus(status) => status.is_retriable(), + etcd_client::Error::TransportError(_) + | etcd_client::Error::IoError(_) + | etcd_client::Error::EndpointError(_) => true, + // A catch-all is required regardless: `etcd_client::Error` has a + // `#[cfg(feature = "tls-openssl")]` variant. Everything else - bad URI, bad + // arguments, bad metadata - is a configuration bug, not a transient failure. + _ => false, + }, ShardManagerError::MigrationError(_) => false, ShardManagerError::IoError(_) => false, ShardManagerError::Internal(_) => false, diff --git a/golem-shard-manager/src/sharding/persistence.rs b/golem-shard-manager/src/sharding/persistence.rs deleted file mode 100644 index 7d66bb81cf..0000000000 --- a/golem-shard-manager/src/sharding/persistence.rs +++ /dev/null @@ -1,189 +0,0 @@ -// Copyright 2024-2026 Golem Cloud -// -// Licensed under the Golem Source License v1.1 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://license.golem.cloud/LICENSE -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use super::error::ShardManagerError; -use super::model::ShardLeaseState; -use anyhow::anyhow; -use async_trait::async_trait; -use conditional_trait_gen::trait_gen; -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}; -use golem_service_base::repo::RepoError; -use sqlx::Row; - -const PERSISTENCE_SVC: &str = "persistence"; - -#[async_trait] -pub trait RoutingTablePersistence: Send + Sync { - async fn write(&self, shard_state: &ShardLeaseState) -> Result<(), ShardManagerError>; - async fn read(&self) -> Result; -} - -pub struct DbRoutingTablePersistence { - pool: DBP, - number_of_shards: usize, -} - -impl DbRoutingTablePersistence { - pub fn new(pool: DBP, number_of_shards: usize) -> Self { - Self { - pool, - number_of_shards, - } - } -} - -#[trait_gen(PostgresPool -> PostgresPool, SqlitePool)] -#[async_trait] -impl RoutingTablePersistence for DbRoutingTablePersistence { - 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") - .execute( - sqlx::query( - "INSERT INTO shard_manager_state (id, state) VALUES (1, $1) \ - ON CONFLICT (id) DO UPDATE SET state = EXCLUDED.state", - ) - .bind(encoded), - ) - .await - .map_err(ShardManagerError::RepoError)?; - - Ok(()) - } - - async fn read(&self) -> Result { - let row = self - .pool - .with_ro(PERSISTENCE_SVC, "read") - .fetch_optional(sqlx::query( - "SELECT state FROM shard_manager_state WHERE id = 1", - )) - .await - .map_err(ShardManagerError::RepoError)?; - - if let Some(row) = row { - let bytes: Vec = row - .try_get("state") - .map_err(|err| RepoError::InternalError(anyhow!(err)))?; - decode_shard_state(&bytes) - } else { - 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/persistence/db.rs b/golem-shard-manager/src/sharding/persistence/db.rs new file mode 100644 index 0000000000..cd9fb91838 --- /dev/null +++ b/golem-shard-manager/src/sharding/persistence/db.rs @@ -0,0 +1,409 @@ +// Copyright 2024-2026 Golem Cloud +// +// Licensed under the Golem Source License v1.1 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://license.golem.cloud/LICENSE +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::{ + ExternalRevision, NO_REVISION, RoutingTablePersistence, check_prev_revision, + check_state_for_write, check_stored_revision, decode_shard_state, +}; +use crate::sharding::error::ShardManagerError; +use crate::sharding::model::ShardLeaseState; +use anyhow::anyhow; +use async_trait::async_trait; +use conditional_trait_gen::trait_gen; +use futures::FutureExt; +use golem_common::serialization::serialize; +use golem_service_base::db::postgres::PostgresPool; +use golem_service_base::db::sqlite::SqlitePool; +use golem_service_base::db::{LabelledPoolApi, Pool, PoolApi}; +use golem_service_base::repo::{Blob, RepoError, SqlDateTime}; +use indoc::indoc; +use sqlx::{QueryBuilder, Row}; +use std::net::IpAddr; +use uuid::Uuid; + +const PERSISTENCE_SVC: &str = "persistence"; + +/// Rows per multi-row `INSERT` into the mirror tables. Six binds per lease row keeps a chunk far +/// below both Postgres' (65535) and SQLite's (32766) bind-parameter limits. +const MIRROR_INSERT_CHUNK_SIZE: usize = 1000; +// Six binds per lease row; SQLite's default SQLITE_MAX_VARIABLE_NUMBER is the smaller limit. +const _: () = assert!(MIRROR_INSERT_CHUNK_SIZE * 6 <= 32766); + +/// Creates the single state row. Succeeds only while the row is absent: `DO NOTHING` turns the +/// primary-key clash into "0 rows affected" instead of a unique-violation error, identically on +/// Postgres and SQLite. +const INSERT_STATE_SQL: &str = indoc! { r#" + INSERT INTO shard_manager_state (id, state, revision) + VALUES (1, $1, $2) + ON CONFLICT (id) DO NOTHING +"#}; + +/// Replaces the single state row. Succeeds only while the row is present and still carries `$3`; +/// an absent row matches nothing and yields 0 rows, which is the same answer etcd gives for a +/// compare against a deleted key. +const UPDATE_STATE_SQL: &str = indoc! { r#" + UPDATE shard_manager_state + SET state = $1, revision = $2 + WHERE id = 1 AND revision = $3 +"#}; + +const SELECT_STATE_SQL: &str = indoc! { r#" + SELECT state, revision + FROM shard_manager_state + WHERE id = 1 +"#}; + +pub struct DbRoutingTablePersistence { + pool: DBP, + number_of_shards: usize, +} + +impl DbRoutingTablePersistence { + pub fn new(pool: DBP, number_of_shards: usize) -> Self { + Self { + pool, + number_of_shards, + } + } +} + +#[trait_gen(PostgresPool -> PostgresPool, SqlitePool)] +#[async_trait] +impl RoutingTablePersistence for DbRoutingTablePersistence { + async fn read(&self) -> Result<(ShardLeaseState, ExternalRevision), ShardManagerError> { + // NOTE: `with_ro` and `with_rw` address the same Postgres pool. If a read replica is + // ever put behind `with_ro`, a stale read here yields a stale revision and every + // subsequent compare-and-swap fails - a livelock, not corruption - and this read has to + // move to `with_rw` at that point. + let row = self + .pool + .with_ro(PERSISTENCE_SVC, "read") + .fetch_optional(sqlx::query(SELECT_STATE_SQL)) + .await + .map_err(ShardManagerError::RepoError)?; + + let Some(row) = row else { + return Ok((ShardLeaseState::new(self.number_of_shards), NO_REVISION)); + }; + + let bytes: Vec = row + .try_get("state") + .map_err(|err| RepoError::InternalError(anyhow!(err)))?; + let revision: ExternalRevision = row + .try_get("revision") + .map_err(|err| RepoError::InternalError(anyhow!(err)))?; + + if revision < 1 { + return Err(ShardManagerError::Internal(format!( + "persisted shard lease state carries revision {revision}, which is reserved for \ + absent state" + ))); + } + + Ok((decode_shard_state(&bytes)?, revision)) + } + + async fn write( + &self, + shard_state: &ShardLeaseState, + prev_revision: ExternalRevision, + ) -> Result { + check_prev_revision(prev_revision)?; + check_state_for_write(shard_state)?; + let next_revision = prev_revision.checked_add(1).ok_or_else(|| { + ShardManagerError::Internal("shard state storage revision overflow".to_string()) + })?; + let encoded = serialize(shard_state).map_err(ShardManagerError::SerializationError)?; + let leases = lease_rows(shard_state); + let assignments = assignment_rows(shard_state)?; + + // One transaction: the compare-and-swap on the blob decides, and the mirror tables follow + // it or are left untouched with it. + let revision = self + .pool + .with_tx_err(PERSISTENCE_SVC, "write", |tx| { + async move { + Self::compare_and_swap_state_in_tx(tx, encoded, prev_revision, next_revision) + .await?; + Self::replace_mirror_rows_in_tx(tx, &leases, &assignments).await?; + Ok::(next_revision) + } + .boxed() + }) + .await?; + + check_stored_revision(revision, prev_revision) + } +} + +#[trait_gen(PostgresPool -> PostgresPool, SqlitePool)] +impl DbRoutingTablePersistence { + async fn compare_and_swap_state_in_tx( + tx: &mut <::LabelledApi as LabelledPoolApi>::LabelledTransaction, + encoded: Vec, + prev_revision: ExternalRevision, + next_revision: ExternalRevision, + ) -> Result<(), ShardManagerError> { + // `prev_revision == NO_REVISION` asserts the row does not exist, which is an + // insert-if-absent, not an update. A single `INSERT ... ON CONFLICT DO UPDATE ... WHERE + // revision = $prev` cannot express it: its INSERT branch is unguarded, so it would happily + // resurrect a row that was deleted underneath a writer holding `prev > NO_REVISION`, while + // etcd refuses the same write. Two statements keep both backends honest. + let query = if prev_revision == NO_REVISION { + sqlx::query(INSERT_STATE_SQL) + .bind(encoded) // $1 + .bind(next_revision) // $2 + } else { + sqlx::query(UPDATE_STATE_SQL) + .bind(encoded) // $1 + .bind(next_revision) // $2 + .bind(prev_revision) // $3 + }; + + let result = tx.execute(query).await?; + if result.rows_affected() == 0 { + return Err(ShardManagerError::ConcurrentModification); + } + Ok(()) + } + + /// Rewrites the mirror tables from scratch. They are a projection of the blob, so replacing + /// them wholesale is both simpler and safer than diffing: nothing can be left behind. + async fn replace_mirror_rows_in_tx( + tx: &mut <::LabelledApi as LabelledPoolApi>::LabelledTransaction, + leases: &[ExecutorLeaseRow], + assignments: &[ShardAssignmentRow], + ) -> Result<(), ShardManagerError> { + // Assignments reference leases, so they go first on delete and last on insert. + tx.execute(sqlx::query("DELETE FROM shard_assignments")) + .await?; + tx.execute(sqlx::query("DELETE FROM executor_leases")) + .await?; + + for chunk in leases.chunks(MIRROR_INSERT_CHUNK_SIZE) { + let mut query = QueryBuilder::<::Db>::new( + "INSERT INTO executor_leases \ + (executor_id, ip, port, granted_at, expires_at, pod_name) ", + ); + query.push_values(chunk, |mut row, lease| { + row.push_bind(lease.executor_id) + .push_bind(lease.ip.clone()) + .push_bind(lease.port) + .push_bind(lease.granted_at.clone()) + .push_bind(lease.expires_at.clone()) + .push_bind(lease.pod_name.clone()); + }); + tx.execute(query.build()).await?; + } + + for chunk in assignments.chunks(MIRROR_INSERT_CHUNK_SIZE) { + let mut query = QueryBuilder::<::Db>::new( + "INSERT INTO shard_assignments (shard_id, executor_id, epoch) ", + ); + query.push_values(chunk, |mut row, assignment| { + row.push_bind(assignment.shard_id) + .push_bind(assignment.executor_id) + .push_bind(assignment.epoch); + }); + tx.execute(query.build()).await?; + } + + Ok(()) + } +} + +/// One `executor_leases` row. +#[derive(Debug, Clone, PartialEq)] +struct ExecutorLeaseRow { + executor_id: Uuid, + ip: Blob, + port: i32, + granted_at: SqlDateTime, + expires_at: SqlDateTime, + pod_name: Option, +} + +/// One `shard_assignments` row. +#[derive(Debug, Clone, PartialEq, Eq)] +struct ShardAssignmentRow { + shard_id: i32, + executor_id: Uuid, + epoch: i64, +} + +fn lease_rows(shard_state: &ShardLeaseState) -> Vec { + shard_state + .executor_leases + .iter() + .map(|(executor_id, lease)| ExecutorLeaseRow { + executor_id: executor_id.0, + ip: Blob::new(lease.addr.ip), + port: i32::from(lease.addr.port), + granted_at: SqlDateTime::new(lease.granted_at), + expires_at: SqlDateTime::new(lease.expires_at), + pod_name: lease.pod_name.clone(), + }) + .collect() +} + +fn assignment_rows( + shard_state: &ShardLeaseState, +) -> Result, ShardManagerError> { + shard_state + .shard_assignments + .iter() + .map(|(shard_id, entry)| { + Ok(ShardAssignmentRow { + shard_id: i32::try_from(shard_id.value()).map_err(|_| { + ShardManagerError::Internal(format!( + "shard id {} does not fit the shard_assignments.shard_id column", + shard_id.value() + )) + })?, + executor_id: entry.executor_id.0, + epoch: i64::try_from(entry.epoch.0).map_err(|_| { + ShardManagerError::Internal(format!( + "shard epoch {} does not fit the shard_assignments.epoch column", + entry.epoch.0 + )) + })?, + }) + }) + .collect() +} + +#[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::ShardId; + use std::net::Ipv4Addr; + use std::time::Duration; + + const TTL: Duration = Duration::from_secs(60); + + fn t0() -> DateTime { + DateTime::from_timestamp(1_700_000_000, 0).unwrap() + } + + fn addr(last_octet: u8, port: u16) -> ExecutorAddr { + ExecutorAddr { + ip: IpAddr::V4(Ipv4Addr::new(10, 0, 0, last_octet)), + port, + } + } + + #[test] + fn mirror_rows_project_the_state() { + let mut shard_state = ShardLeaseState::new(8); + shard_state.add_executor( + ExecutorId(Uuid::from_u128(1)), + addr(1, 9010), + Some("worker-executor-0".to_string()), + t0(), + TTL, + ); + shard_state.add_executor( + ExecutorId(Uuid::from_u128(2)), + addr(2, 9011), + None, + t0(), + TTL, + ); + shard_state.assign_shard(ExecutorId(Uuid::from_u128(1)), ShardId::new(0)); + shard_state.assign_shard(ExecutorId(Uuid::from_u128(2)), ShardId::new(1)); + // moving a shard bumps its epoch + shard_state.assign_shard(ExecutorId(Uuid::from_u128(2)), ShardId::new(0)); + + let leases = lease_rows(&shard_state); + assert_eq!(leases.len(), 2); + assert_eq!(leases[0].executor_id, Uuid::from_u128(1)); + assert_eq!( + *leases[0].ip.value(), + IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)) + ); + assert_eq!(leases[0].port, 9010); + assert_eq!(leases[0].granted_at, SqlDateTime::new(t0())); + assert_eq!( + leases[0].expires_at, + SqlDateTime::new(t0() + chrono::Duration::from_std(TTL).unwrap()) + ); + assert_eq!(leases[0].pod_name.as_deref(), Some("worker-executor-0")); + assert_eq!(leases[1].pod_name, None); + + let assignments = assignment_rows(&shard_state).unwrap(); + assert_eq!( + assignments, + vec![ + ShardAssignmentRow { + shard_id: 0, + executor_id: Uuid::from_u128(2), + epoch: 1, + }, + ShardAssignmentRow { + shard_id: 1, + executor_id: Uuid::from_u128(2), + epoch: 0, + }, + ] + ); + assert_eq!( + shard_state.epoch_for_shard(ShardId::new(0)), + Some(ShardEpoch(1)) + ); + } + + #[test] + fn mirror_rows_refuse_values_the_columns_cannot_hold() { + let executor_id = ExecutorId(Uuid::from_u128(1)); + + let mut shard_state = ShardLeaseState::new(8); + shard_state.add_executor(executor_id, addr(1, 9010), None, t0(), TTL); + shard_state.shard_assignments.insert( + ShardId::new(0), + ShardAssignmentEntry { + executor_id, + epoch: ShardEpoch(u64::MAX), + }, + ); + match assignment_rows(&shard_state) { + Err(ShardManagerError::Internal(msg)) => assert!(msg.contains("epoch"), "{msg}"), + other => panic!("expected Internal, got {other:?}"), + } + + let mut shard_state = ShardLeaseState::new(8); + shard_state.add_executor(executor_id, addr(1, 9010), None, t0(), TTL); + shard_state.shard_assignments.insert( + ShardId::new(i64::from(i32::MAX) + 1), + ShardAssignmentEntry { + executor_id, + epoch: ShardEpoch::initial(), + }, + ); + match assignment_rows(&shard_state) { + Err(ShardManagerError::Internal(msg)) => assert!(msg.contains("shard id"), "{msg}"), + other => panic!("expected Internal, got {other:?}"), + } + } + + #[test] + fn mirror_rows_of_an_empty_state_are_empty() { + let shard_state = ShardLeaseState::new(8); + assert!(lease_rows(&shard_state).is_empty()); + assert!(assignment_rows(&shard_state).unwrap().is_empty()); + } +} diff --git a/golem-shard-manager/src/sharding/persistence/etcd.rs b/golem-shard-manager/src/sharding/persistence/etcd.rs new file mode 100644 index 0000000000..5374a67b76 --- /dev/null +++ b/golem-shard-manager/src/sharding/persistence/etcd.rs @@ -0,0 +1,150 @@ +// Copyright 2024-2026 Golem Cloud +// +// Licensed under the Golem Source License v1.1 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://license.golem.cloud/LICENSE +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::{ + ExternalRevision, NO_REVISION, RoutingTablePersistence, check_prev_revision, + check_state_for_write, check_stored_revision, decode_shard_state, +}; +use crate::config::EtcdConfig; +use crate::sharding::error::ShardManagerError; +use crate::sharding::model::ShardLeaseState; +use async_trait::async_trait; +use etcd_client::{Client, Compare, CompareOp, ConnectOptions, Txn, TxnOp}; +use golem_common::serialization::serialize; +use tracing::info; + +/// Key holding the serialized [`ShardLeaseState`]. +pub const STATE_KEY: &str = "/golem/shard-manager/state"; + +pub struct EtcdRoutingTablePersistence { + client: Client, + number_of_shards: usize, +} + +impl EtcdRoutingTablePersistence { + pub async fn new( + config: &EtcdConfig, + number_of_shards: usize, + ) -> Result { + if config.endpoints.is_empty() { + return Err(ShardManagerError::Internal( + "etcd shard state persistence requires at least one endpoint".to_string(), + )); + } + + // TLS is not configurable, so an `https://` endpoint could only fail at connect time with + // an opaque transport error. Refuse it up front and say why instead. + if let Some(endpoint) = config + .endpoints + .iter() + .find(|endpoint| !endpoint.starts_with("http://")) + { + return Err(ShardManagerError::Internal(format!( + "etcd endpoint {endpoint} is not an http:// URL; TLS is not supported" + ))); + } + + let options = ConnectOptions::new() + .with_connect_timeout(config.connect_timeout) + .with_timeout(config.request_timeout); + + let client = Client::connect(&config.endpoints, Some(options)).await?; + info!( + endpoints = config.endpoints.join(", "), + state_key = STATE_KEY, + "Connected to etcd for shard lease state persistence" + ); + + Ok(Self { + client, + number_of_shards, + }) + } +} + +#[async_trait] +impl RoutingTablePersistence for EtcdRoutingTablePersistence { + async fn read(&self) -> Result<(ShardLeaseState, ExternalRevision), ShardManagerError> { + // `KvClient` is a cheap handle over the shared multiplexed channel; cloning it per call is + // how the generated stubs' `&mut self` is satisfied. + let mut kv = self.client.kv_client(); + let response = kv.get(STATE_KEY, None).await?; + + let Some(kv_pair) = response.kvs().first() else { + return Ok((ShardLeaseState::new(self.number_of_shards), NO_REVISION)); + }; + + // etcd's store starts at revision 1 and every mutation increments it, so a live key cannot + // have mod_revision 0. If it does, the "0 means absent" invariant is broken. + let revision = kv_pair.mod_revision(); + if revision < 1 { + return Err(ShardManagerError::Internal(format!( + "etcd returned key {STATE_KEY} with mod_revision {revision}, which is reserved for \ + absent keys" + ))); + } + + Ok((decode_shard_state(kv_pair.value())?, revision)) + } + + async fn write( + &self, + shard_state: &ShardLeaseState, + prev_revision: ExternalRevision, + ) -> Result { + check_prev_revision(prev_revision)?; + check_state_for_write(shard_state)?; + let encoded = serialize(shard_state).map_err(ShardManagerError::SerializationError)?; + + // etcd reports mod_revision 0 for a key that does not exist, so `prev_revision == + // NO_REVISION` already means exactly "must not exist yet" - unlike the SQL backend, no + // branching is needed here. + // + // `or_else` is deliberately empty. Returning the winning state from a lost compare-and-swap + // would have to travel inside the error variant, which the SQL backend cannot fill in + // without a second query; the caller re-reads instead. + let txn = Txn::new() + .when([Compare::mod_revision( + STATE_KEY, + CompareOp::Equal, + prev_revision, + )]) + .and_then([TxnOp::put(STATE_KEY, encoded, None)]); + + let mut kv = self.client.kv_client(); + let response = kv.txn(txn).await?; + + if !response.succeeded() { + return Err(ShardManagerError::ConcurrentModification); + } + + // A transaction that applied at least one mutation reports the NEW store revision in its + // header, and etcd stamps every key written by that transaction with exactly that + // revision - so this is the mod_revision our PUT produced. + // + // Do NOT replace this with a follow-up GET: another writer can land between the + // transaction and the GET, and we would cache their revision as ours, so our next + // compare-and-swap would compare-equal against their write and silently clobber it. + let revision = response + .header() + .ok_or_else(|| { + ShardManagerError::Internal( + "etcd transaction response carried no header".to_string(), + ) + })? + .revision(); + + check_stored_revision(revision, prev_revision) + } +} diff --git a/golem-shard-manager/src/sharding/persistence/mod.rs b/golem-shard-manager/src/sharding/persistence/mod.rs new file mode 100644 index 0000000000..eb9e17f357 --- /dev/null +++ b/golem-shard-manager/src/sharding/persistence/mod.rs @@ -0,0 +1,259 @@ +// Copyright 2024-2026 Golem Cloud +// +// Licensed under the Golem Source License v1.1 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://license.golem.cloud/LICENSE +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +mod db; +mod etcd; + +pub use db::DbRoutingTablePersistence; +pub use etcd::{EtcdRoutingTablePersistence, STATE_KEY}; + +use super::error::ShardManagerError; +use super::model::ShardLeaseState; +use async_trait::async_trait; +use golem_common::serialization::try_deserialize; + +/// An opaque, backend-assigned version of the persisted [`ShardLeaseState`]. +/// +/// This is a *storage* fencing token, not a domain concept. It is deliberately unrelated to +/// [`super::model::ShardLeaseRevision`], which lives inside the state blob and is bumped by the +/// shard management loop when the routing table changes. The two are expected to drift apart and +/// must never be compared or derived from each other. +/// +/// Guarantees, upheld by every implementation: +/// * [`NO_REVISION`] means "no state is stored". +/// * Any *stored* state carries a revision `>= 1`, so [`NO_REVISION`] is unambiguous. +/// * Successive successful writes return strictly increasing values. +/// +/// The magnitude is meaningless: the SQL backend assigns `previous + 1`, while the etcd backend +/// assigns the cluster-wide etcd revision, which jumps by arbitrary amounts because every +/// unrelated etcd mutation advances it. +/// +/// Monotonicity is guaranteed only for as long as state remains stored. The SQL token is derived +/// from the row, so deleting the row restarts the sequence at 1, whereas etcd's keeps climbing. +/// A writer holding a token from before such a deletion is therefore not fenced by it; only +/// leader election makes that safe. +pub type ExternalRevision = i64; + +/// The revision reported by [`RoutingTablePersistence::read`] when nothing is stored, and the +/// only value [`RoutingTablePersistence::write`] accepts for a write that must create the state. +pub const NO_REVISION: ExternalRevision = 0; + +#[async_trait] +pub trait RoutingTablePersistence: Send + Sync { + /// Loads the persisted shard lease state together with the revision it is stored at. + /// + /// If nothing is stored, returns a freshly initialized [`ShardLeaseState`] paired with + /// [`NO_REVISION`]. The caller cannot distinguish "never written" from "written and then + /// externally deleted", and does not need to: both mean the routing table is rebuilt from + /// scratch as executors register. + /// + /// A stored blob that cannot be decoded, or that violates the state invariants, is an error. + /// It is never silently replaced by a default state - that would drop a live routing table + /// on a transient decoding bug. + async fn read(&self) -> Result<(ShardLeaseState, ExternalRevision), ShardManagerError>; + + /// Stores `shard_state`, but only if the currently stored revision is exactly + /// `prev_revision`. Returns the revision it was stored at. + /// + /// * `prev_revision == NO_REVISION` means **"no state must exist yet"**. The write succeeds + /// only if nothing is stored, and creates it. If state does exist, at any revision, the + /// write fails - it never overwrites. + /// * `prev_revision > NO_REVISION` means **"the stored state must still be the one I read at + /// that revision"**. In particular, a write with `prev_revision > NO_REVISION` against + /// absent state fails; it does not resurrect it. + /// + /// When that condition does not hold, returns [`ShardManagerError::ConcurrentModification`] + /// and stores nothing. Recovery is always the same: discard the in-memory state, + /// [`Self::read`] again, re-derive the intended change against what was read, and write with + /// the revision that came with it. Retrying with the same `prev_revision` can never succeed, + /// which is why the error is reported as non-retriable. + /// + /// The returned revision is always `>= 1` and always strictly greater than `prev_revision`. + async fn write( + &self, + shard_state: &ShardLeaseState, + prev_revision: ExternalRevision, + ) -> Result; +} + +/// Decodes a persisted state blob and refuses to load one that violates the state invariants. +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) +} + +/// Refuses to persist a state that violates [`ShardLeaseState::check_invariants`]. +/// +/// Every backend calls this before touching its store, so an invalid state is rejected +/// identically everywhere and before any I/O - rather than by whichever constraint a backend +/// happens to enforce (the SQL mirror tables' foreign key, which SQLite only checks with +/// `foreign_keys = true`), and rather than poisoning the store for every later +/// [`RoutingTablePersistence::read`]. +fn check_state_for_write(shard_state: &ShardLeaseState) -> Result<(), ShardManagerError> { + shard_state.check_invariants().map_err(|violation| { + ShardManagerError::Internal(format!( + "refusing to persist a shard lease state that violates invariants: {violation}" + )) + }) +} + +/// Rejects a revision a backend claims to have stored at, if it would be indistinguishable from +/// "absent" or would break monotonicity. Both implementations funnel their result through this, +/// so a protocol violation surfaces as an error instead of silently corrupting the fencing chain. +fn check_stored_revision( + revision: ExternalRevision, + prev_revision: ExternalRevision, +) -> Result { + if revision < 1 || revision <= prev_revision { + return Err(ShardManagerError::Internal(format!( + "persistence backend reported revision {revision} after a write guarded on \ + {prev_revision}, which is not a valid successor" + ))); + } + Ok(revision) +} + +/// Rejects a `prev_revision` that cannot have come from this layer. +fn check_prev_revision(prev_revision: ExternalRevision) -> Result<(), ShardManagerError> { + if prev_revision < NO_REVISION { + return Err(ShardManagerError::Internal(format!( + "negative previous revision {prev_revision}" + ))); + } + Ok(()) +} + +#[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 golem_common::serialization::serialize; + 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:?}"), + } + } + + #[test] + fn a_state_violating_invariants_is_refused_for_write() { + 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(), + }, + ); + match check_state_for_write(&shard_state) { + Err(ShardManagerError::Internal(msg)) => { + assert!(msg.contains("violates invariants"), "{msg}"); + } + other => panic!("expected Internal, got {other:?}"), + } + } + + #[test] + fn stored_revision_must_be_a_valid_successor() { + assert_eq!(check_stored_revision(1, NO_REVISION).unwrap(), 1); + assert_eq!(check_stored_revision(9, 4).unwrap(), 9); + // Indistinguishable from "absent". + assert!(check_stored_revision(0, NO_REVISION).is_err()); + // Not monotonic. + assert!(check_stored_revision(4, 4).is_err()); + assert!(check_stored_revision(3, 4).is_err()); + } +} diff --git a/golem-shard-manager/src/sharding/shard_management.rs b/golem-shard-manager/src/sharding/shard_management.rs index 08dd9a905f..552b3fd082 100644 --- a/golem-shard-manager/src/sharding/shard_management.rs +++ b/golem-shard-manager/src/sharding/shard_management.rs @@ -15,7 +15,7 @@ use super::error::ShardManagerError; use super::healthcheck::{HealthCheck, get_unhealthy_executors}; use super::model::{Assignments, ExecutorAddr, ExecutorAddrs, ExecutorId, ShardLeaseState}; -use super::persistence::RoutingTablePersistence; +use super::persistence::{ExternalRevision, RoutingTablePersistence}; use super::rebalancing::Rebalance; use super::worker_executor::{ WorkerExecutorService, assign_shards, revoke_shards, set_shard_assignments, @@ -29,13 +29,16 @@ use std::sync::Arc; use std::time::Duration; use tokio::sync::{Mutex, Notify}; use tokio::task::JoinSet; -use tracing::{Instrument, debug, info, warn}; +use tracing::{Instrument, debug, error, info, warn}; #[derive(Clone)] pub struct ShardManagement { shard_state: Arc>, change: Arc, updates: Arc>, + persistence: Arc, + /// Compare-and-swap token of the persisted state; see [`Self::mutate_and_persist`]. + external_revision: Arc>, } impl ShardManagement { @@ -52,7 +55,7 @@ impl ShardManagement { lease_ttl: Duration, join_set: &mut JoinSet>, ) -> Result { - let shard_state = persistence_service.read().await?; + let (shard_state, external_revision) = persistence_service.read().await?; info!("Initial healthcheck started"); @@ -66,43 +69,33 @@ impl ShardManagement { info!("Initial healthcheck finished"); - let change = Arc::new(Notify::new()); - let updates = Arc::new(Mutex::new(ShardManagementChanges::new( - healthy_executors, - unhealthy_executors, - ))); - let shard_state = Arc::new(RwLock::new(shard_state)); + let shard_management = ShardManagement { + shard_state: Arc::new(RwLock::new(shard_state)), + change: Arc::new(Notify::new()), + updates: Arc::new(Mutex::new(ShardManagementChanges::new( + healthy_executors, + unhealthy_executors, + ))), + persistence: persistence_service, + external_revision: Arc::new(Mutex::new(external_revision)), + }; { - let change = change.clone(); - let updates = updates.clone(); - let shard_state = shard_state.clone(); - + let shard_management = shard_management.clone(); join_set.spawn( async move { - Self::worker( - shard_state, - change, - updates, - persistence_service, - worker_executors, - threshold, - lease_ttl, - ) - .await; - Ok(()) + shard_management + .worker(worker_executors, threshold, lease_ttl) + .await + .map_err(anyhow::Error::from) } .in_current_span(), ); - }; + } - change.notify_one(); + shard_management.change.notify_one(); - Ok(ShardManagement { - shard_state, - change, - updates, - }) + Ok(shard_management) } /// Registers a new executor instance listening at `addr`. @@ -138,20 +131,17 @@ impl ShardManagement { } async fn worker( - shard_state: Arc>, - change: Arc, - updates: Arc>, - persistence_service: Arc, + self, worker_executors: Arc, threshold: f64, lease_ttl: Duration, - ) { + ) -> Result<(), ShardManagerError> { loop { debug!("Shard management loop awaiting changes"); - change.notified().await; + self.change.notified().await; let (new_executors, removed_executors, full_assignment_requests) = - updates.lock().await.reset(); + self.updates.lock().await.reset(); debug!( new_executors = new_executors .values() @@ -161,44 +151,36 @@ impl ShardManagement { full_assignment_requests = full_assignment_requests.iter().join(", "), "Shard management loop woken up", ); - // Getting a write lock while + // The write lock is held while // - registrations and removals are applied to the state and got persisted, // - the rebalance plan is calculated, // 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 full_assignment_executors = apply_executor_changes( - &mut current_shard_state, - new_executors, - removed_executors, - full_assignment_requests, - lease_ttl, - ); + let (mut rebalance, full_assignment_executors, addrs) = self + .mutate_and_persist(|current_shard_state| { + // 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_shard_state(¤t_shard_state, threshold); - let addrs = current_shard_state.executor_addrs(); + let full_assignment_executors = apply_executor_changes( + current_shard_state, + new_executors, + removed_executors, + full_assignment_requests, + lease_ttl, + ); - current_shard_state - .bump_revision() - .expect("Failed to bump shard lease state revision"); - persistence_service - .write(¤t_shard_state) - .await - .expect("Failed to persist shard lease state after executor changes"); + let rebalance = Rebalance::from_shard_state(current_shard_state, threshold); + let addrs = current_shard_state.executor_addrs(); - (rebalance, full_assignment_executors, addrs) - }; + (rebalance, full_assignment_executors, addrs) + }) + .await?; debug!(rebalance=%rebalance, "Applying rebalance plan"); let rebalance_failures = @@ -219,7 +201,7 @@ impl ShardManagement { ); { - let mut updates_guard = updates.lock().await; + let mut updates_guard = self.updates.lock().await; for (executor_id, _) in &rebalance_failures.failed_assignments { if full_assignment_executors.contains(executor_id) { updates_guard.retry_full_assignment(*executor_id); @@ -241,18 +223,15 @@ impl ShardManagement { needs_retry = true; } - 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(&shard_state_snapshot) - .await - .expect("Failed to persist shard lease state after rebalance"); + self.mutate_and_persist(|current_shard_state| { + current_shard_state.apply_rebalance(&rebalance) + }) + .await?; + + // Race-free: the worker task is the only writer of the shard state. Reading the + // snapshot after the persist rather than out of the closure keeps its `revision` + // field consistent with what was stored. + let shard_state_snapshot = self.shard_state.read().await.clone(); let mut full_assignments = Assignments::new(); for executor_id in &full_assignment_executors { @@ -288,7 +267,7 @@ impl ShardManagement { ); { - let mut updates_guard = updates.lock().await; + let mut updates_guard = self.updates.lock().await; for (executor_id, _) in &failed_full_assignments { updates_guard.retry_full_assignment(*executor_id); } @@ -297,7 +276,64 @@ impl ShardManagement { } if needs_retry { - change.notify_one(); + self.change.notify_one(); + } + } + } + + /// Applies `mutate` to the shard lease state and persists the result compare-and-swap style: + /// snapshot, mutate, bump the revision, write guarded on the cached external revision. On any + /// failure the in-memory state is rolled back to the snapshot and the error is returned - the + /// same pattern as the quota service's lease mutations. + /// + /// The write lock is held across the persistence round-trip so that readers of + /// [`Self::current_snapshot`] can never observe a state that was not durably stored and then + /// watch it go backwards. Lock order is `shard_state`, then `external_revision`. Every writer + /// of the persisted state must go through here, which is what makes in-process conflicts + /// impossible. + /// + /// A [`ShardManagerError::ConcurrentModification`] therefore means another shard manager + /// *process* wrote the state. The cached revision is deliberately not refreshed on failure: it + /// is the fencing token, and a writer that lost it must stop, not adopt the winner's and go on. + async fn mutate_and_persist(&self, mutate: F) -> Result + where + F: FnOnce(&mut ShardLeaseState) -> T, + { + let mut current_shard_state = self.shard_state.write().await; + let mut external_revision = self.external_revision.lock().await; + + let snapshot = current_shard_state.clone(); + let prev_external_revision = *external_revision; + let outcome = mutate(&mut current_shard_state); + + let written = match current_shard_state.bump_revision() { + Ok(_) => { + self.persistence + .write(¤t_shard_state, prev_external_revision) + .await + } + Err(err) => Err(err), + }; + + match written { + Ok(new_external_revision) => { + *external_revision = new_external_revision; + Ok(outcome) + } + Err(err) => { + match &err { + ShardManagerError::ConcurrentModification => error!( + prev_external_revision, + "Revision conflict: another shard manager wrote the shard lease state. \ + Rolling back" + ), + other => error!( + error = %other, + "Persisting the shard lease state failed, rolling back" + ), + } + *current_shard_state = snapshot; + Err(err) } } } diff --git a/golem-shard-manager/tests/persistence.rs b/golem-shard-manager/tests/persistence.rs index ee673fcc94..f9b13c82cc 100644 --- a/golem-shard-manager/tests/persistence.rs +++ b/golem-shard-manager/tests/persistence.rs @@ -17,10 +17,14 @@ use chrono::{DateTime, Utc}; use golem_common::config::{DbPostgresConfig, DbSqliteConfig}; use golem_common::model::ShardId; use golem_service_base::migration::{IncludedMigrationsDir, Migrations}; +use golem_service_base::repo::{Blob, SqlDateTime}; +use golem_shard_manager::config::EtcdConfig; use golem_shard_manager::{ - DbRoutingTablePersistence, ExecutorAddr, ExecutorId, RoutingTablePersistence, - ShardLeaseRevision, ShardLeaseState, + DbRoutingTablePersistence, EtcdRoutingTablePersistence, ExecutorAddr, ExecutorId, + ExternalRevision, NO_REVISION, RoutingTablePersistence, STATE_KEY, ShardAssignmentEntry, + ShardEpoch, ShardLeaseRevision, ShardLeaseState, ShardManagerError, }; +use golem_test_framework::components::etcd::docker_etcd::DockerEtcd; use golem_test_framework::components::rdb::docker_postgres::DockerPostgresRdb; use std::net::{IpAddr, Ipv4Addr}; use std::sync::Arc; @@ -30,13 +34,129 @@ use test_r::{define_matrix_dimension, test, test_dep}; use url::Url; use uuid::Uuid; +/// One `executor_leases` row, as a person reading the table would see it. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +struct LeaseRow { + executor_id: Uuid, + ip: IpAddr, + port: i32, + granted_at: DateTime, + expires_at: DateTime, + pod_name: Option, +} + +/// The shape `SELECT executor_id, ip, port, granted_at, expires_at, pod_name` decodes into. +type RawLeaseRow = ( + Uuid, + Blob, + i32, + SqlDateTime, + SqlDateTime, + Option, +); + +/// What the local-mode mirror tables hold, normalized for comparison. +#[derive(Debug, PartialEq, Eq)] +struct MirrorSnapshot { + leases: Vec, + /// `(shard_id, executor_id, epoch)` per `shard_assignments` row, sorted. + assignments: Vec<(i32, Uuid, i64)>, +} + +impl MirrorSnapshot { + fn of(shard_state: &ShardLeaseState) -> Self { + let leases = shard_state + .executor_leases + .iter() + .map(|(id, lease)| LeaseRow { + executor_id: id.0, + ip: lease.addr.ip, + port: i32::from(lease.addr.port), + granted_at: lease.granted_at, + expires_at: lease.expires_at, + pod_name: lease.pod_name.clone(), + }) + .collect(); + let assignments = shard_state + .shard_assignments + .iter() + .map(|(shard_id, entry)| { + ( + i32::try_from(shard_id.value()).unwrap(), + entry.executor_id.0, + i64::try_from(entry.epoch.0).unwrap(), + ) + }) + .collect(); + Self::sorted(leases, assignments) + } + + fn from_rows(leases: Vec, assignments: Vec<(i32, Uuid, i64)>) -> Self { + let leases = leases + .into_iter() + .map( + |(executor_id, ip, port, granted_at, expires_at, pod_name)| LeaseRow { + executor_id, + ip: ip.into_value(), + port, + granted_at: granted_at.into_utc(), + expires_at: expires_at.into_utc(), + pod_name, + }, + ) + .collect(); + Self::sorted(leases, assignments) + } + + fn empty() -> Self { + Self::sorted(Vec::new(), Vec::new()) + } + + fn sorted(mut leases: Vec, mut assignments: Vec<(i32, Uuid, i64)>) -> Self { + leases.sort(); + assignments.sort(); + Self { + leases, + assignments, + } + } +} + const LEASE_TTL: Duration = Duration::from_secs(60); +const NUMBER_OF_SHARDS: usize = 16; +/// A place where shard lease state can be stored. +/// +/// Hands out any number of independent clients over the *same* underlying store, which is what +/// the compare-and-swap tests need: two clients must see each other's writes. +#[async_trait] +trait PersistenceStore: std::fmt::Debug + Send + Sync { + async fn connect(&self) -> Arc; + + /// Writes something unrelated to the shard state on the same backend, so a test can check + /// that activity elsewhere on the backend does not move this store's revision. + async fn unrelated_write(&self); + + /// Whether this backend has the local-mode mirror tables at all. + fn has_mirror(&self) -> bool; + + /// The local-mode mirror tables, or `None` for a backend that has none. + async fn mirror_snapshot(&self) -> Option; +} + +/// Creates isolated stores: two stores never see each other's data. #[async_trait] trait GetRoutingTablePersistence: std::fmt::Debug + Send + Sync { - async fn get_persistence(&self) -> Arc; + async fn new_store(&self) -> Arc; + + /// For the tests that only need a single client over a fresh store. + async fn get_persistence(&self) -> Arc { + self.new_store().await.connect().await + } } +// -- postgres: isolated by a fresh database per store ------------------------------------------ + struct PostgresRoutingTablePersistence { postgres: DockerPostgresRdb, } @@ -47,9 +167,70 @@ impl std::fmt::Debug for PostgresRoutingTablePersistence { } } +#[derive(Debug)] +struct PostgresStore { + config: DbPostgresConfig, +} + +#[async_trait] +impl PersistenceStore for PostgresStore { + async fn connect(&self) -> Arc { + let pool = golem_service_base::db::postgres::PostgresPool::configured(&self.config) + .await + .expect("Cannot create postgres pool"); + + Arc::new(DbRoutingTablePersistence::new(pool, NUMBER_OF_SHARDS)) + } + + async fn unrelated_write(&self) { + use sqlx::Connection; + + let mut conn = sqlx::PgConnection::connect_with(&self.config.connect_options()) + .await + .expect("Cannot connect to postgres"); + sqlx::query("CREATE TABLE IF NOT EXISTS unrelated_writes (id INTEGER)") + .execute(&mut conn) + .await + .expect("Cannot create the unrelated table"); + sqlx::query("INSERT INTO unrelated_writes (id) VALUES (1)") + .execute(&mut conn) + .await + .expect("Cannot perform the unrelated write"); + } + + fn has_mirror(&self) -> bool { + true + } + + async fn mirror_snapshot(&self) -> Option { + use sqlx::Connection; + + let mut conn = sqlx::PgConnection::connect_with(&self.config.connect_options()) + .await + .expect("Cannot connect to the database"); + // One transaction, so both tables are read from the same committed state. + let mut tx = conn.begin().await.expect("Cannot begin a transaction"); + let leases: Vec = sqlx::query_as( + "SELECT executor_id, ip, port, granted_at, expires_at, pod_name FROM executor_leases", + ) + .fetch_all(&mut *tx) + .await + .expect("Cannot read executor_leases"); + let assignments = + sqlx::query_as("SELECT shard_id, executor_id, epoch FROM shard_assignments") + .fetch_all(&mut *tx) + .await + .expect("Cannot read shard_assignments"); + tx.commit() + .await + .expect("Cannot commit the read transaction"); + Some(MirrorSnapshot::from_rows(leases, assignments)) + } +} + #[async_trait] impl GetRoutingTablePersistence for PostgresRoutingTablePersistence { - async fn get_persistence(&self) -> Arc { + async fn new_store(&self) -> Arc { let db_name = format!("shard_{}", Uuid::new_v4().simple()); let admin_pool = sqlx::postgres::PgPoolOptions::new() @@ -63,7 +244,7 @@ impl GetRoutingTablePersistence for PostgresRoutingTablePersistence { .await .expect("Cannot create postgres test database"); - let postgres_config = DbPostgresConfig { + let config = DbPostgresConfig { host: "localhost".to_string(), database: db_name, username: "postgres".to_string(), @@ -78,21 +259,16 @@ impl GetRoutingTablePersistence for PostgresRoutingTablePersistence { let migrations = IncludedMigrationsDir::new(&golem_shard_manager::DB_MIGRATIONS); - golem_service_base::db::postgres::migrate( - &postgres_config, - migrations.postgres_migrations(), - ) - .await - .expect("Cannot apply postgres migrations"); - - let pool = golem_service_base::db::postgres::PostgresPool::configured(&postgres_config) + golem_service_base::db::postgres::migrate(&config, migrations.postgres_migrations()) .await - .expect("Cannot create postgres pool"); + .expect("Cannot apply postgres migrations"); - Arc::new(DbRoutingTablePersistence::new(pool, 16)) + Arc::new(PostgresStore { config }) } } +// -- sqlite: isolated by a fresh file per store ------------------------------------------------ + struct SqliteRoutingTablePersistence { temp_dir: TempDir, } @@ -103,9 +279,70 @@ impl std::fmt::Debug for SqliteRoutingTablePersistence { } } +#[derive(Debug)] +struct SqliteStore { + config: DbSqliteConfig, +} + +#[async_trait] +impl PersistenceStore for SqliteStore { + async fn connect(&self) -> Arc { + let pool = golem_service_base::db::sqlite::SqlitePool::configured(&self.config) + .await + .expect("Cannot create sqlite pool"); + + Arc::new(DbRoutingTablePersistence::new(pool, NUMBER_OF_SHARDS)) + } + + async fn unrelated_write(&self) { + use sqlx::Connection; + + let mut conn = sqlx::SqliteConnection::connect_with(&self.config.connect_options()) + .await + .expect("Cannot connect to sqlite"); + sqlx::query("CREATE TABLE IF NOT EXISTS unrelated_writes (id INTEGER)") + .execute(&mut conn) + .await + .expect("Cannot create the unrelated table"); + sqlx::query("INSERT INTO unrelated_writes (id) VALUES (1)") + .execute(&mut conn) + .await + .expect("Cannot perform the unrelated write"); + } + + fn has_mirror(&self) -> bool { + true + } + + async fn mirror_snapshot(&self) -> Option { + use sqlx::Connection; + + let mut conn = sqlx::SqliteConnection::connect_with(&self.config.connect_options()) + .await + .expect("Cannot connect to the database"); + // One transaction, so both tables are read from the same committed state. + let mut tx = conn.begin().await.expect("Cannot begin a transaction"); + let leases: Vec = sqlx::query_as( + "SELECT executor_id, ip, port, granted_at, expires_at, pod_name FROM executor_leases", + ) + .fetch_all(&mut *tx) + .await + .expect("Cannot read executor_leases"); + let assignments = + sqlx::query_as("SELECT shard_id, executor_id, epoch FROM shard_assignments") + .fetch_all(&mut *tx) + .await + .expect("Cannot read shard_assignments"); + tx.commit() + .await + .expect("Cannot commit the read transaction"); + Some(MirrorSnapshot::from_rows(leases, assignments)) + } +} + #[async_trait] impl GetRoutingTablePersistence for SqliteRoutingTablePersistence { - async fn get_persistence(&self) -> Arc { + async fn new_store(&self) -> Arc { let database_file = self .temp_dir .path() @@ -114,7 +351,7 @@ impl GetRoutingTablePersistence for SqliteRoutingTablePersistence { .expect("tempfile path was not valid unicode") .to_string(); - let sqlite_config = DbSqliteConfig { + let config = DbSqliteConfig { database: database_file, max_connections: 10, foreign_keys: true, @@ -122,15 +359,89 @@ impl GetRoutingTablePersistence for SqliteRoutingTablePersistence { let migrations = IncludedMigrationsDir::new(&golem_shard_manager::DB_MIGRATIONS); - golem_service_base::db::sqlite::migrate(&sqlite_config, migrations.sqlite_migrations()) + golem_service_base::db::sqlite::migrate(&config, migrations.sqlite_migrations()) .await .expect("Cannot apply sqlite migrations"); - let pool = golem_service_base::db::sqlite::SqlitePool::configured(&sqlite_config) + Arc::new(SqliteStore { config }) + } +} + +// -- etcd: isolated by a fresh key prefix per store -------------------------------------------- + +struct EtcdRoutingTablePersistenceFactory { + etcd: DockerEtcd, +} + +impl std::fmt::Debug for EtcdRoutingTablePersistenceFactory { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("EtcdRoutingTablePersistenceFactory") + } +} + +#[derive(Debug)] +struct EtcdStore { + config: EtcdConfig, +} + +impl EtcdStore { + async fn kv(&self) -> etcd_client::KvClient { + etcd_client::Client::connect(&self.config.endpoints, None) .await - .expect("Cannot create sqlite pool"); + .expect("Cannot connect to etcd") + .kv_client() + } +} - Arc::new(DbRoutingTablePersistence::new(pool, 16)) +#[async_trait] +impl PersistenceStore for EtcdStore { + async fn connect(&self) -> Arc { + Arc::new( + EtcdRoutingTablePersistence::new(&self.config, NUMBER_OF_SHARDS) + .await + .expect("Cannot connect to etcd"), + ) + } + + async fn unrelated_write(&self) { + // Any write advances etcd's cluster-wide revision; the state key's mod_revision must not + // follow it. + self.kv() + .await + .put("/golem/test/unrelated", "x", None) + .await + .expect("Cannot perform the unrelated write"); + } + + fn has_mirror(&self) -> bool { + false + } + + async fn mirror_snapshot(&self) -> Option { + // Mirror tables are a local-mode feature; distributed mode holds only the blob. + None + } +} + +#[async_trait] +impl GetRoutingTablePersistence for EtcdRoutingTablePersistenceFactory { + async fn new_store(&self) -> Arc { + // The state key is fixed, so stores on one etcd server cannot be isolated from each other + // the way postgres stores are by database. Instead the server is per test worker - tests + // on a worker run one at a time - and every new store starts by wiping the key. + let store = EtcdStore { + config: EtcdConfig { + endpoints: vec![self.etcd.client_url()], + ..EtcdConfig::default() + }, + }; + store + .kv() + .await + .delete(STATE_KEY, None) + .await + .expect("Cannot wipe the etcd state key"); + Arc::new(store) } } @@ -147,7 +458,14 @@ async fn postgres_persistence() -> Arc { Arc::new(PostgresRoutingTablePersistence { postgres }) } -define_matrix_dimension!(persistence: Arc -> "sqlite", "postgres"); +#[test_dep(scope = PerWorker, tagged_as = "etcd")] +async fn etcd_persistence() -> Arc { + Arc::new(EtcdRoutingTablePersistenceFactory { + etcd: DockerEtcd::new().await, + }) +} + +define_matrix_dimension!(persistence: Arc -> "sqlite", "postgres", "etcd"); #[test] #[tracing::instrument] @@ -155,12 +473,13 @@ async fn read_returns_default_when_empty( #[dimension(persistence)] persistence: &Arc, ) { let persistence = persistence.get_persistence().await; - let shard_state = persistence + let (shard_state, revision) = persistence .read() .await .expect("Reading default shard lease state should succeed"); - assert_eq!(shard_state.number_of_shards, 16); + assert_eq!(revision, NO_REVISION); + assert_eq!(shard_state.number_of_shards, NUMBER_OF_SHARDS); assert_eq!(shard_state.revision, ShardLeaseRevision::INITIAL); assert!(shard_state.shard_assignments.is_empty()); assert!(shard_state.executor_leases.is_empty()); @@ -173,45 +492,454 @@ async fn write_then_read_roundtrip( #[dimension(persistence)] persistence: &Arc, ) { let persistence = persistence.get_persistence().await; - let expected = sample_shard_state(16); + let expected = sample_shard_state(NUMBER_OF_SHARDS); - persistence - .write(&expected) + let written_revision = persistence + .write(&expected, NO_REVISION) .await .expect("Writing routing table should succeed"); + assert!(written_revision > NO_REVISION); - let actual = persistence + let (actual, read_revision) = persistence .read() .await .expect("Reading persisted routing table should succeed"); assert_eq!(actual, expected); + assert_eq!(read_revision, written_revision); } #[test] #[tracing::instrument] -async fn last_write_wins( +async fn sequential_writes_advance_the_revision( #[dimension(persistence)] persistence: &Arc, ) { let persistence = persistence.get_persistence().await; - let first = sample_shard_state(16); - let second = replacement_shard_state(16); + let first = sample_shard_state(NUMBER_OF_SHARDS); + let second = replacement_shard_state(NUMBER_OF_SHARDS); - persistence - .write(&first) + let first_revision = persistence + .write(&first, NO_REVISION) .await .expect("Writing first routing table should succeed"); - persistence - .write(&second) + let second_revision = persistence + .write(&second, first_revision) .await .expect("Writing second routing table should succeed"); - let actual = persistence + // Strictly greater, never `first_revision + 1`: etcd's revision is a cluster-global counter. + assert!(second_revision > first_revision); + + let (actual, revision) = persistence .read() .await .expect("Reading persisted routing table should succeed"); assert_eq!(actual, second); + assert_eq!(revision, second_revision); +} + +#[test] +#[tracing::instrument] +async fn writing_the_same_state_twice_still_advances_the_revision( + #[dimension(persistence)] persistence: &Arc, +) { + // Guards against ever deriving the storage revision from the in-blob ShardLeaseRevision: a + // write that does not change the state must still move the compare-and-swap token, otherwise + // two concurrent such writes would both succeed. + let persistence = persistence.get_persistence().await; + let shard_state = sample_shard_state(NUMBER_OF_SHARDS); + + let first = persistence + .write(&shard_state, NO_REVISION) + .await + .expect("Writing routing table should succeed"); + let second = persistence + .write(&shard_state, first) + .await + .expect("Rewriting the same routing table should succeed"); + + assert!(second > first); +} + +#[test] +#[tracing::instrument] +async fn mirror_tables_follow_the_persisted_state( + #[dimension(persistence)] persistence: &Arc, +) { + let store = persistence.new_store().await; + let ours = store.connect().await; + + let first = sample_shard_state(NUMBER_OF_SHARDS); + let first_revision = ours + .write(&first, NO_REVISION) + .await + .expect("Writing the first routing table should succeed"); + + let snapshot = store.mirror_snapshot().await; + assert_eq!( + snapshot.is_some(), + store.has_mirror(), + "a backend must report a mirror exactly when it has one" + ); + let Some(snapshot) = snapshot else { + // Distributed mode has no mirror tables; nothing to check on this backend. + return; + }; + assert_eq!(snapshot, MirrorSnapshot::of(&first)); + // Orphaned shards are pending, not assigned, and must not show up as assignments. + for pending in &first.pending_rebalance { + assert!( + !snapshot + .assignments + .iter() + .any(|(shard_id, _, _)| i64::from(*shard_id) == pending.value()), + "pending shard {pending} must not appear in shard_assignments" + ); + } + + // A rewrite replaces the mirror wholesale: rows of executors that are gone must not linger. + let second = replacement_shard_state(NUMBER_OF_SHARDS); + let second_revision = ours + .write(&second, first_revision) + .await + .expect("Writing the second routing table should succeed"); + assert_eq!( + store.mirror_snapshot().await.unwrap(), + MirrorSnapshot::of(&second) + ); + + // The mirror rides the blob's compare-and-swap: a rejected write leaves it untouched. + let result = ours.write(&first, first_revision).await; + assert!( + matches!(result, Err(ShardManagerError::ConcurrentModification)), + "expected ConcurrentModification, got {result:?}" + ); + assert_eq!( + store.mirror_snapshot().await.unwrap(), + MirrorSnapshot::of(&second) + ); + let (_, revision) = ours + .read() + .await + .expect("Reading the routing table should succeed"); + assert_eq!(revision, second_revision); + + // ...and empties out with the last lease. + let mut emptied = second.clone(); + let _ = emptied.housekeep(granted_at() + chrono::Duration::from_std(LEASE_TTL).unwrap()); + emptied + .bump_revision() + .expect("revision bump should succeed"); + assert!(emptied.executor_leases.is_empty()); + assert!(emptied.shard_assignments.is_empty()); + ours.write(&emptied, second_revision) + .await + .expect("Writing the emptied routing table should succeed"); + assert_eq!( + store.mirror_snapshot().await.unwrap(), + MirrorSnapshot::empty() + ); +} + +#[test] +#[tracing::instrument] +async fn a_full_size_routing_table_roundtrips_and_is_mirrored( + #[dimension(persistence)] persistence: &Arc, +) { + // The production default is 1024 shards - more than one multi-row INSERT chunk of the mirror + // tables, so this is the path every real local-mode deployment takes on every write. + const SHARDS: usize = 1024; + const EXECUTORS: usize = 4; + let shard_ids: Vec> = (0..EXECUTORS) + .map(|executor| { + (0..SHARDS as i64) + .filter(|shard| *shard as usize % EXECUTORS == executor) + .collect() + }) + .collect(); + let executors: Vec<(ExecutorId, ExecutorAddr, Option<&str>, &[i64])> = (0..EXECUTORS) + .map(|executor| { + ( + self::executor(executor as u128 + 1), + addr(executor as u8 + 1, 9010 + executor as u16), + None, + shard_ids[executor].as_slice(), + ) + }) + .collect(); + let shard_state = shard_state_with_executors(SHARDS, &executors); + assert_eq!(shard_state.shard_assignments.len(), SHARDS); + + let store = persistence.new_store().await; + let ours = store.connect().await; + + let revision = ours + .write(&shard_state, NO_REVISION) + .await + .expect("Writing a full-size routing table should succeed"); + let (actual, actual_revision) = ours + .read() + .await + .expect("Reading a full-size routing table should succeed"); + assert_eq!(actual, shard_state); + assert_eq!(actual_revision, revision); + + if let Some(snapshot) = store.mirror_snapshot().await { + assert_eq!(snapshot.assignments.len(), SHARDS); + assert_eq!(snapshot, MirrorSnapshot::of(&shard_state)); + } +} + +#[test] +#[tracing::instrument] +async fn a_state_that_violates_invariants_is_refused_before_it_is_stored( + #[dimension(persistence)] persistence: &Arc, +) { + // Every backend refuses identically and before any I/O - not by whichever constraint a + // backend happens to enforce (the SQL mirror's foreign key, which SQLite only checks with + // foreign_keys = true), and not by poisoning the store for every later read. + let store = persistence.new_store().await; + let ours = store.connect().await; + + let good = sample_shard_state(NUMBER_OF_SHARDS); + let revision = ours + .write(&good, NO_REVISION) + .await + .expect("Writing the routing table should succeed"); + + let mut bad = good.clone(); + // shard 5 is unassigned in the fixture; executor 42 holds no lease + bad.shard_assignments.insert( + ShardId::new(5), + ShardAssignmentEntry { + executor_id: executor(42), + epoch: ShardEpoch::initial(), + }, + ); + let result = ours.write(&bad, revision).await; + assert!( + matches!(result, Err(ShardManagerError::Internal(_))), + "expected Internal, got {result:?}" + ); + + let (actual, actual_revision) = ours + .read() + .await + .expect("Reading the routing table should succeed"); + assert_eq!(actual, good); + assert_eq!(actual_revision, revision); + if let Some(snapshot) = store.mirror_snapshot().await { + assert_eq!(snapshot, MirrorSnapshot::of(&good)); + } +} + +#[test] +#[tracing::instrument] +async fn housekeep_expiry_reclamation_roundtrips( + #[dimension(persistence)] persistence: &Arc, +) { + // The state shape `housekeep` produces - leases gone, their shards moved to + // `pending_rebalance`, `shard_epochs` retained as the high-water mark - has to survive a + // round-trip on every backend. + let persistence = persistence.get_persistence().await; + let mut shard_state = sample_shard_state(NUMBER_OF_SHARDS); + let before = shard_state.executor_count(); + + let expired = + shard_state.housekeep(granted_at() + chrono::Duration::from_std(LEASE_TTL).unwrap()); + assert_eq!( + expired.len(), + before, + "every lease in the fixture should have expired at granted_at + ttl" + ); + assert!(shard_state.executor_leases.is_empty()); + assert!(!shard_state.pending_rebalance.is_empty()); + assert!(shard_state.shard_assignments.is_empty()); + shard_state + .bump_revision() + .expect("revision bump should succeed"); + + let revision = persistence + .write(&shard_state, NO_REVISION) + .await + .expect("Writing the reclaimed state should succeed"); + + let (actual, actual_revision) = persistence + .read() + .await + .expect("Reading the reclaimed state should succeed"); + + assert_eq!(actual, shard_state); + assert_eq!(actual_revision, revision); + // The epoch high-water mark must outlive the leases, or a reassigned shard could reuse an + // epoch a fenced zombie still holds. + assert!(!actual.shard_epochs.is_empty()); +} + +#[test] +#[tracing::instrument] +async fn two_clients_over_the_same_store_see_each_others_writes( + #[dimension(persistence)] persistence: &Arc, +) { + // A sanity check on the fixture itself: if `connect` handed out isolated stores, the + // compare-and-swap tests below would pass vacuously. + let store = persistence.new_store().await; + let first = store.connect().await; + let second = store.connect().await; + + let expected = sample_shard_state(NUMBER_OF_SHARDS); + let revision = first + .write(&expected, NO_REVISION) + .await + .expect("Writing routing table should succeed"); + + let (actual, actual_revision) = second + .read() + .await + .expect("Reading persisted routing table should succeed"); + + assert_eq!(actual, expected); + assert_eq!(actual_revision, revision); +} + +#[test] +#[tracing::instrument] +async fn write_with_no_revision_when_already_present_is_rejected( + #[dimension(persistence)] persistence: &Arc, +) { + // A cold-started rival that believes the store is empty must not overwrite live state. + let store = persistence.new_store().await; + let first = store.connect().await; + let second = store.connect().await; + + first + .write(&sample_shard_state(NUMBER_OF_SHARDS), NO_REVISION) + .await + .expect("Writing the initial routing table should succeed"); + + let result = second + .write(&replacement_shard_state(NUMBER_OF_SHARDS), NO_REVISION) + .await; + + assert!( + matches!(result, Err(ShardManagerError::ConcurrentModification)), + "expected ConcurrentModification, got {result:?}" + ); +} + +#[test] +#[tracing::instrument] +async fn write_with_a_superseded_revision_is_rejected( + #[dimension(persistence)] persistence: &Arc, +) { + let store = persistence.new_store().await; + let first = store.connect().await; + let second = store.connect().await; + + let initial = sample_shard_state(NUMBER_OF_SHARDS); + let winner = replacement_shard_state(NUMBER_OF_SHARDS); + + let stale_revision = first + .write(&initial, NO_REVISION) + .await + .expect("Writing the initial routing table should succeed"); + let winning_revision = first + .write(&winner, stale_revision) + .await + .expect("Writing the second routing table should succeed"); + + let result = second.write(&initial, stale_revision).await; + assert!( + matches!(result, Err(ShardManagerError::ConcurrentModification)), + "expected ConcurrentModification, got {result:?}" + ); + + // The loser's payload must never become visible. + let (actual, revision) = second + .read() + .await + .expect("Reading persisted routing table should succeed"); + assert_eq!(actual, winner); + assert_eq!(revision, winning_revision); +} + +#[test] +#[tracing::instrument] +async fn write_with_a_revision_when_nothing_is_stored_is_rejected( + #[dimension(persistence)] persistence: &Arc, +) { + // A revision against absent state must fail rather than resurrect it. This is the case a + // guarded upsert gets wrong on SQL, because its INSERT branch is unguarded. + let persistence = persistence.get_persistence().await; + + let result = persistence + .write(&sample_shard_state(NUMBER_OF_SHARDS), 12345) + .await; + + assert!( + matches!(result, Err(ShardManagerError::ConcurrentModification)), + "expected ConcurrentModification, got {result:?}" + ); + + let (_, revision) = persistence + .read() + .await + .expect("Reading shard lease state should succeed"); + assert_eq!(revision, NO_REVISION); +} + +#[test] +#[tracing::instrument] +async fn conflict_does_not_advance_the_stored_revision( + #[dimension(persistence)] persistence: &Arc, +) { + let persistence = persistence.get_persistence().await; + let stored = sample_shard_state(NUMBER_OF_SHARDS); + + let revision = persistence + .write(&stored, NO_REVISION) + .await + .expect("Writing routing table should succeed"); + + let result = persistence + .write(&replacement_shard_state(NUMBER_OF_SHARDS), NO_REVISION) + .await; + assert!( + matches!(result, Err(ShardManagerError::ConcurrentModification)), + "expected ConcurrentModification, got {result:?}" + ); + + let (actual, actual_revision) = persistence + .read() + .await + .expect("Reading persisted routing table should succeed"); + assert_eq!(actual, stored); + assert_eq!(actual_revision, revision); +} + +#[test] +#[tracing::instrument] +async fn an_unrelated_write_on_the_backend_does_not_move_our_revision( + #[dimension(persistence)] persistence: &Arc, +) { + // The revision must belong to the key/row, not to the backend as a whole. This is the one + // test that catches returning etcd's cluster-global header revision from `read`. + let store = persistence.new_store().await; + let ours = store.connect().await; + + let revision = ours + .write(&sample_shard_state(NUMBER_OF_SHARDS), NO_REVISION) + .await + .expect("Writing our routing table should succeed"); + + store.unrelated_write().await; + + let (_, our_revision): (ShardLeaseState, ExternalRevision) = ours + .read() + .await + .expect("Reading our routing table should succeed"); + assert_eq!(our_revision, revision); } fn granted_at() -> DateTime { diff --git a/golem-shard-manager/tests/shard_management.rs b/golem-shard-manager/tests/shard_management.rs index 9d7cdd91de..5c69d27fa5 100644 --- a/golem-shard-manager/tests/shard_management.rs +++ b/golem-shard-manager/tests/shard_management.rs @@ -16,10 +16,11 @@ use async_trait::async_trait; use chrono::{DateTime, Utc}; use golem_common::model::{Pod, ShardId}; use golem_shard_manager::{ - ExecutorAddr, ExecutorId, HealthCheck, HealthCheckError, RoutingTablePersistence, ShardEpoch, - ShardLeaseState, ShardManagement, ShardManagerError, WorkerExecutorService, + ExecutorAddr, ExecutorId, ExternalRevision, HealthCheck, HealthCheckError, NO_REVISION, + RoutingTablePersistence, ShardEpoch, ShardLeaseState, ShardManagement, ShardManagerError, + WorkerExecutorService, }; -use std::collections::{BTreeSet, HashMap}; +use std::collections::{BTreeSet, HashMap, VecDeque}; use std::net::{IpAddr, Ipv4Addr}; use std::sync::Arc; use std::time::Duration; @@ -31,35 +32,92 @@ use uuid::Uuid; const LEASE_TTL: Duration = Duration::from_secs(60); +#[derive(Debug)] +struct TestStore { + shard_state: ShardLeaseState, + revision: ExternalRevision, +} + +/// An in-memory [`RoutingTablePersistence`] with the same compare-and-swap semantics as the real +/// backends, plus a hook for failing writes. #[derive(Clone, Debug)] struct TestPersistence { - shard_state: Arc>, - writes: Arc>>, + store: Arc>, + /// `(prev_revision, written_state)` for every accepted write. + writes: Arc>>, + /// Every `write` call, accepted or not. The gap between this and `writes` is how a test sees + /// that a write was rejected. + attempts: Arc>, + /// Per-write script: a `None` entry lets a write through, a `Some(err)` entry fails it before + /// it reaches the store. Writes past the end of the script always succeed. + injected: Arc>>>, } impl TestPersistence { + /// Seeds a store that already holds state. Revision 1 rather than `NO_REVISION`, because the + /// real backends can never hold state at revision 0 - that value means "nothing stored". fn new(initial: ShardLeaseState) -> Self { + Self::at_revision(initial, 1) + } + + fn at_revision(initial: ShardLeaseState, revision: ExternalRevision) -> Self { Self { - shard_state: Arc::new(Mutex::new(initial)), + store: Arc::new(Mutex::new(TestStore { + shard_state: initial, + revision, + })), writes: Arc::new(Mutex::new(Vec::new())), + attempts: Arc::new(Mutex::new(0)), + injected: Arc::new(Mutex::new(VecDeque::new())), } } async fn latest(&self) -> ShardLeaseState { - self.shard_state.lock().await.clone() + self.store.lock().await.shard_state.clone() + } + + /// Scripts the outcome of the next writes. `vec![None, Some(err)]` fails only the second one. + async fn fail_writes(&self, script: Vec>) { + *self.injected.lock().await = script.into(); + } + + async fn write_count(&self) -> usize { + self.writes.lock().await.len() + } + + async fn attempt_count(&self) -> usize { + *self.attempts.lock().await } } #[async_trait] impl RoutingTablePersistence for TestPersistence { - 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 write( + &self, + shard_state: &ShardLeaseState, + prev_revision: ExternalRevision, + ) -> Result { + *self.attempts.lock().await += 1; + if let Some(err) = self.injected.lock().await.pop_front().flatten() { + return Err(err); + } + + let mut store = self.store.lock().await; + if prev_revision != store.revision { + return Err(ShardManagerError::ConcurrentModification); + } + store.revision += 1; + store.shard_state = shard_state.clone(); + self.writes + .lock() + .await + .push((prev_revision, shard_state.clone())); + Ok(store.revision) } - async fn read(&self) -> Result { - Ok(self.shard_state.lock().await.clone()) + async fn read(&self) -> Result<(ShardLeaseState, ExternalRevision), ShardManagerError> { + let store = self.store.lock().await; + Ok((store.shard_state.clone(), store.revision)) } } @@ -282,10 +340,28 @@ async fn new_shard_management( .expect("failed to create shard management"); tokio::time::sleep(Duration::from_millis(50)).await; + // The startup pass persists exactly twice (executor changes, then the applied rebalance). + // Wait for both before handing the fixture over, so a test that scripts write failures cannot + // have one of them swallowed by a startup write that had not landed yet. + wait_for_writes(&persistence, 2).await; (shard_management, persistence, join_set) } +/// Waits until the loop has performed at least `count` accepted writes. +async fn wait_for_writes(persistence: &TestPersistence, count: usize) { + let start = Instant::now(); + while persistence.write_count().await < count { + if start.elapsed() > Duration::from_secs(5) { + panic!( + "timed out waiting for {count} writes, saw {}", + persistence.write_count().await + ); + } + tokio::time::sleep(Duration::from_millis(10)).await; + } +} + #[test] // On shard-manager restart, live executors are reset to the routing table. async fn shard_manager_restart_clears_stale_executor_shards() { @@ -632,7 +708,7 @@ async fn same_address_reregistration_transfers_shards_and_reconciles() { 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() { + for (_, written) in persistence.writes.lock().await.iter() { assert!( written.get_unassigned_shards().is_empty(), "shards became unassigned during re-registration: {written}" @@ -641,3 +717,139 @@ async fn same_address_reregistration_transfers_shards_and_reconciles() { join_set.abort_all(); } + +#[test] +// Snapshot-and-rollback: a failed persist restores the in-memory state, and the loop +// does not carry on against a store it can no longer trust - the task ends with the error, so +// the process restarts and re-reads. A revision conflict in particular means another shard +// manager is writing the state; the cached revision is deliberately not refreshed, because it is +// the fencing token and the loser of it has to stop. +async fn a_persistence_failure_rolls_back_the_state_and_stops_the_loop() { + let existing_pod = pod(1, 9000); + let new_pod = pod(2, 9001); + let worker_executors = Arc::new(TestWorkerExecutors::default()); + + let (shard_management, persistence, mut join_set) = new_shard_management( + shard_state_with_executors( + 4, + vec![( + executor(1), + existing_pod, + "worker-executor-0", + &[0, 1, 2, 3], + )], + ), + worker_executors.clone(), + ) + .await; + let before = shard_management.current_snapshot().await; + let persisted_before = persistence.latest().await; + + persistence + .fail_writes(vec![Some(ShardManagerError::ConcurrentModification)]) + .await; + let new_executor = shard_management + .register_executor( + ExecutorAddr::from(new_pod), + Some("worker-executor-1".into()), + ) + .await; + + let outcome = tokio::time::timeout(Duration::from_secs(5), join_set.join_next()) + .await + .expect("the shard management loop should have stopped") + .expect("the loop task should exist") + .expect("the loop task should not panic"); + assert!( + outcome.is_err(), + "the loop must end with the persistence error, got {outcome:?}" + ); + + // Rolled back: the registration reached neither memory nor the store, and nothing was pushed + // to the executor that never made it in. + let after = shard_management.current_snapshot().await; + assert!(!after.has_executor(new_executor)); + assert_eq!(after, before); + assert_eq!(persistence.latest().await, persisted_before); + assert!(worker_executors.local_assignment(new_pod).await.is_empty()); +} + +#[test] +// First boot through the loop: nothing is stored, so the very first write must be guarded on +// NO_REVISION. Getting this wrong is rejected by both real backends, not silently tolerated. +async fn the_first_write_on_an_empty_store_uses_no_revision() { + let worker_executors = Arc::new(TestWorkerExecutors::default()); + let persistence = TestPersistence::at_revision(ShardLeaseState::new(4), NO_REVISION); + let health_check = Arc::new(TestHealthCheck::all_healthy()); + let mut join_set = JoinSet::new(); + + let _shard_management = ShardManagement::new( + Arc::new(persistence.clone()), + worker_executors, + health_check, + 0.0, + LEASE_TTL, + &mut join_set, + ) + .await + .expect("failed to create shard management"); + + wait_for_writes(&persistence, 1).await; + + let writes = persistence.writes.lock().await; + let (first_prev_revision, _) = writes.first().expect("the loop should have written once"); + assert_eq!(*first_prev_revision, NO_REVISION); + drop(writes); + + join_set.abort_all(); +} + +#[test] +// The loop must write with the revision it read at startup, not with a hardcoded one. +async fn the_first_write_uses_the_revision_read_at_startup() { + let existing_pod = pod(1, 9000); + let worker_executors = Arc::new(TestWorkerExecutors::default()); + let shard_state = shard_state_with_executors( + 1, + vec![(executor(1), existing_pod, "worker-executor-0", &[0])], + ); + + let persistence = TestPersistence::at_revision(shard_state, 7); + let health_check = Arc::new(TestHealthCheck::all_healthy()); + let mut join_set = JoinSet::new(); + + let _shard_management = ShardManagement::new( + Arc::new(persistence.clone()), + worker_executors, + health_check, + 0.0, + LEASE_TTL, + &mut join_set, + ) + .await + .expect("failed to create shard management"); + + wait_for_writes(&persistence, 2).await; + + let prev_revisions: Vec = persistence + .writes + .lock() + .await + .iter() + .map(|(prev, _)| *prev) + .collect(); + + // The first write is guarded on the revision `read()` returned at startup, the second on what + // the first write returned. + assert_eq!(prev_revisions, vec![7, 8]); + + // ...and both were accepted first time. This pins the cached revision being advanced on + // success: without it the second write of the startup pass would conflict and stop the loop. + assert_eq!( + persistence.attempt_count().await, + 2, + "a conflict-free startup pass must not need a retry" + ); + + join_set.abort_all(); +} diff --git a/golem-test-framework/src/components/etcd/docker_etcd.rs b/golem-test-framework/src/components/etcd/docker_etcd.rs new file mode 100644 index 0000000000..f68294037d --- /dev/null +++ b/golem-test-framework/src/components/etcd/docker_etcd.rs @@ -0,0 +1,138 @@ +// Copyright 2024-2026 Golem Cloud +// +// Licensed under the Golem Source License v1.1 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://license.golem.cloud/LICENSE +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::components::docker::ContainerHandle; +use std::fmt::{Debug, Formatter}; +use std::time::{Duration, Instant}; +use testcontainers::core::{IntoContainerPort, WaitFor}; +use testcontainers::runners::AsyncRunner; +use testcontainers::{GenericImage, ImageExt}; +use tracing::{error, info}; + +/// A single-node etcd server started in Docker. +/// +/// The client port is exposed on a random host port; use [`client_url`](DockerEtcd::client_url) +/// to get the `http://127.0.0.1:` URL for the service under test. TLS is not configured. +pub struct DockerEtcd { + _container: ContainerHandle, + public_port: u16, +} + +impl DockerEtcd { + const CLIENT_PORT: u16 = 2379; + const PEER_PORT: u16 = 2380; + const DEFAULT_IMAGE_NAME: &'static str = "gcr.io/etcd-development/etcd"; + const DEFAULT_IMAGE_TAG: &'static str = "v3.5.17"; + + pub async fn new() -> Self { + Self::new_with_image(Self::DEFAULT_IMAGE_NAME, Self::DEFAULT_IMAGE_TAG).await + } + + pub async fn new_with_image(image: &str, tag: &str) -> Self { + info!("Starting etcd container ({image}:{tag})"); + + let client_port = Self::CLIENT_PORT; + let client_urls = format!("http://0.0.0.0:{client_port}"); + let peer_urls = format!("http://0.0.0.0:{}", Self::PEER_PORT); + + let container = tryhard::retry_fn(move || { + // The official image's default entrypoint binds the client port to the container's + // own localhost, so the mapped host port would refuse connections. Bind 0.0.0.0. + let cmd = vec![ + "/usr/local/bin/etcd".to_string(), + "--name".to_string(), + "golem-test-etcd".to_string(), + "--data-dir".to_string(), + "/etcd-data".to_string(), + "--listen-client-urls".to_string(), + client_urls.clone(), + "--advertise-client-urls".to_string(), + client_urls.clone(), + "--listen-peer-urls".to_string(), + peer_urls.clone(), + "--initial-advertise-peer-urls".to_string(), + peer_urls.clone(), + "--initial-cluster".to_string(), + format!("golem-test-etcd={peer_urls}"), + "--initial-cluster-token".to_string(), + "golem-test-etcd".to_string(), + "--initial-cluster-state".to_string(), + "new".to_string(), + "--log-level".to_string(), + "info".to_string(), + ]; + + GenericImage::new(image, tag) + .with_exposed_port(client_port.tcp()) + // etcd's zap logger writes to stderr. + .with_wait_for(WaitFor::message_on_stderr("ready to serve client requests")) + .with_cmd(cmd) + .start() + }) + .retries(5) + .exponential_backoff(Duration::from_millis(10)) + .max_delay(Duration::from_secs(10)) + .await + .expect("Failed to start etcd container"); + + let public_port = container + .get_host_port_ipv4(client_port) + .await + .expect("Failed to get etcd host port"); + + // The log line alone is not enough: the port forward can be established slightly after it. + etcd_wait_for_startup("127.0.0.1", public_port, Duration::from_secs(60)).await; + + info!("etcd container started on port {public_port}"); + + Self { + _container: ContainerHandle::new(container), + public_port, + } + } + + /// Returns the client URL, e.g. `http://127.0.0.1:2379`. + pub fn client_url(&self) -> String { + format!("http://127.0.0.1:{}", self.public_port) + } +} + +impl Debug for DockerEtcd { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "DockerEtcd(port={})", self.public_port) + } +} + +async fn etcd_wait_for_startup(host: &str, port: u16, timeout: Duration) { + info!( + "Waiting for etcd client port on {host}:{port} (timeout {}s)", + timeout.as_secs() + ); + let start = Instant::now(); + loop { + match tokio::net::TcpStream::connect(format!("{host}:{port}")).await { + Ok(_) => { + info!("etcd client port {port} is accepting connections"); + return; + } + Err(e) => { + if start.elapsed() > timeout { + error!("etcd {host}:{port} did not become ready: {e}"); + panic!("etcd {host}:{port} did not become ready within the timeout"); + } + } + } + tokio::time::sleep(Duration::from_secs(2)).await; + } +} diff --git a/golem-test-framework/src/components/etcd/mod.rs b/golem-test-framework/src/components/etcd/mod.rs new file mode 100644 index 0000000000..38448f2259 --- /dev/null +++ b/golem-test-framework/src/components/etcd/mod.rs @@ -0,0 +1,15 @@ +// Copyright 2024-2026 Golem Cloud +// +// Licensed under the Golem Source License v1.1 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://license.golem.cloud/LICENSE +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +pub mod docker_etcd; diff --git a/golem-test-framework/src/components/mod.rs b/golem-test-framework/src/components/mod.rs index edc12b8a25..9e92146e5c 100644 --- a/golem-test-framework/src/components/mod.rs +++ b/golem-test-framework/src/components/mod.rs @@ -31,6 +31,7 @@ pub mod blob_storage; pub mod component_compilation_service; mod docker; mod dynamic_span; +pub mod etcd; pub mod jaeger; pub mod otel_collector; pub mod rdb; diff --git a/golem-test-framework/src/components/shard_manager/mod.rs b/golem-test-framework/src/components/shard_manager/mod.rs index 67b10051bf..565043fe46 100644 --- a/golem-test-framework/src/components/shard_manager/mod.rs +++ b/golem-test-framework/src/components/shard_manager/mod.rs @@ -108,8 +108,19 @@ async fn env_vars( registry_service.grpc_port().to_string(), ) .with_all( + // The shard manager's SQL settings live under `persistence`, not `db`: the same key + // also selects etcd in distributed mode. The sub-keys are identical to DbConfig's. rdb.info() - .env("golem_shard_manager", rdb_private_connection), + .env("golem_shard_manager", rdb_private_connection) + .into_iter() + .map(|(key, value)| { + let key = key + .strip_prefix("GOLEM__DB__") + .map(|rest| format!("GOLEM__PERSISTENCE__{rest}")) + .unwrap_or(key); + (key, value) + }) + .collect(), ) .with_optional_otlp("shard_manager", otlp); diff --git a/integration-tests/tests/sharding.rs b/integration-tests/tests/sharding.rs index f6415722b3..619c913e65 100644 --- a/integration-tests/tests/sharding.rs +++ b/integration-tests/tests/sharding.rs @@ -695,6 +695,14 @@ mod tests { let _ = sqlx::query("DELETE FROM golem_shard_manager.quota_resources") .execute(&pool) .await; + // The local-mode mirror of the shard state; assignments first, because of + // the foreign key. + let _ = sqlx::query("DELETE FROM golem_shard_manager.shard_assignments") + .execute(&pool) + .await; + let _ = sqlx::query("DELETE FROM golem_shard_manager.executor_leases") + .execute(&pool) + .await; pool.close().await; } DbInfo::Mysql(_) => { diff --git a/local-run/start.sh b/local-run/start.sh index 187b0ca187..2387d597a1 100644 --- a/local-run/start.sh +++ b/local-run/start.sh @@ -114,8 +114,8 @@ pushd "${GOLEM_DIR}/golem-shard-manager" || exit RUST_LOG=info,h2=warn,hyper=warn,tower=warn \ GOLEM__HTTP_PORT=${SHARD_MANAGER_HTTP_PORT} \ GOLEM__GRPC__PORT=${SHARD_MANAGER_GRPC_PORT} \ -GOLEM__DB__TYPE="Sqlite" \ -GOLEM__DB__CONFIG__DATABASE="../local-run/data/shard-manager/golem_shard_manager.sqlite" \ +GOLEM__PERSISTENCE__TYPE="Sqlite" \ +GOLEM__PERSISTENCE__CONFIG__DATABASE="../local-run/data/shard-manager/golem_shard_manager.sqlite" \ GOLEM__REGISTRY_SERVICE__HOST="localhost" \ GOLEM__REGISTRY_SERVICE__PORT=${REGISTRY_SERVICE_GRPC_PORT} \ ../target/debug/golem-shard-manager & From 49221257e151b2398e4277545c9ee48a8702c448 Mon Sep 17 00:00:00 2001 From: Aditya Salunkhe Date: Tue, 1 Sep 2026 01:51:28 +0530 Subject: [PATCH 2/6] Reject the removed db config key, bound persistence round-trips, and close two test gaps --- Makefile.toml | 2 +- .../postgres/004_shard_state_revision.sql | 31 +--- .../sqlite/004_shard_state_revision.sql | 31 +--- golem-shard-manager/src/config.rs | 63 ++++++- golem-shard-manager/src/server.rs | 6 +- .../src/sharding/persistence/db.rs | 11 +- .../src/sharding/persistence/etcd.rs | 26 +-- .../src/sharding/persistence/mod.rs | 48 ++---- .../src/sharding/shard_management.rs | 34 +++- golem-shard-manager/tests/persistence.rs | 22 ++- golem-shard-manager/tests/shard_management.rs | 157 +++++++++++++++++- 11 files changed, 297 insertions(+), 134 deletions(-) diff --git a/Makefile.toml b/Makefile.toml index 497b60c1b3..39e65a6891 100644 --- a/Makefile.toml +++ b/Makefile.toml @@ -681,7 +681,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 +cargo-test-r run --package golem-shard-manager --test integration -- --nocapture --test-threads=1 --report-time $JUNIT_OPTS ''' [tasks.integration-tests-group6] diff --git a/golem-shard-manager/db/migration/postgres/004_shard_state_revision.sql b/golem-shard-manager/db/migration/postgres/004_shard_state_revision.sql index c9223acc56..0c5d04c648 100644 --- a/golem-shard-manager/db/migration/postgres/004_shard_state_revision.sql +++ b/golem-shard-manager/db/migration/postgres/004_shard_state_revision.sql @@ -1,17 +1,6 @@ --- The persisted shard lease state is now written with a compare-and-swap guard, and mirrored into --- queryable tables. --- --- `revision` is a storage-level fencing token. It is incremented by exactly one on every --- successful write and is deliberately UNRELATED to the domain-level ShardLeaseRevision inside --- `state`: that one starts at 0 and is only bumped when the routing table meaningfully changes. --- Never derive this column from it. --- --- Revision 0 is reserved to mean "no state stored", so a stored row always carries >= 1. The --- column therefore has NO DEFAULT: a write that forgets to bind it must fail loudly. --- --- The table is recreated rather than altered: any row written before this migration has no --- meaningful revision, and the state rebuilds itself from executor registrations (same reasoning --- as 003). Recreating also avoids SQLite's requirement that ADD COLUMN supply a non-null default. +-- `revision` is a storage-level fencing token for compare-and-swap writes, unrelated to the +-- ShardLeaseRevision inside `state`; never derive one from the other. Revision 0 means "no state +-- stored", so a stored row always carries >= 1 and the column has no default. DROP TABLE shard_manager_state; CREATE TABLE shard_manager_state @@ -21,18 +10,8 @@ CREATE TABLE shard_manager_state revision BIGINT NOT NULL ); --- Local-mode mirror of the state blob, for inspection with plain SQL (`executor_leases` is the --- counterpart of the quota system's `quota_leases`). The blob is the source of truth: both tables --- are rewritten wholesale in the same transaction as every write the shard manager makes, and it --- never reads them back. Not mirrored: `pending_rebalance`, `shard_epochs` (the per-shard epoch --- high-water marks, which outlive leases) and the in-blob `ShardLeaseRevision`. --- --- Anyone clearing the state by hand (`DELETE FROM shard_manager_state`) must clear these two --- tables as well, or they keep describing leases that no longer exist until the next write. --- --- The shard manager refuses to persist a state whose assignments reference an executor without a --- lease, so the foreign key below is belt-and-braces. Postgres always enforces it; SQLite only --- with `foreign_keys = true` (the default is `false`). +-- A mirror of the state blob for inspection with plain SQL; rewritten wholesale on every write. +-- Clearing shard_manager_state by hand means clearing these two as well. CREATE TABLE executor_leases ( executor_id UUID NOT NULL, diff --git a/golem-shard-manager/db/migration/sqlite/004_shard_state_revision.sql b/golem-shard-manager/db/migration/sqlite/004_shard_state_revision.sql index 55f1468e68..0b7b7abb0a 100644 --- a/golem-shard-manager/db/migration/sqlite/004_shard_state_revision.sql +++ b/golem-shard-manager/db/migration/sqlite/004_shard_state_revision.sql @@ -1,17 +1,6 @@ --- The persisted shard lease state is now written with a compare-and-swap guard, and mirrored into --- queryable tables. --- --- `revision` is a storage-level fencing token. It is incremented by exactly one on every --- successful write and is deliberately UNRELATED to the domain-level ShardLeaseRevision inside --- `state`: that one starts at 0 and is only bumped when the routing table meaningfully changes. --- Never derive this column from it. --- --- Revision 0 is reserved to mean "no state stored", so a stored row always carries >= 1. The --- column therefore has NO DEFAULT: a write that forgets to bind it must fail loudly. --- --- The table is recreated rather than altered: any row written before this migration has no --- meaningful revision, and the state rebuilds itself from executor registrations (same reasoning --- as 003). Recreating also avoids SQLite's requirement that ADD COLUMN supply a non-null default. +-- `revision` is a storage-level fencing token for compare-and-swap writes, unrelated to the +-- ShardLeaseRevision inside `state`; never derive one from the other. Revision 0 means "no state +-- stored", so a stored row always carries >= 1 and the column has no default. DROP TABLE shard_manager_state; CREATE TABLE shard_manager_state @@ -21,18 +10,8 @@ CREATE TABLE shard_manager_state revision BIGINT NOT NULL ); --- Local-mode mirror of the state blob, for inspection with plain SQL (`executor_leases` is the --- counterpart of the quota system's `quota_leases`). The blob is the source of truth: both tables --- are rewritten wholesale in the same transaction as every write the shard manager makes, and it --- never reads them back. Not mirrored: `pending_rebalance`, `shard_epochs` (the per-shard epoch --- high-water marks, which outlive leases) and the in-blob `ShardLeaseRevision`. --- --- Anyone clearing the state by hand (`DELETE FROM shard_manager_state`) must clear these two --- tables as well, or they keep describing leases that no longer exist until the next write. --- --- The shard manager refuses to persist a state whose assignments reference an executor without a --- lease, so the foreign key below is belt-and-braces. Postgres always enforces it; SQLite only --- with `foreign_keys = true` (the default is `false`). +-- A mirror of the state blob for inspection with plain SQL; rewritten wholesale on every write. +-- Clearing shard_manager_state by hand means clearing these two as well. CREATE TABLE executor_leases ( executor_id UUID NOT NULL, diff --git a/golem-shard-manager/src/config.rs b/golem-shard-manager/src/config.rs index 76baba4557..3579705c6c 100644 --- a/golem-shard-manager/src/config.rs +++ b/golem-shard-manager/src/config.rs @@ -477,14 +477,75 @@ pub fn make_config_loader() -> ConfigLoader { ConfigLoader::new_with_examples(Path::new("config/shard-manager.toml")) } +/// Environment variables that configured the shard manager's database before `db` became +/// [`PersistenceConfig`]. +const LEGACY_DB_ENV_VAR_PREFIX: &str = "GOLEM__DB__"; + +fn legacy_db_env_vars>(names: I) -> Vec { + let mut found: Vec = names + .into_iter() + .filter(|name| name.starts_with(LEGACY_DB_ENV_VAR_PREFIX)) + .collect(); + found.sort(); + found +} + +/// Fails if the environment still configures the shard manager through the removed `db` key. +/// +/// figment layers the environment over the defaults and serde then discards unknown keys, so +/// without this a deployment that was not updated starts *successfully* on the default SQLite +/// database, with an empty routing table and a quota ledger that resets on every restart. +pub fn reject_legacy_db_env_vars() -> Result<(), String> { + let found = legacy_db_env_vars(std::env::vars().map(|(name, _)| name)); + if found.is_empty() { + return Ok(()); + } + + Err(format!( + "The shard manager's `db` configuration was replaced by `persistence`, which also selects \ + which backend holds the shard lease state. The following environment variables are no \ + longer read, and ignoring them would start the shard manager on the default SQLite \ + database with an empty routing table:\n {}\nRename them to `GOLEM__PERSISTENCE__*` (for \ + example `GOLEM__DB__TYPE` becomes `GOLEM__PERSISTENCE__TYPE`).", + found.join("\n ") + )) +} + #[cfg(test)] mod tests { use test_r::test; - use crate::config::make_config_loader; + use crate::config::{legacy_db_env_vars, make_config_loader}; #[test] pub fn config_is_loadable() { let _ = make_config_loader().load().expect("Failed to load config"); } + + #[test] + pub fn legacy_db_env_vars_are_detected() { + let found = legacy_db_env_vars( + [ + "GOLEM__DB__CONFIG__HOST", + "GOLEM__PERSISTENCE__CONFIG__HOST", + "GOLEM__DB__TYPE", + "GOLEM__HTTP_PORT", + "PATH", + "SOMETHING__GOLEM__DB__TYPE", + "GOLEM__DBX__TYPE", + ] + .map(str::to_string), + ); + + assert_eq!(found, vec!["GOLEM__DB__CONFIG__HOST", "GOLEM__DB__TYPE"]); + } + + #[test] + pub fn an_environment_without_the_legacy_key_is_accepted() { + let found = legacy_db_env_vars( + ["GOLEM__PERSISTENCE__TYPE", "GOLEM__HTTP_PORT"].map(str::to_string), + ); + + assert!(found.is_empty(), "unexpected legacy variables: {found:?}"); + } } diff --git a/golem-shard-manager/src/server.rs b/golem-shard-manager/src/server.rs index 56b14b10be..be29b938ee 100644 --- a/golem-shard-manager/src/server.rs +++ b/golem-shard-manager/src/server.rs @@ -14,7 +14,9 @@ use golem_common::SafeDisplay; use golem_common::tracing::init_tracing_with_default_env_filter; -use golem_shard_manager::config::{ShardManagerConfig, make_config_loader}; +use golem_shard_manager::config::{ + ShardManagerConfig, make_config_loader, reject_legacy_db_env_vars, +}; use prometheus::default_registry; use tokio::task::JoinSet; use tracing::info; @@ -29,6 +31,8 @@ fn main() -> Result<(), anyhow::Error> { init_tracing_with_default_env_filter(&config.tracing); info!("Using configuration:\n{}", config.to_safe_string_indented()); + reject_legacy_db_env_vars().map_err(|err| anyhow::anyhow!(err))?; + let registry = default_registry().clone(); tokio::runtime::Builder::new_multi_thread() diff --git a/golem-shard-manager/src/sharding/persistence/db.rs b/golem-shard-manager/src/sharding/persistence/db.rs index cd9fb91838..9a78c478ea 100644 --- a/golem-shard-manager/src/sharding/persistence/db.rs +++ b/golem-shard-manager/src/sharding/persistence/db.rs @@ -34,10 +34,9 @@ use uuid::Uuid; const PERSISTENCE_SVC: &str = "persistence"; -/// Rows per multi-row `INSERT` into the mirror tables. Six binds per lease row keeps a chunk far -/// below both Postgres' (65535) and SQLite's (32766) bind-parameter limits. +/// Rows per multi-row `INSERT` into the mirror tables. Six binds per lease row, against SQLite's +/// 32766 bind-parameter limit - the smaller of the two dialects'. const MIRROR_INSERT_CHUNK_SIZE: usize = 1000; -// Six binds per lease row; SQLite's default SQLITE_MAX_VARIABLE_NUMBER is the smaller limit. const _: () = assert!(MIRROR_INSERT_CHUNK_SIZE * 6 <= 32766); /// Creates the single state row. Succeeds only while the row is absent: `DO NOTHING` turns the @@ -128,8 +127,8 @@ impl RoutingTablePersistence for DbRoutingTablePersistence { let leases = lease_rows(shard_state); let assignments = assignment_rows(shard_state)?; - // One transaction: the compare-and-swap on the blob decides, and the mirror tables follow - // it or are left untouched with it. + // One transaction, so the mirror tables follow the compare-and-swap or, if it is rejected, + // are left untouched with it. let revision = self .pool .with_tx_err(PERSISTENCE_SVC, "write", |tx| { @@ -223,7 +222,6 @@ impl DbRoutingTablePersistence { } } -/// One `executor_leases` row. #[derive(Debug, Clone, PartialEq)] struct ExecutorLeaseRow { executor_id: Uuid, @@ -234,7 +232,6 @@ struct ExecutorLeaseRow { pod_name: Option, } -/// One `shard_assignments` row. #[derive(Debug, Clone, PartialEq, Eq)] struct ShardAssignmentRow { shard_id: i32, diff --git a/golem-shard-manager/src/sharding/persistence/etcd.rs b/golem-shard-manager/src/sharding/persistence/etcd.rs index 5374a67b76..700b7e05fa 100644 --- a/golem-shard-manager/src/sharding/persistence/etcd.rs +++ b/golem-shard-manager/src/sharding/persistence/etcd.rs @@ -43,15 +43,15 @@ impl EtcdRoutingTablePersistence { )); } - // TLS is not configurable, so an `https://` endpoint could only fail at connect time with - // an opaque transport error. Refuse it up front and say why instead. + // Only plain `http://` works: TLS is not configurable, and anything else - including a + // scheme-less `host:port` - would otherwise fail at connect time with an opaque error. if let Some(endpoint) = config .endpoints .iter() .find(|endpoint| !endpoint.starts_with("http://")) { return Err(ShardManagerError::Internal(format!( - "etcd endpoint {endpoint} is not an http:// URL; TLS is not supported" + "etcd endpoint {endpoint} must start with http:// (TLS is not supported)" ))); } @@ -76,8 +76,6 @@ impl EtcdRoutingTablePersistence { #[async_trait] impl RoutingTablePersistence for EtcdRoutingTablePersistence { async fn read(&self) -> Result<(ShardLeaseState, ExternalRevision), ShardManagerError> { - // `KvClient` is a cheap handle over the shared multiplexed channel; cloning it per call is - // how the generated stubs' `&mut self` is satisfied. let mut kv = self.client.kv_client(); let response = kv.get(STATE_KEY, None).await?; @@ -107,13 +105,10 @@ impl RoutingTablePersistence for EtcdRoutingTablePersistence { check_state_for_write(shard_state)?; let encoded = serialize(shard_state).map_err(ShardManagerError::SerializationError)?; - // etcd reports mod_revision 0 for a key that does not exist, so `prev_revision == - // NO_REVISION` already means exactly "must not exist yet" - unlike the SQL backend, no - // branching is needed here. - // - // `or_else` is deliberately empty. Returning the winning state from a lost compare-and-swap - // would have to travel inside the error variant, which the SQL backend cannot fill in - // without a second query; the caller re-reads instead. + // etcd reports mod_revision 0 for a key that does not exist, so comparing it equal to + // `prev_revision == NO_REVISION` already means "the key must not exist yet" - create-only + // semantics, without the separate INSERT statement the SQL backend needs for the same + // guarantee. let txn = Txn::new() .when([Compare::mod_revision( STATE_KEY, @@ -129,13 +124,6 @@ impl RoutingTablePersistence for EtcdRoutingTablePersistence { return Err(ShardManagerError::ConcurrentModification); } - // A transaction that applied at least one mutation reports the NEW store revision in its - // header, and etcd stamps every key written by that transaction with exactly that - // revision - so this is the mod_revision our PUT produced. - // - // Do NOT replace this with a follow-up GET: another writer can land between the - // transaction and the GET, and we would cache their revision as ours, so our next - // compare-and-swap would compare-equal against their write and silently clobber it. let revision = response .header() .ok_or_else(|| { diff --git a/golem-shard-manager/src/sharding/persistence/mod.rs b/golem-shard-manager/src/sharding/persistence/mod.rs index eb9e17f357..66fb9fb2ed 100644 --- a/golem-shard-manager/src/sharding/persistence/mod.rs +++ b/golem-shard-manager/src/sharding/persistence/mod.rs @@ -25,24 +25,17 @@ use golem_common::serialization::try_deserialize; /// An opaque, backend-assigned version of the persisted [`ShardLeaseState`]. /// -/// This is a *storage* fencing token, not a domain concept. It is deliberately unrelated to -/// [`super::model::ShardLeaseRevision`], which lives inside the state blob and is bumped by the -/// shard management loop when the routing table changes. The two are expected to drift apart and -/// must never be compared or derived from each other. +/// A *storage* fencing token, deliberately unrelated to [`super::model::ShardLeaseRevision`], +/// which lives inside the blob and is bumped when the routing table changes. The two drift apart +/// by design and must never be compared or derived from each other. /// -/// Guarantees, upheld by every implementation: -/// * [`NO_REVISION`] means "no state is stored". -/// * Any *stored* state carries a revision `>= 1`, so [`NO_REVISION`] is unambiguous. -/// * Successive successful writes return strictly increasing values. +/// * [`NO_REVISION`] means "no state is stored"; any stored state carries `>= 1`. +/// * Successive successful writes return strictly increasing values. The magnitude is meaningless: +/// the SQL backend assigns `previous + 1`, etcd assigns its cluster-wide revision, which jumps. /// -/// The magnitude is meaningless: the SQL backend assigns `previous + 1`, while the etcd backend -/// assigns the cluster-wide etcd revision, which jumps by arbitrary amounts because every -/// unrelated etcd mutation advances it. -/// -/// Monotonicity is guaranteed only for as long as state remains stored. The SQL token is derived -/// from the row, so deleting the row restarts the sequence at 1, whereas etcd's keeps climbing. -/// A writer holding a token from before such a deletion is therefore not fenced by it; only -/// leader election makes that safe. +/// Monotonicity holds only while state remains stored: the SQL token is derived from the row, so +/// deleting it restarts the sequence at 1 while etcd's keeps climbing. A writer holding a token +/// from before such a deletion is therefore not fenced by it. pub type ExternalRevision = i64; /// The revision reported by [`RoutingTablePersistence::read`] when nothing is stored, and the @@ -53,11 +46,6 @@ pub const NO_REVISION: ExternalRevision = 0; pub trait RoutingTablePersistence: Send + Sync { /// Loads the persisted shard lease state together with the revision it is stored at. /// - /// If nothing is stored, returns a freshly initialized [`ShardLeaseState`] paired with - /// [`NO_REVISION`]. The caller cannot distinguish "never written" from "written and then - /// externally deleted", and does not need to: both mean the routing table is rebuilt from - /// scratch as executors register. - /// /// A stored blob that cannot be decoded, or that violates the state invariants, is an error. /// It is never silently replaced by a default state - that would drop a live routing table /// on a transient decoding bug. @@ -73,11 +61,9 @@ pub trait RoutingTablePersistence: Send + Sync { /// that revision"**. In particular, a write with `prev_revision > NO_REVISION` against /// absent state fails; it does not resurrect it. /// - /// When that condition does not hold, returns [`ShardManagerError::ConcurrentModification`] - /// and stores nothing. Recovery is always the same: discard the in-memory state, - /// [`Self::read`] again, re-derive the intended change against what was read, and write with - /// the revision that came with it. Retrying with the same `prev_revision` can never succeed, - /// which is why the error is reported as non-retriable. + /// Otherwise returns [`ShardManagerError::ConcurrentModification`] having stored nothing. + /// Recovery is always re-read, re-derive, write; retrying the same `prev_revision` can never + /// succeed, which is why the error is non-retriable. /// /// The returned revision is always `>= 1` and always strictly greater than `prev_revision`. async fn write( @@ -107,11 +93,8 @@ fn decode_shard_state(bytes: &[u8]) -> Result Result<(), ShardManagerError> { shard_state.check_invariants().map_err(|violation| { ShardManagerError::Internal(format!( @@ -121,8 +104,7 @@ fn check_state_for_write(shard_state: &ShardLeaseState) -> Result<(), ShardManag } /// Rejects a revision a backend claims to have stored at, if it would be indistinguishable from -/// "absent" or would break monotonicity. Both implementations funnel their result through this, -/// so a protocol violation surfaces as an error instead of silently corrupting the fencing chain. +/// "absent" or would break monotonicity. fn check_stored_revision( revision: ExternalRevision, prev_revision: ExternalRevision, diff --git a/golem-shard-manager/src/sharding/shard_management.rs b/golem-shard-manager/src/sharding/shard_management.rs index 552b3fd082..4f2bb6da8d 100644 --- a/golem-shard-manager/src/sharding/shard_management.rs +++ b/golem-shard-manager/src/sharding/shard_management.rs @@ -29,8 +29,14 @@ use std::sync::Arc; use std::time::Duration; use tokio::sync::{Mutex, Notify}; use tokio::task::JoinSet; +use tokio::time::timeout; use tracing::{Instrument, debug, error, info, warn}; +/// Bounds a persistence round-trip, so a wedged backend cannot hold the shard state lock forever, +/// leaving every [`ShardManagement::current_snapshot`] reader waiting while the fail-stop that +/// should end the process never runs. +const PERSISTENCE_TIMEOUT: Duration = Duration::from_secs(30); + #[derive(Clone)] pub struct ShardManagement { shard_state: Arc>, @@ -55,7 +61,15 @@ impl ShardManagement { lease_ttl: Duration, join_set: &mut JoinSet>, ) -> Result { - let (shard_state, external_revision) = persistence_service.read().await?; + let (shard_state, external_revision) = + match timeout(PERSISTENCE_TIMEOUT, persistence_service.read()).await { + Ok(read) => read?, + Err(_) => { + return Err(ShardManagerError::Internal(format!( + "reading the shard lease state timed out after {PERSISTENCE_TIMEOUT:?}" + ))); + } + }; info!("Initial healthcheck started"); @@ -228,9 +242,8 @@ impl ShardManagement { }) .await?; - // Race-free: the worker task is the only writer of the shard state. Reading the - // snapshot after the persist rather than out of the closure keeps its `revision` - // field consistent with what was stored. + // Read after the persist rather than out of the closure, so the snapshot's + // `revision` field is consistent with what was stored. let shard_state_snapshot = self.shard_state.read().await.clone(); let mut full_assignments = Assignments::new(); @@ -308,9 +321,16 @@ impl ShardManagement { let written = match current_shard_state.bump_revision() { Ok(_) => { - self.persistence - .write(¤t_shard_state, prev_external_revision) - .await + let write = self + .persistence + .write(¤t_shard_state, prev_external_revision); + + match timeout(PERSISTENCE_TIMEOUT, write).await { + Ok(written) => written, + Err(_) => Err(ShardManagerError::Internal(format!( + "persisting the shard lease state timed out after {PERSISTENCE_TIMEOUT:?}" + ))), + } } Err(err) => Err(err), }; diff --git a/golem-shard-manager/tests/persistence.rs b/golem-shard-manager/tests/persistence.rs index f9b13c82cc..7577d67b9c 100644 --- a/golem-shard-manager/tests/persistence.rs +++ b/golem-shard-manager/tests/persistence.rs @@ -367,7 +367,7 @@ impl GetRoutingTablePersistence for SqliteRoutingTablePersistence { } } -// -- etcd: isolated by a fresh key prefix per store -------------------------------------------- +// -- etcd: isolated by one server per test worker, plus a wipe per store ----------------------- struct EtcdRoutingTablePersistenceFactory { etcd: DockerEtcd, @@ -587,16 +587,6 @@ async fn mirror_tables_follow_the_persisted_state( return; }; assert_eq!(snapshot, MirrorSnapshot::of(&first)); - // Orphaned shards are pending, not assigned, and must not show up as assignments. - for pending in &first.pending_rebalance { - assert!( - !snapshot - .assignments - .iter() - .any(|(shard_id, _, _)| i64::from(*shard_id) == pending.value()), - "pending shard {pending} must not appear in shard_assignments" - ); - } // A rewrite replaces the mirror wholesale: rows of executors that are gone must not linger. let second = replacement_shard_state(NUMBER_OF_SHARDS); @@ -685,7 +675,15 @@ async fn a_full_size_routing_table_roundtrips_and_is_mirrored( assert_eq!(actual, shard_state); assert_eq!(actual_revision, revision); - if let Some(snapshot) = store.mirror_snapshot().await { + let snapshot = store.mirror_snapshot().await; + assert_eq!( + snapshot.is_some(), + store.has_mirror(), + "a backend must report a mirror exactly when it has one" + ); + // The only coverage of writes spanning more than one MIRROR_INSERT_CHUNK_SIZE chunk, which + // every production-sized write does, so it must not become skippable by accident. + if let Some(snapshot) = snapshot { assert_eq!(snapshot.assignments.len(), SHARDS); assert_eq!(snapshot, MirrorSnapshot::of(&shard_state)); } diff --git a/golem-shard-manager/tests/shard_management.rs b/golem-shard-manager/tests/shard_management.rs index 5c69d27fa5..ecba36ecd2 100644 --- a/golem-shard-manager/tests/shard_management.rs +++ b/golem-shard-manager/tests/shard_management.rs @@ -25,7 +25,7 @@ use std::net::{IpAddr, Ipv4Addr}; use std::sync::Arc; use std::time::Duration; use test_r::test; -use tokio::sync::Mutex; +use tokio::sync::{Mutex, oneshot}; use tokio::task::JoinSet; use tokio::time::Instant; use uuid::Uuid; @@ -51,6 +51,14 @@ struct TestPersistence { /// Per-write script: a `None` entry lets a write through, a `Some(err)` entry fails it before /// it reaches the store. Writes past the end of the script always succeed. injected: Arc>>>, + gate: Arc>>, +} + +/// Suspends a single write: `entered` fires when it is reached, then it waits for `release`. +#[derive(Debug)] +struct WriteGate { + entered: oneshot::Sender<()>, + release: oneshot::Receiver<()>, } impl TestPersistence { @@ -69,6 +77,7 @@ impl TestPersistence { writes: Arc::new(Mutex::new(Vec::new())), attempts: Arc::new(Mutex::new(0)), injected: Arc::new(Mutex::new(VecDeque::new())), + gate: Arc::new(Mutex::new(None)), } } @@ -81,6 +90,18 @@ impl TestPersistence { *self.injected.lock().await = script.into(); } + /// Suspends the next write. Returns a receiver that resolves once that write is in flight, + /// and the sender that lets it finish. + async fn block_next_write(&self) -> (oneshot::Receiver<()>, oneshot::Sender<()>) { + let (entered_tx, entered_rx) = oneshot::channel(); + let (release_tx, release_rx) = oneshot::channel(); + *self.gate.lock().await = Some(WriteGate { + entered: entered_tx, + release: release_rx, + }); + (entered_rx, release_tx) + } + async fn write_count(&self) -> usize { self.writes.lock().await.len() } @@ -102,6 +123,13 @@ impl RoutingTablePersistence for TestPersistence { return Err(err); } + // Taken out of the mutex before awaiting, so the gate never holds a lock. + let gate = self.gate.lock().await.take(); + if let Some(gate) = gate { + let _ = gate.entered.send(()); + let _ = gate.release.await; + } + let mut store = self.store.lock().await; if prev_revision != store.revision { return Err(ShardManagerError::ConcurrentModification); @@ -774,6 +802,133 @@ async fn a_persistence_failure_rolls_back_the_state_and_stops_the_loop() { assert!(worker_executors.local_assignment(new_pod).await.is_empty()); } +#[test] +// The second persist of a pass runs *after* the rebalance reached the executors over gRPC, so +// rolling it back leaves them holding shards the routing table no longer records - and a rebalance +// planned from the rolled-back state then sees a balanced table and plans nothing, forever. Ending +// the task is what saves it: the process restarts and re-sends every authoritative assignment. +async fn a_persistence_failure_after_the_rebalance_was_executed_stops_the_loop() { + let existing_pod = pod(1, 9000); + let new_pod = pod(2, 9001); + let worker_executors = Arc::new(TestWorkerExecutors::default()); + + let (shard_management, persistence, mut join_set) = new_shard_management( + shard_state_with_executors( + 4, + vec![( + executor(1), + existing_pod, + "worker-executor-0", + &[0, 1, 2, 3], + )], + ), + worker_executors.clone(), + ) + .await; + + // Let the registration persist, then fail the persist of the applied rebalance. + persistence + .fail_writes(vec![None, Some(ShardManagerError::ConcurrentModification)]) + .await; + let new_executor = shard_management + .register_executor( + ExecutorAddr::from(new_pod), + Some("worker-executor-1".into()), + ) + .await; + + let outcome = tokio::time::timeout(Duration::from_secs(5), join_set.join_next()) + .await + .expect("the shard management loop should have stopped") + .expect("the loop task should exist") + .expect("the loop task should not panic"); + assert!( + outcome.is_err(), + "the loop must end with the persistence error, got {outcome:?}" + ); + + // The first persist landed, so the registration survived both in memory and in the store. + let after = shard_management.current_snapshot().await; + assert!( + after.has_executor(new_executor), + "the registration was persisted by the first write and must not be rolled back" + ); + assert_eq!( + after, + persistence.latest().await, + "the in-memory state must match the last state that was actually stored" + ); + + // ... but the rebalance it triggered was rolled back after the executor was told about it. + // That divergence is why the loop must stop rather than carry on. + assert!( + shards_at(&after, new_pod).is_empty(), + "the rolled-back rebalance must not appear in the routing table" + ); + assert!( + !worker_executors.local_assignment(new_pod).await.is_empty(), + "the rebalance should have reached the executor before the persist failed - without that, \ + this test is not exercising the second persist site" + ); +} + +#[test] +// The write lock is held across the persistence round-trip, so nobody can read a state that is +// still being stored and that a failed write is about to roll back. Releasing it early would let +// `GetRoutingTable` hand out a routing table that never reached the store, then go backwards. +async fn readers_cannot_observe_a_state_that_is_still_being_persisted() { + let existing_pod = pod(1, 9000); + let new_pod = pod(2, 9001); + let worker_executors = Arc::new(TestWorkerExecutors::default()); + + let (shard_management, persistence, _join_set) = new_shard_management( + shard_state_with_executors( + 4, + vec![( + executor(1), + existing_pod, + "worker-executor-0", + &[0, 1, 2, 3], + )], + ), + worker_executors.clone(), + ) + .await; + + let (entered, release) = persistence.block_next_write().await; + + let _ = shard_management + .register_executor( + ExecutorAddr::from(new_pod), + Some("worker-executor-1".into()), + ) + .await; + + // Signalled from inside `write`, which the loop calls while holding the guard. + tokio::time::timeout(Duration::from_secs(5), entered) + .await + .expect("the loop should have reached a write") + .expect("the gate should have been signalled"); + + let blocked = tokio::time::timeout( + Duration::from_millis(200), + shard_management.current_snapshot(), + ) + .await; + assert!( + blocked.is_err(), + "current_snapshot() must not be able to read a state that is still being persisted" + ); + + release.send(()).expect("the write should still be waiting"); + + wait_for_writes(&persistence, 4).await; + let after = tokio::time::timeout(Duration::from_secs(5), shard_management.current_snapshot()) + .await + .expect("the read lock must be free once the persist completed"); + assert_eq!(after, persistence.latest().await); +} + #[test] // First boot through the loop: nothing is stored, so the very first write must be guarded on // NO_REVISION. Getting this wrong is rejected by both real backends, not silently tolerated. From e8e731b7dc2cbfe39097abbe7780edbe0a90942b Mon Sep 17 00:00:00 2001 From: Aditya Salunkhe Date: Tue, 1 Sep 2026 23:07:25 +0530 Subject: [PATCH 3/6] Fail quota operations in distributed mode instead of silently discarding them --- golem-shard-manager/src/lib.rs | 13 ++- golem-shard-manager/src/quota/mod.rs | 2 +- golem-shard-manager/src/quota/quota_repo.rs | 121 +++++++++++++++++++- 3 files changed, 125 insertions(+), 11 deletions(-) diff --git a/golem-shard-manager/src/lib.rs b/golem-shard-manager/src/lib.rs index 0ba71cd1fc..f3ac41e837 100644 --- a/golem-shard-manager/src/lib.rs +++ b/golem-shard-manager/src/lib.rs @@ -23,7 +23,9 @@ use self::grpc::ShardManagerServiceImpl; #[cfg(feature = "kubernetes")] use crate::config::HealthCheckK8sConfig; use crate::config::{HealthCheckMode, PersistenceConfig}; -use crate::quota::{DbQuotaRepo, GrpcResourceDefinitionFetcher, InMemoryQuotaRepo, QuotaService}; +use crate::quota::{ + DbQuotaRepo, GrpcResourceDefinitionFetcher, QuotaService, UnavailableQuotaRepo, +}; use crate::registry_event_subscriber::ShardManagerRegistryInvalidationHandler; use crate::sharding::healthcheck::GrpcHealthCheck; use crate::sharding::worker_executor::WorkerExecutorServiceDefault; @@ -145,10 +147,9 @@ pub async fn run( ) } PersistenceConfig::Etcd(etcd) => { - // Distributed mode. Quota state has no durable store in this mode: it lives in - // the quota service's memory and is rebuilt as executors re-acquire their leases - // after a restart. - // InMemoryQuotaRepo is a placeholder for now + // Distributed mode. The shard lease state is durable in etcd, but the quota + // tables have not moved there and there is no SQL pool here to hold them, so + // quota operations fail rather than silently succeeding against nothing. ( Arc::new( EtcdRoutingTablePersistence::new( @@ -157,7 +158,7 @@ pub async fn run( ) .await?, ), - Arc::new(InMemoryQuotaRepo), + Arc::new(UnavailableQuotaRepo), ) } } diff --git a/golem-shard-manager/src/quota/mod.rs b/golem-shard-manager/src/quota/mod.rs index 4e19754157..ee526e9620 100644 --- a/golem-shard-manager/src/quota/mod.rs +++ b/golem-shard-manager/src/quota/mod.rs @@ -20,6 +20,6 @@ mod quota_service_tests; mod quota_state; pub mod resource_definition_fetcher; -pub use quota_repo::{DbQuotaRepo, InMemoryQuotaRepo, QuotaRepo}; +pub use quota_repo::{DbQuotaRepo, QuotaRepo, UnavailableQuotaRepo}; pub use quota_service::{QuotaError, QuotaService}; pub use resource_definition_fetcher::{GrpcResourceDefinitionFetcher, ResourceDefinitionFetcher}; diff --git a/golem-shard-manager/src/quota/quota_repo.rs b/golem-shard-manager/src/quota/quota_repo.rs index 2f99240ec2..1458d2ea09 100644 --- a/golem-shard-manager/src/quota/quota_repo.rs +++ b/golem-shard-manager/src/quota/quota_repo.rs @@ -100,14 +100,127 @@ pub trait QuotaRepo: Send + Sync { ) -> Result<(), QuotaRepoError>; } -/// A [`QuotaRepo`] that persists nothing. +/// A [`QuotaRepo`] that refuses to record anything. /// -/// Quota state then lives only in the quota service's memory: it starts empty and is rebuilt as -/// executors acquire their leases. Used in distributed (etcd) mode, which has no durable quota -/// repository, and by the quota service unit tests. +/// Distributed mode has no quota repository yet: the shard lease state moved to etcd, but the +/// quota tables have not, and there is no SQL pool in that mode to hold them. Rather than accept +/// writes and drop them - which would hand every executor the full budget while reporting +/// success - every write fails and says why. +/// +/// The reads return empty because that is the truth: nothing is stored. Failing them instead +/// would abort startup in `QuotaService::restore_state` and take the shard lease state, which +/// *is* durable in this mode, down with it. +#[derive(Debug, Default)] +pub struct UnavailableQuotaRepo; + +impl UnavailableQuotaRepo { + fn unavailable(operation: &str) -> Result { + Err(QuotaRepoError::InternalError(anyhow::anyhow!( + "cannot {operation}: quota state has no durable store when the shard manager is \ + configured with etcd persistence, so quota leases cannot be granted or tracked" + ))) + } +} + +#[async_trait] +impl QuotaRepo for UnavailableQuotaRepo { + async fn save_lease_change( + &self, + _resource: &QuotaResourceRecord, + _previous_resource_revision: i64, + _lease: &QuotaLeaseRecord, + _expired_pods: &[(Blob, i32)], + ) -> Result<(), QuotaRepoError> { + Self::unavailable("record a quota lease change") + } + + async fn save_lease_release( + &self, + _resource: &QuotaResourceRecord, + _previous_resource_revision: i64, + _pod_ip: Blob, + _pod_port: i32, + ) -> Result<(), QuotaRepoError> { + Self::unavailable("record a quota lease release") + } + + async fn save_resource( + &self, + _record: &QuotaResourceRecord, + _previous_revision: i64, + ) -> Result<(), QuotaRepoError> { + Self::unavailable("record a quota resource") + } + + async fn delete_resource_and_leases( + &self, + _resource_definition_id: ResourceDefinitionId, + ) -> Result<(), QuotaRepoError> { + Self::unavailable("delete a quota resource and its leases") + } + + async fn delete_leases_for_resource( + &self, + _resource_definition_id: ResourceDefinitionId, + ) -> Result<(), QuotaRepoError> { + Self::unavailable("delete the leases of a quota resource") + } + + /// Empty, not an error: see the type's documentation. + async fn get_all_resources(&self) -> Result, QuotaRepoError> { + Ok(Vec::new()) + } + + /// Empty, not an error: see the type's documentation. + async fn get_all_leases(&self) -> Result, QuotaRepoError> { + Ok(Vec::new()) + } +} + +#[cfg(test)] +mod unavailable_quota_repo_tests { + use test_r::test; + + use super::{QuotaRepo, UnavailableQuotaRepo}; + use golem_common::model::quota::ResourceDefinitionId; + + #[test] + async fn writes_are_refused_rather_than_silently_dropped() { + let repo = UnavailableQuotaRepo; + let id = ResourceDefinitionId(uuid::Uuid::new_v4()); + + assert!(repo.delete_resource_and_leases(id).await.is_err()); + assert!(repo.delete_leases_for_resource(id).await.is_err()); + } + + #[test] + // Reads must stay infallible: `QuotaService::restore_state` propagates their errors, so + // failing them would abort startup and take down the shard lease state, which *is* durable + // in this mode. Making every method fail uniformly looks tidier and breaks etcd mode. + async fn reads_report_no_state_instead_of_failing() { + let repo = UnavailableQuotaRepo; + + assert!( + repo.get_all_resources() + .await + .expect("reads must not fail") + .is_empty() + ); + assert!( + repo.get_all_leases() + .await + .expect("reads must not fail") + .is_empty() + ); + } +} + +/// A [`QuotaRepo`] that persists nothing, for the quota service unit tests. +#[cfg(test)] #[derive(Debug, Default)] pub struct InMemoryQuotaRepo; +#[cfg(test)] #[async_trait] impl QuotaRepo for InMemoryQuotaRepo { async fn save_lease_change( From f0ea20e00c6cf47a4011d2afa2c906cb52a403d7 Mon Sep 17 00:00:00 2001 From: Aditya Salunkhe Date: Wed, 2 Sep 2026 00:05:06 +0530 Subject: [PATCH 4/6] Install protoc in the cross container and fix five review findings --- .agents/skills/db-migration-scripts/SKILL.md | 6 +- Cross.toml | 4 +- golem-shard-manager/src/config.rs | 19 +- golem-shard-manager/src/quota/quota_repo.rs | 196 +++++++++--------- .../src/quota/quota_service_tests.rs | 61 +++++- 5 files changed, 180 insertions(+), 106 deletions(-) diff --git a/.agents/skills/db-migration-scripts/SKILL.md b/.agents/skills/db-migration-scripts/SKILL.md index 7fec14dde5..df3f9c0d15 100644 --- a/.agents/skills/db-migration-scripts/SKILL.md +++ b/.agents/skills/db-migration-scripts/SKILL.md @@ -30,7 +30,7 @@ root's query layer is shared: | Migration root | Query layer | Convention | |---|---|---| | `golem-registry-service/db/migration/` | one impl expanded per backend by `#[trait_gen(PostgresPool -> PostgresPool, SqlitePool)]` | converge by default | -| `golem-shard-manager/db/migration/` | same `trait_gen` pattern (`src/quota/quota_repo.rs`, `src/sharding/persistence.rs`) | converge by default | +| `golem-shard-manager/db/migration/` | same `trait_gen` pattern (`src/quota/quota_repo.rs`, `src/sharding/persistence/db.rs`) | converge by default | | `golem-worker-executor/db/migration/{indexed,keyvalue,scheduler}/` | separate `postgres.rs` / `sqlite.rs` under `src/storage//` (alongside `multi_sqlite.rs`, `redis.rs`, `memory.rs`) | diverge where the engine calls for it | ### Shared-query roots @@ -109,8 +109,8 @@ Examples of the relevant integration coverage: - Registry repository tests initialize PostgreSQL and SQLite from `golem-registry-service/db/migration/`. -- `golem-shard-manager`'s `persistence` tests run each test through both `sqlite` and `postgres` - matrix dimensions. +- `golem-shard-manager`'s `persistence` tests run each test through the `sqlite`, `postgres` and + `etcd` matrix dimensions. - Worker-executor indexed/key-value storage tests include SQLite and PostgreSQL dimensions; select the storage test that exercises the changed schema. diff --git a/Cross.toml b/Cross.toml index 93412a8ff6..0f727d99e8 100644 --- a/Cross.toml +++ b/Cross.toml @@ -1,7 +1,9 @@ [target.aarch64-unknown-linux-gnu] pre-build = [ "dpkg --add-architecture $CROSS_DEB_ARCH", - "apt-get update && apt-get --assume-yes install pkg-config libssl-dev:$CROSS_DEB_ARCH libssl-dev unzip curl", + # protobuf-compiler is a build-time tool for the host, not a target library: etcd-client's + # build script shells out to protoc. The runner's protoc is outside this container. + "apt-get update && apt-get --assume-yes install pkg-config libssl-dev:$CROSS_DEB_ARCH libssl-dev unzip curl protobuf-compiler", ] env.passthrough = [ "OPENSSL_LIB_DIR=/usr/lib/aarch64-linux-gnu", diff --git a/golem-shard-manager/src/config.rs b/golem-shard-manager/src/config.rs index 3579705c6c..bde50d0af6 100644 --- a/golem-shard-manager/src/config.rs +++ b/golem-shard-manager/src/config.rs @@ -221,18 +221,29 @@ pub struct EtcdConfig { /// fails to deserialize: /// `GOLEM__PERSISTENCE__CONFIG__ENDPOINTS=["http://a:2379","http://b:2379"]` pub endpoints: Vec, - #[serde(with = "humantime_serde")] + /// Defaulted, so that selecting etcd by environment variable does not also require setting + /// every timeout: figment merges `GOLEM__PERSISTENCE__CONFIG__*` over the *default* variant's + /// map, which is SQLite's, so these two would otherwise be missing rather than inherited. + #[serde(with = "humantime_serde", default = "default_etcd_connect_timeout")] pub connect_timeout: Duration, - #[serde(with = "humantime_serde")] + #[serde(with = "humantime_serde", default = "default_etcd_request_timeout")] pub request_timeout: Duration, } +fn default_etcd_connect_timeout() -> Duration { + Duration::from_secs(10) +} + +fn default_etcd_request_timeout() -> Duration { + Duration::from_secs(5) +} + impl Default for EtcdConfig { fn default() -> Self { Self { endpoints: vec!["http://localhost:2379".to_string()], - connect_timeout: Duration::from_secs(10), - request_timeout: Duration::from_secs(5), + connect_timeout: default_etcd_connect_timeout(), + request_timeout: default_etcd_request_timeout(), } } } diff --git a/golem-shard-manager/src/quota/quota_repo.rs b/golem-shard-manager/src/quota/quota_repo.rs index 1458d2ea09..8ccc67761c 100644 --- a/golem-shard-manager/src/quota/quota_repo.rs +++ b/golem-shard-manager/src/quota/quota_repo.rs @@ -177,103 +177,6 @@ impl QuotaRepo for UnavailableQuotaRepo { } } -#[cfg(test)] -mod unavailable_quota_repo_tests { - use test_r::test; - - use super::{QuotaRepo, UnavailableQuotaRepo}; - use golem_common::model::quota::ResourceDefinitionId; - - #[test] - async fn writes_are_refused_rather_than_silently_dropped() { - let repo = UnavailableQuotaRepo; - let id = ResourceDefinitionId(uuid::Uuid::new_v4()); - - assert!(repo.delete_resource_and_leases(id).await.is_err()); - assert!(repo.delete_leases_for_resource(id).await.is_err()); - } - - #[test] - // Reads must stay infallible: `QuotaService::restore_state` propagates their errors, so - // failing them would abort startup and take down the shard lease state, which *is* durable - // in this mode. Making every method fail uniformly looks tidier and breaks etcd mode. - async fn reads_report_no_state_instead_of_failing() { - let repo = UnavailableQuotaRepo; - - assert!( - repo.get_all_resources() - .await - .expect("reads must not fail") - .is_empty() - ); - assert!( - repo.get_all_leases() - .await - .expect("reads must not fail") - .is_empty() - ); - } -} - -/// A [`QuotaRepo`] that persists nothing, for the quota service unit tests. -#[cfg(test)] -#[derive(Debug, Default)] -pub struct InMemoryQuotaRepo; - -#[cfg(test)] -#[async_trait] -impl QuotaRepo for InMemoryQuotaRepo { - async fn save_lease_change( - &self, - _resource: &QuotaResourceRecord, - _previous_resource_revision: i64, - _lease: &QuotaLeaseRecord, - _expired_pods: &[(Blob, i32)], - ) -> Result<(), QuotaRepoError> { - Ok(()) - } - - async fn save_lease_release( - &self, - _resource: &QuotaResourceRecord, - _previous_resource_revision: i64, - _pod_ip: Blob, - _pod_port: i32, - ) -> Result<(), QuotaRepoError> { - Ok(()) - } - - async fn save_resource( - &self, - _record: &QuotaResourceRecord, - _previous_revision: i64, - ) -> Result<(), QuotaRepoError> { - Ok(()) - } - - async fn delete_resource_and_leases( - &self, - _resource_definition_id: ResourceDefinitionId, - ) -> Result<(), QuotaRepoError> { - Ok(()) - } - - async fn get_all_resources(&self) -> Result, QuotaRepoError> { - Ok(Vec::new()) - } - - async fn get_all_leases(&self) -> Result, QuotaRepoError> { - Ok(Vec::new()) - } - - async fn delete_leases_for_resource( - &self, - _resource_definition_id: ResourceDefinitionId, - ) -> Result<(), QuotaRepoError> { - Ok(()) - } -} - static SPAN_NAME: &str = "quota repository"; pub struct LoggedQuotaRepo { @@ -629,3 +532,102 @@ impl DbQuotaRepo { Ok(()) } } + +#[cfg(test)] +mod unavailable_quota_repo_tests { + use test_r::test; + + use super::*; + use chrono::Utc; + use golem_common::model::environment::EnvironmentId; + use golem_common::model::quota::{ + EnforcementAction, ResourceCapacityLimit, ResourceDefinitionRevision, ResourceLimit, + ResourceName, + }; + use std::net::Ipv4Addr; + + fn a_resource() -> QuotaResourceRecord { + let now = SqlDateTime::new(Utc::now()); + QuotaResourceRecord { + resource_definition_id: Uuid::new_v4(), + revision: 1, + definition: Blob::new(ResourceDefinition { + id: ResourceDefinitionId(Uuid::new_v4()), + revision: ResourceDefinitionRevision::INITIAL, + environment_id: EnvironmentId(Uuid::new_v4()), + name: ResourceName("tokens".to_string()), + limit: ResourceLimit::Capacity(ResourceCapacityLimit { value: 100 }), + enforcement_action: EnforcementAction::Reject, + unit: "token".to_string(), + units: "tokens".to_string(), + }), + remaining: NumericU64::new(100), + last_refilled_at: now.clone(), + last_refreshed_at: now, + } + } + + fn a_lease() -> QuotaLeaseRecord { + let now = SqlDateTime::new(Utc::now()); + QuotaLeaseRecord { + resource_definition_id: Uuid::new_v4(), + pod_ip: Blob::new(IpAddr::V4(Ipv4Addr::LOCALHOST)), + pod_port: 9000, + epoch: NumericU64::new(1), + allocated: NumericU64::new(10), + granted_at: now.clone(), + expires_at: now, + pending_reservations: Blob::new(Vec::new()), + } + } + + #[test] + // Every write, not just the cheap ones. `save_lease_change` is the path that grants budget, + // and it is the specific silent success this type exists to prevent - a test that skipped it + // would stay green if the method were reverted to `Ok(())`. + async fn writes_are_refused_rather_than_silently_dropped() { + let repo = UnavailableQuotaRepo; + let id = ResourceDefinitionId(Uuid::new_v4()); + + assert!( + repo.save_lease_change(&a_resource(), 1, &a_lease(), &[]) + .await + .is_err(), + "granting a quota lease must fail rather than report success against nothing" + ); + assert!( + repo.save_lease_release( + &a_resource(), + 1, + Blob::new(IpAddr::V4(Ipv4Addr::LOCALHOST)), + 9000 + ) + .await + .is_err() + ); + assert!(repo.save_resource(&a_resource(), 1).await.is_err()); + assert!(repo.delete_resource_and_leases(id).await.is_err()); + assert!(repo.delete_leases_for_resource(id).await.is_err()); + } + + #[test] + // Reads must stay infallible: `QuotaService::restore_state` propagates their errors, so + // failing them would abort startup and take down the shard lease state, which *is* durable + // in this mode. Making every method fail uniformly looks tidier and breaks etcd mode. + async fn reads_report_no_state_instead_of_failing() { + let repo = UnavailableQuotaRepo; + + assert!( + repo.get_all_resources() + .await + .expect("reads must not fail") + .is_empty() + ); + assert!( + repo.get_all_leases() + .await + .expect("reads must not fail") + .is_empty() + ); + } +} diff --git a/golem-shard-manager/src/quota/quota_service_tests.rs b/golem-shard-manager/src/quota/quota_service_tests.rs index 197b044fee..e731f2ee74 100644 --- a/golem-shard-manager/src/quota/quota_service_tests.rs +++ b/golem-shard-manager/src/quota/quota_service_tests.rs @@ -13,7 +13,7 @@ // limitations under the License. use super::quota_lease::QuotaLease; -use super::quota_repo::{InMemoryQuotaRepo, QuotaRepo}; +use super::quota_repo::{QuotaLeaseRecord, QuotaRepo, QuotaRepoError, QuotaResourceRecord}; use super::quota_service::{QuotaError, QuotaService}; use super::resource_definition_fetcher::{FetchError, ResourceDefinitionFetcher}; use crate::config::QuotaServiceConfig; @@ -26,6 +26,7 @@ use golem_common::model::quota::{ ResourceDefinitionId, ResourceDefinitionRevision, ResourceLimit, ResourceName, ResourceRateLimit, TimePeriod, }; +use golem_service_base::repo::Blob; use std::collections::HashMap; use std::net::{IpAddr, Ipv4Addr}; use std::sync::Arc; @@ -33,6 +34,64 @@ use std::time::Duration; use test_r::test; use tokio::sync::RwLock; +/// A [`QuotaRepo`] that persists nothing, so these tests exercise the service and state +/// machinery without a database. +#[derive(Debug, Default)] +struct InMemoryQuotaRepo; + +#[async_trait] +impl QuotaRepo for InMemoryQuotaRepo { + async fn save_lease_change( + &self, + _resource: &QuotaResourceRecord, + _previous_resource_revision: i64, + _lease: &QuotaLeaseRecord, + _expired_pods: &[(Blob, i32)], + ) -> Result<(), QuotaRepoError> { + Ok(()) + } + + async fn save_lease_release( + &self, + _resource: &QuotaResourceRecord, + _previous_resource_revision: i64, + _pod_ip: Blob, + _pod_port: i32, + ) -> Result<(), QuotaRepoError> { + Ok(()) + } + + async fn save_resource( + &self, + _record: &QuotaResourceRecord, + _previous_revision: i64, + ) -> Result<(), QuotaRepoError> { + Ok(()) + } + + async fn delete_resource_and_leases( + &self, + _resource_definition_id: ResourceDefinitionId, + ) -> Result<(), QuotaRepoError> { + Ok(()) + } + + async fn get_all_resources(&self) -> Result, QuotaRepoError> { + Ok(Vec::new()) + } + + async fn get_all_leases(&self) -> Result, QuotaRepoError> { + Ok(Vec::new()) + } + + async fn delete_leases_for_resource( + &self, + _resource_definition_id: ResourceDefinitionId, + ) -> Result<(), QuotaRepoError> { + Ok(()) + } +} + fn test_repo() -> Arc { Arc::new(InMemoryQuotaRepo) } From c67596f04a444b596d3e8a01e6472a5146c9dcc2 Mon Sep 17 00:00:00 2001 From: Aditya Salunkhe Date: Wed, 2 Sep 2026 16:35:16 +0530 Subject: [PATCH 5/6] fixes --- AGENTS.md | 3 + golem-shard-manager/src/server.rs | 6 +- .../src/sharding/persistence/etcd.rs | 4 +- golem-shard-manager/tests/shard_management.rs | 56 +++++++++++++++++++ 4 files changed, 66 insertions(+), 3 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 9aa7c21f3b..3bbdbc794c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -25,6 +25,9 @@ explicitly revised, **do not perform backward-compatibility work**. - **cargo-make**: Latest version (`cargo install --force cargo-make`) - **cargo-test-r**: Match the version CI installs, pinned in `.github/actions/restore-binaries/action.yml` (`cargo install --force --locked cargo-test-r@`) +- **protoc**: Required to build `golem-shard-manager`, and therefore the `golem` binary. The + `etcd-client` crate's build script compiles its protos with `tonic-prost-build`, which shells out + to the protobuf compiler; the workspace's own protos use `protox` and need no binary - **redis-server**: Required by worker-executor, service-integration, and CLI integration tests that use spawned test dependencies; not required by every unit test - **docker**: Required by tests that use containerized dependencies diff --git a/golem-shard-manager/src/server.rs b/golem-shard-manager/src/server.rs index be29b938ee..52329aa286 100644 --- a/golem-shard-manager/src/server.rs +++ b/golem-shard-manager/src/server.rs @@ -22,6 +22,10 @@ use tokio::task::JoinSet; use tracing::info; fn main() -> Result<(), anyhow::Error> { + // Before the configuration is loaded at all, so that `--dump-config` cannot print a config + // that silently ignores a deployment's legacy settings. + reject_legacy_db_env_vars().map_err(|err| anyhow::anyhow!(err))?; + match make_config_loader().load_or_dump_config() { Some(config) => { rustls::crypto::ring::default_provider() @@ -31,8 +35,6 @@ fn main() -> Result<(), anyhow::Error> { init_tracing_with_default_env_filter(&config.tracing); info!("Using configuration:\n{}", config.to_safe_string_indented()); - reject_legacy_db_env_vars().map_err(|err| anyhow::anyhow!(err))?; - let registry = default_registry().clone(); tokio::runtime::Builder::new_multi_thread() diff --git a/golem-shard-manager/src/sharding/persistence/etcd.rs b/golem-shard-manager/src/sharding/persistence/etcd.rs index 700b7e05fa..2dc3af57ed 100644 --- a/golem-shard-manager/src/sharding/persistence/etcd.rs +++ b/golem-shard-manager/src/sharding/persistence/etcd.rs @@ -59,11 +59,13 @@ impl EtcdRoutingTablePersistence { .with_connect_timeout(config.connect_timeout) .with_timeout(config.request_timeout); + // The client connects lazily, on its first request, so nothing is known about the + // endpoints' reachability yet; the startup read is what first finds out. let client = Client::connect(&config.endpoints, Some(options)).await?; info!( endpoints = config.endpoints.join(", "), state_key = STATE_KEY, - "Connected to etcd for shard lease state persistence" + "Configured the etcd client for shard lease state persistence" ); Ok(Self { diff --git a/golem-shard-manager/tests/shard_management.rs b/golem-shard-manager/tests/shard_management.rs index ecba36ecd2..aaf701d471 100644 --- a/golem-shard-manager/tests/shard_management.rs +++ b/golem-shard-manager/tests/shard_management.rs @@ -1008,3 +1008,59 @@ async fn the_first_write_uses_the_revision_read_at_startup() { join_set.abort_all(); } + +#[test] +// The rollback arm is not specific to a revision conflict. A backend failure that says nothing +// about revisions leaves the store in an unknown state - the write may or may not have landed - +// so it is rolled back and stops the loop exactly like a conflict does, and the loop surfaces the +// backend's own error rather than a generic one. +async fn a_non_conflict_persistence_error_is_rolled_back_the_same_way() { + let existing_pod = pod(1, 9000); + let new_pod = pod(2, 9001); + let worker_executors = Arc::new(TestWorkerExecutors::default()); + + let (shard_management, persistence, mut join_set) = new_shard_management( + shard_state_with_executors( + 4, + vec![( + executor(1), + existing_pod, + "worker-executor-0", + &[0, 1, 2, 3], + )], + ), + worker_executors.clone(), + ) + .await; + let before = shard_management.current_snapshot().await; + let persisted_before = persistence.latest().await; + + persistence + .fail_writes(vec![Some(ShardManagerError::Internal( + "injected backend failure".into(), + ))]) + .await; + let new_executor = shard_management + .register_executor( + ExecutorAddr::from(new_pod), + Some("worker-executor-1".into()), + ) + .await; + + let outcome = tokio::time::timeout(Duration::from_secs(5), join_set.join_next()) + .await + .expect("the shard management loop should have stopped") + .expect("the loop task should exist") + .expect("the loop task should not panic"); + let err = outcome.expect_err("the loop must end with the persistence error"); + assert!( + err.to_string().contains("injected backend failure"), + "the loop should surface the backend's own error, got: {err:#}" + ); + + let after = shard_management.current_snapshot().await; + assert!(!after.has_executor(new_executor)); + assert_eq!(after, before); + assert_eq!(persistence.latest().await, persisted_before); + assert!(worker_executors.local_assignment(new_pod).await.is_empty()); +} From 5c6b538e44712fd98fe9743832f57758526dd302 Mon Sep 17 00:00:00 2001 From: Aditya Salunkhe Date: Wed, 2 Sep 2026 17:50:04 +0530 Subject: [PATCH 6/6] Sequence the persistence tests in code instead of test-threads flag --- Makefile.toml | 2 +- golem-shard-manager/tests/lib.rs | 7 ++++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/Makefile.toml b/Makefile.toml index 39e65a6891..497b60c1b3 100644 --- a/Makefile.toml +++ b/Makefile.toml @@ -681,7 +681,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 --test-threads=1 --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/tests/lib.rs b/golem-shard-manager/tests/lib.rs index abb8e476d6..cb4fd5617f 100644 --- a/golem-shard-manager/tests/lib.rs +++ b/golem-shard-manager/tests/lib.rs @@ -16,10 +16,15 @@ mod persistence; mod shard_management; use golem_common::tracing::{TracingConfig, init_tracing_with_default_debug_env_filter}; -use test_r::test_dep; +use test_r::{sequential_suite, test_dep}; test_r::enable!(); +// The etcd dimension shares one server per worker and the fixed `STATE_KEY`, and every store wipes +// that key when it connects, so two persistence tests running at once would see each other's +// writes as revision conflicts. +sequential_suite!(persistence); + #[derive(Debug)] pub struct Tracing;