Add throttle middleware. (#76)

* Limit the number of concurrent connections to the API service.
* Add api.maxconns configuration parameter.

Author: Mike Davis

pkg/api/router.go

@@ -23,6 +23,7 @@ import (
 	"github.com/optimizely/sidedoor/pkg/api/handlers"
 	"github.com/optimizely/sidedoor/pkg/api/middleware"
 	"github.com/optimizely/sidedoor/pkg/optimizely"
+	"github.com/spf13/viper"
 
 	"github.com/go-chi/chi"
 	chimw "github.com/go-chi/chi/middleware"
@@ -31,6 +32,7 @@ import (
 
 // RouterOptions defines the configuration parameters for Router.
 type RouterOptions struct {
+	maxConns     int
 	middleware   middleware.OptlyMiddleware
 	featureAPI   handlers.FeatureAPI
 	userEventAPI handlers.UserEventAPI
@@ -40,6 +42,7 @@ type RouterOptions struct {
 // NewDefaultRouter creates a new router with the default backing optimizely.Cache
 func NewDefaultRouter(optlyCache optimizely.Cache) http.Handler {
 	spec := &RouterOptions{
+		maxConns:     viper.GetInt("api.maxconns"),
 		middleware:   &middleware.CachedOptlyMiddleware{Cache: optlyCache},
 		featureAPI:   new(handlers.FeatureHandler),
 		userEventAPI: new(handlers.UserEventHandler),
@@ -53,6 +56,11 @@ func NewDefaultRouter(optlyCache optimizely.Cache) http.Handler {
 func NewRouter(opt *RouterOptions) *chi.Mux {
 	r := chi.NewRouter()
 
+	if opt.maxConns > 0 {
+		// Note this is NOT a rate limiter, but a concurrency threshold
+		r.Use(chimw.Throttle(opt.maxConns))
+	}
+
 	r.Use(middleware.SetTime)
 	r.Use(render.SetContentType(render.ContentTypeJSON), middleware.SetRequestID)
 

pkg/api/router_test.go

@@ -124,6 +124,7 @@ func (suite *RouterTestSuite) SetupTest() {
 		suite.tc = testClient
 
 		opts := &RouterOptions{
+			maxConns:     1,
 			featureAPI:   new(MockFeatureAPI),
 			userEventAPI: new(MockUserEventAPI),
 			userAPI:      new(MockUserAPI),
@@ -255,6 +256,30 @@ func (suite *RouterTestSuite) TestRemoveForcedVariation() {
 	suite.assertValid(rec, expected)
 }
 
+func (suite *RouterTestSuite) TestThrottleConfig() {
+	req := httptest.NewRequest("GET", "/throttled", nil)
+
+	wg1 := sync.WaitGroup{}
+	wg1.Add(1)
+	suite.mux.Get("/throttled", func(w http.ResponseWriter, r *http.Request) {
+		wg1.Wait()
+	})
+
+	wg2 := sync.WaitGroup{}
+	wg2.Add(1)
+	go func() {
+		wg2.Done()
+		rec := httptest.NewRecorder()
+		suite.mux.ServeHTTP(rec, req)
+	}()
+	wg2.Wait()
+
+	rec := httptest.NewRecorder()
+	suite.mux.ServeHTTP(rec, req)
+	suite.Equal(http.StatusServiceUnavailable, rec.Code)
+	wg1.Done()
+}
+
 func (suite *RouterTestSuite) assertValid(rec *httptest.ResponseRecorder, expected map[string]string) {
 	suite.Equal(http.StatusOK, rec.Code)