postprocess: scan SSE lines linearly, and skip chunks without the field

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

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,8 @@ 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))
+-- 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

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: <frames> content frames of a few tokens each, the
+-- finish frame, the usage frame (with actual_usage) and [DONE]. Each upstream chunk carries
+-- <frames_per_chunk> 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 <rounds> 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)))

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 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))
+
+print(fail == 0 and "sse scan + fast path: PASS" or ("sse scan + fast path: FAIL=" .. fail))
+os.exit(fail == 0 and 0 or 1)