refactor

Author: linkdata

dialer_test.go

@@ -34,10 +34,10 @@ func TestDialer_Dial(t *testing.T) {
 			Dial: d2.Dial,
 		},
 	}
-	<-Ticker.TickCh()
+	<-DefaultTicker.Ch()
 	resp1, err1 := client1.Get(srv.URL)
 	resp2, err2 := client2.Get(srv.URL)
-	<-Ticker.TickCh()
+	<-DefaultTicker.Ch()
 
 	if err1 != nil {
 		t.Fatal(err1)

limiter.go

@@ -10,19 +10,18 @@ type DialContextFn func(ctx context.Context, network string, address string) (ne
 var DefaultNetDialer = &net.Dialer{}
 
 type Limiter struct {
+	*Ticker
 	Reads  *Operation
 	Writes *Operation
 }
 
-// NewLimiter returns a new limiter. If you provide limits, the first will set
+// NewLimiter returns a new limiter from DefaultTicker.
+// If you provide limits, the first will set
 // both read and write limits, the second will set the write limit.
 //
-// To stop the limiter and free it's resources, call Stop.
+// To stop the Limiter and free it's resources, call Stop.
 func NewLimiter(limits ...int64) *Limiter {
-	return &Limiter{
-		Reads:  NewOperation(limits, 0),
-		Writes: NewOperation(limits, 1),
-	}
+	return DefaultTicker.NewLimiter(limits...)
 }
 
 // Stop stops the Limiter and frees any resources. Reads and writes on

limiter_test.go

@@ -21,7 +21,7 @@ func TestLimiter_Stop(t *testing.T) {
 		t.Error(n)
 	}
 	l.Stop()
-	<-Ticker.TickCh()
+	<-DefaultTicker.Ch()
 
 	n, err = l.Reads.io(r.Read, buf)
 	if err != io.EOF {

operation.go

@@ -71,7 +71,7 @@ func (op *Operation) run(ch chan<- int) {
 				todo = max(1, int(limit/secparts))
 				batch = min(batchsize, todo)
 			}
-			tickCh := Ticker.TickCh()
+			tickCh := DefaultTicker.Ch()
 
 		partialsecond:
 			for {

operation_test.go

@@ -103,18 +103,18 @@ func TestOperation_read_rate_low(t *testing.T) {
 	buf := make([]byte, 1001)
 
 	var tickCount atomic.Int32
-	oldOnTick := Ticker.GetOnTick()
-	defer Ticker.SetOnTick(oldOnTick)
-	<-Ticker.TickCh()
-	Ticker.SetOnTick(func() { tickCount.Add(1) })
+	oldOnTick := DefaultTicker.GetOnTick()
+	defer DefaultTicker.SetOnTick(oldOnTick)
+	<-DefaultTicker.Ch()
+	DefaultTicker.SetOnTick(func() { tickCount.Add(1) })
 
-	<-Ticker.TickCh()
+	<-DefaultTicker.Ch()
 	// should read in batches of 1000/secparts (=100) bytes
 	now := time.Now()
 	n, err := l.Reads.io(r.Read, buf)
 	elapsed := time.Since(now)
 	rate := l.Reads.Rate.Load()
-	<-Ticker.TickCh()
+	<-DefaultTicker.Ch()
 
 	if n := tickCount.Load(); n < 11 || n > 13 {
 		t.Error(n)

operationticker.go

@@ -0,0 +1,81 @@
+package bwlimit
+
+import (
+	"sync"
+	"time"
+)
+
+// A Ticker synchronizes rate calculation among multiple Limiters and
+// provides the on-tick callback.
+type Ticker struct {
+	mu sync.Mutex
+	ch chan struct{}
+	fn func()
+}
+
+var DefaultTicker *Ticker
+
+// NewLimiter returns a new Limiter using this Ticker.
+// If you provide limits, the first will set
+// both read and write limits, the second will set the write limit.
+//
+// To stop the limiter and free it's resources, call Stop.
+func (ot *Ticker) NewLimiter(limits ...int64) *Limiter {
+	return &Limiter{
+		Ticker: ot,
+		Reads:  NewOperation(limits, 0),
+		Writes: NewOperation(limits, 1),
+	}
+}
+
+func (ot *Ticker) SetOnTick(fn func()) {
+	ot.mu.Lock()
+	ot.fn = fn
+	ot.mu.Unlock()
+}
+
+func (ot *Ticker) GetOnTick() (fn func()) {
+	ot.mu.Lock()
+	fn = ot.fn
+	ot.mu.Unlock()
+	return
+}
+
+func (ot *Ticker) Ch() (ch <-chan struct{}) {
+	ot.mu.Lock()
+	ch = ot.ch
+	ot.mu.Unlock()
+	return
+}
+
+func (ot *Ticker) run() {
+	defer close(ot.ch)
+
+	tckr := time.NewTicker(interval)
+	defer tckr.Stop()
+
+	for range tckr.C {
+		ot.mu.Lock()
+		oldCh := ot.ch
+		ot.ch = make(chan struct{})
+		ot.mu.Unlock()
+		close(oldCh)
+	}
+}
+
+func (ot *Ticker) runOnTick() {
+	for {
+		<-ot.Ch()
+		if fn := ot.GetOnTick(); fn != nil {
+			fn()
+		}
+	}
+}
+
+func init() {
+	DefaultTicker = &Ticker{
+		ch: make(chan struct{}),
+	}
+	go DefaultTicker.run()
+	go DefaultTicker.runOnTick()
+}