Citus Shard Imbalance and Rebalance Guard

id: citus-shard-imbalance-rebalance-guard
namespace: company.team

description: |
  Measure how a Citus cluster spreads its distributed shards over the workers that accept them.
  Report a worker that holds nothing, which is what a freshly added node looks like, plus a byte
  skew above a threshold. Print the exact moves Citus would make as a dry run, then optionally
  run the rebalance behind an approval, poll it to completion and read the placement again,
  failing the run if the skew did not improve or a worker is still empty.

inputs:
  - id: jdbc_url
    type: STRING
    displayName: Coordinator JDBC URL
    description: The Citus coordinator. Worker nodes are read from its metadata.
    defaults: "jdbc:postgresql://localhost:5432/citus"

  - id: skew_threshold
    type: FLOAT
    displayName: Byte skew that counts as imbalanced
    description: Largest worker bytes divided by smallest worker bytes. 1.25 means
      the hottest worker may hold a quarter more than the coldest before the
      cluster is called imbalanced.
    defaults: 1.25

  - id: min_total_mb
    type: INT
    displayName: Floor on total shard bytes
    description: Below this, skew is not worth moving data for. An empty worker is
      still reported under the floor.
    defaults: 10

  - id: rebalance_strategy
    type: SELECT
    displayName: Rebalance strategy
    description: Citus built-in strategies. by_disk_size balances bytes,
      by_shard_count balances the number of shards.
    values:
      - by_disk_size
      - by_shard_count
    defaults: by_disk_size

  - id: shard_transfer_mode
    type: SELECT
    displayName: Shard transfer mode
    description: >-
      auto uses logical replication when the shard has a replica identity, which
      needs wal_level=logical on the workers. block_writes copies the shard with
      writes blocked on it and works on any wal_level. force_logical always uses
      logical replication.
    values:
      - auto
      - block_writes
      - force_logical
    defaults: auto

  - id: max_shard_moves
    type: INT
    displayName: Maximum shard moves
    description: Cap on the moves in one rebalance, so a large cluster can be
      levelled over several windows.
    defaults: 1000000

  - id: rebalance_timeout
    type: DURATION
    displayName: Rebalance timeout
    description: How long to poll citus_rebalance_status before failing the run. The
      rebalance keeps going in the background, so a timeout is a reporting
      limit, not a cancel.
    defaults: PT30M

  - id: auto_remediate
    type: BOOL
    displayName: Auto remediate
    description: Allow the flow to start a rebalance. Off means report only.
    defaults: false

  - id: dry_run
    type: BOOL
    displayName: Dry run
    description: Must be false, together with auto_remediate, before any shard moves.
    defaults: true

variables:
  placement_sql: |
    WITH node AS (
      SELECT n.groupid, n.nodename, n.nodeport, n.shouldhaveshards
        FROM pg_dist_node n
       WHERE n.noderole = 'primary' AND n.isactive AND n.groupid <> 0
    ), shard_size AS (
      SELECT shard_id, max(size) AS bytes FROM citus_shard_sizes() GROUP BY shard_id
    ), rebalanceable AS (
      SELECT logicalrelid, (repmodel <> 't' AND colocationid <> 0) AS movable FROM pg_dist_partition
    ), placed AS (
      SELECT p.groupid,
             count(*) FILTER (WHERE r.movable) AS shard_count,
             count(*) FILTER (WHERE NOT r.movable) AS pinned_shards,
             coalesce(sum(z.bytes) FILTER (WHERE r.movable), 0) AS bytes
        FROM pg_dist_placement p
        JOIN pg_dist_shard s ON s.shardid = p.shardid
        JOIN rebalanceable r ON r.logicalrelid = s.logicalrelid
        LEFT JOIN shard_size z ON z.shard_id = p.shardid
       WHERE p.shardstate = 1
       GROUP BY p.groupid
    ), per_node AS (
      SELECT n.nodename, n.nodeport, n.shouldhaveshards,
             coalesce(pl.shard_count, 0) AS shard_count,
             coalesce(pl.pinned_shards, 0) AS pinned_shards,
             coalesce(pl.bytes, 0) AS bytes
        FROM node n LEFT JOIN placed pl ON pl.groupid = n.groupid
    )

  assess_sql: |
    {{ render(vars.placement_sql) }}, agg AS (
      SELECT count(*) AS worker_count,
             count(*) FILTER (WHERE shouldhaveshards) AS target_count,
             coalesce(sum(bytes), 0) AS total_bytes,
             coalesce(sum(shard_count), 0) AS total_shards,
             coalesce(sum(pinned_shards), 0) AS pinned_shards,
             coalesce(max(bytes) FILTER (WHERE shouldhaveshards), 0) AS max_bytes,
             coalesce(min(bytes) FILTER (WHERE shouldhaveshards), 0) AS min_bytes,
             count(*) FILTER (WHERE shouldhaveshards AND shard_count = 0) AS empty_count
        FROM per_node
    ), job AS (
      SELECT job_id, state::text AS state FROM citus_rebalance_status()
       WHERE state::text IN ('scheduled', 'running', 'cancelling', 'failing')
       ORDER BY job_id DESC LIMIT 1
    )
    SELECT
      CASE
        WHEN (SELECT count(*) FROM job) > 0 THEN 'REBALANCE_IN_PROGRESS'
        WHEN a.target_count < 2 THEN 'OK'
        WHEN a.empty_count > 0 AND a.total_bytes > 0 THEN 'EMPTY_WORKER'
        WHEN a.total_bytes >= {{ inputs.min_total_mb }}::bigint * 1048576
             AND a.max_bytes::numeric / greatest(a.min_bytes, 1) > {{ inputs.skew_threshold }} THEN 'IMBALANCED'
        ELSE 'OK'
      END AS state,
      a.worker_count, a.target_count, a.total_shards, a.pinned_shards, a.empty_count,
      a.total_bytes, pg_size_pretty(a.total_bytes) AS total_size,
      a.max_bytes, pg_size_pretty(a.max_bytes) AS max_size,
      a.min_bytes, pg_size_pretty(a.min_bytes) AS min_size,
      least(a.max_bytes::numeric * 10000 / greatest(a.min_bytes, 1), 100000000)::bigint AS skew_bp,
      CASE WHEN a.min_bytes = 0 THEN 'unbounded (a worker holds nothing)'
           ELSE to_char(a.max_bytes::numeric / a.min_bytes, 'FM999999990.00') || 'x' END AS skew_ratio,
      to_char(100.0 * a.max_bytes / greatest(a.total_bytes, 1), 'FM990.0') AS max_share_pct,
      to_char(100.0 / greatest(a.target_count, 1), 'FM990.0') AS fair_share_pct,
      coalesce((SELECT string_agg(nodename || ':' || nodeport, ', ' ORDER BY nodename)
                  FROM per_node WHERE shouldhaveshards AND shard_count = 0), '') AS empty_workers,
      coalesce((SELECT nodename || ':' || nodeport FROM per_node WHERE shouldhaveshards
                 ORDER BY bytes DESC, nodename LIMIT 1), '') AS hottest_worker,
      coalesce((SELECT nodename || ':' || nodeport FROM per_node WHERE shouldhaveshards
                 ORDER BY bytes ASC, nodename LIMIT 1), '') AS coldest_worker,
      coalesce((SELECT job_id FROM job), 0) AS running_job_id,
      coalesce((SELECT state FROM job), '') AS running_job_state,
      CASE WHEN a.target_count < 2
             THEN 'Only ' || a.target_count || ' worker accepts shards, so there is nothing to balance.'
           WHEN a.empty_count = 0 AND a.total_bytes < {{ inputs.min_total_mb }}::bigint * 1048576
             THEN 'Total shard bytes are under the {{ inputs.min_total_mb }} MB floor, so skew is ignored.'
           ELSE '' END AS note
      FROM agg a

  distribution_sql: |
    {{ render(vars.placement_sql) }}
    SELECT nodename || ':' || nodeport AS worker, shouldhaveshards AS accepts_shards,
           shard_count, pinned_shards, bytes, pg_size_pretty(bytes) AS size,
           to_char(100.0 * bytes / greatest((SELECT sum(bytes) FROM per_node), 1), 'FM990.0') AS share_pct
      FROM per_node ORDER BY bytes DESC, nodename

tasks:
  - id: distribution
    type: io.kestra.plugin.jdbc.postgresql.Query
    description: >-
      Read-only. One row per active primary worker, built from
      pg_dist_placement, pg_dist_shard, pg_dist_partition and citus_shard_sizes
      so a worker holding nothing still gets a row. Reference table and Citus
      local table shards are counted separately as pinned, because the
      rebalancer never moves them.
    url: "{{ inputs.jdbc_url }}"
    username: "{{ secret('POSTGRES_USERNAME') }}"
    password: "{{ secret('POSTGRES_PASSWORD') }}"
    fetchType: FETCH
    sql: "{{ render(vars.distribution_sql) }}"

  - id: assess
    type: io.kestra.plugin.jdbc.postgresql.Query
    description: >-
      Read-only. Classify the cluster as REBALANCE_IN_PROGRESS, EMPTY_WORKER,
      IMBALANCED or OK, and return the numbers behind the verdict.
    url: "{{ inputs.jdbc_url }}"
    username: "{{ secret('POSTGRES_USERNAME') }}"
    password: "{{ secret('POSTGRES_PASSWORD') }}"
    fetchType: FETCH_ONE
    sql: "{{ render(vars.assess_sql) }}"

  - id: report
    type: io.kestra.plugin.core.log.Log
    message: >-
      Citus placement on {{ inputs.jdbc_url }} is {{ outputs.assess.row.state
      }}. {{ outputs.assess.row.total_shards }} movable shards holding {{
      outputs.assess.row.total_size }} over {{ outputs.assess.row.target_count
      }} of {{ outputs.assess.row.worker_count }} active workers. Hottest {{
      outputs.assess.row.hottest_worker }} at {{ outputs.assess.row.max_size }}
      ({{ outputs.assess.row.max_share_pct }}% of bytes against a fair share of
      {{ outputs.assess.row.fair_share_pct }}%), coldest {{
      outputs.assess.row.coldest_worker }} at {{ outputs.assess.row.min_size }},
      skew {{ outputs.assess.row.skew_ratio }} against a threshold of {{
      inputs.skew_threshold }}x. {{ outputs.assess.row.empty_workers == '' ? ''
      : 'Workers holding no movable shard: ' ~ outputs.assess.row.empty_workers
      ~ '. ' }}{{ outputs.assess.row.running_job_id > 0 ? 'Rebalance job ' ~
      outputs.assess.row.running_job_id ~ ' is ' ~
      outputs.assess.row.running_job_state ~ '. ' : '' }}{{
      outputs.assess.row.note }} {% for r in outputs.distribution.rows %}[{{
      r.worker }} shards={{ r.shard_count }} {{ r.size }} {{ r.share_pct }}%{{
      r.pinned_shards > 0 ? ' pinned=' ~ r.pinned_shards : '' }}{{
      r.accepts_shards ? '' : ' draining' }}] {% endfor %}

  - id: route
    type: io.kestra.plugin.core.flow.If
    condition: "{{ outputs.assess.row.state != 'OK' }}"
    then:
      - id: alert
        type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
        url: "{{ secret('SLACK_WEBHOOK_URL') }}"
        payload: |
          {{ {'text': 'Citus ' ~ inputs.jdbc_url ~ ' is ' ~ outputs.assess.row.state ~ '. ' ~ outputs.assess.row.total_shards ~ ' movable shards holding ' ~ outputs.assess.row.total_size ~ ' over ' ~ outputs.assess.row.target_count ~ ' workers. Hottest ' ~ outputs.assess.row.hottest_worker ~ ' at ' ~ outputs.assess.row.max_size ~ ' (' ~ outputs.assess.row.max_share_pct ~ '% of bytes), coldest ' ~ outputs.assess.row.coldest_worker ~ ' at ' ~ outputs.assess.row.min_size ~ ', skew ' ~ outputs.assess.row.skew_ratio ~ '.' ~ (outputs.assess.row.empty_workers == '' ? '' : ' Holding nothing: ' ~ outputs.assess.row.empty_workers ~ '.') ~ (inputs.auto_remediate and not inputs.dry_run ? ' Rebalance pending approval.' : ' No change applied (auto_remediate=' ~ inputs.auto_remediate ~ ', dry_run=' ~ inputs.dry_run ~ ').')} | toJson }}

      - id: already_running
        type: io.kestra.plugin.core.flow.If
        condition: "{{ outputs.assess.row.state == 'REBALANCE_IN_PROGRESS' }}"
        then:
          - id: leave_it_alone
            type: io.kestra.plugin.core.log.Log
            message: >-
              Rebalance job {{ outputs.assess.row.running_job_id }} is {{
              outputs.assess.row.running_job_state }}. Citus runs one rebalance
              at a time, so this run reports and stops. Re-run once the job
              finishes.
        else:
          - id: plan
            type: io.kestra.plugin.jdbc.postgresql.Query
            description: >-
              The dry run. get_rebalance_table_shards_plan returns the moves
              Citus would make and changes nothing. Its own shard_size column is
              always 0, so the real size comes from citus_shard_sizes.
            url: "{{ inputs.jdbc_url }}"
            username: "{{ secret('POSTGRES_USERNAME') }}"
            password: "{{ secret('POSTGRES_PASSWORD') }}"
            fetchType: FETCH
            sql: >-
              SELECT p.table_name::text AS table_name, p.shardid,
                     p.sourcename || ':' || p.sourceport AS source,
                     p.targetname || ':' || p.targetport AS target,
                     coalesce(z.bytes, 0) AS bytes, pg_size_pretty(coalesce(z.bytes, 0)) AS size
                FROM get_rebalance_table_shards_plan(NULL, NULL, {{ inputs.max_shard_moves }},
                       '{}'::bigint[], false, '{{ inputs.rebalance_strategy }}') p
                LEFT JOIN (SELECT shard_id, max(size) AS bytes FROM citus_shard_sizes() GROUP BY shard_id) z
                  ON z.shard_id = p.shardid
               ORDER BY coalesce(z.bytes, 0) DESC, p.shardid

          - id: log_plan
            type: io.kestra.plugin.core.log.Log
            message: >-
              Planned moves with strategy {{ inputs.rebalance_strategy }}: {{
              outputs.plan.size }} shards, {{ outputs.plan.rows |
              jq('map(.bytes) | add // 0 | . / 1048576 | . * 10 | floor | . /
              10') | first }} MB. {% for r in outputs.plan.rows %}{{
              r.table_name }} shard {{ r.shardid }} ({{ r.size }}) {{ r.source
              }} to {{ r.target }}. {% endfor %}

          - id: remediate_gate
            type: io.kestra.plugin.core.flow.If
            condition: "{{ inputs.auto_remediate and not inputs.dry_run and
              outputs.plan.size > 0 }}"
            then:
              - id: approve
                type: io.kestra.plugin.core.flow.Pause
                description: Resume the execution to run the moves listed by log_plan. Kill it
                  to cancel. A shard move copies data between workers, so it is
                  never applied without this.

              - id: start_rebalance
                type: io.kestra.plugin.jdbc.postgresql.Query
                description: citus_rebalance_start queues the moves as a background job and
                  returns its id straight away.
                url: "{{ inputs.jdbc_url }}"
                username: "{{ secret('POSTGRES_USERNAME') }}"
                password: "{{ secret('POSTGRES_PASSWORD') }}"
                fetchType: FETCH_ONE
                sql: >-
                  SELECT coalesce(citus_rebalance_start('{{
                  inputs.rebalance_strategy }}', false,
                    '{{ inputs.shard_transfer_mode }}'), 0) AS job_id

              - id: wait_rebalance
                type: io.kestra.plugin.core.flow.LoopUntil
                description: Poll the job every 10 seconds until Citus reports it finished,
                  failed or cancelled.
                condition: "{{ outputs.poll_rebalance.row.terminal }}"
                checkFrequency:
                  interval: PT10S
                  maxDuration: "{{ inputs.rebalance_timeout }}"
                failOnMaxReached: true
                tasks:
                  - id: poll_rebalance
                    type: io.kestra.plugin.jdbc.postgresql.Query
                    url: "{{ inputs.jdbc_url }}"
                    username: "{{ secret('POSTGRES_USERNAME') }}"
                    password: "{{ secret('POSTGRES_PASSWORD') }}"
                    fetchType: FETCH_ONE
                    sql: >-
                      WITH j AS (
                        SELECT state::text AS state, details FROM citus_rebalance_status()
                         WHERE job_id = {{ outputs.start_rebalance.row.job_id }}
                      ) SELECT coalesce((SELECT state FROM j), 'missing') AS
                      state,
                             coalesce((SELECT state FROM j), 'missing')
                               NOT IN ('scheduled', 'running', 'cancelling', 'failing') AS terminal,
                             coalesce((SELECT details::text FROM j), 'none') AS details

              - id: job_result
                type: io.kestra.plugin.core.flow.If
                condition: "{{ outputs.poll_rebalance.row.state != 'finished' }}"
                then:
                  - id: job_failed
                    type: io.kestra.plugin.core.execution.Fail
                    errorMessage: >-
                      Rebalance job {{ outputs.start_rebalance.row.job_id }}
                      ended as {{ outputs.poll_rebalance.row.state }} with
                      details {{ outputs.poll_rebalance.row.details }}.
                      citus_rebalance_status only reports the newest rebalance
                      job, so a state of missing means another rebalance
                      replaced this one. Read pg_dist_background_task on the
                      coordinator.

              - id: redistribution
                type: io.kestra.plugin.jdbc.postgresql.Query
                description: Read the placement again, after the moves.
                url: "{{ inputs.jdbc_url }}"
                username: "{{ secret('POSTGRES_USERNAME') }}"
                password: "{{ secret('POSTGRES_PASSWORD') }}"
                fetchType: FETCH
                sql: "{{ render(vars.distribution_sql) }}"

              - id: reassess
                type: io.kestra.plugin.jdbc.postgresql.Query
                description: >-
                  The same assessment query, wrapped so the before numbers
                  decide whether the move helped. improved is false unless every
                  worker that accepts shards now holds at least one and the skew
                  either dropped or reached OK.
                url: "{{ inputs.jdbc_url }}"
                username: "{{ secret('POSTGRES_USERNAME') }}"
                password: "{{ secret('POSTGRES_PASSWORD') }}"
                fetchType: FETCH_ONE
                sql: >-
                  WITH after AS (
                    {{ render(vars.assess_sql) }}
                  ) SELECT a.*, {{ outputs.assess.row.skew_bp }} AS
                  before_skew_bp,
                         {{ outputs.assess.row.empty_count }} AS before_empty_count,
                         (a.empty_count = 0
                          AND (a.skew_bp < {{ outputs.assess.row.skew_bp }} OR a.state = 'OK')) AS improved
                    FROM after a

              - id: verify
                type: io.kestra.plugin.core.flow.If
                condition: "{{ outputs.reassess.row.improved }}"
                then:
                  - id: rebalanced
                    type: io.kestra.plugin.core.log.Log
                    message: >-
                      Rebalance job {{ outputs.start_rebalance.row.job_id }}
                      moved {{ outputs.plan.size }} shards. Skew {{
                      outputs.assess.row.skew_ratio }} before, {{
                      outputs.reassess.row.skew_ratio }} after. Hottest worker
                      share {{ outputs.assess.row.max_share_pct }}% before, {{
                      outputs.reassess.row.max_share_pct }}% after against a
                      fair share of {{ outputs.reassess.row.fair_share_pct }}%.
                      Workers holding nothing: {{ outputs.assess.row.empty_count
                      }} before, {{ outputs.reassess.row.empty_count }} after.
                      Cluster is now {{ outputs.reassess.row.state }}. {{
                      outputs.reassess.row.state == 'OK' ? '' : 'Bytes are still
                      skewed past skew_threshold, which happens when
                      citus.rebalancer_by_disk_size_base_cost, 100MB by default,
                      dominates the real shard sizes and turns by_disk_size into
                      a shard count balance. ' }}{% for r in
                      outputs.redistribution.rows %}[{{ r.worker }} shards={{
                      r.shard_count }} {{ r.size }} {{ r.share_pct }}%] {%
                      endfor %}
                else:
                  - id: remediation_failed
                    type: io.kestra.plugin.core.execution.Fail
                    errorMessage: >-
                      Rebalance job {{ outputs.start_rebalance.row.job_id }}
                      reported finished, but the placement did not improve. Skew
                      {{ outputs.assess.row.skew_ratio }} before, {{
                      outputs.reassess.row.skew_ratio }} after. Workers holding
                      no movable shard: {{ outputs.assess.row.empty_count }}
                      before, {{ outputs.reassess.row.empty_count }} after.
                      Cluster is {{ outputs.reassess.row.state }}. Check
                      pg_dist_background_task for skipped moves.
            else:
              - id: reported_only
                type: io.kestra.plugin.core.log.Log
                message: >-
                  {{ outputs.plan.size }} shard moves reported, none applied. {{
                  outputs.plan.size > 0 ? 'auto_remediate=' ~
                  inputs.auto_remediate ~ ', dry_run=' ~ inputs.dry_run ~ '.' :
                  'Citus itself plans no move, so there is nothing to apply. The
                  ' ~ inputs.rebalance_strategy ~ ' strategy finds this
                  placement acceptable while the byte skew is ' ~
                  outputs.assess.row.skew_ratio ~ '. by_disk_size adds
                  citus.rebalancer_by_disk_size_base_cost, 100MB by default, to
                  every shard group before it compares nodes, so a cluster whose
                  shards are smaller than that balances on shard count instead
                  of bytes. Lower that setting on the coordinator. Raise
                  skew_threshold instead if this placement should read as OK.'
                  }}
    else:
      - id: balanced
        type: io.kestra.plugin.core.log.Log
        message: >-
          Placement is within {{ inputs.skew_threshold }}x across {{
          outputs.assess.row.target_count }} workers, and every worker that
          accepts shards holds at least one. {{ outputs.assess.row.note }}

errors:
  - id: alert_flow_failure
    type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
    url: "{{ secret('SLACK_WEBHOOK_URL') }}"
    payload: |
      {{ {'text': 'Citus shard imbalance guard failed on ' ~ inputs.jdbc_url ~ ' (execution ' ~ execution.id ~ '). Check the logs.'} | toJson }}

triggers:
  - id: daily
    type: io.kestra.plugin.core.trigger.Schedule
    description: Shipped disabled and always in safe mode.
    cron: "0 2 * * *"
    disabled: true
    inputs:
      auto_remediate: false
      dry_run: true