From 936a916e784f46b9edb9dd98797c67b6d81b5732 Mon Sep 17 00:00:00 2001
From: Harshavardhana <harsha@minio.io>
Date: Sun, 26 Jun 2016 03:18:07 -0700
Subject: [PATCH] server: Add connection rate limiter for server. (#1977)

---
 globals.go            |  5 ++++
 rate-limit-handler.go | 63 +++++++++++++++++++++++++++++++++++++++++++
 routers.go            |  2 ++
 server-main.go        |  8 ++++++
 4 files changed, 78 insertions(+)
 create mode 100644 rate-limit-handler.go

diff --git a/globals.go b/globals.go
index 523b535e8..0ffb61a6f 100644
--- a/globals.go
+++ b/globals.go
@@ -39,6 +39,11 @@ var (
 	globalQuiet = false // Quiet flag set via command line
 	globalTrace = false // Trace flag set via environment setting.
 	// Add new global flags here.
+
+	// Maximum connections handled per
+	// server, defaults to 0 (unlimited).
+	globalMaxConn = 0
+	// Add new variable global values here.
 )
 
 // global colors.
diff --git a/rate-limit-handler.go b/rate-limit-handler.go
new file mode 100644
index 000000000..2972be213
--- /dev/null
+++ b/rate-limit-handler.go
@@ -0,0 +1,63 @@
+/*
+ * Minio Cloud Storage, (C) 2016 Minio, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package main
+
+import (
+	"net/http"
+	"sync"
+)
+
+// rateLimit - represents datatype of the functionality implemented to
+// limit the number of concurrent http requests.
+type rateLimit struct {
+	handler     http.Handler
+	rqueue      chan struct{}
+	releaseOnce sync.Once
+}
+
+// acquire and release implement a way to send and receive from the
+// channel this is in-turn used to rate limit incoming connections in
+// ServeHTTP() http.Handler method.
+func (c *rateLimit) acquire() { c.rqueue <- struct{}{} }
+func (c *rateLimit) release() { <-c.rqueue }
+
+// ServeHTTP is an http.Handler ServeHTTP method, implemented to rate
+// limit incoming HTTP requests.
+func (c *rateLimit) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	// Acquire the connection if queue is not full, otherwise
+	// code path waits here until the previous case is true.
+	c.acquire()
+
+	// Serves the request.
+	c.handler.ServeHTTP(w, r)
+
+	// Release by draining the channel once.
+	c.releaseOnce.Do(c.release)
+}
+
+// setRateLimitHandler limits the number of concurrent http requests based on MINIO_MAXCONN.
+func setRateLimitHandler(handler http.Handler) http.Handler {
+	if globalMaxConn == 0 {
+		return handler
+	} // else proceed to rate limiting.
+
+	// For max connection limit of > '0' we initialize rate limit handler.
+	return &rateLimit{
+		handler: handler,
+		rqueue:  make(chan struct{}, globalMaxConn),
+	}
+}
diff --git a/routers.go b/routers.go
index 78c742a1f..72e866425 100644
--- a/routers.go
+++ b/routers.go
@@ -70,6 +70,8 @@ func configureServerHandler(srvCmdConfig serverCmdConfig) http.Handler {
 	// List of some generic handlers which are applied for all
 	// incoming requests.
 	var handlerFns = []HandlerFunc{
+		// Limits the number of concurrent http requests.
+		setRateLimitHandler,
 		// Redirect some pre-defined browser request paths to a static
 		// location prefix.
 		setBrowserRedirectHandler,
diff --git a/server-main.go b/server-main.go
index fd72e2119..eef11f01d 100644
--- a/server-main.go
+++ b/server-main.go
@@ -133,6 +133,14 @@ func initServerConfig(c *cli.Context) {
 	err := serverConfig.Save()
 	fatalIf(err, "Unable to save config.")
 
+	// Fetch max conn limit from environment variable.
+	if maxConnStr := os.Getenv("MINIO_MAXCONN"); maxConnStr != "" {
+		// We need to parse to its integer value.
+		var err error
+		globalMaxConn, err = strconv.Atoi(maxConnStr)
+		fatalIf(err, "Unable to convert MINIO_MAXCONN=%s environment variable into its integer value.", maxConnStr)
+	}
+
 	// Fetch access keys from environment variables if any and update the config.
 	accessKey := os.Getenv("MINIO_ACCESS_KEY")
 	secretKey := os.Getenv("MINIO_SECRET_KEY")