Add rate limit

pool.go

@@ -20,9 +20,9 @@ type Pool interface {
 }
 
 type NonTimeLimitedPool struct {
-	input   chan Task
-	output  chan interface{}
-	wg      *sync.WaitGroup
+	input  chan Task
+	output chan interface{}
+	wg     *sync.WaitGroup
 }
 
 func (p NonTimeLimitedPool) End() {
@@ -73,12 +73,15 @@ func (p *NonTimeLimitedPool) worker() {
 	}
 }
 
-func (p *NonTimeLimitedPool) AddWorker() {
+func (p NonTimeLimitedPool) AddWorker() {
 	p.wg.Add(1)
 	go p.worker()
 }
 
 func NewPool(options PoolOptions) Pool {
+	if options.Duration != 0 && options.PerDuration != 0 {
+		return NewTimeLimitedPool(options)
+	}
 	return NewNonTimeLimitedPool(options)
 }
 
@@ -94,4 +97,4 @@ func NewNonTimeLimitedPool(options PoolOptions) NonTimeLimitedPool {
 		p.AddWorker()
 	}
 	return p
-}
+}

pool_test.go

@@ -3,7 +3,6 @@ package cartel
 import (
 	"runtime"
 	"testing"
-	"time"
 )
 
 func Test_Create_Starts_Correct_Number_Of_Goroutines(t *testing.T) {
@@ -50,7 +49,7 @@ func Test_Returns_Output(t *testing.T) {
 
 	values := p.GetOutput()
 
-	if expected, actual := 1, len(values);expected != actual {
+	if expected, actual := 1, len(values); expected != actual {
 		t.Errorf("expected %v  but got %v ", expected, actual)
 	}
 	value := values[0]
@@ -88,11 +87,3 @@ type TestTask struct {
 func (tt TestTask) Execute() interface{} {
 	return tt.Name
 }
-
-type TestTimeTask struct {
-	Name string
-}
-
-func (ttt TestTimeTask) Execute() interface{} {
-	return time.Now()
-}

time_limited_pool.go

@@ -1,19 +1,20 @@
 package cartel
 
 import (
+	"runtime"
 	"sync"
 	"time"
 )
 
 type TimeLimitedPool struct {
-	input     chan Task
-	output  chan interface{}
-	wg      *sync.WaitGroup
-	available map[int]bool
-	options PoolOptions
+	input       chan Task
+	output      chan interface{}
+	wg          *sync.WaitGroup
+	options     PoolOptions
+	lastChecked time.Time
+	allowace    float64
 }
 
-
 func (p TimeLimitedPool) End() {
 	close(p.input)
 	p.wg.Wait()
@@ -41,32 +42,52 @@ func (p TimeLimitedPool) GetOutput() []interface{} {
 
 func (p *TimeLimitedPool) worker() {
 	for {
+		p.rateLimit()
 		t, ok := <-p.input
+
 		if !ok {
 			p.wg.Done()
 			break
 		}
+
 		v := t.Execute()
+		p.lastChecked = time.Now()
+
 		p.output <- v
 		runtime.GC()
 	}
 }
 
-func (p *TimeLimitedPool) AddWorker() {
+func (p *TimeLimitedPool) rateLimit() {
+	lastChecked := p.lastChecked
+	now := time.Now()
+
+	p.lastChecked = now
+	diff := float64(now.Unix() - lastChecked.Unix())
+	floatRate := float64(p.options.PerDuration)
+	floatDuration := float64(p.options.Duration.Seconds())
+
+	p.allowace = p.allowace + (diff * (floatRate / floatDuration))
+	if p.allowace < 1.0 {
+		time.Sleep(p.options.Duration)
+	}
+}
+
+func (p TimeLimitedPool) AddWorker() {
 	p.wg.Add(1)
 	go p.worker()
 }
+
 func NewTimeLimitedPool(options PoolOptions) TimeLimitedPool {
 
 	jobs := make(chan Task, 100)
 	results := make(chan interface{}, 100)
-	available := map[int]bool{}
 
 	var wg sync.WaitGroup
-	p := TimeLimitedPool{jobs, results, &wg, available, options}
+	p := TimeLimitedPool{jobs, results, &wg, options, time.Now(), 0.0}
 
 	for w := 1; w <= options.Size; w++ {
 		p.AddWorker()
 	}
 	return p
-}
+}

time_limited_pool_test.go

@@ -1 +1,41 @@
 package cartel
+
+import (
+	"testing"
+	"time"
+)
+
+func Test_Runs_One_Per_Second(t *testing.T) {
+
+	p := NewPool(
+		PoolOptions{
+			Size:        1,
+			PerDuration: 1,
+			Duration:    time.Second,
+		},
+	)
+	p.Do(TestTimeTask{})
+	p.Do(TestTimeTask{})
+	p.End()
+
+	values := p.GetOutput()
+
+	if actual := len(values); actual != 2 {
+		t.Errorf("expected 2 items but got %v", actual)
+	}
+
+	firstTime := values[0].(time.Time)
+	secondTime := values[1].(time.Time)
+
+	d := secondTime.Sub(firstTime)
+	if d.Seconds() < 1 {
+		t.Errorf("expected 1 seconds but got %v - %v %v", d.Seconds(), firstTime, secondTime)
+	}
+}
+
+type TestTimeTask struct {
+}
+
+func (ttt TestTimeTask) Execute() interface{} {
+	return time.Now()
+}