V2bX/common/limiter/rate.go
2022-06-02 01:35:41 +08:00

32 lines
583 B
Go

package limiter
import (
"io"
"github.com/juju/ratelimit"
"github.com/xtls/xray-core/common"
"github.com/xtls/xray-core/common/buf"
)
type Writer struct {
writer buf.Writer
limiter *ratelimit.Bucket
w io.Writer
}
func (l *Limiter) RateWriter(writer buf.Writer, limiter *ratelimit.Bucket) buf.Writer {
return &Writer{
writer: writer,
limiter: limiter,
}
}
func (w *Writer) Close() error {
return common.Close(w.writer)
}
func (w *Writer) WriteMultiBuffer(mb buf.MultiBuffer) error {
w.limiter.Wait(int64(mb.Len()))
return w.writer.WriteMultiBuffer(mb)
}