Skip to content

feat: ADD SHARD command plus lookup_query for fixed shard mapping - #1300

Draft
rlittlefield wants to merge 35 commits into
pgdogdev:mainfrom
rlittlefield:add-shard
Draft

feat: ADD SHARD command plus lookup_query for fixed shard mapping#1300
rlittlefield wants to merge 35 commits into
pgdogdev:mainfrom
rlittlefield:add-shard

Conversation

@rlittlefield

Copy link
Copy Markdown

ADD SHARD: orchestrated shard provisioning with WAL catch-up

Base: fixed shard placement

This branch includes the work for the lookup_query (hashed), and also adds support
for fixed shard mapping based on lookup_query:

  • lookup_query on [[sharded_tables]] routes queries by asking the
    database. pgdog runs the query with the sharding key as $1, lazily and
    per key, and caches the result. The cache and query timeout are tuned
    with sharding_lookup_cache_size and sharding_lookup_timeout in
    [general].
  • lookup_result = "shard" means the query returns the 0-based shard
    number itself. The application owns placement, for example a permanent
    shard_id column on its tenants table. Placement is stored, not derived
    from a hash, so the shard count can change without moving any rows.
[[sharded_tables]]
database = "prod"
column = "tenant_id"
data_type = "bigint"
lookup_query = "SELECT shard_id FROM tenants WHERE id = $1"
lookup_result = "shard"

Stored placement is what makes growing a cluster in place safe. This PR
automates the growing.

What this PR adds

An ADD SHARD admin command (and a pgdog add-shard CLI) that grows a
cluster by one shard end to end while production traffic keeps flowing.

You declare each future shard once, in its final shape. The
provisioning flag keeps it out of the serving topology until its cutover
activates it:

[[databases]]
name = "prod"        # The database gaining a shard: this entry joins
                     # the existing [[databases]] entries for "prod".
host = "10.0.0.3"    # The new server.
shard = 2            # The shard number it will serve. Shards are added
                     # in order: with 0 and 1 serving, 2 goes first.
provisioning = true  # Keeps the entry out of the serving topology until
                     # the cutover. Declare as many future shards as you
                     # like (2, 3, 4...); each command names one.
-- Provision shard 2 of the "prod" database, from its provisioning
-- entry: sync all DDL, copy the omnisharded data, stream WAL until
-- caught up. Then park at "awaiting cutover".
psql admin -c "ADD SHARD prod 2"

-- Cut over: pause omni writes, drain to zero, activate the shard,
-- resume. (Also addressable by the task id from SHOW TASKS:
-- CUTOVER <id>.)
psql admin -c "CUTOVER SHARD prod 2"

-- Or both in one step, without the confirmation checkpoint.
psql admin -c "ADD SHARD prod 2 AUTO"

You create an empty database on the new server, deploy the entry, and
run the command. The database's existing schema_admin user covers the
new shard. pgdog does the rest:

  1. Guards. The task refuses to run unless every sharded table is
    placement-stable, meaning lookup_result = "shard" or an explicit
    mapping. Hash-routed tables need a real reshard, because adding a shard
    would move their rows. The task also refuses when the named shard has
    no provisioning entry, when it isn't the next shard number (declared
    shards are added in order, one ADD SHARD each), when the new shard
    isn't empty, when another topology change for the database is already
    running, and when another pgdog instance holds the provisioning lock.
    That lock is a session-scoped advisory lock taken on the new shard
    itself, so a crashed holder releases it automatically.
  2. Schema sync. The existing pg_dump-based schema sync provisions all
    DDL onto the new shard.
  3. Omni data. Omnisharded tables are snapshot-copied from shard 0
    using the resharding data-sync machinery, then WAL-streamed until
    caught up. A publication scoped to the omnisharded tables is created
    automatically and dropped when the task ends. Pass one explicitly to
    manage it yourself. A database with no omnisharded tables skips this
    phase and the cutover's write pause entirely: the task syncs schema,
    parks, and activates.
  4. Cutover. Only omnisharded writes pause. Sharded traffic and all
    reads flow throughout. The task drains replication to zero, flips the
    provisioning flag off in the running config, reloads, refreshes
    pgdog.config on every shard, persists the files when
    cutover_save_config is on, and resumes omni writes. The pause is the
    drain plus coordination, typically around a second.

Progress is visible in SHOW TASKS (validating, syncing schema, syncing
data, replicating, awaiting cutover, draining, swapping topology).
STOP_TASK aborts cleanly any time before the swap: slots and the
auto-created publication are dropped, and the topology never changed. The
swap is the point of no return. Everything after it is best-effort, and
each failure warns with its recovery step.

Sequence

sequenceDiagram
    autonumber
    actor Op as You
    participant App as Application
    participant PD as pgdog running ADD SHARD
    participant Fleet as Other pgdog instances
    participant S0 as Shard 0
    participant NEW as New shard
    participant ALL as All shards

    Note over Op,NEW: MANUAL - create empty DB, deploy the provisioning entry, RELOAD

    Op->>PD: ADD SHARD prod 2
    PD->>PD: guards - placement stable, shard 2 is next, new shard empty
    PD->>NEW: take the cross-instance advisory lock
    PD->>S0: create publication for omni tables
    PD->>S0: pg_dump schema only
    PD->>NEW: apply all DDL with idempotent restore
    PD->>S0: create durable replication slot at L0
    loop each omni table
        PD->>S0: temp slot with snapshot, COPY
        PD->>NEW: rows consistent at Lt after L0
    end
    PD->>S0: start replication from L0
    loop until lag under threshold
        S0-->>PD: WAL omni changes
        PD->>NEW: apply with upsert and per table LSN dedup
        App->>ALL: traffic flows normally throughout
    end
    Note over PD: status awaiting cutover - WAL retained on shard 0 while parked

    Op->>PD: CUTOVER SHARD prod 2
    PD->>PD: pause local omni writes
    PD->>NEW: publish armed state
    Fleet-->>NEW: see armed, pause omni writes, ack
    Note over App,Fleet: omni writes park everywhere - reads and sharded traffic keep flowing
    S0-->>PD: in flight omni writes drain to zero lag
    PD->>S0: stop stream and drop slots
    PD->>PD: POINT OF NO RETURN - provisioning flag flips off, registry reloads with 3 shards
    PD->>PD: resume local omni writes, persist config if enabled
    PD->>ALL: refresh pgdog.config rows on every shard
    PD->>NEW: publish activated state
    Fleet-->>NEW: see activated, activate shard 2, resume omni writes, ack
    PD->>S0: drop publication, also on new shard
    App->>ALL: omni writes now reach all three shards
    Note over Op,App: MANUAL - remove the provisioning flag from the config source, app assigns tenants to shard 2
Loading

Config as state

The provisioning = true flag makes the config declarative rather than
stateful. The entry already describes the end state, and activation only
deletes one line. When the config comes from a manifest or template
(GitOps), you remove the flag at the source at your leisure.
cutover_save_config covers file-managed setups.

The database records the truth either way: the cutover writes
(shard = N, shards = N + 1) into the new shard's own pgdog.config, and
convergence reads it back at startup (before the listener opens) and after
every RELOAD. A restart with a stale manifest re-activates the shard
before serving a single query.

Multiple pgdog instances

The cutover coordinates a fleet, using the new shard itself as the shared
medium. Every instance that has the entry can reach it, and it serves
nothing else.

  • Every instance heartbeats pgdog.instances on shard 0. That table is
    the fleet registry, installed by SETUP SCHEMA and by the topology
    tasks. Graceful shutdowns (including SIGTERM) deregister the instance,
    so rolling restarts don't disturb a cutover.
  • Every instance that sees provisioning entries runs an agent for the
    next declared shard. The agent registers on that shard and polls a
    state row there once a second.
  • The cutover refuses if a live instance hasn't registered on the new
    shard, naming it in the error. That instance is running a config without
    the entry, and the fix is deploying the config everywhere first.
  • The protocol is a two-phase handshake with acks. The coordinator arms
    the omni-write barrier locally, publishes armed, and waits for every
    peer's ack. It drains to zero, swaps, resumes its own writes, and
    publishes activated. Peers activate, resume writes, and ack. Aborts
    publish released.
  • Fail-safety is structural. The coordinator refreshes the armed row every
    5 seconds for as long as it holds the fleet's barriers, and agents
    resume writes only when that heartbeat goes silent for 30 seconds, which
    means the coordinator actually died. A dead but recently seen peer times
    out the arm phase and the task re-parks for a retry. An instance that
    missed the activation keeps its omni writes parked until it learns the
    outcome, and a crashed one converges at restart before serving. There is
    no path that writes with a stale topology.
  • A solo instance coordinates with nobody and behaves exactly as before.

Failure and crash behavior

Everything before the swap aborts cleanly and reruns from scratch.
Per-table copy slots are temporary. The durable slot and the auto-created
publication are dropped by cleanup guards. The omni-write barrier and the
provisioning cluster (whose session holds the advisory lock) are drop
guards, so even a hard-aborted cancellation releases them. After the swap,
a crash before the pgdog.config refresh is recovered with SETUP SCHEMA.
Once the marker is written, restarts converge on their own. There is no
rollback after the swap, because once lookups assign a tenant to the new
shard, removing it would strand data.

Notes for review

  • The drain honors cutover_timeout and cutover_timeout_action. With
    abort, the task releases the barriers and re-parks at "awaiting
    cutover" instead of failing, so the cutover can be retried.
  • While parked, the durable slot on shard 0 retains WAL. It's visible in
    SHOW REPLICATION_SLOTS, so don't park indefinitely.
  • Coordination timing constants (agent poll 1s, arm ack timeout 10s,
    liveness window 15s, heartbeat 5s) are code constants for now. Promoting
    them to config is a reasonable follow-up if fleets need it.
  • Other possible follow-ups: NOTIFY/LISTEN via the existing pub/sub
    listener to replace the agents' poll, and a SHOW INSTANCES command
    over the registry.
  • Integration suite: integration/add_shard/run.sh covers a parked
    cutover with a live omni writer, auto cutover, abort mid-task, the
    cross-instance lock, convergence after a restart with a stale flag, a
    two-instance coordinated cutover, a dead-peer refusal, and a
    schema-only run without omnisharded tables.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant