github.com/grafana/pyroscope@v1.18.0/pkg/util/ratelimit/writer.go (about)

     1  package ratelimit
     2  
     3  import (
     4  	"io"
     5  )
     6  
     7  type Writer struct {
     8  	w io.Writer
     9  	l *Limiter
    10  }
    11  
    12  func NewWriter(w io.Writer, l *Limiter) *Writer {
    13  	return &Writer{w: w, l: l}
    14  }
    15  
    16  func (rw *Writer) Write(p []byte) (int, error) {
    17  	var total int
    18  	for len(p) > 0 {
    19  		n := len(p)
    20  		if n > int(rw.l.rate) {
    21  			n = int(rw.l.rate)
    22  		}
    23  		rw.l.Wait(n)
    24  		written, err := rw.w.Write(p[:n])
    25  		total += written
    26  		p = p[written:]
    27  		if err != nil {
    28  			return total, err
    29  		}
    30  	}
    31  	return total, nil
    32  }