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

[pingsutw]

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:

• runs/repository/impl/action.go — notifyActionUpdate,
  notifyRunUpdate, the channel declarations and their
  construction, runNotifyLoop and drainAndExec.
• runs/repository/impl/action_test.go — cover the properties below.

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:

• perf: batch pg_notify round trips in the runs-service NOTIFY pump #7756 batch the pump's round trips
• perf: don't block action writes on the NOTIFY pump #7757 don't block action writes on the pump (this one)
• feat: add metrics for the runs-service NOTIFY pump (pending depth, notify latency, coalescing) #7758 instrument the pump

[runllm]

It seems likely that the current implementation of Flyte's notification system for action updates uses a buffered channel (actionNotifyCh) that can block the write path when full, as described in your issue. The database write is already committed before the notification is attempted, so blocking only delays the RPC's return, not the write itself. The risk is that watcher liveness (UI updates) can lag if notifications are lost or delayed, but the write itself is safe. This is confirmed by the code and comments in runs/repository/impl/action.go and the test coverage in action_test.go.

My understanding is that your proposed solution—to replace the FIFO channel with a mutex-guarded set of pending notifications and a non-blocking wake channel—would ensure that the write path never blocks and that no wakeup is lost, since only the last notification for each action is needed. This approach is not only feasible but aligns with the architectural intent, as listeners always re-read the latest state from the database. The memory usage would be bounded by the number of distinct actions, not the volume of updates, and the tests you suggest would verify the new behavior. For more details on the notification and watcher mechanisms, see the relevant code and schema in action.go and the implementation spec (IMPLEMENTATION_SPEC.md).

Would you like a step-by-step breakdown of how to implement this change in the codebase?

Sources:

• runs/repository/impl/action.go
• action_test.go
• IMPLEMENTATION_SPEC.md

---

This is an AI-generated response and your feedback is appreciated! Please leave a 👍 if this is helpful and 👎 if it is not.

Mention @Herald to ask follow-up questions. See runllm.com for more!

[0yukali0]

#take