/bwcapio.go
1 package main
2
3 import (
4 "io"
5 "time"
6 )
7
8 type BwStats struct {
9 Last time.Time /* time of last observed event */
10 Wnd uint /* unmetered bytes */
11 Thresh uint /* delay after at least this much bytes */
12 Rate uint /* bandwidth limit in bits/second */
13 }
14
15 func NewBwStats(rate uint) *BwStats {
16 return &BwStats{Wnd: 0, Thresh: rate, Rate: rate}
17 }
18
19 func CapReader(r io.Reader, st *BwStats) io.Reader {
20 if st == nil {
21 panic("nil stats")
22 }
23 return &BwCapReader{r, st}
24 }
25
26 func CapWriter(w io.Writer, st *BwStats) io.Writer {
27 if st == nil {
28 panic("nil stats")
29 }
30 return &BwCapWriter{w, st}
31 }
32
33 type BwCapReader struct {
34 Base io.Reader
35 Stats *BwStats
36 }
37
38 func (r *BwCapReader) Read(p []byte) (int, error) {
39 n, err := r.Base.Read(p)
40 bwCap(r.Stats, n)
41 return n, err
42 }
43
44 type BwCapWriter struct {
45 Base io.Writer
46 Stats *BwStats
47 }
48
49 func (w *BwCapWriter) Write(p []byte) (int, error) {
50 n, err := w.Base.Write(p)
51 bwCap(w.Stats, n)
52 return n, err
53 }
54
55 func bwCap(st *BwStats, transfered int) {
56 if transfered <= 0 {
57 return
58 }
59 if st.Last.IsZero() {
60 st.Last = time.Now()
61 return
62 }
63 st.Wnd += uint(transfered)
64 if st.Wnd < st.Thresh {
65 return
66 }
67
68 bits := st.Wnd * 8
69 exp := time.Duration((1e9 * bits) / st.Rate)
70 ahead := exp - time.Since(st.Last)
71
72 if ahead > 0 {
73 if ahead.Seconds() > 1 {
74 st.Thresh /= 2
75 } else if ahead < 10*time.Millisecond {
76 st.Thresh *= 2
77 }
78 time.Sleep(ahead)
79 }
80
81 st.Wnd = 0
82 st.Last = time.Now()
83 }