flyteorg / flyteorg/flyte

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

Open
#7,756 1 comment 0 reactions 0 assignees View on GitHub
enhancement flyte2 help wanted scale
Dominant language
Go
Stars
7.5k
Forks
886
Avg merge
1d 14h
Merged PRs (30d)
120

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`](https://github.com/flyteorg/flyte/blob/a2d7e82be425372e6d94252932b31800b34992f3/runs/repository/impl/action.go#L1078-L1090) **already drains** the channel into a batch — it just issues one
`execNotify` round trip per payload instead of one per batch:

```go
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**:

```sql
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`](https://github.com/flyteorg/flyte/blob/a2d7e82be425372e6d94252932b31800b34992f3/runs/repository/impl/action.go#L901) 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`](https://github.com/flyteorg/flyte/blob/a2d7e82be425372e6d94252932b31800b34992f3/runs/repository/impl/action.go#L117) 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`](https://github.com/flyteorg/flyte/blob/a2d7e82be425372e6d94252932b31800b34992f3/runs/repository/impl/action.go#L1062-L1076) currently discards a payload on any error:

```go
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.go` — [`drainAndExec`](https://github.com/flyteorg/flyte/blob/a2d7e82be425372e6d94252932b31800b34992f3/runs/repository/impl/action.go#L1078-L1090) / [`execNotify`](https://github.com/flyteorg/flyte/blob/a2d7e82be425372e6d94252932b31800b34992f3/runs/repository/impl/action.go#L1062-L1077)
inside [`runNotifyLoop`](https://github.com/flyteorg/flyte/blob/a2d7e82be425372e6d94252932b31800b34992f3/runs/repository/impl/action.go#L1039).
- `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`](https://github.com/flyteorg/flyte/blob/a2d7e82be425372e6d94252932b31800b34992f3/runs/repository/impl/action.go#L659),
[`WatchAllRunUpdates`](https://github.com/flyteorg/flyte/blob/a2d7e82be425372e6d94252932b31800b34992f3/runs/repository/impl/action.go#L699), [`WatchAllActionUpdates`](https://github.com/flyteorg/flyte/blob/a2d7e82be425372e6d94252932b31800b34992f3/runs/repository/impl/action.go#L751) and
[`WatchActionUpdates`](https://github.com/flyteorg/flyte/blob/a2d7e82be425372e6d94252932b31800b34992f3/runs/repository/impl/action.go#L818) 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](https://github.com/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`](https://github.com/flyteorg/flyte/blob/a2d7e82be425372e6d94252932b31800b34992f3/runs/repository/impl/action.go#L1039) — 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](https://github.com/flyteorg/flyte/blob/main/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

Contributor guide

Open the contributing guide

Research direction

Start at runNotifyLoop, drainAndExec, and execNotify in runs/repository/impl/action.go, then read the related tests in runs/repository/impl/action_test.go. Run go test ./runs/repository/impl/... with the configured local PostgreSQL instance. Done means drained payloads are sent in chunked batches with retries, all logical notifications remain delivered, and tests verify fewer round trips.

Written by the indexing model from the issue text.

Assessment

Tech stack
go, postgresql
Domain
backend, databases, performance, testing-qa
Issue type
Refactor
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.