]> Sergey Matveev's repositories - btrtrc.git/blobdiff - ratelimitreader.go
Add download rate limiting
[btrtrc.git] / ratelimitreader.go
diff --git a/ratelimitreader.go b/ratelimitreader.go
new file mode 100644 (file)
index 0000000..236efd4
--- /dev/null
@@ -0,0 +1,25 @@
+package torrent
+
+import (
+       "context"
+       "io"
+       "time"
+
+       "golang.org/x/time/rate"
+)
+
+type rateLimitedReader struct {
+       l *rate.Limiter
+       r io.Reader
+}
+
+func (me rateLimitedReader) Read(b []byte) (n int, err error) {
+       if err := me.l.WaitN(context.Background(), 1); err != nil {
+               panic(err)
+       }
+       n, err = me.r.Read(b)
+       if !me.l.ReserveN(time.Now(), n-1).OK() {
+               panic(n - 1)
+       }
+       return
+}