File: io.go

package info (click to toggle)
golang-github-tsenart-tb 0.0~git20151208.0.19f4c3d-2
  • links: PTS, VCS
  • area: main
  • in suites: bullseye, buster, sid, stretch
  • size: 92 kB
  • ctags: 51
  • sloc: makefile: 2
file content (35 lines) | stat: -rw-r--r-- 807 bytes parent folder | download
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
26
27
28
29
30
31
32
33
34
35
package io

import (
	"github.com/tsenart/tb"
	"io"
	"time"
)

// NewThrottledWriter is an io.Writer wrapping another io.Writer with
// byte rate throttling, flushing block bytes at a time.
func NewThrottledWriter(rate, block int64, w io.Writer) io.Writer {
	return &throttledWriter{rate, block, w, tb.NewBucket(rate, -1)}
}

type throttledWriter struct {
	rate, block int64
	w           io.Writer
	b           *tb.Bucket
}

func (tw *throttledWriter) Write(p []byte) (n int, err error) {
	for wr := 0; wr < len(p); {
		var got int64
		for got < tw.block {
			if got += tw.b.Take(tw.block - got); got != tw.block {
				time.Sleep(time.Duration((1e9 / tw.rate) * (tw.block - got)))
			}
		}
		if n, err = tw.w.Write(p[wr : wr+int(got)]); err != nil {
			return wr, err
		}
		wr += n
	}
	return len(p), nil
}