]> Sergey Matveev's repositories - btrtrc.git/blob - ratelimitreader.go
gorond
[btrtrc.git] / ratelimitreader.go
1 package torrent
2
3 import (
4         "context"
5         "fmt"
6         "io"
7         "time"
8
9         "golang.org/x/time/rate"
10 )
11
12 type rateLimitedReader struct {
13         l *rate.Limiter
14         r io.Reader
15
16         // This is the time of the last Read's reservation.
17         lastRead time.Time
18 }
19
20 func (me *rateLimitedReader) Read(b []byte) (n int, err error) {
21         const oldStyle = false // Retained for future reference.
22         if oldStyle {
23                 // Wait until we can read at all.
24                 if err := me.l.WaitN(context.Background(), 1); err != nil {
25                         panic(err)
26                 }
27                 // Limit the read to within the burst.
28                 if me.l.Limit() != rate.Inf && len(b) > me.l.Burst() {
29                         b = b[:me.l.Burst()]
30                 }
31                 n, err = me.r.Read(b)
32                 // Pay the piper.
33                 now := time.Now()
34                 me.lastRead = now
35                 if !me.l.ReserveN(now, n-1).OK() {
36                         panic(fmt.Sprintf("burst exceeded?: %d", n-1))
37                 }
38         } else {
39                 // Limit the read to within the burst.
40                 if me.l.Limit() != rate.Inf && len(b) > me.l.Burst() {
41                         b = b[:me.l.Burst()]
42                 }
43                 n, err = me.r.Read(b)
44                 now := time.Now()
45                 r := me.l.ReserveN(now, n)
46                 if !r.OK() {
47                         panic(n)
48                 }
49                 me.lastRead = now
50                 time.Sleep(r.Delay())
51         }
52         return
53 }