4 Star 23 Fork 15

Gitee 极速下载/goreplay

加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
此仓库是为了提升国内下载速度的镜像仓库,每日同步一次。 原始仓库: https://github.com/buger/goreplay
克隆/下载
limiter.go 2.72 KB
一键复制 编辑 原始数据 按行查看 历史
package goreplay
import (
"fmt"
"io"
"math/rand"
"strconv"
"strings"
"time"
)
// Limiter is a wrapper for input or output plugin which adds rate limiting
type Limiter struct {
plugin interface{}
limit int
isPercent bool
currentRPS int
currentTime int64
}
func parseLimitOptions(options string) (limit int, isPercent bool) {
if n := strings.Index(options, "%"); n > 0 {
limit, _ = strconv.Atoi(options[:n])
isPercent = true
} else {
limit, _ = strconv.Atoi(options)
isPercent = false
}
return
}
func newLimiterExceptions(l *Limiter) {
if !l.isPercent {
return
}
speedFactor := float64(l.limit) / float64(100)
// FileInput、KafkaInput have its own rate limiting. Unlike other inputs we not just dropping requests, we can slow down or speed up request emittion.
switch input := l.plugin.(type) {
case *FileInput:
input.speedFactor = speedFactor
case *KafkaInput:
input.speedFactor = speedFactor
}
}
// NewLimiter constructor for Limiter, accepts plugin and options
// `options` allow to sprcify relatve or absolute limiting
func NewLimiter(plugin interface{}, options string) PluginReadWriter {
l := new(Limiter)
l.limit, l.isPercent = parseLimitOptions(options)
l.plugin = plugin
l.currentTime = time.Now().UnixNano()
newLimiterExceptions(l)
return l
}
func (l *Limiter) isLimitedExceptions() bool {
if !l.isPercent {
return false
}
// Fileinput、Kafkainput have its own limiting algorithm
switch l.plugin.(type) {
case *FileInput:
return true
case *KafkaInput:
return true
default:
return false
}
}
func (l *Limiter) isLimited() bool {
if l.isLimitedExceptions() {
return false
}
if l.isPercent {
return l.limit <= rand.Intn(100)
}
if (time.Now().UnixNano() - l.currentTime) > time.Second.Nanoseconds() {
l.currentTime = time.Now().UnixNano()
l.currentRPS = 0
}
if l.currentRPS >= l.limit {
return true
}
l.currentRPS++
return false
}
// PluginWrite writes message to this plugin
func (l *Limiter) PluginWrite(msg *Message) (n int, err error) {
if l.isLimited() {
return 0, nil
}
if w, ok := l.plugin.(PluginWriter); ok {
return w.PluginWrite(msg)
}
// avoid further writing
return 0, io.ErrClosedPipe
}
// PluginRead reads message from this plugin
func (l *Limiter) PluginRead() (msg *Message, err error) {
if r, ok := l.plugin.(PluginReader); ok {
msg, err = r.PluginRead()
} else {
// avoid further reading
return nil, io.ErrClosedPipe
}
if l.isLimited() {
return nil, nil
}
return
}
func (l *Limiter) String() string {
return fmt.Sprintf("Limiting %s to: %d (isPercent: %v)", l.plugin, l.limit, l.isPercent)
}
// Close closes the resources.
func (l *Limiter) Close() error {
if fi, ok := l.plugin.(io.Closer); ok {
fi.Close()
}
return nil
}
马建仓 AI 助手
尝试更多
代码解读
代码找茬
代码优化
Go
1
https://gitee.com/mirrors/goreplay.git
git@gitee.com:mirrors/goreplay.git
mirrors
goreplay
goreplay
master

搜索帮助

D67c1975 1850385 1daf7b77 1850385