nghttpx: Add rate limit options

src/shrpx.cc

@@ -400,10 +400,22 @@ void fill_default_config()
   mod_config()->downstream_http_proxy_host = 0;
   mod_config()->downstream_http_proxy_port = 0;
   mod_config()->downstream_http_proxy_addrlen = 0;
-  mod_config()->rate_limit_cfg =
+  mod_config()->rate_limit_cfg = nullptr;
-    ev_token_bucket_cfg_new(1024*1024, 4*1024*1024,
+  mod_config()->read_rate = 1024*1024;
-                            EV_RATE_LIMIT_MAX, EV_RATE_LIMIT_MAX,
+  mod_config()->read_burst = 4*1024*1024;
-                            nullptr);
+  mod_config()->write_rate = 0;
+  mod_config()->write_burst = 0;
+}
+} // namespace
+
+namespace {
+size_t get_rate_limit(size_t rate_limit)
+{
+  if(rate_limit == 0) {
+    return EV_RATE_LIMIT_MAX;
+  } else {
+    return rate_limit;
+  }
 }
 } // namespace
 
@@ -461,6 +473,29 @@ void print_help(std::ostream& out)
       << "                       Set the number of worker threads.\n"
       << "                       Default: "
       << get_config()->num_worker << "\n"
+      << "    --read-rate=<RATE> Set maximum average read rate on frontend\n"
+      << "                       connection. Setting 0 to this option means\n"
+      << "                       read rate is unlimited.\n"
+      << "                       Default: "
+      << get_config()->read_rate << "\n"
+      << "    --read-burst=<SIZE>\n"
+      << "                       Set maximum read burst size on frontend\n"
+      << "                       connection. Setting 0 to this option means\n"
+      << "                       read burst size is unlimited.\n"
+      << "                       Default: "
+      << get_config()->read_burst << "\n"
+      << "    --write-rate=<RATE>\n"
+      << "                       Set maximum average write rate on frontend\n"
+      << "                       connection. Setting 0 to this option means\n"
+      << "                       write rate is unlimited.\n"
+      << "                       Default: "
+      << get_config()->write_rate << "\n"
+      << "    --write-burst=<SIZE>\n"
+      << "                       Set maximum write burst size on frontend\n"
+      << "                       connection. Setting 0 to this option means\n"
+      << "                       write burst size is unlimited.\n"
+      << "                       Default: "
+      << get_config()->write_burst << "\n"
       << "\n"
       << "  Timeout:\n"
       << "    --frontend-spdy-read-timeout=<SEC>\n"
@@ -671,6 +706,10 @@ int main(int argc, char **argv)
       {"backend-tls-sni-field", required_argument, &flag, 31},
       {"honor-cipher-order", no_argument, &flag, 32},
       {"dh-param-file", required_argument, &flag, 33},
+      {"read-rate", required_argument, &flag, 34},
+      {"read-burst", required_argument, &flag, 35},
+      {"write-rate", required_argument, &flag, 36},
+      {"write-burst", required_argument, &flag, 37},
       {nullptr, 0, nullptr, 0 }
     };
     int option_index = 0;
@@ -854,6 +893,22 @@ int main(int argc, char **argv)
         // --dh-param-file
         cmdcfgs.push_back(std::make_pair(SHRPX_OPT_DH_PARAM_FILE, optarg));
         break;
+      case 34:
+        // --read-rate
+        cmdcfgs.push_back(std::make_pair(SHRPX_OPT_READ_RATE, optarg));
+        break;
+      case 35:
+        // --read-burst
+        cmdcfgs.push_back(std::make_pair(SHRPX_OPT_READ_BURST, optarg));
+        break;
+      case 36:
+        // --write-rate
+        cmdcfgs.push_back(std::make_pair(SHRPX_OPT_WRITE_RATE, optarg));
+        break;
+      case 37:
+        // --write-burst
+        cmdcfgs.push_back(std::make_pair(SHRPX_OPT_WRITE_BURST, optarg));
+        break;
       default:
         break;
       }
@@ -980,6 +1035,13 @@ int main(int argc, char **argv)
     mod_config()->use_syslog = true;
   }
 
+  mod_config()->rate_limit_cfg = ev_token_bucket_cfg_new
+    (get_rate_limit(get_config()->read_rate),
+     get_rate_limit(get_config()->read_burst),
+     get_rate_limit(get_config()->write_rate),
+     get_rate_limit(get_config()->write_burst),
+     nullptr);
+
   struct sigaction act;
   memset(&act, 0, sizeof(struct sigaction));
   act.sa_handler = SIG_IGN;

src/shrpx_config.cc

@@ -95,6 +95,10 @@ const char SHRPX_OPT_CACERT[] = "cacert";
 const char SHRPX_OPT_BACKEND_IPV4[] = "backend-ipv4";
 const char SHRPX_OPT_BACKEND_IPV6[] = "backend-ipv6";
 const char SHRPX_OPT_BACKEND_HTTP_PROXY_URI[] = "backend-http-proxy-uri";
+const char SHRPX_OPT_READ_RATE[] = "read-rate";
+const char SHRPX_OPT_READ_BURST[] = "read-burst";
+const char SHRPX_OPT_WRITE_RATE[] = "write-rate";
+const char SHRPX_OPT_WRITE_BURST[] = "write-burst";
 
 namespace {
 Config *config = 0;
@@ -373,6 +377,14 @@ int parse_config(const char *opt, const char *optarg)
       LOG(ERROR) << "Could not parse backend-http-proxy-uri";
         return -1;
     }
+  } else if(util::strieq(opt, SHRPX_OPT_READ_RATE)) {
+    mod_config()->read_rate = strtoul(optarg, nullptr, 10);
+  } else if(util::strieq(opt, SHRPX_OPT_READ_BURST)) {
+    mod_config()->read_burst = strtoul(optarg, nullptr, 10);
+  } else if(util::strieq(opt, SHRPX_OPT_WRITE_RATE)) {
+    mod_config()->write_rate = strtoul(optarg, nullptr, 10);
+  } else if(util::strieq(opt, SHRPX_OPT_WRITE_BURST)) {
+    mod_config()->write_burst = strtoul(optarg, nullptr, 10);
   } else if(util::strieq(opt, "conf")) {
     LOG(WARNING) << "conf is ignored";
   } else {

src/shrpx_config.h

@@ -86,6 +86,10 @@ extern const char SHRPX_OPT_BACKEND_IPV4[];
 extern const char SHRPX_OPT_BACKEND_IPV6[];
 extern const char SHRPX_OPT_BACKEND_HTTP_PROXY_URI[];
 extern const char SHRPX_OPT_BACKEND_TLS_SNI_FIELD[];
+extern const char SHRPX_OPT_READ_RATE[];
+extern const char SHRPX_OPT_READ_BURST[];
+extern const char SHRPX_OPT_WRITE_RATE[];
+extern const char SHRPX_OPT_WRITE_BURST[];
 
 union sockaddr_union {
   sockaddr sa;
@@ -170,6 +174,10 @@ struct Config {
   size_t downstream_http_proxy_addrlen;
   // Rate limit configuration
   ev_token_bucket_cfg *rate_limit_cfg;
+  size_t read_rate;
+  size_t read_burst;
+  size_t write_rate;
+  size_t write_burst;
 };
 
 const Config* get_config();