vendor/github.com/shazow/rateio/writer.go


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

package rateio

import "io"

type writer struct {
	io.Writer
	Limiter
}

// Write writes the contents of p into the buffer.
// Returns ErrRateExceeded error if our specified read is exceeded.
func (w *writer) Write(p []byte) (n int, err error) {
	n, err = w.Writer.Write(p)
	if err != nil {
		return
	}

	err = w.Limiter.Count(n)
	return
}

// NewWriter proxies an io.Writer but keeps track of bytes read based on our Limiter.
func NewWriter(w io.Writer, limiter Limiter) io.Writer {
	return &writer{w, limiter}
}