40 lines
745 B
Go
40 lines
745 B
Go
package rate
|
|
|
|
import (
|
|
"errors"
|
|
|
|
"github.com/juju/ratelimit"
|
|
)
|
|
|
|
type Limit struct {
|
|
*ratelimit.Bucket
|
|
}
|
|
|
|
func New(opts ...Option) (*Limit, error) {
|
|
opt := new(options)
|
|
for _, o := range opts {
|
|
o(opt)
|
|
}
|
|
|
|
res := new(Limit)
|
|
if opt.fillInterval > 0 && opt.capacity > 0 && opt.quantum > 0 {
|
|
res.Bucket = ratelimit.NewBucketWithQuantum(opt.fillInterval, opt.capacity, opt.quantum)
|
|
|
|
return res, nil
|
|
}
|
|
|
|
if opt.fillInterval > 0 && opt.capacity > 0 {
|
|
res.Bucket = ratelimit.NewBucket(opt.fillInterval, opt.capacity)
|
|
|
|
return res, nil
|
|
}
|
|
|
|
if opt.rate > 0 && opt.capacity > 0 {
|
|
res.Bucket = ratelimit.NewBucketWithRate(opt.rate, opt.capacity)
|
|
|
|
return res, nil
|
|
}
|
|
|
|
return nil, errors.New("options error can not find func to init")
|
|
}
|