From c9846c8d4179246ade8f5cf33ce94f7cfdbda244 Mon Sep 17 00:00:00 2001
From: An Tran <tkan145@gmail.com>
Date: Fri, 18 Oct 2024 15:45:26 +1000
Subject: [PATCH] Introduce APICAST_LUA_SOCKET_KEEPALIVE_REQUESTS env var

Under highload, APIcast keepalive connection could cause unbalance
traffic to backend-listenr.

This PR add a new environment variable to limit the number of request
a single keepalive connection can handle. Once the limit is reached
APIcast will close the connection and open a new one.
---
 CHANGELOG.md                        |  1 +
 doc/parameters.md                   | 10 ++++++++++
 gateway/src/resty/resolver/http.lua | 22 ++++++++++++++++++++++
 3 files changed, 33 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 19f5506ff..3a06aedef 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -33,6 +33,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 - Added the `APICAST_PROXY_BUFFER_SIZE` variable to allow configuration of the buffer size for handling response from the proxied servers. [PR #1473](https://github.com/3scale/APIcast/pull/1473), [THREESCALE-8410](https://issues.redhat.com/browse/THREESCALE-8410)
 
 - Added the `APICAST_HTTPS_VERIFY_CLIENT` variable to allow configuration of the `ssl_verify_client` directive. [PR #1491](https://github.com/3scale/APIcast/pull/1491) [THREESCALE-10156](https://issues.redhat.com/browse/THREESCALE-10156)
+- Add `APICAST_LUA_SOCKET_KEEPALIVE_REQUESTS` to limit the number of requests a single keepalive socket can handle [PR #1496](https://github.com/3scale/APIcast/pull/1496) [THREESCALE-11321](https://issues.redhat.com/browse/THREESCALE-11321)
 
 ## [3.15.0] 2024-04-04
 
diff --git a/doc/parameters.md b/doc/parameters.md
index 3565d9556..ab1e0f4fe 100644
--- a/doc/parameters.md
+++ b/doc/parameters.md
@@ -487,6 +487,16 @@ connections.
 By default Gateway does not enable it, and the keepalive timeout on nginx is set
 to [75 seconds](http://nginx.org/en/docs/http/ngx_http_core_module.html#keepalive_timeout)
 
+### `APICAST_LUA_SOCKET_KEEPALIVE_REQUESTS`
+
+**Value:** positive integers
+**Example:** "1"
+
+Sets the maximum number of requests that one keepalive connection can serve.
+After reaching the limit, the connection closes.
+
+NOTE: This value affects connections opened by APIcast and will not have any
+impact on requests proxied via APIcast.
 
 ### `APICAST_CACHE_STATUS_CODES`
 
diff --git a/gateway/src/resty/resolver/http.lua b/gateway/src/resty/resolver/http.lua
index 4f646ffbc..82e6e2789 100644
--- a/gateway/src/resty/resolver/http.lua
+++ b/gateway/src/resty/resolver/http.lua
@@ -5,6 +5,9 @@ local url_helper = require('resty.url_helper')
 local format = string.format
 
 local setmetatable = setmetatable
+local resty_env = require 'resty.env'
+local tonumber = tonumber
+local keepalive_request = resty_env.get('APICAST_LUA_SOCKET_KEEPALIVE_REQUESTS')
 
 local _M = setmetatable({}, { __index = resty_http })
 
@@ -85,4 +88,23 @@ function _M.connect(self, options, ...)
   return ok, err
 end
 
+function _M:set_keepalive()
+  if keepalive_request then
+    local count, err = resty_http.get_reused_times(self)
+    if err then
+      return nil, err
+    end
+    if count >= tonumber(keepalive_request) then
+      resty_http.close(self)
+      return true
+    end
+  end
+
+  local ok, err = resty_http.set_keepalive(self)
+  if not ok then
+    return nil, err
+  end
+  return true
+end
+
 return _M