fix high limit flow

Author: linkdata

operation.go

@@ -30,35 +30,41 @@ func NewOperation(ctx context.Context, reader bool) (op *Operation) {
 func (op *Operation) run(ctx context.Context, ch chan<- struct{}) {
 	defer close(ch)
 	seccount := 0
-	now := time.Now()
+	counts := make([]int32, secparts)
+	tckr := time.NewTicker(interval)
+	defer tckr.Stop()
 
 	for {
-		if elapsed := time.Since(now); elapsed > 0 {
-			now = now.Add(elapsed)
-
-			if limit := op.Limit.Load(); limit > 0 {
-				todo := max(1, limit/secparts)
-				batch := min(1024, todo)
-				op.batch.Store(batch)
-				for todo >= batch && time.Since(now) < (interval-(interval/10)) {
+		if limit := op.Limit.Load(); limit > 0 {
+			todo := max(1, limit/secparts)
+			batch := min(1024, todo)
+			op.batch.Store(batch)
+		drive:
+			for {
+				select {
+				case <-ctx.Done():
+					return
+				case ch <- struct{}{}:
+					todo -= batch
 					todo += op.avail.Swap(0)
-					select {
-					case <-ctx.Done():
-						return
-					case ch <- struct{}{}:
-						todo -= batch
-					default:
-						time.Sleep(interval / 10)
+					if todo < batch {
+						<-tckr.C
+						break drive
 					}
+				case <-tckr.C:
+					break drive
 				}
 			}
-
-			time.Sleep(interval - time.Since(now))
+			counts[seccount] = op.count.Swap(0)
 			seccount++
 			if seccount >= secparts {
 				seccount = 0
-				op.Rate.Store(op.count.Swap(0))
 			}
+			var rate int32
+			for i := 0; i < secparts; i++ {
+				rate += counts[i]
+			}
+			op.Rate.Store(rate)
 		}
 	}
 }

operation_test.go

@@ -62,7 +62,7 @@ func TestOperation_io_read_limit(t *testing.T) {
 	}
 }
 
-func TestOperation_read_rate(t *testing.T) {
+func TestOperation_read_rate_low(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
 	defer cancel()
 	l := NewLimiter(ctx)
@@ -89,6 +89,36 @@ func TestOperation_read_rate(t *testing.T) {
 	}
 }
 
+func TestOperation_read_rate_high(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
+	defer cancel()
+	l := NewLimiter(ctx)
+
+	const numbytes = 10 * 1000000
+	l.Reads.Limit.Store(numbytes)
+
+	now := time.Now()
+	r := bytes.NewReader(make([]byte, numbytes*2))
+	var numread int
+
+	for numread < numbytes {
+		buf := make([]byte, 1000)
+		n, err := l.Reads.io(r.Read, buf)
+		numread += n
+		if err != nil {
+			t.Error(numread, n)
+			t.Fatal(err)
+		}
+	}
+
+	if elapsed := time.Since(now); elapsed < time.Millisecond*900 || elapsed > time.Millisecond*1100 {
+		t.Error(elapsed)
+	}
+	if rate := int(l.Reads.Rate.Load()); rate < numbytes*0.9 || rate > numbytes {
+		t.Error(rate)
+	}
+}
+
 func TestOperation_write_rate(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
 	defer cancel()
@@ -106,10 +136,14 @@ func TestOperation_write_rate(t *testing.T) {
 		t.Error(err)
 	}
 
-	for l.Writes.Rate.Load() == 0 && ctx.Err() == nil {
-		time.Sleep(time.Second / 100)
+	rate := l.Writes.Rate.Load()
+	for rate == 0 && ctx.Err() == nil {
+		rate = l.Writes.Rate.Load()
+	}
+	if ctx.Err() != nil {
+		t.Fatal(ctx.Err())
 	}
-	if rate := l.Writes.Rate.Load(); rate != 10 {
+	if rate != 10 {
 		t.Error(rate)
 	}
 }