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

Author: AlinsRan

apisix/plugin.lua

@@ -1471,52 +1471,45 @@ end
 --   can detect client disconnection. Defaults to false (async flush).
 -- @return boolean, string|nil Always returns (ok, err). On success returns true.
 --   On flush failure or print failure returns false, err.
-function _M.lua_response_filter(api_ctx, headers, body, wait)
+function _M.lua_response_filter(api_ctx, headers, body, no_flush, wait)
     local plugins = api_ctx.plugins
-    if not plugins or #plugins == 0 then
-        -- if there is no any plugin, just print the original body to downstream
-        local ok, err = ngx_print(body)
-        if not ok then
-            return false, err
-        end
-        ok, err = ngx_flush(wait == true)
-        if not ok then
-            return false, err
-        end
-        return true
-    end
-    for i = 1, #plugins, 2 do
-        local phase_func = plugins[i]["lua_body_filter"]
-        if phase_func then
-            local conf = plugins[i + 1]
-            if not meta_filter(api_ctx, plugins[i]["name"], conf)then
-                goto CONTINUE
-            end
-
-            run_meta_pre_function(conf, api_ctx, plugins[i]["name"])
-            local code, new_body = phase_func(conf, api_ctx, headers, body)
-            if code then
-                if code ~= ngx_ok then
-                    ngx.status = code
+    if plugins and #plugins > 0 then
+        for i = 1, #plugins, 2 do
+            local phase_func = plugins[i]["lua_body_filter"]
+            if phase_func then
+                local conf = plugins[i + 1]
+                if not meta_filter(api_ctx, plugins[i]["name"], conf)then
+                    goto CONTINUE
                 end
 
-                ngx_print(new_body)
-                ngx_exit(ngx_ok)
-            end
-            if new_body then
-                body = new_body
+                run_meta_pre_function(conf, api_ctx, plugins[i]["name"])
+                local code, new_body = phase_func(conf, api_ctx, headers, body)
+                if code then
+                    if code ~= ngx_ok then
+                        ngx.status = code
+                    end
+
+                    ngx_print(new_body)
+                    ngx_exit(ngx_ok)
+                end
+                if new_body then
+                    body = new_body
+                end
             end
-        end
 
-        ::CONTINUE::
+            ::CONTINUE::
+        end
     end
     local ok, err = ngx_print(body)
     if not ok then
         return false, err
     end
-    ok, err = ngx_flush(wait == true)
-    if not ok then
-        return false, err
+    if not no_flush then
+        core.log.debug("lua_response_filter: flushing chunk to client")
+        ok, err = ngx_flush(wait == true)
+        if not ok then
+            return false, err
+        end
     end
     return true
 end

apisix/plugins/ai-providers/base.lua

@@ -37,8 +37,8 @@ local protocols = require("apisix.plugins.ai-protocols")
 local deep_merge = require("apisix.plugins.ai-proxy.merge").deep_merge
 local ngx = ngx
 local ngx_now = ngx.now
-local tonumber = tonumber
 local require = require
+local tonumber = tonumber
 
 local table = table
 local pairs = pairs
@@ -420,7 +420,12 @@ function _M.parse_streaming_response(self, ctx, res, target_proto, converter, co
     local body_reader = res.body_reader
     local contents = {}
     local sse_state = { is_first = true }
-    local sse_buf = ""
+    -- SSE framing buffer: accumulate chunks with table.insert to avoid
+    -- allocating a new string on every append; reset to {remainder} after
+    -- 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 = {""}
     -- Track whether any output was sent to the client.
     -- When a converter is active but the upstream returns a different SSE format,
     -- all events may be skipped and no output produced, leaving the response
@@ -436,59 +441,146 @@ function _M.parse_streaming_response(self, ctx, res, target_proto, converter, co
     end
     local bytes_read = 0
 
+    -- streaming_flush_interval_ms controls both flush strategy and the thread:
+    --   == 0          : no thread; lua_response_filter flushes synchronously
+    --                   per chunk via ngx.flush(true), guaranteeing immediate
+    --                   client delivery.
+    --   >  0 (default: 10): background thread calls ngx.flush(false) every N ms;
+    --                   lua_response_filter skips per-chunk flush for maximum
+    --                   throughput. Useful when the upstream bursts multiple
+    --                   tokens at once.
+    local flush_interval_ms = conf and conf.streaming_flush_interval_ms or 0
+    -- async_flush: true when the interval thread is responsible for flushing
+    local async_flush = flush_interval_ms > 0
+    -- needs_flush is set to true immediately after dispatching a chunk so the
+    -- thread always flushes exactly the data that has been written.  Cleared
+    -- before ngx.flush() so any new chunks written during the flush yield are
+    -- picked up on the next interval rather than silently dropped.
+    local needs_flush = false
+    local flush_thread
+    local flush_err
+    if async_flush then
+        local interval_s = flush_interval_ms / 1000
+        local spawn_err
+        flush_thread, spawn_err = ngx.thread.spawn(function()
+            while true do
+                ngx.sleep(interval_s)
+                if needs_flush then
+                    needs_flush = false
+                    local ok, err = ngx.flush(false)
+                    if not ok then
+                        flush_err = err
+                        return
+                    end
+                    core.log.debug("ai-proxy: flush_thread periodic flush")
+                end
+            end
+        end)
+        if not flush_thread then
+            core.log.error("failed to spawn flush thread: ", spawn_err)
+            async_flush = false
+        end
+    end
+
     local function abort_on_disconnect(flush_err)
         core.log.info("client disconnected during AI streaming, ",
                       "aborting upstream read: ", flush_err)
+        if flush_thread then
+            ngx.thread.kill(flush_thread)
+            flush_thread = nil
+        end
         if res._httpc then
             res._httpc:close()
             res._httpc = nil
         end
         res._upstream_bytes = bytes_read
+        ctx.var.apisix_upstream_response_time = math.floor(
+            (ngx_now() - ctx.llm_request_start_time) * 1000)
         ctx.var.llm_request_done = true
     end
 
+    -- Use a local flag instead of reading ctx.var on every chunk.
+    local first_token_set = false
+
     while true do
+        if flush_err then
+            abort_on_disconnect(flush_err)
+            return
+        end
+
         local chunk, err = body_reader()
-        ctx.var.apisix_upstream_response_time = math.floor((ngx_now() -
-                                         ctx.llm_request_start_time) * 1000)
         if err then
+            ctx.var.apisix_upstream_response_time = math.floor(
+                (ngx_now() - ctx.llm_request_start_time) * 1000)
             core.log.warn("failed to read response chunk: ", err)
             res._upstream_bytes = bytes_read
+            if flush_thread then
+                ngx.thread.kill(flush_thread)
+                flush_thread = nil
+            end
             return transport_http.handle_error(err)
         end
         if not chunk then
-            if #sse_buf > 0 then
+            local sse_rem = table.concat(sse_parts)
+            if #sse_rem > 0 then
                 core.log.warn("dropping incomplete stream frame at EOF, size: ",
-                              #sse_buf)
+                              #sse_rem)
             end
 
             res._upstream_bytes = bytes_read
+            ctx.var.apisix_upstream_response_time = math.floor(
+                (ngx_now() - ctx.llm_request_start_time) * 1000)
             if converter and not output_sent then
+                if flush_thread then
+                    ngx.thread.kill(flush_thread)
+                end
                 local msg = "streaming response completed without producing "
                             .. "any output; the upstream likely returned a "
                             .. "different stream format than the converter expects"
                 core.log.error(msg)
                 return 502, msg
             end
+            -- Final sync flush: ensure the last async-queued bytes reach the client.
+            -- flush_err means client already disconnected; skip to avoid a noisy log.
+            if flush_thread then
+                ngx.thread.kill(flush_thread)
+                flush_thread = nil
+            end
+            if not flush_err then
+                ngx.flush(true)
+            end
             return
         end
 
         bytes_read = bytes_read + #chunk
 
-        if ctx.var.llm_time_to_first_token == "0" then
+        if not first_token_set then
             ctx.var.llm_time_to_first_token = math.floor(
-                                            (ngx_now() - ctx.llm_request_start_time) * 1000)
+                (ngx_now() - ctx.llm_request_start_time) * 1000)
+            first_token_set = true
+        end
+
+        -- Skip table.concat when there is no carry-over remainder (common case).
+        -- sse_parts is reset to {remainder} after each iteration; when the
+        -- previous chunk ended on a boundary, sse_parts[1] == "" and we can
+        -- hand the new chunk directly to decode_buf without allocating a concat.
+        local candidate
+        if sse_parts[1] == "" then
+            candidate = chunk
+        else
+            sse_parts[#sse_parts + 1] = chunk
+            candidate = table.concat(sse_parts)
         end
 
-        sse_buf = sse_buf .. chunk
-        local complete, remainder = framing.split_buf(sse_buf)
+        -- One-pass split + decode: finds all complete SSE events and the
+        -- trailing remainder in a single forward scan (no PCRE, no double scan).
+        local events, remainder = framing.decode_buf(candidate)
         local max_remainder = framing.max_remainder or 1024 * 1024
         if #remainder > max_remainder then
             core.log.warn("stream remainder exceeded ", max_remainder, " bytes, resetting")
             remainder = ""
         end
-        sse_buf = remainder
-        local events = complete ~= "" and framing.decode(complete) or {}
+        sse_parts = {remainder}
         ctx.llm_response_contents_in_chunk = {}
         local converted_chunks = {}
 
@@ -532,24 +624,33 @@ function _M.parse_streaming_response(self, ctx, res, target_proto, converter, co
             ::CONTINUE::
         end
 
-        -- Output: converter events or passthrough raw chunk.
-        -- Pass wait=true for synchronous flush so we can detect client disconnection.
+        -- Dispatch chunk downstream.  Plugins run per-chunk so body_filter
+        -- hooks (e.g. content moderation) receive every SSE event individually.
+        -- no_flush=true when the interval thread handles flushing; otherwise
+        -- no_flush=nil + wait=true for synchronous per-chunk delivery guarantee.
+        local no_flush = async_flush or nil
         if converter then
             for _, c in ipairs(converted_chunks) do
-                local ok, flush_err = plugin.lua_response_filter(ctx, res.headers, c, true)
-                output_sent = true
+                local ok, flush_err = plugin.lua_response_filter(
+                    ctx, res.headers, c, no_flush, true)
                 if not ok then
                     abort_on_disconnect(flush_err)
                     return
                 end
+                output_sent = true
             end
         else
-            local ok, flush_err = plugin.lua_response_filter(ctx, res.headers, chunk, true)
-            output_sent = true
+            local ok, flush_err = plugin.lua_response_filter(
+                ctx, res.headers, chunk, no_flush, true)
             if not ok then
                 abort_on_disconnect(flush_err)
                 return
             end
+            output_sent = true
+        end
+        -- Let the interval flush thread know there is unflushed output.
+        if async_flush then
+            needs_flush = true
         end
 
         -- Enforce runaway-upstream safeguards after processing the chunk.
@@ -561,6 +662,10 @@ function _M.parse_streaming_response(self, ctx, res, target_proto, converter, co
             limit_hit = "max_response_bytes"
         end
         if limit_hit then
+            if flush_thread then
+                ngx.thread.kill(flush_thread)
+                flush_thread = nil
+            end
             local duration_ms = math.floor((ngx_now() -
                                             ctx.llm_request_start_time) * 1000)
             core.log.warn("aborting AI stream: ", limit_hit, " exceeded;",
@@ -574,6 +679,7 @@ function _M.parse_streaming_response(self, ctx, res, target_proto, converter, co
             end
             -- Signal downstream filters (e.g. moderation plugins that defer
             -- work until request completion) that no more content is coming.
+            ctx.var.apisix_upstream_response_time = duration_ms
             ctx.var.llm_request_done = true
             res._upstream_bytes = bytes_read
             if output_sent then
@@ -602,6 +708,7 @@ function _M.parse_streaming_response(self, ctx, res, target_proto, converter, co
         -- backpressure, or time out stalled streams. See #13256 for a proper
         -- solution.
         ngx.sleep(0)
+
     end
 end
 

apisix/plugins/ai-proxy/schema.lua

@@ -267,6 +267,16 @@ _M.ai_proxy_schema = {
             description = "keepalive timeout in milliseconds",
         },
         keepalive_pool = {type = "integer", minimum = 1, default = 30},
+        streaming_flush_interval_ms = {
+            type = "integer",
+            minimum = 0,
+            default = 10,
+            description = "A background thread flushes the output buffer every N "
+                       .. "milliseconds (async flush). Useful when the upstream bursts "
+                       .. "multiple tokens at once and you need to bound client latency. "
+                       .. "Set to 0 to disable the background thread and flush each "
+                       .. "chunk synchronously inline.",
+        },
         ssl_verify = {type = "boolean", default = true },
         override = override_schema,
     },
@@ -353,6 +363,16 @@ _M.ai_proxy_multi_schema = {
             description = "keepalive timeout in milliseconds",
         },
         keepalive_pool = {type = "integer", minimum = 1, default = 30},
+        streaming_flush_interval_ms = {
+            type = "integer",
+            minimum = 0,
+            default = 10,
+            description = "A background thread flushes the output buffer every N "
+                       .. "milliseconds (async flush). Useful when the upstream bursts "
+                       .. "multiple tokens at once and you need to bound client latency. "
+                       .. "Set to 0 to disable the background thread and flush each "
+                       .. "chunk synchronously inline.",
+        },
         ssl_verify = {type = "boolean", default = true },
     },
     required = {"instances"},

apisix/plugins/ai-transport/aws-eventstream.lua

@@ -289,4 +289,15 @@ function _M.decode(buf)
 end
 
 
+--- Combined split + decode in one call, matching the interface expected by
+-- ai-providers/base.lua.
+-- @param buf string  Accumulated bytes from the upstream socket.
+-- @return table  Array of decoded event tables (same as decode()).
+-- @return string  Trailing bytes that did not form a complete frame.
+function _M.decode_buf(buf)
+    local complete, remainder = _M.split_buf(buf)
+    return _M.decode(complete), remainder
+end
+
+
 return _M