perf(http_proxy): stream the response to avoid memory bloat

Author: tkan145

gateway/src/apicast/http_proxy.lua

@@ -13,7 +13,7 @@ local file_reader = require("resty.file").file_reader
 local file_size = require("resty.file").file_size
 local client_body_reader = require("resty.http.request_reader").get_client_body_reader
 local send_response = require("resty.http.response_writer").send_response
-local concat = table.concat
+local proxy_response = require("resty.http.response_writer").proxy_response
 
 local _M = { }
 
@@ -163,13 +163,18 @@ local function forward_https_request(proxy_uri, uri, proxy_opts)
 
     if res then
         if opts.request_unbuffered and raw then
-            local bytes, err = send_response(sock, res, DEFAULT_CHUNKSIZE)
-            if not bytes then
+            err = send_response(sock, res, DEFAULT_CHUNKSIZE)
+            if err then
                 ngx.log(ngx.ERR, "failed to send response: ", err)
-                return sock:send("HTTP/1.1 502 Bad Gateway")
+                sock:close()
+              return ngx.exit(ngx.HTTP_BAD_GATEWAY)
             end
         else
-            httpc:proxy_response(res)
+            err = proxy_response(res, DEFAULT_CHUNKSIZE)
+            if err then
+              ngx.log(ngx.ERR, 'failed to proxy request to: ', proxy_uri, ' err : ', err)
+              return ngx.exit(ngx.HTTP_BAD_GATEWAY)
+            end
             httpc:set_keepalive()
         end
     else

gateway/src/resty/http/response_writer.lua

@@ -3,6 +3,8 @@ local str_lower = string.lower
 local insert = table.insert
 local concat = table.concat
 
+local ngx       = ngx
+
 local _M = {
 }
 
@@ -31,6 +33,45 @@ local function send(socket, data)
     return socket:send(data)
 end
 
+local function send_chunk(chunk)
+    if not chunk then
+        return nil
+    end
+
+    local ok, err = ngx.print(chunk)
+    if not ok then
+        return "output response failed: " .. (err or "")
+    end
+
+    return nil
+end
+
+-- forward_body reads chunks from a body_reader and passes them to the callback
+-- function cb.
+-- cb(chunk) should return a true on success, or nil/false, err on failure.
+local function forward_body(reader, cb, chunksize)
+  if not reader then
+      return "no body reader"
+  end
+
+  local buffer_size = chunksize or 65536
+
+  repeat
+      local buffer, read_err, send_err
+      buffer, read_err = reader(buffer_size)
+      if read_err then
+          return "failed to read response body: " .. read_err
+      end
+
+      if buffer then
+          send_err = cb(buffer)
+          if send_err then
+              return "failed to send response body: " .. (send_err or "unknown")
+          end
+      end
+  until not buffer
+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)
@@ -42,7 +83,7 @@ function _M.send_response(sock, response, chunksize)
     end
 
     if not sock then
-        return nil, "socket not initialized yet"
+        return "socket not initialized yet"
     end
 
     -- Build status line + headers into a single buffer to minimize send() calls
@@ -62,32 +103,31 @@ function _M.send_response(sock, response, chunksize)
 
     local bytes, err = sock:send(concat(buf))
     if not bytes then
-        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"
+        return "failed to send headers, err: " .. (err or "unknown")
     end
 
-    repeat
-        local chunk, read_err
-
-        chunk, read_err = reader(chunksize)
-        if read_err then
-            return nil, "failed to read response body, err: " .. (err or "unknown")
+    return forward_body(response.body_reader, function(chunk)
+        bytes, err = send(sock, chunk)
+        if not bytes then
+            return "failed to send response body, err: " .. (err or "unknown")
         end
+    end, chunksize)
+end
 
-        if chunk then
-            bytes, err = send(sock, chunk)
-            if not bytes then
-                return nil, "failed to send response body, err: " .. (err or "unknown")
-            end
-        end
-    until not chunk
+function _M.proxy_response(res, chunksize)
+  if not res then
+      ngx.log(ngx.ERR, "no response provided")
+      return
+  end
+
+  ngx.status = res.status
+  for k, v in pairs(res.headers) do
+    if not HOP_BY_HOP_HEADERS[str_lower(k)] then
+        ngx.header[k] = v
+     end
+  end
 
-    return true, nil
+  return forward_body(res.body_reader, send_chunk, chunksize)
 end
 
 return _M

spec/http_proxy_spec.lua

@@ -20,6 +20,8 @@ describe('http_proxy', function()
 
       local resty_http_proxy = require 'resty.http.proxy'
       stub(resty_http_proxy, 'new', function() return httpc end)
+      local http_writer = require 'resty.http.response_writer'
+      stub(http_writer, 'proxy_response')
     end
 
     before_each(function()

spec/resty/http/response_writer_spec.lua

@@ -36,14 +36,12 @@ describe('resty.http.response_writer', function()
   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)
+      local err = response_writer.send_response(mock_sock, nil)
       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)
+      local err = response_writer.send_response(nil, make_response())
       assert.equal("socket not initialized yet", err)
     end)
 
@@ -145,18 +143,16 @@ describe('resty.http.response_writer', function()
 
     it('returns true on success', function()
       local response = make_response()
-      local ok, err = response_writer.send_response(mock_sock, response)
+      local 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)
+      local err = response_writer.send_response(mock_sock, response)
 
-      assert.is_nil(ok)
       assert.truthy(string.find(err, "failed to send headers"))
     end)
 
@@ -171,9 +167,8 @@ describe('resty.http.response_writer', function()
       end
 
       local response = make_response({ chunks = { "hello" } })
-      local ok, err = response_writer.send_response(mock_sock, response)
+      local err = response_writer.send_response(mock_sock, response)
 
-      assert.is_nil(ok)
       assert.truthy(string.find(err, "failed to send response body"))
     end)
 
@@ -186,9 +181,8 @@ describe('resty.http.response_writer', function()
           return nil, "read error"
         end
       }
-      local ok, err = response_writer.send_response(mock_sock, response)
+      local err = response_writer.send_response(mock_sock, response)
 
-      assert.is_nil(ok)
       assert.truthy(string.find(err, "failed to read response body"))
     end)
 
@@ -198,9 +192,8 @@ describe('resty.http.response_writer', function()
         reason = "OK",
         headers = {},
       }
-      local ok, err = response_writer.send_response(mock_sock, response)
+      local err = response_writer.send_response(mock_sock, response)
 
-      assert.is_nil(ok)
       assert.equal("no body reader", err)
     end)
 
@@ -238,10 +231,168 @@ describe('resty.http.response_writer', function()
 
     it('handles empty body', function()
       local response = make_response({ chunks = {} })
-      local ok, err = response_writer.send_response(mock_sock, response)
+      local err = response_writer.send_response(mock_sock, response)
+
+      assert.is_nil(err)
+    end)
+  end)
+
+  describe('.proxy_response', function()
+    local printed_data
+    local headers_set
+
+    before_each(function()
+      printed_data = {}
+      headers_set = {}
+
+      ngx.status = nil
+      ngx.header = setmetatable({}, {
+        __newindex = function(_, k, v)
+          headers_set[k] = v
+        end
+      })
+
+      stub(ngx, 'print', function(data)
+        table.insert(printed_data, data)
+        return true
+      end)
+      stub(ngx, 'flush', function() return true end)
+    end)
+
+    it('returns nil when no response is provided', function()
+      local ok = response_writer.proxy_response(nil)
+      assert.is_nil(ok)
+    end)
+
+    it('sets ngx.status from response', function()
+      local res = make_response({ status = 201 })
+      local err = response_writer.proxy_response(res)
+
+      assert.is_nil(err)
+      assert.equal(201, ngx.status)
+    end)
+
+    it('sets response headers on ngx.header', function()
+      local res = make_response({
+        headers = {
+          ["Content-Type"] = "application/json",
+          ["X-Request-Id"] = "abc123",
+        }
+      })
+      local err = response_writer.proxy_response(res)
+
+      assert.is_nil(err)
+      assert.equal("application/json", headers_set["Content-Type"])
+      assert.equal("abc123", headers_set["X-Request-Id"])
+    end)
+
+    it('filters hop-by-hop headers', function()
+      local res = 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"] = "kept",
+        }
+      })
+      local err = response_writer.proxy_response(res)
+
+      assert.is_nil(err)
+      assert.is_nil(headers_set["Connection"])
+      assert.is_nil(headers_set["Keep-Alive"])
+      assert.is_nil(headers_set["Transfer-Encoding"])
+      assert.is_nil(headers_set["Proxy-Authenticate"])
+      assert.is_nil(headers_set["Proxy-Authorization"])
+      assert.is_nil(headers_set["TE"])
+      assert.is_nil(headers_set["Trailers"])
+      assert.is_nil(headers_set["Upgrade"])
+      assert.is_nil(headers_set["Content-Length"])
+      assert.equal("kept", headers_set["X-Custom"])
+    end)
+
+    it('filters hop-by-hop headers case-insensitively', function()
+      local res = make_response({
+        headers = {
+          ["CONNECTION"] = "close",
+          ["TRANSFER-ENCODING"] = "chunked",
+        }
+      })
+      local err = response_writer.proxy_response(res)
+
+      assert.is_nil(err)
+      assert.is_nil(headers_set["CONNECTION"])
+      assert.is_nil(headers_set["TRANSFER-ENCODING"])
+    end)
+
+    it('prints and flushes body chunks', function()
+      local res = make_response({ chunks = { "chunk1", "chunk2" } })
+      local err = response_writer.proxy_response(res)
 
-      assert.is_true(ok)
       assert.is_nil(err)
+      assert.stub(ngx.print).was_called(2)
+      assert.stub(ngx.print).was_called_with("chunk1")
+      assert.stub(ngx.print).was_called_with("chunk2")
+      -- assert.stub(ngx.flush).was_called(2)
+    end)
+
+    it('returns true on success', function()
+      -- local res = make_response({ chunks = { "data" } })
+      -- local ok, err = response_writer.proxy_response(res)
+
+      -- assert.is_true(ok)
+      -- assert.is_nil(err)
+    end)
+
+    it('handles empty body', function()
+      local res = make_response({ chunks = {} })
+      local err = response_writer.proxy_response(res)
+
+      assert.is_nil(err)
+      assert.stub(ngx.print).was_not_called()
+    end)
+
+    it('returns error when ngx.print fails', function()
+      ngx.print:revert()
+      stub(ngx, 'print', function() return nil, "broken pipe" end)
+
+      local res = make_response({ chunks = { "data" } })
+      local err = response_writer.proxy_response(res)
+
+      assert.truthy(string.find(err, "output response failed"))
+    end)
+
+    it('returns error when body reader fails', function()
+      local res = {
+        status = 200,
+        headers = {},
+        body_reader = function()
+          return nil, "read timeout"
+        end
+      }
+      local err = response_writer.proxy_response(res)
+
+      assert.truthy(string.find(err, "failed to read response body"))
+    end)
+
+    it('passes chunksize to body_reader', function()
+      local received_chunksize
+      local res = {
+        status = 200,
+        headers = {},
+        body_reader = function(size)
+          received_chunksize = size
+          return nil
+        end
+      }
+      response_writer.proxy_response(res, 2048)
+
+      assert.equal(2048, received_chunksize)
     end)
   end)
 end)