perf(response_writer): minimize socket:send() calls

Author: tkan145

gateway/src/resty/http/response_writer.lua

@@ -1,5 +1,7 @@
 local fmt = string.format
 local str_lower = string.lower
+local insert = table.insert
+local concat = table.concat
 
 local _M = {
 }
@@ -32,7 +34,6 @@ end
 -- write_response writes response body reader to sock in the HTTP/1.x server response format,
 -- The connection is closed if send() fails or when returning a non-zero
 function _M.send_response(sock, response, chunksize)
-    local bytes, err
     chunksize = chunksize or 65536
 
     if not response then
@@ -44,33 +45,32 @@ function _M.send_response(sock, response, chunksize)
         return nil, "socket not initialized yet"
     end
 
-    -- Status line
-    -- TODO: get HTTP version from request
-    local status = fmt("HTTP/%d.%d %03d %s\r\n", 1, 1, response.status, response.reason)
-    bytes, err = send(sock, status)
-    if not bytes then
-        return nil, "failed to send status line, err: " .. (err or "unknown")
-    end
+    -- Build status line + headers into a single buffer to minimize send() calls
+    local buf = {
+        fmt("HTTP/1.1 %03d %s\r\n", response.status, response.reason)
+    }
 
      -- Filter out hop-by-hop headeres
      for k, v in pairs(response.headers) do
         if not HOP_BY_HOP_HEADERS[str_lower(k)] then
-            local header = fmt("%s: %s\r\n", k, v)
-            bytes, err = sock:send(header)
-            if not bytes then
-              return nil, "failed to send status line, err: " .. (err or "unknown")
-            end
+          insert(buf, k .. ": " .. v .. cr_lf)
         end
      end
 
     -- End-of-header
-    bytes, err = send(sock, cr_lf)
+    insert(buf, cr_lf)
+
+    local bytes, err = sock:send(concat(buf))
     if not bytes then
-        return nil, "failed to send status line, err: " .. (err or "unknown")
+        return nil, "failed to send headers, err: " .. (err or "unknown")
     end
 
     -- Write body
     local reader = response.body_reader
+    if not reader then
+      return nil, "no body reader"
+    end
+
     repeat
         local chunk, read_err
 

spec/resty/http/response_writer_spec.lua

@@ -0,0 +1,248 @@
+local response_writer = require('resty.http.response_writer')
+
+
+describe('resty.http.response_writer', function()
+
+  local mock_sock
+  local sent_data
+
+  local function make_response(opts)
+    opts = opts or {}
+    local chunks = opts.chunks or { "hello" }
+    local idx = 0
+    return {
+      status = opts.status or 200,
+      reason = opts.reason or "OK",
+      headers = opts.headers or {},
+      body_reader = function()
+        idx = idx + 1
+        return chunks[idx]
+      end
+    }
+  end
+
+  before_each(function()
+    sent_data = {}
+    mock_sock = {
+      send = function(_, data)
+        table.insert(sent_data, data)
+        return #data, nil
+      end
+    }
+
+    stub(ngx, 'log')
+  end)
+
+  describe('.send_response', function()
+
+    it('returns nil when no response is provided', function()
+      local ok, err = response_writer.send_response(mock_sock, nil)
+      assert.is_nil(ok)
+      assert.is_nil(err)
+    end)
+
+    it('returns error when no socket is provided', function()
+      local ok, err = response_writer.send_response(nil, make_response())
+      assert.is_nil(ok)
+      assert.equal("socket not initialized yet", err)
+    end)
+
+    it('sends status line in first send', function()
+      local response = make_response({ status = 200, reason = "OK" })
+      response_writer.send_response(mock_sock, response)
+
+      assert.truthy(string.find(sent_data[1], "^HTTP/1.1 200 OK\r\n"))
+    end)
+
+    it('formats different status codes', function()
+      local response = make_response({ status = 404, reason = "Not Found" })
+      response_writer.send_response(mock_sock, response)
+
+      assert.truthy(string.find(sent_data[1], "^HTTP/1.1 404 Not Found\r\n"))
+    end)
+
+    it('batches status line and headers in a single send', function()
+      local response = make_response({
+        headers = { ["Content-Type"] = "text/plain", ["X-Custom"] = "value" }
+      })
+      response_writer.send_response(mock_sock, response)
+
+      -- First send contains status line + headers + end-of-header CRLF
+      local header_block = sent_data[1]
+      assert.truthy(string.find(header_block, "^HTTP/1.1"))
+      assert.truthy(string.find(header_block, "Content%-Type: text/plain\r\n"))
+      assert.truthy(string.find(header_block, "X%-Custom: value\r\n"))
+      assert.truthy(string.find(header_block, "\r\n\r\n$"))
+    end)
+
+    it('filters hop-by-hop headers', function()
+      local response = make_response({
+        headers = {
+          ["Connection"] = "keep-alive",
+          ["Keep-Alive"] = "timeout=5",
+          ["Transfer-Encoding"] = "chunked",
+          ["Proxy-Authenticate"] = "Basic",
+          ["Proxy-Authorization"] = "Basic abc",
+          ["TE"] = "trailers",
+          ["Trailers"] = "Expires",
+          ["Upgrade"] = "websocket",
+          ["Content-Length"] = "5",
+          ["X-Custom"] = "value",
+        }
+      })
+      response_writer.send_response(mock_sock, response)
+
+      local all_sent = table.concat(sent_data)
+      assert.falsy(string.find(all_sent, "Connection:"))
+      assert.falsy(string.find(all_sent, "Keep%-Alive:"))
+      assert.falsy(string.find(all_sent, "Transfer%-Encoding:"))
+      assert.falsy(string.find(all_sent, "Proxy%-Authenticate:"))
+      assert.falsy(string.find(all_sent, "Proxy%-Authorization:"))
+      assert.falsy(string.find(all_sent, "TE:"))
+      assert.falsy(string.find(all_sent, "Trailers:"))
+      assert.falsy(string.find(all_sent, "Upgrade:"))
+      assert.falsy(string.find(all_sent, "Content%-Length:"))
+      assert.truthy(string.find(all_sent, "X%-Custom: value"))
+    end)
+
+    it('filters hop-by-hop headers case-insensitively', function()
+      local response = make_response({
+        headers = {
+          ["CONNECTION"] = "close",
+          ["KEEP-ALIVE"] = "timeout=5",
+        }
+      })
+      response_writer.send_response(mock_sock, response)
+
+      local all_sent = table.concat(sent_data)
+      assert.falsy(string.find(all_sent, "CONNECTION:"))
+      assert.falsy(string.find(all_sent, "KEEP%-ALIVE:"))
+    end)
+
+    it('sends end-of-header marker', function()
+      local response = make_response({ headers = {} })
+      response_writer.send_response(mock_sock, response)
+
+      assert.truthy(string.find(sent_data[1], "\r\n\r\n$"))
+    end)
+
+    it('sends body chunks after headers', function()
+      local response = make_response({ chunks = { "hello world" } })
+      response_writer.send_response(mock_sock, response)
+
+      -- sent_data[1] is headers, sent_data[2] is body chunk
+      assert.equal("hello world", sent_data[2])
+    end)
+
+    it('sends multi-chunk body', function()
+      local response = make_response({ chunks = { "chunk1", "chunk2", "chunk3" } })
+      response_writer.send_response(mock_sock, response)
+
+      assert.equal("chunk1", sent_data[2])
+      assert.equal("chunk2", sent_data[3])
+      assert.equal("chunk3", sent_data[4])
+    end)
+
+    it('returns true on success', function()
+      local response = make_response()
+      local ok, err = response_writer.send_response(mock_sock, response)
+
+      assert.is_true(ok)
+      assert.is_nil(err)
+    end)
+
+    it('returns error when headers send fails', function()
+      mock_sock.send = function() return nil, "closed" end
+      local response = make_response()
+      local ok, err = response_writer.send_response(mock_sock, response)
+
+      assert.is_nil(ok)
+      assert.truthy(string.find(err, "failed to send headers"))
+    end)
+
+    it('returns error when body send fails', function()
+      local call_count = 0
+      mock_sock.send = function(_, data)
+        call_count = call_count + 1
+        if call_count == 1 then
+          return #data, nil  -- headers succeed
+        end
+        return nil, "closed"  -- body chunk fails
+      end
+
+      local response = make_response({ chunks = { "hello" } })
+      local ok, err = response_writer.send_response(mock_sock, response)
+
+      assert.is_nil(ok)
+      assert.truthy(string.find(err, "failed to send response body"))
+    end)
+
+    it('returns error when body reader fails', function()
+      local response = {
+        status = 200,
+        reason = "OK",
+        headers = {},
+        body_reader = function()
+          return nil, "read error"
+        end
+      }
+      local ok, err = response_writer.send_response(mock_sock, response)
+
+      assert.is_nil(ok)
+      assert.truthy(string.find(err, "failed to read response body"))
+    end)
+
+    it('returns error when response has no body_reader', function()
+      local response = {
+        status = 200,
+        reason = "OK",
+        headers = {},
+      }
+      local ok, err = response_writer.send_response(mock_sock, response)
+
+      assert.is_nil(ok)
+      assert.equal("no body reader", err)
+    end)
+
+    it('passes chunksize to body_reader', function()
+      local received_chunksize
+      local response = {
+        status = 200,
+        reason = "OK",
+        headers = {},
+        body_reader = function(size)
+          received_chunksize = size
+          return nil
+        end
+      }
+      response_writer.send_response(mock_sock, response, 1024)
+
+      assert.equal(1024, received_chunksize)
+    end)
+
+    it('uses default chunksize of 65536', function()
+      local received_chunksize
+      local response = {
+        status = 200,
+        reason = "OK",
+        headers = {},
+        body_reader = function(size)
+          received_chunksize = size
+          return nil
+        end
+      }
+      response_writer.send_response(mock_sock, response)
+
+      assert.equal(65536, received_chunksize)
+    end)
+
+    it('handles empty body', function()
+      local response = make_response({ chunks = {} })
+      local ok, err = response_writer.send_response(mock_sock, response)
+
+      assert.is_true(ok)
+      assert.is_nil(err)
+    end)
+  end)
+end)
+