27 lines
430 B
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)
|
|
}
|