flyteorg/flyte

perf: don't block action writes on the NOTIFY pump

Open

#7,757 opened on Aug 3, 2026

 (1 comment) (0 reactions) (0 assignees)Python (378 forks)batch import
enhancementflyte2help wantedscale

Repository metrics

Stars
 (3,705 stars)
PR merge metrics
 (Avg merge 3d 8h) (114 merged PRs in 30d)

Description

Why

A PostgreSQL NOTIFY in the runs-service is a wakeup hint — the payload carries only an action identifier, and every listener re-reads state from the database when it wakes. Losing one is recoverable. But today, sending one can block a write RPC indefinitely.

notifyActionUpdate does a blocking channel send:

select {
case r.actionNotifyCh <- payload:
case <-ctx.Done():
    logger.Errorf(ctx, "Action NOTIFY send cancelled for %s: %v", payload, ctx.Err())
}

The channel is buffered at 256 and drained by a single goroutine that round-trips to Postgres once per payload — a ceiling of roughly 240 notifications/second (measured: 4.1 ms per pg_notify round trip from the same VPC, versus 32 µs for the pg_notify call itself). When demand exceeds that, the buffer fills and the send blocks — inside the RPC handler, because UpdateActionPhase calls it right after the row is written.

What is not at risk: the database write

Worth stating up front, because it is the first question everyone asks. The status update is already committed before the notification is attempted — every one of the seven notify call sites runs after its DB statement returns, and UpdateActionPhase additionally gates on rowsAffected > 0:

result, err := r.db.ExecContext(ctx, queryBuilder.String(), args...)  // committed here
if err != nil { return err }                                         // no notify on failure
rowsAffected, err := result.RowsAffected()
if rowsAffected > 0 {
    r.notifyActionUpdate(ctx, actionID)   // only after the row is written
}

These are autocommit single statements — there is no transaction spanning the write and the notification. So blocking on the channel does not make the write more likely to succeed; it only delays the RPC's return after the write already succeeded.

What is at risk is watcher liveness: WatchActionUpdates / WatchAllActionUpdates take their wakeup from this notification and re-read state from the database when woken. Lose a wakeup and a stream — and therefore the phase a user sees in the UI — can lag until something else touches that action. So the fix must not block the write path and must not lose a wakeup. The design below does both; see "What to change".

The result is that watch fan-out latency becomes write latency. Measured on a dev cluster under a 100,000-action benchmark:

RPC rate p50 p95 mean
InternalRunService.UpdateActionStatus 101.8/s 394 ms 3085 ms 780 ms
InternalRunService.RecordActionEvents 25.6/s 652 ms 2418 ms 800 ms
ActionsService.Enqueue 57.5/s 1351 ms 375 ms

Meanwhile the database was idle (WriteLatency 0.9 ms, 14 connections open out of 100, all in ClientRead) and the pod used 1.34 of 4 cores with 0.09% CPU throttling. About 78 requests were in flight (Little's law: 101.8 rps × 0.78 s) waiting on neither CPU nor Postgres.

Batching the pump's round trips (see the companion issue) raises the ceiling, but the structural problem remains: a durable write should not wait on a best-effort notification. As long as the send can block, any future slowdown in the notify path silently becomes user-visible write latency.

What to change

The fix must satisfy two constraints at once: never block the write path, and never lose a wakeup — a lost wakeup means a run sits at a stale phase in the UI until something else happens to that action.

Both are achievable, because the notification payload is an identity (project/domain/run/action), not a state. Listeners re-read from the database when woken. So collapsing five notifications for the same action into one is not a partial delivery — the single wakeup still causes a read of the latest state. What must never be lost is the last wakeup for a given action.

That points at replacing the FIFO channel with a coalescing pending set:

// today:
actionNotifyCh chan string          // 256 FIFO slots, blocking send

// proposed:
pendingActions map[string]struct{}  // mutex-guarded; key = the payload
pendingCh      chan struct{}        // capacity 1, non-blocking "something is pending" signal

(Name the signal channel to match the surrounding style — pendingCh alongside pendingActions reads as "what is pending" plus "a nudge that something is pending". The existing fields use the ...Ch suffix, so wakeCh/signalCh would also be fine; just avoid a bare wake.)

  • Writer (notifyActionUpdate): take the lock, insert the key, unlock, then a non-blocking send on pendingCh. It cannot block and cannot discard a key. A signal that is already buffered needs no second one — the map is the queue, the channel is only a nudge.
  • Pump (runNotifyLoop): on each signal, swap the map out under the lock, then emit the whole set — ideally as one batched statement, which is exactly what the companion batching issue does. The two changes compose.

Delivery: keep the key until the NOTIFY succeeds

"If the writer doesn't block, how do we know the NOTIFY was actually sent?" — worth answering directly, because blocking does not answer it today either. execNotify discards the payload on any failure:

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
}

The blocking send only guarantees that the payload reached a 256-slot buffer. Once the pump takes it, a connection blip drops it silently. So delivery is best-effort today, and a run can already sit at a stale phase in the UI after a transient database hiccup.

The pending set fixes this, and that is the main reason to prefer it over a bigger buffer:

  1. The pump swaps the pending set out into a batch under the lock.
  2. It attempts the batched pg_notify.
  3. On success the keys are gone.
  4. On failure it merges the keys back into the pending set, reconnects, and retries with backoff.

Step 4 is safe because it is a set: re-adding a key is idempotent, so a retry cannot duplicate or double-count, and any new updates for the same action that arrived meanwhile merge into the same entry. A FIFO would force a choice between dropping and unbounded re-queueing. The result is at-least-once delivery for the life of the process — the right target, since a notification is an idempotent wakeup and a duplicate costs one redundant re-read.

What about a runs-service restart? In-memory pending entries are lost, but that does not strand the UI, because every new watch stream starts from the database rather than from notifications. WatchActions subscribes first, then sends a full snapshot paged out of the actions table, and only then streams deltas:

// Start watching for updates from DB first to prevent event miss
go s.repo.ActionRepo().WatchAllActionUpdates(ctx, runID, updatesCh, errsCh)
...
if err := s.listAndSendAllActions(ctx, runID, rsm, stream); err != nil {   // full snapshot

So a restart breaks the streams, clients reconnect, and each reconnect re-reads authoritative state that was committed long before any notification was attempted. Losing in-flight wakeups costs a reconnect window, not a permanently stale run. This is also why the notification path can be best-effort at all: it is an optimization over polling, and Postgres is the source of truth.

Why this is better than a bigger buffer or a drop policy:

today drop when full pending set
Blocks the write path yes no no
Can lose a wakeup yes, on ctx.Done() yes, under pressure no
Queue bounded by 256 updates 256 updates distinct actions

The memory bound is the key property: the set can never hold more entries than there are distinct actions with an unsent notification, no matter how fast updates arrive. At 65,000 live actions that is a few MB.

It also removes a loss path that exists today: the current select discards the notification on ctx.Done(), so a client disconnecting mid-request can lose a wakeup that other watchers needed. A non-blocking send has no reason to consult the request context at all.

Files:

Ordering note: notifyRunUpdate needs the same treatment, and the two sets can share one pump.

Outcome

  • UpdateActionPhase never blocks on notification delivery, however slow or stalled the pump is
  • No wakeup is lost: for every action that was updated, at least one notification is delivered after its last update
  • A failed pg_notify retries rather than dropping the payload — keys stay pending until the statement succeeds (strictly better than today's behaviour)
  • Repeated notifications for the same action between drains collapse into one
  • Request-context cancellation no longer discards a notification
  • Memory is bounded by distinct pending actions rather than by update volume
  • Tests cover: pump stalled → writer returns immediately and the row is still written; N updates to one action between drains → one notification; every updated action is represented in the delivered set

Getting started

  • Reproduce: generate ~100 action-updates/second against a dev cluster and watch rpc_server_duration_milliseconds for UpdateActionStatus. The flyteorg/benchmark suite does this with uv run scripts/v2/swarm.py --k 25 --n 2000.
  • Relevant code: notifyActionUpdate, then the pump in runNotifyLoop, then the listener side in processNotifications.
  • 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
  • #7757 don't block action writes on the pump (this one)
  • #7758 instrument the pump

Contributor guide