flyteorg/flyte

perf: batch pg_notify round trips in the runs-service NOTIFY pump

Ouverte

#7 756 ouverte le 3 août 2026

 (1 commentaire) (0 réaction) (0 personne assignée)Python (378 forks)batch import
enhancementflyte2help wantedscale

Métriques du dépôt

Stars
 (3 705 étoiles)
Métriques de merge PR
 (Merge moyen 3j 8h) (114 PRs mergées en 30 j)

Description

Why

The runs-service sends one PostgreSQL NOTIFY per action update, and each one costs a full network round trip on a single dedicated connection. That serial round trip — not Postgres, not CPU — is what limits action-update throughput today.

Measured on a dev cluster during a 100,000-action benchmark:

RPC rate p50 p95 mean
InternalRunService.UpdateActionStatus 101.8/s 394 ms 3085 ms 780 ms
EventsProxyService.Record 25.6/s 653 ms 2418 ms 801 ms
InternalRunService.RecordActionEvents 25.6/s 652 ms 2418 ms 800 ms
ActionsService.Enqueue 57.5/s 1351 ms 375 ms
InternalRunService.RecordAction 52.9/s 326 ms 588 ms 288 ms

The mean is 780 ms, so this is the whole distribution being slow, not a tail. And it is not the usual suspects:

  • Not Postgres — RDS WriteLatency 0.9 ms, CPU 58–72%, and only 14 database connections open against a pool max of 100. Every connection sat in ClientRead, i.e. the database was waiting on the application.
  • Not CPU — the pod used 1.34 of its 4 cores with CFS throttling at 0.09% of periods.

By Little's law, 101.8 rps × 0.78 s ≈ 78 requests in flight — with an idle CPU and an idle connection pool. Those requests are blocked on something inside the process.

They are blocked on the NOTIFY pump. pg_stat_activity shows the single dedicated connection:

 pid  | state | wait_event | query
 7330 | idle  | ClientRead | SELECT pg_notify($1, $2)     <- the pump, one connection

Timing pg_notify from a pod in the same VPC shows where the cost actually is:

total per notify
1,000 × pg_notify inside one server-side DO loop 32.3 ms 32 µs
1 × pg_notify as its own statement 4.1 ms 4.1 ms

pg_notify itself costs 32 µs; the other ~4.07 ms is pure round trip. Since the pump pays that per payload, serially, its ceiling is ~240 notifications/second on an idle database — and it degrades as the database gets busier. Demand at the measured rates (UpdateActionStatus alone at 101.8/s, plus RecordAction at 52.9/s, plus six other call sites) lands at or above that ceiling.

The user-visible effect: a swarm of 50 concurrent runs × 2,000 tasks (100k actions) finished only 37/50 runs before timing out at 6,014 s, while the executor sat at 63% of its memory limit with zero restarts. The deployment is throughput-bound here, not memory-bound.

What to change

All in runs/repository/impl/action.go.

drainAndExec already drains the channel into a batch — it just issues one execNotify round trip per payload instead of one per batch:

drainAndExec := func(channel, firstPayload string, ch <-chan string) {
    execNotify(channel, firstPayload)        // <- one round trip
    for {
        select {
        case payload, ok := <-ch:
            if !ok { return }
            execNotify(channel, payload)     // <- one round trip each
        default:
            return
        }
    }
}

Send the whole drained batch in one multi-call statement:

SELECT pg_notify($1,$2), pg_notify($1,$3), pg_notify($1,$4), ...

The channel name is the same for every call in a batch, so it binds once as $1 and each payload adds one parameter. This keeps one notification per payload, which means the listener side is completely untouched — no payload framing, no splitting, no changes to processNotifications or to any of the Watch* consumers.

Two constraints to respect when building the statement:

  • Bind-parameter limit. PostgreSQL allows at most 65,535 parameters per statement, so a batch of N payloads uses N+1. Chunk well below that — a cap in the low thousands keeps the statement string small and parse time negligible. (InsertEvents already chunks for the same reason and is a good model.)
  • Payload size is per notification, not per batch — the 8000-byte NOTIFY limit applies to each individual payload, and these payloads are short identifiers, so it is not a concern here.

A useful property worth knowing: PostgreSQL documents that when the same channel is signalled multiple times with identical payloads inside one transaction, only one notification is delivered. Since a batched statement is a single implicit transaction, duplicate payloads within a batch collapse for free — the pump gets a degree of coalescing without writing any.

Retry the batch — don't widen the existing drop

Batching changes the blast radius of a failure, so it has to come with a retry. execNotify currently discards a payload on any error:

if conn == nil {
    logger.Errorf(ctx, "No NOTIFY connection available, dropping %s notification", channel)
    return                                    // payload lost
}
if _, err := conn.ExecContext(ctx, "SELECT pg_notify($1, $2)", channel, payload); err != nil {
    logger.Errorf(ctx, "Failed to NOTIFY %s: %v", channel, err)
    if isConnError(err) { reconnect() }
    // payload lost — no retry
}

Today a connection blip loses one notification. Batch 500 payloads into one statement and the same blip loses all 500 — 500 runs stuck at a stale phase in the UI until something else touches them. That is a real regression if batching lands on its own.

So: hold the drained payloads in a slice until the statement succeeds. On failure, reconnect and retry the same batch with a backoff instead of dropping it. Duplicate delivery is harmless — a notification is an idempotent wakeup that causes a re-read — so at-least-once is the right target and worth far more than avoiding a redundant wakeup.

If #7757 (the coalescing pending set) lands first, this becomes even simpler: merge the failed batch's keys back into the pending set, which is idempotent by construction. The two issues converge on the same retry story, so whoever goes second should reuse what the first built rather than adding a parallel mechanism.

Files:

  • runs/repository/impl/action.godrainAndExec / execNotify inside runNotifyLoop.
  • runs/repository/impl/action_test.go — add coverage that N queued updates produce N logical notifications with far fewer round trips.

The listener side must keep working unchanged: WatchRunUpdates, WatchAllRunUpdates, WatchAllActionUpdates and WatchActionUpdates all consume these notifications.

Outcome

  • A drained batch of N payloads costs one database round trip, not N
  • Every payload still reaches listeners — existing Watch* streams behave identically, with no changes needed on the listener side
  • Batches are chunked to stay well under PostgreSQL's 65,535 bind-parameter limit
  • A failed batch is retried, not dropped — batching must not turn one lost notification into N
  • A test covers the batching path (N queued notifications → 1 round trip, all delivered)
  • Ideally: a before/after number for UpdateActionStatus p95 under load

Getting started

  • Reproduce: run a wide fan-out against a dev cluster and watch rpc_server_duration_milliseconds for UpdateActionStatus. The flyteorg/benchmark suite does this — uv run scripts/v2/swarm.py --k 25 --n 2000 generates ~100 action-updates/second.
  • See the ceiling yourself: against any Postgres, compare DO $$ BEGIN FOR i IN 1..1000 LOOP PERFORM pg_notify('probe','x'); END LOOP; END $$; (one round trip) with 1,000 separate SELECT pg_notify('probe','x'); statements.
  • Relevant code: start at runNotifyLoop — the pump is about 70 lines.
  • How to test: go test ./runs/repository/impl/... — needs a local Postgres on port 15432 (user/password: postgres, db flyte_runs_test; see testDbConfig in action_test.go).
  • Setup: CONTRIBUTING.md

Related — three issues on the same bottleneck:

  • #7756 batch the pump's round trips (this one)
  • #7757 don't block action writes on the pump
  • #7758 instrument the pump

Guide contributeur