Skip to content

perf(ai-proxy): optimize SSE decoder - remove PCRE, add decode_buf, fix comment lines#13391

Merged
AlinsRan merged 17 commits into
apache:masterfrom
AlinsRan:perf/ai-sse-optimize
May 22, 2026
Merged

perf(ai-proxy): optimize SSE decoder - remove PCRE, add decode_buf, fix comment lines#13391
AlinsRan merged 17 commits into
apache:masterfrom
AlinsRan:perf/ai-sse-optimize

Conversation

@AlinsRan

@AlinsRan AlinsRan commented May 19, 2026

Copy link
Copy Markdown
Contributor

What

Optimize the AI proxy SSE streaming hot path: fewer allocations, less CPU per chunk, and a new interval-flush mode for bursty LLM streams.

Changes

# Change Approximate Benefit
1 Replace ngx_re.split (PCRE) with string.find in SSE decoder Eliminates PCRE JIT compilation on every chunk (~1–5 µs/chunk saved)
2 New decode_buf(buf) — combines split_buf + decode in one forward pass ~50% fewer byte-scan iterations per chunk on the hot path
3 Fix SSE comment lines (:) being emitted as empty events Correctness fix; providers using : keep-alive pings no longer produce spurious events
4 Replace buf = buf .. chunk with a 2-element table + table.concat O(n) → O(1) allocation per chunk; a 100-chunk stream drops from ~250 KB of copies to ~5 KB
5 Add streaming_flush_interval_ms config option (default 10 ms) Reduces ngx.flush syscalls by ~95% on bursty streams (100 chunks → 2–4 flushes at 10 ms)
6 Add no_flush param to lua_response_filter Enables background-thread flush ownership (required by #5)

streaming_flush_interval_ms example

{
  "provider": "openai",
  "streaming_flush_interval_ms": 10
}
Value Behavior
10 (default) Background thread calls ngx.flush(false) every 10 ms — batches output on bursty streams
> 0 (e.g. 50) Background thread flushes every 50 ms; lua_response_filter skips ngx.flush per chunk
0 Background thread disabled; each chunk flushed synchronously via ngx.flush(true)

The background thread only flushes when there is new data (via a needs_flush flag), so empty intervals cost nothing.

…ix comment lines

- Replace ngx_re.split (PCRE) in sse.decode() with a pure Lua forward scan.
  Avoids regex compilation and PCRE overhead on every streaming chunk.
- Add sse.decode_buf(buf) that combines split_buf + decode into a single
  O(n) forward pass, halving the number of scans per chunk in hot paths.
- Fix SSE comment lines (lines starting with ':') being incorrectly treated
  as fields. Per the SSE spec, comment lines must be silently ignored; the
  previous code set has_field=true for them, producing spurious empty events.
- Add test: SSE comment lines are correctly discarded by decode_buf.

Co-authored-by: Copilot <[email protected]>
@dosubot dosubot Bot added size:L This PR changes 100-499 lines, ignoring generated files. performance generate flamegraph for the current PR labels May 19, 2026
…path

- base.lua: replace sse_buf string concatenation with sse_parts table to
  avoid O(n) string allocation on every chunk; use table.insert + concat
  only when needed. Also adds ngx.sleep(0) yield at end of each iteration
  so health checks and concurrent requests are not starved under bursty
  SSE upstreams.

- base.lua: add streaming_flush_interval_ms support - when > 0, a
  background ngx.thread flushes the output buffer every N milliseconds
  instead of per-chunk. Useful when the upstream bursts multiple tokens
  at once and you want to bound client latency without synchronous flush
  overhead on every chunk. Thread spawn failure falls back to sync flush.
  flush_thread is killed on all exit paths (EOF, error, 502, limits).

- plugin.lua: add no_flush parameter to lua_response_filter(). When true,
  ngx.flush is skipped so the background interval thread can own flushing.
  Backward compatible - existing callers passing (ctx, headers, body, wait)
  are unaffected.

- schema.lua: add streaming_flush_interval_ms field (integer, min=0,
  default=0) to both ai_proxy_schema and ai_proxy_multi_schema.

- t/plugin/ai-proxy-flush.t: E2E tests for flush interval behavior,
  using error_log / no_error_log to assert whether the background flush
  thread fires. Also covers schema validation of negative values.
@dosubot dosubot Bot added size:XL This PR changes 500-999 lines, ignoring generated files. and removed size:L This PR changes 100-499 lines, ignoring generated files. labels May 20, 2026
Foo Bar and others added 10 commits May 20, 2026 09:32
…blank-line chunks

decode() previously broke callers that pass full response bodies or
per-chunk bodies without a terminal blank line. The old ngx.re-based
decoder treated the remaining content as a complete event; the new
implementation discarded it.

Fix: when no boundary is found for the remaining content, parse it as a
final event via parse_raw_event instead of silently dropping it.
…eaming_flush_interval_ms docs

Co-authored-by: Copilot <[email protected]>
Replace debug log grep approach with a hook-based counter:
- patch plugin_mod.lua_response_filter in extra_init_worker_by_lua to
  count flush calls per request in ngx.ctx._flush_count
- log the count via serverless-post-function in the log phase
- TEST 2 asserts 'flush count: 5' (one flush per SSE chunk)

Co-authored-by: Copilot <[email protected]>
… mode)

TEST 3 route now includes serverless-post-function log hook.
TEST 4 adds 'flush count: 0' assertion to confirm sync flush is
skipped when background flush_thread is active.

Co-authored-by: Copilot <[email protected]>
…ate flush count

Instead of wrapping lua_response_filter or using rawset(ngx,'flush') in
init_by_lua (which doesn't survive worker fork), use debug.setupvalue in
init_worker to replace the ngx_flush upvalue captured by lua_response_filter.

This directly counts actual ngx.flush calls:
- TEST 2 (interval_ms=0): 5 per-chunk flushes
- TEST 4 (interval_ms=50): flush_thread handles flushing, no per-chunk flushes

Co-authored-by: Copilot <[email protected]>
- Deduplicate print+flush code in lua_response_filter by inverting
  the early-return guard and merging both paths into one
- Add debug log in lua_response_filter flush path to prove per-chunk
  sync flush behavior in tests
- Simplify ai-proxy-flush tests: replace ngx.flush hook mechanism
  with log-based assertions (grep_error_log for exact counts)

Co-authored-by: Copilot <[email protected]>
base.lua parse_streaming_response calls framing.decode_buf() on
every chunk. sse.lua already had this method but aws-eventstream.lua
did not, causing a nil-call 500 error on all Bedrock streaming requests.

Co-authored-by: Copilot <[email protected]>
…m decode_buf docstring

- Restore the original upstream WORKAROUND comment for ngx.sleep(0)
- Update aws-eventstream.lua decode_buf docstring to match EE

Co-authored-by: Copilot <[email protected]>
Foo Bar and others added 2 commits May 21, 2026 10:12
- Check ngx.flush(false) return in flush_thread; propagate error via
  shared flush_err upvalue so the main loop can abort cleanly
- Check flush_err at the top of the streaming loop and call
  abort_on_disconnect to stop upstream reads on client disconnect
- Skip final ngx.flush(true) when flush_err is set to suppress noisy
  'final flush failed' log on client disconnect
- Change streaming_flush_interval_ms default from 0 to 10 ms
- Add TEST 3/4 to ai-proxy-client-disconnect.t covering async flush
  disconnect scenario
- Minor ai-proxy-flush.t cleanups (explicit ms=0 in TEST 1, TEST 6 title)

Co-authored-by: Copilot <[email protected]>
-- each split so the table never grows beyond two elements.
-- Initialized with "" so the fast-path (sse_parts[1] == "") activates
-- immediately on the first chunk, avoiding an unnecessary table.concat.
local sse_parts = {""}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

very nice 👍🏼

nic-6443
nic-6443 previously approved these changes May 21, 2026
@membphis

Copy link
Copy Markdown
Member

@AlinsRan

Changing the default to 10ms looks fine, but a few places still need to be updated so the behavior is documented consistently:

The PR description still says 0 (default); it should say the default is 10ms, and 0 disables background interval flushing.
The comment in base.lua still says == 0 (default), but the schema default is now 10.
The configuration docs should explicitly state what 0 means: streaming_flush_interval_ms = 0 uses per-chunk synchronous flushing, i.e. ngx.flush(true); values > 0 use the background thread to flush periodically with ngx.flush(false).
The tests currently cover explicit 0 and 50; it would be good to add one case where the field is omitted and the 10ms default is used.

… default

- Update base.lua comment: clarify default is 10ms, not 0; describe both
  modes (0 = per-chunk ngx.flush(true), >0 = background ngx.flush(false))
- Update EN/ZH docs: explicitly document 0 behavior vs >0 behavior
- Add TEST 7/8 to ai-proxy-flush.t: omit streaming_flush_interval_ms and
  verify the 10ms default activates the background flush thread

Co-authored-by: Copilot <[email protected]>
Foo Bar and others added 2 commits May 21, 2026 14:30
Replace route-setup + streaming request pair with a single check_schema
test that verifies the JSON Schema default (10) is injected into the conf
table when streaming_flush_interval_ms is omitted.

Co-authored-by: Copilot <[email protected]>
@AlinsRan
AlinsRan merged commit c1fc58f into apache:master May 22, 2026
25 of 26 checks passed
@AlinsRan
AlinsRan deleted the perf/ai-sse-optimize branch May 22, 2026 08:17
wistefan pushed a commit to wistefan/apisix that referenced this pull request Jun 16, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

performance generate flamegraph for the current PR size:XL This PR changes 500-999 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants