mfproxy/throttle.go

27 lines
430 B
Go

package main
import (
"context"
"io"
"golang.org/x/time/rate"
)
type throttledWriter struct {
ctx context.Context
w io.Writer
limiter *rate.Limiter
}
func (tw throttledWriter) Write(b []byte) (int, error) {
if tw.limiter != nil {
if block := tw.limiter.Burst(); len(b) > block {
b = b[:block]
}
if err := tw.limiter.WaitN(tw.ctx, len(b)); err != nil {
return 0, err
}
}
return tw.w.Write(b)
}