Skip to content

Commit

Permalink
feat: support global concurrent limiter
Browse files Browse the repository at this point in the history
  • Loading branch information
vicanso committed Aug 6, 2020
1 parent 20a9d2f commit 95b5ae9
Show file tree
Hide file tree
Showing 3 changed files with 88 additions and 1 deletion.
37 changes: 37 additions & 0 deletions docs/middlewares.md
Original file line number Diff line number Diff line change
Expand Up @@ -190,6 +190,43 @@ func main() {
```


## global concurrent limiter

全局的并发请求限制,可以用于控制应用的并发请求量。

**Example**
```go
package main

import (
"bytes"
"sync"
"time"

"github.com/vicanso/elton"
"github.com/vicanso/elton/middleware"
)

func main() {

e := elton.New()
e.Use(middleware.NewGlobalConcurrentLimiter(middleware.GlobalConcurrentLimiterConfig{
Max: 1000,
}))

e.POST("/login", func(c *elton.Context) (err error) {
time.Sleep(3 * time.Second)
c.BodyBuffer = bytes.NewBufferString("hello world")
return
})

err := e.ListenAndServe(":3000")
if err != nil {
panic(err)
}
}
```

## concurrent limiter

并发请求限制,可以通过指定请求的参数,如IP、query的字段或者body等获取,限制同时并发性的提交请求,主要用于避免相同的请求多次提交。指定的Key分为以下几种:
Expand Down
28 changes: 27 additions & 1 deletion middleware/concurrent_limiter.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ import (
"errors"
"net/http"
"strings"
"sync/atomic"

"github.com/tidwall/gjson"
"github.com/vicanso/elton"
Expand All @@ -39,6 +40,12 @@ var (
Message: "submit too frequently",
Category: ErrConcurrentLimiterCategory,
}
// ErrTooManyRequests too many request
ErrTooManyRequests = &hes.Error{
StatusCode: http.StatusTooManyRequests,
Message: "Too Many Requests",
Category: ErrConcurrentLimiterCategory,
}
ErrNotAllowEmpty = &hes.Error{
StatusCode: http.StatusBadRequest,
Message: "empty value is not allowed",
Expand Down Expand Up @@ -78,9 +85,14 @@ type (
Body bool
IP bool
}
// GlobalConcurrentLimiterConfig
GlobalConcurrentLimiterConfig struct {
Skipper elton.Skipper
Max uint32
}
)

// New create a concurrent limiter middleware
// NewConcurrentLimiter create a concurrent limiter middleware
func NewConcurrentLimiter(config ConcurrentLimiterConfig) elton.Handler {

if config.Lock == nil {
Expand Down Expand Up @@ -177,3 +189,17 @@ func NewConcurrentLimiter(config ConcurrentLimiterConfig) elton.Handler {
return c.Next()
}
}

// NewGlobalConcurrentLimiter create a new global concurrent limiter
func NewGlobalConcurrentLimiter(config GlobalConcurrentLimiterConfig) elton.Handler {
var count uint32
return func(c *elton.Context) (err error) {
value := atomic.AddUint32(&count, 1)
defer atomic.AddUint32(&count, ^uint32(0))
if value >= config.Max {
err = ErrTooManyRequests
return
}
return c.Next()
}
}
24 changes: 24 additions & 0 deletions middleware/concurrent_limiter_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -136,3 +136,27 @@ func TestConcurrentLimiter(t *testing.T) {
assert.Equal(ErrNotAllowEmpty, err)
})
}

func TestGlobalConcurrentLimiter(t *testing.T) {
assert := assert.New(t)
fn := NewGlobalConcurrentLimiter(GlobalConcurrentLimiterConfig{
Max: 1,
})
req := httptest.NewRequest("POST", "/users/login?type=1", nil)
resp := httptest.NewRecorder()
c := elton.NewContext(resp, req)
err := fn(c)
assert.Equal(ErrTooManyRequests, err)

fn = NewGlobalConcurrentLimiter(GlobalConcurrentLimiterConfig{
Max: 2,
})
done := false
c.Next = func() error {
done = true
return nil
}
err = fn(c)
assert.Nil(err)
assert.True(done)
}

0 comments on commit 95b5ae9

Please sign in to comment.