This is an automated email from the ASF dual-hosted git repository.

cmcfarlen pushed a commit to branch 10.0.x
in repository https://gitbox.apache.org/repos/asf/trafficserver.git

commit 95ca54586d96b4c60431748ae9de10998f28a7c5
Author: Leif Hedstrom <[email protected]>
AuthorDate: Thu Apr 4 08:07:02 2024 -0600

    Support rate limit on active connections instead of txn (#11172)
    
    (cherry picked from commit 5cb31cd6273f23e0e29e8fe4e2f78c1dc2d3a025)
---
 doc/admin-guide/plugins/rate_limit.en.rst      |  9 ++++++
 plugins/experimental/rate_limit/rate_limit.cc  | 24 +++++++++++---
 plugins/experimental/rate_limit/txn_limiter.cc | 44 +++++++++++++++++++-------
 plugins/experimental/rate_limit/txn_limiter.h  | 14 ++++++--
 4 files changed, 73 insertions(+), 18 deletions(-)

diff --git a/doc/admin-guide/plugins/rate_limit.en.rst 
b/doc/admin-guide/plugins/rate_limit.en.rst
index 8e0f8c8cbd..1b1b85af8a 100644
--- a/doc/admin-guide/plugins/rate_limit.en.rst
+++ b/doc/admin-guide/plugins/rate_limit.en.rst
@@ -101,6 +101,15 @@ are available:
    noting that in the latter exampe, the non-standard scheme and port led to
    ":8080" being appended to the string.
 
+.. option:: --conntrack
+   This flag tells the limiter that rather than limiting the number of active 
transactions,
+   it should limit the number of active connections. This allows an 
established connection
+   to make any number of transactions, but limits the number of connections 
that can be
+   active at any one time.
+
+   Note that it's highly recommended that you keep a very low ``keep-alive`` 
timeout for the
+   connections that are using this rate limiter.
+
 Global Plugin
 -------------
 
diff --git a/plugins/experimental/rate_limit/rate_limit.cc 
b/plugins/experimental/rate_limit/rate_limit.cc
index f7ba630eb2..f67541165f 100644
--- a/plugins/experimental/rate_limit/rate_limit.cc
+++ b/plugins/experimental/rate_limit/rate_limit.cc
@@ -95,8 +95,8 @@ TSRemapNewInstance(int argc, char *argv[], void **ih, char * 
/* errbuf ATS_UNUSE
   limiter->initialize(argc, const_cast<const char **>(argv));
   *ih = static_cast<void *>(limiter);
 
-  Dbg(dbg_ctl, "Added active_in limiter rule (limit=%u, queue=%u, 
max-age=%ldms, error=%u)", limiter->limit(), limiter->max_queue(),
-      static_cast<long>(limiter->max_age().count()), limiter->error());
+  Dbg(dbg_ctl, "Added active_in limiter rule (limit=%u, queue=%u, 
max-age=%ldms, error=%u, conntrack=%s)", limiter->limit(),
+      limiter->max_queue(), static_cast<long>(limiter->max_age().count()), 
limiter->error(), limiter->conntrack() ? "yes" : "no");
 
   return TS_SUCCESS;
 }
@@ -110,6 +110,17 @@ TSRemapDoRemap(void *ih, TSHttpTxn txnp, 
TSRemapRequestInfo *rri)
   auto *limiter = static_cast<TxnRateLimiter *>(ih);
 
   if (limiter) {
+    TSHttpSsn ssnp = TSHttpTxnSsnGet(txnp);
+
+    if (limiter->conntrack()) {
+      int count = TSHttpSsnTransactionCount(ssnp);
+
+      if (count > 1) { // The first transaction is the connect, so we need to 
have at least 2 to be "established"
+        Dbg(dbg_ctl, "Allowing an established connection to pass through, 
txn=%d", count);
+        return TSREMAP_NO_REMAP;
+      }
+    }
+
     if (!limiter->reserve()) {
       if (!limiter->max_queue() || limiter->full()) {
         // We are running at limit, and the queue has reached max capacity, 
give back an error and be done.
@@ -121,8 +132,13 @@ TSRemapDoRemap(void *ih, TSHttpTxn txnp, 
TSRemapRequestInfo *rri)
         Dbg(dbg_ctl, "Adding rate limiting hook, we are at capacity");
       }
     } else {
-      limiter->setupTxnCont(txnp, TS_HTTP_TXN_CLOSE_HOOK);
-      Dbg(dbg_ctl, "Adding txn-close hook, we're not at capacity");
+      if (limiter->conntrack()) {
+        limiter->setupSsnCont(ssnp);
+        Dbg(dbg_ctl, "Adding ssn-close hook, we're not at capacity");
+      } else {
+        limiter->setupTxnCont(txnp, TS_HTTP_TXN_CLOSE_HOOK);
+        Dbg(dbg_ctl, "Adding txn-close hook, we're not at capacity");
+      }
     }
   }
 
diff --git a/plugins/experimental/rate_limit/txn_limiter.cc 
b/plugins/experimental/rate_limit/txn_limiter.cc
index 507c402e0a..294d843efe 100644
--- a/plugins/experimental/rate_limit/txn_limiter.cc
+++ b/plugins/experimental/rate_limit/txn_limiter.cc
@@ -38,6 +38,13 @@ txn_limit_cont(TSCont cont, TSEvent event, void *edata)
     return TS_EVENT_CONTINUE;
     break;
 
+  case TS_EVENT_HTTP_SSN_CLOSE:
+    limiter->free();
+    TSContDestroy(cont); // We are done with this continuation now
+    TSHttpSsnReenable(static_cast<TSHttpSsn>(edata), TS_EVENT_HTTP_CONTINUE);
+    return TS_EVENT_NONE;
+    break;
+
   case TS_EVENT_HTTP_POST_REMAP:
     limiter->push(static_cast<TSHttpTxn>(edata), cont);
     limiter->incrementMetric(RATE_LIMITER_METRIC_QUEUED);
@@ -75,8 +82,8 @@ txn_queue_cont(TSCont cont, TSEvent event, void *edata)
     Dbg(dbg_ctl, "Enabling queued txn after %ldms", 
static_cast<long>(delay.count()));
     // Since this was a delayed transaction, we need to add the TXN_CLOSE hook 
to free the slot when done
     TSHttpTxnHookAdd(txnp, TS_HTTP_TXN_CLOSE_HOOK, contp);
-    TSHttpTxnReenable(txnp, TS_EVENT_HTTP_CONTINUE);
     limiter->incrementMetric(RATE_LIMITER_METRIC_RESUMED);
+    TSHttpTxnReenable(txnp, TS_EVENT_HTTP_CONTINUE);
   }
 
   // Kill any queued txns if they are too old
@@ -92,8 +99,8 @@ txn_queue_cont(TSCont cont, TSEvent event, void *edata)
       Dbg(dbg_ctl, "Queued TXN is too old (%ldms), erroring out", 
static_cast<long>(age.count()));
       TSHttpTxnStatusSet(txnp, static_cast<TSHttpStatus>(limiter->error()));
       TSHttpTxnHookAdd(txnp, TS_HTTP_SEND_RESPONSE_HDR_HOOK, contp);
-      TSHttpTxnReenable(txnp, TS_EVENT_HTTP_ERROR);
       limiter->incrementMetric(RATE_LIMITER_METRIC_EXPIRED);
+      TSHttpTxnReenable(txnp, TS_EVENT_HTTP_ERROR);
     }
   }
 
@@ -107,16 +114,17 @@ bool
 TxnRateLimiter::initialize(int argc, const char *argv[])
 {
   static const struct option longopt[] = {
-    {const_cast<char *>("limit"),  required_argument, nullptr, 'l' },
-    {const_cast<char *>("queue"),  required_argument, nullptr, 'q' },
-    {const_cast<char *>("error"),  required_argument, nullptr, 'e' },
-    {const_cast<char *>("retry"),  required_argument, nullptr, 'r' },
-    {const_cast<char *>("header"), required_argument, nullptr, 'h' },
-    {const_cast<char *>("maxage"), required_argument, nullptr, 'm' },
-    {const_cast<char *>("prefix"), required_argument, nullptr, 'p' },
-    {const_cast<char *>("tag"),    required_argument, nullptr, 't' },
+    {const_cast<char *>("limit"),     required_argument, nullptr, 'l' },
+    {const_cast<char *>("queue"),     required_argument, nullptr, 'q' },
+    {const_cast<char *>("error"),     required_argument, nullptr, 'e' },
+    {const_cast<char *>("retry"),     required_argument, nullptr, 'r' },
+    {const_cast<char *>("header"),    required_argument, nullptr, 'h' },
+    {const_cast<char *>("maxage"),    required_argument, nullptr, 'm' },
+    {const_cast<char *>("prefix"),    required_argument, nullptr, 'p' },
+    {const_cast<char *>("tag"),       required_argument, nullptr, 't' },
+    {const_cast<char *>("conntrack"), no_argument,       nullptr, 'c' },
  // EOF
-    {nullptr,                      no_argument,       nullptr, '\0'},
+    {nullptr,                         no_argument,       nullptr, '\0'},
   };
   optind             = 1;
   std::string prefix = RATE_LIMITER_METRIC_PREFIX;
@@ -150,6 +158,9 @@ TxnRateLimiter::initialize(int argc, const char *argv[])
     case 't':
       tag = optarg;
       break;
+    case 'c':
+      this->_conntrack = true;
+      break;
     }
     if (opt == -1) {
       break;
@@ -180,3 +191,14 @@ TxnRateLimiter::setupTxnCont(TSHttpTxn txnp, TSHttpHookID 
hook)
   TSContDataSet(cont, this);
   TSHttpTxnHookAdd(txnp, hook, cont);
 }
+
+// This only needs the TS_HTTP_SSN_CLOSE_HOOK, for now at least, so not passed 
as argument.
+void
+TxnRateLimiter::setupSsnCont(TSHttpSsn ssnp)
+{
+  TSCont cont = TSContCreate(txn_limit_cont, nullptr);
+  TSReleaseAssert(cont);
+
+  TSContDataSet(cont, this);
+  TSHttpSsnHookAdd(ssnp, TS_HTTP_SSN_CLOSE_HOOK, cont);
+}
diff --git a/plugins/experimental/rate_limit/txn_limiter.h 
b/plugins/experimental/rate_limit/txn_limiter.h
index 6c7a651226..434a66ecbd 100644
--- a/plugins/experimental/rate_limit/txn_limiter.h
+++ b/plugins/experimental/rate_limit/txn_limiter.h
@@ -37,6 +37,7 @@ public:
   }
 
   void setupTxnCont(TSHttpTxn txnp, TSHttpHookID hook);
+  void setupSsnCont(TSHttpSsn ssnp);
   bool initialize(int argc, const char *argv[]);
 
   const std::string &
@@ -57,10 +58,17 @@ public:
     return _retry;
   }
 
+  bool
+  conntrack() const
+  {
+    return _conntrack;
+  }
+
 private:
-  std::string _header = "";  // Header to put the latency metrics in, e.g. 
@RateLimit-Delay
-  unsigned _error     = 429; // Error code when we decide not to allow a txn 
to be processed (e.g. queue full)
-  unsigned _retry     = 0;   // If > 0, we will also send a Retry-After: 
header with this retry value
+  std::string _header = "";    // Header to put the latency metrics in, e.g. 
@RateLimit-Delay
+  unsigned _error     = 429;   // Error code when we decide not to allow a txn 
to be processed (e.g. queue full)
+  unsigned _retry     = 0;     // If > 0, we will also send a Retry-After: 
header with this retry value
+  bool _conntrack     = false; // If true, we will track connections and limit 
based on that instead of transactions
 
   TSCont _queue_cont = nullptr; // Continuation processing the queue 
periodically
   TSAction _action   = nullptr; // The action associated with the queue 
continuation, needed to shut it down

Reply via email to