Skip to content

Commit 7132fb7

Browse files
author
Mike Davis
authored
Add throttle middleware. (#76)
* Limit the number of concurrent connections to the API service. * Add api.maxconns configuration parameter.
1 parent 9c23077 commit 7132fb7

File tree

2 files changed

+33
-0
lines changed

2 files changed

+33
-0
lines changed

pkg/api/router.go

+8
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ import (
2323
"github.com/optimizely/sidedoor/pkg/api/handlers"
2424
"github.com/optimizely/sidedoor/pkg/api/middleware"
2525
"github.com/optimizely/sidedoor/pkg/optimizely"
26+
"github.com/spf13/viper"
2627

2728
"github.com/go-chi/chi"
2829
chimw "github.com/go-chi/chi/middleware"
@@ -31,6 +32,7 @@ import (
3132

3233
// RouterOptions defines the configuration parameters for Router.
3334
type RouterOptions struct {
35+
maxConns int
3436
middleware middleware.OptlyMiddleware
3537
featureAPI handlers.FeatureAPI
3638
userEventAPI handlers.UserEventAPI
@@ -40,6 +42,7 @@ type RouterOptions struct {
4042
// NewDefaultRouter creates a new router with the default backing optimizely.Cache
4143
func NewDefaultRouter(optlyCache optimizely.Cache) http.Handler {
4244
spec := &RouterOptions{
45+
maxConns: viper.GetInt("api.maxconns"),
4346
middleware: &middleware.CachedOptlyMiddleware{Cache: optlyCache},
4447
featureAPI: new(handlers.FeatureHandler),
4548
userEventAPI: new(handlers.UserEventHandler),
@@ -53,6 +56,11 @@ func NewDefaultRouter(optlyCache optimizely.Cache) http.Handler {
5356
func NewRouter(opt *RouterOptions) *chi.Mux {
5457
r := chi.NewRouter()
5558

59+
if opt.maxConns > 0 {
60+
// Note this is NOT a rate limiter, but a concurrency threshold
61+
r.Use(chimw.Throttle(opt.maxConns))
62+
}
63+
5664
r.Use(middleware.SetTime)
5765
r.Use(render.SetContentType(render.ContentTypeJSON), middleware.SetRequestID)
5866

pkg/api/router_test.go

+25
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,7 @@ func (suite *RouterTestSuite) SetupTest() {
124124
suite.tc = testClient
125125

126126
opts := &RouterOptions{
127+
maxConns: 1,
127128
featureAPI: new(MockFeatureAPI),
128129
userEventAPI: new(MockUserEventAPI),
129130
userAPI: new(MockUserAPI),
@@ -255,6 +256,30 @@ func (suite *RouterTestSuite) TestRemoveForcedVariation() {
255256
suite.assertValid(rec, expected)
256257
}
257258

259+
func (suite *RouterTestSuite) TestThrottleConfig() {
260+
req := httptest.NewRequest("GET", "/throttled", nil)
261+
262+
wg1 := sync.WaitGroup{}
263+
wg1.Add(1)
264+
suite.mux.Get("/throttled", func(w http.ResponseWriter, r *http.Request) {
265+
wg1.Wait()
266+
})
267+
268+
wg2 := sync.WaitGroup{}
269+
wg2.Add(1)
270+
go func() {
271+
wg2.Done()
272+
rec := httptest.NewRecorder()
273+
suite.mux.ServeHTTP(rec, req)
274+
}()
275+
wg2.Wait()
276+
277+
rec := httptest.NewRecorder()
278+
suite.mux.ServeHTTP(rec, req)
279+
suite.Equal(http.StatusServiceUnavailable, rec.Code)
280+
wg1.Done()
281+
}
282+
258283
func (suite *RouterTestSuite) assertValid(rec *httptest.ResponseRecorder, expected map[string]string) {
259284
suite.Equal(http.StatusOK, rec.Code)
260285

0 commit comments

Comments
 (0)