Skip to content

Make the shard manager's persistence compare-and-swap based across local and distributed backends - #3790

Merged
Aditya1404Sal merged 7 commits into
golemcloud:mainfrom
Aditya1404Sal:shard-manager/ticket2-cas-persistence
Sep 2, 2026
Merged

Make the shard manager's persistence compare-and-swap based across local and distributed backends#3790
Aditya1404Sal merged 7 commits into
golemcloud:mainfrom
Aditya1404Sal:shard-manager/ticket2-cas-persistence

Conversation

@Aditya1404Sal

@Aditya1404Sal Aditya1404Sal commented Aug 29, 2026

Copy link
Copy Markdown
Contributor

Resolves GOL-446

Summary

Ticket 2 of the shard manager redesign. Today RoutingTablePersistence::write is an unconditional last-write-wins upsert: two shard managers writing the same state clobber each other and nothing detects it. Ticket 3 adds leader election, which is only sound if the persisted state carries a fencing token.

This PR delivers that substrate:

  1. Compare-and-swap on every write, via an opaque, store-assigned revision.
  2. A second backend, etcd, behind the existing RoutingTablePersistence trait, so distributed deployments do not keep shard state in a per-service SQL database. The backend is selected at startup from config; ShardManagement only ever sees the trait.
  3. Defined recovery in place of the four .expect(...) panics in the shard management loop: snapshot, write, and on any failure roll back and stop.

Blast radius is golem-shard-manager plus one new test component in golem-test-framework. No proto, executor, or golem-common change. Local mode (Postgres/SQLite) keeps working exactly as before apart from the write becoming conditional.

What changed

The trait (sharding/persistence/mod.rs). read() returns (ShardLeaseState, ExternalRevision); write(&state, prev_revision) stores only if the current revision is exactly prev_revision and returns the new one. NO_REVISION = 0 means "nothing is stored", so a write at 0 is create-only and a stored state always carries >= 1. A stale token yields ConcurrentModification, which is non-retriable: retrying the same prev_revision can never succeed.

The actor (sharding/shard_management.rs). ShardManagement carries external_revision: Arc<Mutex<ExternalRevision>>, seeded from the startup read. Every mutation goes through mutate_and_persist: snapshot, mutate, bump the domain revision, write guarded on the cached token, and on any error restore the snapshot and return Err. That ends the worker task and the process, which restarts and re-reads. The write lock is held across the round-trip so GetRoutingTable can never hand out a state that then rolls back, and every round-trip is bounded at 30s so a wedged store cannot hold that lock forever.

SQL backend (sharding/persistence/db.rs, migration 004). shard_manager_state gains revision BIGINT NOT NULL. The write is one transaction: the compare-and-swap statement, then a wholesale rewrite of the executor_leases and shard_assignments mirror tables. The mirror is a projection of the blob for inspection with plain SQL; the blob stays the source of truth, and a rejected CAS never touches the mirror.

etcd backend (sharding/persistence/etcd.rs). Fixed STATE_KEY = "/golem/shard-manager/state". Read returns the key's mod_revision; write is a transaction comparing mod_revision(STATE_KEY) == prev_revision and putting on success. Endpoints must be http://; auth and TLS are out of scope and are refused with a message naming the gap.

Config (config.rs). persistence: PersistenceConfig { Postgres, Sqlite, Etcd } replaces db: DbConfig. Env keys are GOLEM__PERSISTENCE__TYPE and GOLEM__PERSISTENCE__CONFIG__*. The standalone binary refuses to start if any GOLEM__DB__* variable is still set, because the rename otherwise failed open: figment layered the env over defaults, serde dropped the unknown keys, and a deployment that had not been updated started successfully on the default SQLite file with an empty routing table.

Wiring (lib.rs). Postgres and Sqlite run migrations, build the pool, and construct DbRoutingTablePersistence and DbQuotaRepo, unchanged. Etcd constructs EtcdRoutingTablePersistence and UnavailableQuotaRepo, which fails quota writes loudly and returns nothing for reads: quota has no durable store in distributed mode until ticket 7.

Tests. tests/persistence.rs runs a shared suite across a sqlite / postgres / etcd matrix (read on empty store, round-trip, stale-revision conflict, housekeep expiry, mirror consistency, deleted-row refusal). tests/shard_management.rs covers rollback on both persist sites, the lock held across the round-trip, first boot at NO_REVISION, and the non-conflict rollback arm. DockerEtcd is a new test-framework component. integration-tests-group5 now runs with --test-threads=1, since the etcd dimension shares one container and one fixed key.

CI. etcd-client's build script needs protoc; it is installed in the shared setup-rust action and in the cross container for the aarch64 build.

Deviations from the ticket, and why

  1. "Replace DbRoutingTablePersistence with a trait-based layer." The trait already existed and DbRoutingTablePersistence was already one implementation behind Arc<dyn RoutingTablePersistence>. The real work was changing the trait's signature to carry a token, adding the etcd implementation, and adding config selection.

  2. The SQL write is two statements, not the ticket's single UPDATE … WHERE revision = $prev. A lone UPDATE cannot insert on first boot. The obvious alternative, a guarded upsert, has an unguarded INSERT branch: write(state, prev = 5) against a row that was deleted underneath the writer would silently resurrect it, while etcd refuses the same write. One trait cannot have two behaviours, so the branch happens in Rust: prev == 0 runs INSERT … ON CONFLICT (id) DO NOTHING, prev > 0 runs the guarded UPDATE, and rows_affected == 0 means conflict in both. The deleted-row case is real: integration-tests/tests/sharding.rs deletes that row.

  3. Migration 004 drops and recreates the table instead of ALTER TABLE … ADD COLUMN revision BIGINT NOT NULL DEFAULT 0. A DEFAULT 0 row would be indistinguishable from "absent", since 0 is the sentinel, and the read path rejects a stored revision below 1. 003 had already emptied the table, and SQLite's ADD COLUMN rules would otherwise leave the two dialects with different schemas. Consequence: a row written by a ticket-1 build is dropped at upgrade. No released version carries migration 003, so no released deployment is affected. After the upgrade the routing table is empty until every executor process restarts; that belongs in the release notes.

  4. etcd returns TxnResponse::header().revision(), not the ticket's follow-up GET. A transaction that applied a mutation reports the new store revision, and etcd stamps every key it wrote with exactly that value. The GET is a wasted round-trip and a race: a rival write landing in between would be cached as ours, and the next CAS would compare equal against their write and overwrite it.

  5. etcd-client = "0.19.0", not "0.14". 0.19 matches the workspace's tonic/prost/http/tower versions, so it adds one package to the lockfile and reuses the already-compiled gRPC stack. 0.14 pulls tonic 0.12 and duplicates the whole tree.

  6. PersistenceConfig uses the repo's serde form, #[serde(tag = "type", content = "config")] with PascalCase tags (Postgres, Sqlite, Etcd) and the concrete DbPostgresConfig / DbSqliteConfig payloads, rather than the ticket's rename_all = "snake_case" over a single DbConfig. This matches HealthCheckMode and the previous db config, so generated TOML and env files keep their existing shape.

  7. Rollback covers every failure, not only ConcurrentModification. The ticket's snippet bumps the revision with ? and only rolls back around the write. Here the bump, the write, a backend error and the 30s timeout all restore the snapshot. Nothing is retried on conflict: the cached token is deliberately not refreshed, because a loser that re-reads and writes again is not fenced. Compare-and-swap is only a fence if the loser stops.

  8. Two revisions coexist and are never derived from each other. The in-blob ShardLeaseRevision is domain-owned and increments by one; ExternalRevision is store-owned and opaque. Tying them would break CAS for any write that does not bump the domain revision, which ticket 4's lease renewals will be.

Decisions the ticket deferred and this PR honours as deferrals: no leader election (ticket 3), quota in etcd mode (ticket 7), etcd auth and TLS (with ticket 3).

Verification

  • cargo-test-r run -p golem-shard-manager --lib: 85 passed.
  • cargo-test-r run -p golem-shard-manager --test integration -- --test-threads=1: 56 passed across the sqlite/postgres/etcd matrix.
  • clippy with default and kubernetes features, cargo fmt, cargo make check-configs: clean.
  • Mutation-checked rather than assumed: weakening the second persist site's .await? to .await.ok() turns a_persistence_failure_after_the_rebalance_was_executed_stops_the_loop red; releasing the state lock before the round-trip turns readers_cannot_observe_a_state_that_is_still_being_persisted red; skipping the rollback for non-conflict errors turns a_non_conflict_persistence_error_is_rolled_back_the_same_way red while the conflict test stays green.
  • Smoke tested on both backends: revision climbing in SQL across passes, mod_revision advancing under etcdctl. A revision conflict ends in a clean process exit, not a crash to debug.

Reviewed with a seven-dimension adversarial pass against the ticket text; no blocker or high findings. The low findings that were worth fixing landed in the fixes commit. The merge commit resolves a generated-config conflict with main by regenerating the file; against this branch it differs only in whitespace.

Known limits, handed forward

  • Run a single shard manager until ticket 3. A lost CAS ends the process with no backoff, so two managers during a rolling upgrade kill each other one pod per pass. deploy.mdx documents single-instance operation.
  • The SQL token is not a fence across a row deletion. It is prev + 1, so after DELETE FROM shard_manager_state revisions restart at 1 and a stale cached token can eventually compare equal. etcd's cluster revision cannot do this. Ticket 3's soundness argument depends on the fence, so on SQL this is a precondition to state, not an implementation detail.
  • etcd MVCC growth is unguarded. Every write stores the whole blob as a new revision. Without --auto-compaction-retention the backend eventually hits its quota and goes read-only, which with fail-stop is a crash loop. Needs a deployment note when etcd mode becomes real.
  • STATE_KEY is a constant, so two deployments cannot share one etcd cluster.
  • Cancellation safety is undefined. Dropping the worker mid-write leaves the state mutated and the token stale. Unreachable today; ticket 3 makes leadership own the task's lifetime.
  • A registration acknowledged in a pass whose persist fails is dropped until that executor restarts. Pre-existing behaviour (the old code panicked at the same point); ticket 4's lease renewal is the natural fix.
  • Rolling upgrade: an old binary's blind INSERT hits the NOT NULL revision column and panics. Intended fail-fast for a breaking release; belongs in the release notes.
  • number_of_shards config/state mismatch is unvalidated (pre-existing). Ticket 3 adds a startup guard.

@netlify

netlify Bot commented Aug 29, 2026

Copy link
Copy Markdown

Deploy Preview for golemcloud canceled.

Name Link
🔨 Latest commit 5c6b538
🔍 Latest deploy log https://app.netlify.com/projects/golemcloud/deploys/6a98147f5e4cc90008011d05

@Aditya1404Sal
Aditya1404Sal force-pushed the shard-manager/ticket2-cas-persistence branch from b8a5958 to d68a3b3 Compare August 29, 2026 14:57
@Aditya1404Sal

Copy link
Copy Markdown
Contributor Author

@unblocked review this PR

@unblocked unblocked Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

1 issue found.

In reply to #3790 (comment)

About Unblocked

Unblocked has been set up to automatically review your team's pull requests to identify genuine bugs and issues.

📖 Documentation — Learn more in our docs.

💬 Ask questions — Mention @unblocked to request a review or summary, or ask follow-up questions.

👍 Give feedback — React to comments with 👍 or 👎 to help us improve.

⚙️ Customize — Adjust settings in your preferences.

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 ==

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The comment on line 108 is truncated mid-sentence: // etcd reports mod_revision 0 for a key that does not exist, so \prev_revision ==`. It should explain that prev_revision == NO_REVISION(0) means the key must not exist, and theCompareOp::Equalagainst 0 enforces create-only semantics — mirroring the explanation in thedb.rsbackend'scompare_and_swap_state_in_tx`.

@Aditya1404Sal
Aditya1404Sal force-pushed the shard-manager/ticket2-cas-persistence branch from ad6b5ec to 58f812e Compare September 1, 2026 08:42
@Aditya1404Sal
Aditya1404Sal force-pushed the shard-manager/ticket2-cas-persistence branch from 58f812e to 4922125 Compare September 1, 2026 14:59
@Aditya1404Sal

Copy link
Copy Markdown
Contributor Author

Breaking: the db config key is now persistence

ShardManagerConfig's db field is replaced by persistence, which selects the backend that holds the shard lease state as well as the SQL connection details. In TOML, [db] / [db.config] become [persistence] / [persistence.config]; in the environment, GOLEM__DB__TYPE becomes GOLEM__PERSISTENCE__TYPE and GOLEM__DB__CONFIG__HOST becomes GOLEM__PERSISTENCE__CONFIG__HOST, and so on for every key. The in-repo consumers are updated — docker-examples/published-postgres/compose.yaml, local-run/start.sh, and the test framework — but any deployment configured outside this repository needs the same rename.

Local mode is otherwise unchanged in shape and improved in behaviour. Postgres and SQLite still hold both the shard lease state and the quota tables in one database from one pool; the only functional difference is that the state write is now a compare-and-swap instead of a blind last-write-wins upsert, so a conflicting write is detected rather than silently applied. Distributed (etcd) mode is the additive part — a new persistence.type = "Etcd" arm that nothing selects by default and that is not deployable until leader election lands in ticket 3.

Because figment layers environment variables over defaults and serde discards unknown keys, a shard manager that still had GOLEM__DB__* set would previously have started successfully on the default SQLite file with an empty routing table, silently ignoring the intended Postgres configuration. That is now rejected at startup with an error naming the offending variables and the key to rename them to. The check runs only in the standalone golem-shard-manager binary, since the combined golem binary builds this configuration in code and GOLEM__DB__* there still belongs to other services.

@Aditya1404Sal
Aditya1404Sal marked this pull request as ready for review September 1, 2026 19:12
@Aditya1404Sal
Aditya1404Sal requested a review from a team September 1, 2026 19:12
Comment thread Makefile.toml Outdated
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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do we need this? It will probably slow down CI significantly.
Also if it is really needed, note that test-r has support for marking suites (modules) as sequential from code, which is better than "remembering" to pass the thread count config in cli.

@vigoo vigoo left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just one concern about sequential test execution. Once that's resolved good to go

@Aditya1404Sal

Copy link
Copy Markdown
Contributor Author

Just one concern about sequential test execution. Once that's resolved good to go

It's needed for the etcd dimension only: the fixture shares one etcd server per worker on the fixed STATE_KEY, and every store wipes that key on connect, so concurrent persistence tests see each other's writes as revision conflicts. Without any sequencing 10 of the 14 etcd tests fail.

Moved it into code as sequential_suite!(persistence) and dropped the CLI flag; the loop tests in shard_management stay parallel. CI cost is nil, the whole binary runs in ~6.5s.

@Aditya1404Sal
Aditya1404Sal merged commit 5bb9763 into golemcloud:main Sep 2, 2026
56 checks passed
@github-actions github-actions Bot locked and limited conversation to collaborators Sep 2, 2026
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants