refactor

Author: linkdata

dialer_test.go

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

limiter_test.go

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

operation.go

@@ -12,20 +12,22 @@ const interval = time.Second / secparts
 const batchsize = 4096
 
 type Operation struct {
-	Limit  atomic.Int64 // bandwith limit in bytes/sec
-	Rate   atomic.Int64 // current rate in bytes/sec
-	Count  atomic.Int64 // number of bytes seen
-	avail  atomic.Int64
-	count  atomic.Int64
-	ch     <-chan int
-	reader bool
-	mu     sync.Mutex // protects following
-	stopCh chan struct{}
+	*Ticker              // Ticker we belong to
+	Limit   atomic.Int64 // bandwith limit in bytes/sec
+	Rate    atomic.Int64 // current rate in bytes/sec
+	Count   atomic.Int64 // number of bytes seen
+	avail   atomic.Int64
+	count   atomic.Int64
+	ch      <-chan int
+	reader  bool
+	mu      sync.Mutex // protects following
+	stopCh  chan struct{}
 }
 
-func NewOperation(limits []int64, idx int) (op *Operation) {
+func NewOperation(t *Ticker, limits []int64, idx int) (op *Operation) {
 	ch := make(chan int)
 	op = &Operation{
+		Ticker: t,
 		ch:     ch,
 		stopCh: make(chan struct{}),
 		reader: idx == 0,
@@ -71,7 +73,7 @@ func (op *Operation) run(ch chan<- int) {
 				todo = max(1, int(limit/secparts))
 				batch = min(batchsize, todo)
 			}
-			tickCh := DefaultTicker.Ch()
+			waitCh := op.WaitCh()
 
 		partialsecond:
 			for {
@@ -82,10 +84,10 @@ func (op *Operation) run(ch chan<- int) {
 					todo -= batch
 					todo += int(op.avail.Swap(0))
 					if todo < batch {
-						<-tickCh
+						<-waitCh
 						break partialsecond
 					}
-				case <-tickCh:
+				case <-waitCh:
 					break partialsecond
 				}
 			}

operation_test.go

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

ticker.go

@@ -5,12 +5,10 @@ import (
 	"time"
 )
 
-// A Ticker synchronizes rate calculation among multiple Limiters and
-// provides the on-tick callback.
+// A Ticker synchronizes rate calculation among multiple Limiters.
 type Ticker struct {
 	mu sync.Mutex
 	ch chan struct{}
-	fn func()
 }
 
 var DefaultTicker *Ticker
@@ -20,28 +18,17 @@ var DefaultTicker *Ticker
 // 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 {
+func (ot *Ticker) NewLimiter(limits ...int64) (l *Limiter) {
 	return &Limiter{
 		Ticker: ot,
-		Reads:  NewOperation(limits, 0),
-		Writes: NewOperation(limits, 1),
+		Reads:  NewOperation(ot, limits, 0),
+		Writes: NewOperation(ot, 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{}) {
+// WaitCh returns a channel that will close when the current rate limit
+// time slice runs out.
+func (ot *Ticker) WaitCh() (ch <-chan struct{}) {
 	ot.mu.Lock()
 	ch = ot.ch
 	ot.mu.Unlock()
@@ -63,19 +50,9 @@ func (ot *Ticker) run() {
 	}
 }
 
-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()
 }