handle mid-stream unlimiting

Author: linkdata

operation.go

@@ -109,29 +109,35 @@ func (op *Operation) run(ch chan<- int64) {
 }
 
 func (op *Operation) io(fn func([]byte) (int, error), b []byte) (n int, err error) {
-	if op.Limit.Load() < 1 {
-		n, err = fn(b)
-		op.count.Add(int64(n))
-		return
-	}
+outer:
 	for len(b) > 0 && err == nil {
-		batch, ok := <-op.ch
-		err = io.EOF
-		if ok {
-			var done int
-			todo := min(int64(len(b)), batch)
-			done, err = fn(b[:todo])
-			op.avail.Add(batch - int64(done))
-			if done > 0 {
-				op.count.Add(int64(done))
-				n += int(done)
-				b = b[done:]
-			}
-			if op.reader && int64(done) < todo {
-				break
+		var done int
+		if op.Limit.Load() < 1 {
+			done, err = fn(b)
+			n += done
+			op.count.Add(int64(done))
+			return
+		}
+		select {
+		case batch, ok := <-op.ch:
+			err = io.EOF
+			if ok {
+				todo := min(int64(len(b)), batch)
+				done, err = fn(b[:todo])
+				op.avail.Add(batch - int64(done))
+				if done > 0 {
+					op.count.Add(int64(done))
+					n += int(done)
+					b = b[done:]
+				}
+				if op.reader && int64(done) < todo {
+					break outer
+				}
 			}
+		case <-op.WaitCh():
 		}
 	}
+
 	if op.reader && n > 0 && err == io.EOF {
 		err = nil
 	}

operation_test.go

@@ -96,6 +96,39 @@ func TestOperation_io_read_limit(t *testing.T) {
 	}
 }
 
+func TestOperation_io_read_limit_change_to_unlimited(t *testing.T) {
+	l := NewLimiter(100)
+	defer l.Stop()
+
+	r := bytes.NewReader(make([]byte, 1000))
+	done := make(chan struct{})
+	var n int
+	var err error
+
+	go func() {
+		n, err = l.Reads.io(r.Read, make([]byte, 1000))
+		close(done)
+	}()
+
+	// Let the read enter limited mode before switching to unlimited.
+	<-l.WaitCh()
+	<-l.WaitCh()
+	l.Reads.Limit.Store(0)
+
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatal("read hung after limit changed to unlimited")
+	}
+
+	if err != nil {
+		t.Fatal(err)
+	}
+	if n != 1000 {
+		t.Fatal(n)
+	}
+}
+
 func TestOperation_read_rate_low(t *testing.T) {
 	l := NewLimiter(1000)
 	defer l.Stop()