mirror of
https://github.com/VictoriaMetrics/VictoriaMetrics.git
synced 2024-12-23 00:46:23 +01:00
38 lines
791 B
Go
38 lines
791 B
Go
|
package limiter
|
||
|
|
||
|
import (
|
||
|
"io"
|
||
|
)
|
||
|
|
||
|
// NewWriteLimiter creates a new WriteLimiter object
|
||
|
// for the give writer and Limiter.
|
||
|
func NewWriteLimiter(w io.Writer, limiter *Limiter) *WriteLimiter {
|
||
|
return &WriteLimiter{
|
||
|
writer: w,
|
||
|
limiter: limiter,
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// WriteLimiter limits the amount of bytes written
|
||
|
// per second via Write() method.
|
||
|
// Must be created via NewWriteLimiter.
|
||
|
type WriteLimiter struct {
|
||
|
writer io.Writer
|
||
|
limiter *Limiter
|
||
|
}
|
||
|
|
||
|
// Close implements io.Closer
|
||
|
// also calls Close for wrapped io.WriteCloser
|
||
|
func (wl *WriteLimiter) Close() error {
|
||
|
if c, ok := wl.writer.(io.Closer); ok {
|
||
|
return c.Close()
|
||
|
}
|
||
|
return nil
|
||
|
}
|
||
|
|
||
|
// Write implements io.Writer
|
||
|
func (wl *WriteLimiter) Write(p []byte) (n int, err error) {
|
||
|
wl.limiter.Register(len(p))
|
||
|
return wl.writer.Write(p)
|
||
|
}
|