From 1fa7fe7ad7fefaee4cb8aab79bffed39c1716c46 Mon Sep 17 00:00:00 2001 From: zhanghao Date: Tue, 10 Nov 2026 18:22:22 +0800 Subject: [PATCH 1/3] postprocess: scan SSE lines linearly, and skip chunks without the field With strip_actual_usage on, every SSE frame cost ~8 us of CPU, though only the final usage frame is ever rewritten. ~7 us of it was the line pattern ^data:%s*(.-)%s*\r?\n?$: its lazy capture re-tries the rest of the line at every character, so a 216-byte frame paid for ~200 attempts. - sse_payload(): the same result as that pattern, from one linear scan (prefix check, then whitespace trimmed from both ends; "whitespace" is exactly Lua's %s). Used by both handlers. - json_rewrite_handler(rewrite_json, needle): a stream chunk that does not contain the needle is passed on without being split into lines, except a line it ends in the middle of, which is held back with the next chunk -- the needle may straddle the boundary. strip_actual_usage passes "actual_usage"; vllm_format, which touches every frame, passes none. Measured on one core (test/bench_postprocess.lua, 2000 content frames of ~216 bytes, 1 frame per chunk), ns per frame: off 21 -> 21 strip_actual_usage 8038 -> 270 vllm_format 14614 -> 5316 test/utest_postprocess_sse_scan.lua: - sse_payload agrees with the pattern on 20025 lines (edge cases and random); - 500 random chunkings: a stream without the field comes back byte-identical every time, and a stream with it is stripped every time; - a needle split exactly across two chunks is still stripped. --- lua/postprocess.lua | 38 ++++++++++-- test/bench_postprocess.lua | 82 +++++++++++++++++++++++++ test/utest_postprocess_sse_scan.lua | 93 +++++++++++++++++++++++++++++ 3 files changed, 209 insertions(+), 4 deletions(-) create mode 100644 test/bench_postprocess.lua create mode 100644 test/utest_postprocess_sse_scan.lua diff --git a/lua/postprocess.lua b/lua/postprocess.lua index 9392e27..5bb7d4b 100644 --- a/lua/postprocess.lua +++ b/lua/postprocess.lua @@ -290,22 +290,45 @@ local function vllm_rewrite_json(payload) return enc or payload -- fail open on an encode error end +local byte = string.byte + +-- %s in a Lua pattern: space, \t, \n, \v, \f, \r +local function is_space(b) return b == 32 or (b >= 9 and b <= 13) end + +-- The payload of one SSE line: what follows "data:", without the whitespace around it; nil for +-- any other line. Same result as line:match("^data:%s*(.-)%s*\r?\n?$"), in one linear scan -- +-- that pattern's lazy capture re-tries the rest of the line at every character, which made it +-- most of the handler's cost (~7 us of ~8 us on a 216-byte frame). +local function sse_payload(line) + if byte(line, 1) ~= 100 or line:sub(1, 5) ~= "data:" then return nil end -- 100 = "d" + local s, e = 6, #line + while s <= e and is_space(byte(line, s)) do s = s + 1 end + while e >= s and is_space(byte(line, e)) do e = e - 1 end + return line:sub(s, e) +end +M.sse_payload = sse_payload -- exposed for unit tests + -- Rewrite one SSE data frame with rewrite_json. Leave unchanged frames byte-identical. local function rewrite_sse_line(line, rewrite_json) - local payload = line:match("^data:%s*(.-)%s*\r?\n?$") + local payload = sse_payload(line) if not payload or payload == "" or payload == "[DONE]" then return line -- comment/heartbeat/terminator/non-data end local rewritten = rewrite_json(payload) if rewritten == payload then return line end - return "data: " .. rewritten .. (line:match("\r?\n$") or "") + local n, tail = #line, "" + if byte(line, n) == 10 then tail = (n > 1 and byte(line, n - 1) == 13) and "\r\n" or "\n" end + return "data: " .. rewritten .. tail end -- A handler that rewrites OpenAI-style responses one JSON document at a time: -- each SSE data frame, or a whole non-stream JSON body (capped and buffered until -- eof). rewrite_json(payload) -> payload' returns its input unchanged (the same -- string) when there is nothing to do, so untouched frames stay byte-identical. -local function json_rewrite_handler(rewrite_json) +-- needle (optional): a plain substring every frame the handler would change contains. A stream +-- chunk without it is passed on without splitting it into lines -- except a line it ends in the +-- middle of, which is held back with the next chunk, since the needle may straddle the boundary. +local function json_rewrite_handler(rewrite_json, needle) return { buffered = false, rewrites_length = true, @@ -342,6 +365,13 @@ local function json_rewrite_handler(rewrite_json) if not ctx.pp_sse then return nil end -- other media types pass through local pending = (ctx.pp_leftover or "") .. (chunk or "") ctx.pp_leftover = nil + if needle and not pending:find(needle, 1, true) then + if eof or byte(pending, -1) == 10 then return pending end + local last, p = 0, pending:find("\n", 1, true) + while p do last = p; p = pending:find("\n", p + 1, true) end + ctx.pp_leftover = pending:sub(last + 1) -- the partial last line + return pending:sub(1, last) + end local out, pos = {}, 1 while true do local nl = pending:find("\n", pos, true) @@ -379,6 +409,6 @@ local function strip_actual_usage_json(payload) end M.strip_actual_usage_json = strip_actual_usage_json -- exposed for unit tests -M.register("strip_actual_usage", json_rewrite_handler(strip_actual_usage_json)) +M.register("strip_actual_usage", json_rewrite_handler(strip_actual_usage_json, '"actual_usage"')) return M diff --git a/test/bench_postprocess.lua b/test/bench_postprocess.lua new file mode 100644 index 0000000..23cb923 --- /dev/null +++ b/test/bench_postprocess.lua @@ -0,0 +1,82 @@ +-- Micro-benchmark of lua/postprocess.lua's body filter: the cost per SSE chunk of each handler, against +-- the same stream with post-processing off (the default). Pure CPU in one worker; no network. +-- resty test/bench_postprocess.lua [lua/postprocess.lua] [frames=2000] [rounds=3] [frames_per_chunk=1] +-- The stream mimics an sglang chat completion: content frames of a few tokens each, the +-- finish frame, the usage frame (with actual_usage) and [DONE]. Each upstream chunk carries +-- frames, as the engine's flushes do. Reports ns per frame: each case runs the +-- stream until at least 2 s have passed (clocks under resty are coarse: os.clock ticks in 10 ms, so a +-- single 10 ms pass reads as 0 or one tick), best of such batches. +local ngx_real = ngx -- run() swaps in a fake ngx; keep the real one for the clock +local pp = assert(loadfile(arg[1] or "lua/postprocess.lua"))() +local FRAMES, ROUNDS, PER = tonumber(arg[2] or 2000), tonumber(arg[3] or 3), tonumber(arg[4] or 1) + +local function frame(s) return "data: " .. s .. "\n\n" end +local frames = {} +for i = 1, FRAMES do + frames[#frames + 1] = frame(string.format( + '{"id":"chatcmpl-b3a1f0c2d4e5","object":"chat.completion.chunk","created":1794300000,"model":"kimi-k3",' .. + '"choices":[{"index":0,"delta":{"content":"token %d of the answer"},"logprobs":null,"finish_reason":null}]}', i)) +end +frames[#frames + 1] = frame('{"id":"chatcmpl-b3a1f0c2d4e5","object":"chat.completion.chunk","created":1794300000,' .. + '"model":"kimi-k3","choices":[{"index":0,"delta":{},"logprobs":null,"finish_reason":"stop"}]}') +frames[#frames + 1] = frame('{"id":"chatcmpl-b3a1f0c2d4e5","object":"chat.completion.chunk","created":1794300000,' .. + '"model":"kimi-k3","choices":[],"usage":{"prompt_tokens":4096,"completion_tokens":2000,"total_tokens":6096,' .. + '"prompt_tokens_details":{"cached_tokens":2048,"cache_write_tokens":1792}},"actual_usage":{"prompt_tokens":4096,' .. + '"completion_tokens":2000,"total_tokens":6096,"prompt_tokens_details":{"cached_tokens":3072}}}') +frames[#frames + 1] = "data: [DONE]\n\n" +local chunks = {} +for i = 1, #frames, PER do chunks[#chunks + 1] = table.concat(frames, "", i, math.min(i + PER - 1, #frames)) end + +local function run(name) -- one pass of the whole stream through the real header/body entry points + local ctx = {} + _G.ngx = { ctx = ctx, status = 200, header = { content_type = "text/event-stream" }, arg = {}, + ERR = "error", log = function() end } + local opts = { postprocess = name } + _G.do_postprocess_header(opts) + local bytes = 0 + for i, c in ipairs(chunks) do + ngx.arg[1], ngx.arg[2] = c, i == #chunks + _G.do_postprocess_body(opts) + bytes = bytes + #ngx.arg[1] + end + return bytes +end + +local function now() ngx_real.update_time(); return ngx_real.now() end +local function bench(name) -- seconds per stream pass + local best = math.huge + for _ = 1, ROUNDS do + local n, t0 = 0, now() + repeat run(name); n = n + 1 until n % 10 == 0 and now() - t0 >= 2 + local per = (now() - t0) / n + if per < best then best = per end + end + return best +end + +for _ = 1, 3 do run(nil); run("strip_actual_usage"); run("vllm_format") end -- warm the JIT +local base = bench(nil) +print(string.format("stream: %d frames, %d chunks (%d frames/chunk), %d bytes", #frames, #chunks, PER, run(nil))) +for _, name in ipairs({ "off", "strip_actual_usage", "vllm_format" }) do + local t = name == "off" and base or bench(name) + print(string.format(" %-20s %8.0f ns/frame +%7.0f ns/frame vs off %6.2f ms per stream", + name, t / #frames * 1e9, (t - base) / #frames * 1e9, t * 1e3)) +end + +-- Where strip_actual_usage's time goes: the per-line pattern that rewrite_sse_line runs on every frame, +-- against a plain substring test on the same lines. +local lines = {} +for _, f in ipairs(frames) do lines[#lines + 1] = f:match("^[^\n]*\n") end +local function per_line(fn) + local best = math.huge + for _ = 1, ROUNDS do + local n, t0 = 0, now() + repeat for _, l in ipairs(lines) do fn(l) end; n = n + 1 until n % 10 == 0 and now() - t0 >= 1 + local per = (now() - t0) / (n * #lines) + if per < best then best = per end + end + return best * 1e9 +end +print(string.format("per line (%d bytes avg):", math.floor(#table.concat(lines) / #lines))) +print(string.format(" match ^data:%%s*(.-)%%s*\\r?\\n?$ %8.0f ns", per_line(function(l) return l:match("^data:%s*(.-)%s*\r?\n?$") end))) +print(string.format(" find '\"actual_usage\"' (plain) %8.0f ns", per_line(function(l) return l:find('"actual_usage"', 1, true) end))) diff --git a/test/utest_postprocess_sse_scan.lua b/test/utest_postprocess_sse_scan.lua new file mode 100644 index 0000000..7925e85 --- /dev/null +++ b/test/utest_postprocess_sse_scan.lua @@ -0,0 +1,93 @@ +-- Run: resty test/utest_postprocess_sse_scan.lua [lua/postprocess.lua] +-- The SSE line scan and the per-chunk fast path of lua/postprocess.lua: +-- 1. sse_payload(line) gives exactly what the pattern it replaced gave, on hand-picked edge cases and +-- on random lines; +-- 2. strip_actual_usage passes a stream without the field through byte-identical however the chunks +-- are cut, and still strips it however the chunks are cut (the needle may straddle a boundary). +local pp = assert(loadfile(arg[1] or "lua/postprocess.lua"))() +local fail = 0 +local function check(ok, msg) if not ok then fail = fail + 1; print(" FAIL " .. msg) end end + +-- 1. sse_payload == the old pattern +local function old(line) return line:match("^data:%s*(.-)%s*\r?\n?$") end +local cases = { + "data: {\"a\":1}\n", "data:{\"a\":1}\n", "data: {\"a\":1} \r\n", "data: {\"a\":1}", "data:\n", "data:", + "data: \n", "data: [DONE]\n", "data:\t[DONE]\r\n", ": keep-alive\n", "event: message\n", "\n", "", "d", + "dat", "data", " data: x\n", "DATA: x\n", "data: x\r", "data: x\r\r\n", "data: \v\fx\v\f\n", "data: a b\n", + "data: x\n\n", "datax: y\n", "data:\r\n", +} +math.randomseed(20261110) +local alphabet = { "d", "a", "t", ":", " ", "\t", "\r", "\n", "\v", "\f", "{", "}", "x", "\"", "data:", "data: " } +for _ = 1, 20000 do + local parts = {} + for _ = 1, math.random(0, 12) do parts[#parts + 1] = alphabet[math.random(#alphabet)] end + cases[#cases + 1] = table.concat(parts) +end +local mismatches = 0 +for _, l in ipairs(cases) do + local a, b = old(l), pp.sse_payload(l) + if a ~= b then + mismatches = mismatches + 1 + if mismatches <= 5 then check(false, string.format("sse_payload(%q) = %q, pattern gives %q", l, tostring(b), tostring(a))) end + end +end +check(mismatches == 0, mismatches .. " mismatches with the old pattern") +print(string.format("sse_payload vs pattern: %d lines, %d mismatches", #cases, mismatches)) + +-- 2. the fast path, through the real header/body entry points +local function run(name, pieces) + local ctx = {} + local prev = _G.ngx + _G.ngx = { ctx = ctx, status = 200, header = { content_type = "text/event-stream" }, arg = {}, + ERR = "error", log = function() end } + local opts = { postprocess = name } + _G.do_postprocess_header(opts) + local out = {} + for i, p in ipairs(pieces) do + ngx.arg[1], ngx.arg[2] = p, i == #pieces + _G.do_postprocess_body(opts) + out[#out + 1] = ngx.arg[1] + end + _G.ngx = prev + return table.concat(out) +end +local function cut(s) -- random chunking, including empty chunks + local out, i = {}, 1 + while i <= #s do + local n = math.random(0, 40) + out[#out + 1] = s:sub(i, i + n - 1); i = i + n + end + out[#out + 1] = "" + return out +end +local content = {} +for i = 1, 50 do + content[#content + 1] = 'data: {"id":"c","choices":[{"index":0,"delta":{"content":"t' .. i .. ' a/b"},"finish_reason":null}]}\n\n' +end +content = table.concat(content) +local usage = 'data: {"id":"c","choices":[],"usage":{"prompt_tokens":10,"completion_tokens":50},' .. + '"actual_usage":{"prompt_tokens":10,"prompt_tokens_details":{"cached_tokens":8}}}\n\n' +local without = content .. 'data: {"id":"c","choices":[],"usage":{"prompt_tokens":10}}\n\ndata: [DONE]\n\n' +local with = content .. usage .. "data: [DONE]\n\n" +local unchanged, stripped = 0, 0 +for _ = 1, 500 do + if run("strip_actual_usage", cut(without)) == without then unchanged = unchanged + 1 end + local o = run("strip_actual_usage", cut(with)) + if not o:find("actual_usage", 1, true) and o:sub(1, #content) == content + and o:find('"choices":[]', 1, true) and o:sub(-14) == "data: [DONE]\n\n" then + stripped = stripped + 1 + end +end +check(unchanged == 500, "stream without the field changed in " .. (500 - unchanged) .. "/500 random chunkings") +check(stripped == 500, "field not cleanly stripped in " .. (500 - stripped) .. "/500 random chunkings") +-- the needle split exactly across two chunks +local at = with:find("actual_usage", 1, true) +local split = run("strip_actual_usage", { with:sub(1, at + 5), with:sub(at + 6) }) +check(not split:find("actual_usage", 1, true), "needle split across two chunks was not stripped") +-- vllm_format (no needle) still rewrites on the scanned payload: the "/" in content comes back escaped +local v = run("vllm_format", cut(content .. 'data: {"choices":[],"system_fingerprint":"fp"}\n\n')) +check(not v:find("system_fingerprint", 1, true), "vllm_format did not drop system_fingerprint") +print(string.format("fast path: %d/500 unchanged, %d/500 stripped", unchanged, stripped)) + +print(fail == 0 and "sse scan + fast path: PASS" or ("sse scan + fast path: FAIL=" .. fail)) +os.exit(fail == 0 and 0 or 1) From 602fcea8bffb7609630f7113a411fd34ad8c975f Mon Sep 17 00:00:00 2001 From: zhanghao Date: Tue, 10 Nov 2026 19:01:43 +0800 Subject: [PATCH 2/3] README: say exactly when post-processing holds stream bytes back "Streams are not held" was not quite true: a line a chunk ends in the middle of waits for the next chunk (it always has; the line splitter needs the whole line). Complete lines still go out with the chunk they arrived in. --- README.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index da0d6f5..f9854cf 100644 --- a/README.md +++ b/README.md @@ -268,7 +268,9 @@ How it behaves: re-encoded, so its key order and `/` escaping can differ from the backend's. - `Content-Length` is dropped on the route's 2xx responses (they are sent chunked), and a non-stream JSON body is held until it is complete (up to 8 MiB; a larger - one passes through unchanged). Streams are not held: frames go out as they arrive. + one passes through unchanged). A stream is not held: each complete line goes out + with the chunk it arrived in, and only a line a chunk ends in the middle of waits + for the next chunk. - Fail-open: a frame that does not parse, or a handler error, leaves the original bytes. - The bodylog sees the raw upstream response (it runs first). The listener keeps `usage` and `actual_usage` in the record's `resp_meta`, so a stripped From c1923f9f51140d9de97283c84f7b8269753592c9 Mon Sep 17 00:00:00 2001 From: zhanghao Date: Tue, 10 Nov 2026 19:52:29 +0800 Subject: [PATCH 3/3] postprocess: note the needle's plain-key assumption; fix a test comment Review follow-ups: say that the strip_actual_usage fast path assumes the key is not escaped in the JSON (a backend that escaped it would have the field passed through, as with post-processing off), and stop claiming a check the vllm_format case in utest_postprocess_sse_scan.lua does not make. --- lua/postprocess.lua | 2 ++ test/utest_postprocess_sse_scan.lua | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/lua/postprocess.lua b/lua/postprocess.lua index 5bb7d4b..27b6b73 100644 --- a/lua/postprocess.lua +++ b/lua/postprocess.lua @@ -409,6 +409,8 @@ local function strip_actual_usage_json(payload) end M.strip_actual_usage_json = strip_actual_usage_json -- exposed for unit tests +-- The needle assumes the key is written plainly. A backend that escaped it in the JSON (e.g. +-- "actual\u005fusage") would have the field passed through -- as with post-processing off. M.register("strip_actual_usage", json_rewrite_handler(strip_actual_usage_json, '"actual_usage"')) return M diff --git a/test/utest_postprocess_sse_scan.lua b/test/utest_postprocess_sse_scan.lua index 7925e85..5d78482 100644 --- a/test/utest_postprocess_sse_scan.lua +++ b/test/utest_postprocess_sse_scan.lua @@ -84,7 +84,7 @@ check(stripped == 500, "field not cleanly stripped in " .. (500 - stripped) .. " local at = with:find("actual_usage", 1, true) local split = run("strip_actual_usage", { with:sub(1, at + 5), with:sub(at + 6) }) check(not split:find("actual_usage", 1, true), "needle split across two chunks was not stripped") --- vllm_format (no needle) still rewrites on the scanned payload: the "/" in content comes back escaped +-- vllm_format (no needle) still rewrites frames found by the line scan local v = run("vllm_format", cut(content .. 'data: {"choices":[],"system_fingerprint":"fp"}\n\n')) check(not v:find("system_fingerprint", 1, true), "vllm_format did not drop system_fingerprint") print(string.format("fast path: %d/500 unchanged, %d/500 stripped", unchanged, stripped))