show correct counts immediately after Stop

Author: linkdata

limiter_test.go

@@ -42,3 +42,23 @@ func TestLimiter_double_Wrap(t *testing.T) {
 		t.Error(d1, d2)
 	}
 }
+
+func TestLimiter_Stop_flushesCount(t *testing.T) {
+	l := NewLimiter(0)
+
+	r := bytes.NewReader(make([]byte, 10))
+	buf := make([]byte, 10)
+	n, err := l.Reads.io(r.Read, buf)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if n != 10 {
+		t.Fatal(n)
+	}
+
+	l.Stop()
+
+	if got := l.Reads.Count.Load(); got != int64(n) {
+		t.Fatalf("got %d want %d", got, n)
+	}
+}

operation.go

@@ -19,6 +19,7 @@ type Operation struct {
 	avail   atomic.Int64
 	count   atomic.Int64
 	ch      <-chan int64
+	doneCh  chan struct{}
 	reader  bool
 	mu      sync.Mutex // protects following
 	stopCh  chan struct{}
@@ -30,6 +31,7 @@ func NewOperation(t *Ticker, limits []int64, idx int) (op *Operation) {
 		Ticker: t,
 		ch:     ch,
 		stopCh: make(chan struct{}),
+		doneCh: make(chan struct{}),
 		reader: idx == 0,
 	}
 	var limit int64
@@ -52,10 +54,15 @@ func (op *Operation) Stop() {
 	if ch != nil {
 		close(ch)
 	}
+	<-op.doneCh
 }
 
 func (op *Operation) run(ch chan<- int64) {
-	defer close(ch)
+	defer func() {
+		close(ch)
+		op.Count.Add(op.count.Swap(0))
+		close(op.doneCh)
+	}()
 
 	op.mu.Lock()
 	stopCh := op.stopCh