add TestTicker_ConcurrentStopStress

Author: linkdata

ticker_test.go

@@ -3,6 +3,7 @@ package bwlimit
 import (
 	"bytes"
 	"io"
+	"sync"
 	"testing"
 	"testing/synctest"
 	"time"
@@ -59,3 +60,77 @@ func TestTicker_Stop_unblocksLimitedOperation(t *testing.T) {
 		}
 	})
 }
+
+func TestTicker_ConcurrentStopStress(t *testing.T) {
+	const (
+		iterations = 25
+		limiters   = 16
+		payload    = 64 * 1024
+	)
+
+	for iter := 0; iter < iterations; iter++ {
+		synctest.Test(t, func(t *testing.T) {
+			ticker := NewTicker()
+			ls := make([]*Limiter, limiters)
+			ioDone := make(chan struct{}, limiters*2)
+
+			for i := 0; i < limiters; i++ {
+				l := ticker.NewLimiter(1000, 1000)
+				ls[i] = l
+
+				go func(l *Limiter) {
+					_, _ = l.Reads.io(bytes.NewReader(make([]byte, payload)).Read, make([]byte, payload))
+					ioDone <- struct{}{}
+				}(l)
+
+				go func(l *Limiter) {
+					_, _ = l.Writes.io(io.Discard.Write, make([]byte, payload))
+					ioDone <- struct{}{}
+				}(l)
+			}
+
+			// Ensure operations have entered limited mode.
+			<-ticker.WaitCh()
+
+			var stopWG sync.WaitGroup
+			stopWG.Add(limiters + 1)
+
+			for i, l := range ls {
+				l := l
+				go func() {
+					defer stopWG.Done()
+					if i%2 == 0 {
+						<-ticker.WaitCh()
+					}
+					l.Stop()
+				}()
+			}
+
+			go func() {
+				defer stopWG.Done()
+				<-ticker.WaitCh()
+				ticker.Stop()
+			}()
+
+			stopped := make(chan struct{})
+			go func() {
+				stopWG.Wait()
+				close(stopped)
+			}()
+
+			select {
+			case <-stopped:
+			case <-time.After(10 * time.Second):
+				t.Fatal("timeout waiting for concurrent limiter/ticker Stop calls")
+			}
+
+			for i := 0; i < limiters*2; i++ {
+				select {
+				case <-ioDone:
+				case <-time.After(10 * time.Second):
+					t.Fatal("timeout waiting for active IO goroutine to return")
+				}
+			}
+		})
+	}
+}