Add unlimited rate

bwlimit.go

@@ -2,17 +2,19 @@ package bwlimit
 
 import (
 	"io"
+	"math"
 	"sync"
 	"sync/atomic"
 	"time"
 )
 
 const (
 	DEFAULT_TAKT_PER_SECOND = 100
+	RATE_UNLIMITED          = 0
 )
 
 // Create new bandwidth limiter.
-// rateLimit: maximum bandwidth in bytes per second.
+// rateLimit: maximum bandwidth in bytes per second. 0 for unlimited.
 // blocking: block `Read`/`Write` until next takt time.
 func NewBwlimit(rateLimit int, blocking bool) Bwlimit {
 	bw := Bwlimit{
@@ -180,6 +182,10 @@ func (bw *Bwlimit) takt() {
 
 // Returns allowed windows size for streamId.
 func (bw *Bwlimit) currentWindow(streamId uint64) int {
+	if bw.rateLimit == RATE_UNLIMITED {
+		return math.MaxInt32
+	}
+
 	bw.streamMutex.Lock()
 	defer bw.streamMutex.Unlock()
 
@@ -192,6 +198,10 @@ func (bw *Bwlimit) currentWindow(streamId uint64) int {
 
 // Consume bandwidth window for streamId.
 func (bw *Bwlimit) consumeWindow(streamId uint64, consumeBytes int) {
+	if bw.rateLimit == RATE_UNLIMITED {
+		return
+	}
+
 	bw.streamMutex.Lock()
 	defer bw.streamMutex.Unlock()
 

bwlimit_test.go

@@ -9,6 +9,19 @@ import (
 	"time"
 )
 
+func TestBwlimit_Unlimited(t *testing.T) {
+	bw := NewBwlimit(RATE_UNLIMITED, true)
+	seq := make([]byte, 1000)
+	r := bytes.NewReader(seq)
+	wr := bw.Reader(r)
+	buf := make([]byte, 1000)
+
+	n, err := wr.Read(buf)
+	if n != 1000 || err != nil {
+		t.Errorf("Invalid rate n[%d] or err[%s]", n, err)
+	}
+}
+
 func TestBwlimit_SetBlocking(t *testing.T) {
 	bw := NewBwlimit(1000, true)
 	bw.SetTaktPerSecond(10)