gopkg/pond
2
0
Fork 0
Code Pull Requests Activity

Initial commit

Browse Source
This commit is contained in:
adurantecredify
2020-03-24 15:49:58 -03:00
parent 8ab3413ad2
commit 7b9b533e58
5 changed files with 720 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
go.mod
+9
View File
@@ -0,0 +1,9 @@
module github.com/alitto/pond
go 1.14
require (
github.com/gammazero/workerpool v0.0.0-20200311205957-7b00833861c6
github.com/panjf2000/ants/v2 v2.3.1
github.com/stretchr/testify v1.5.1
)
go.sum
+23
View File
@@ -0,0 +1,23 @@
github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/gammazero/deque v0.0.0-20200227231300-1e9af0e52b46 h1:iX4+rD9Fjdx8SkmSO/O5WAIX/j79ll3kuqv5VdYt9J8=
github.com/gammazero/deque v0.0.0-20200227231300-1e9af0e52b46/go.mod h1:D90+MBHVc9Sk1lJAbEVgws0eYEurY4mv2TDso3Nxh3w=
github.com/gammazero/workerpool v0.0.0-20200311205957-7b00833861c6 h1:1Cy/haf7XO4OyrkGid0Wq5CMluIErbvDptVAt8UTy38=
github.com/gammazero/workerpool v0.0.0-20200311205957-7b00833861c6/go.mod h1:/XWO2YAUUpPi3smDlFBl0vpX0JHwUomDM/oRMwRmnSs=
github.com/panjf2000/ants v1.3.0 h1:8pQ+8leaLc9lys2viEEr8md0U4RN6uOSUCE9bOYjQ9M=
github.com/panjf2000/ants/v2 v2.3.1 h1:9iOZHO5XlSO1Gs5K7x06uDFy8bkicWlhOKGh/TufAZg=
github.com/panjf2000/ants/v2 v2.3.1/go.mod h1:LtwNaBX6OeF5qRtQlaeGndalVwJlS2ueur7uwoAHbPA=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
github.com/stretchr/testify v1.5.1 h1:nOGnQDM7FYENwehXlg/kFVnos3rEvtKTjRvOWSzb6H4=
github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/yaml.v2 v2.2.2 h1:ZCJp+EgiOT7lHqUV2J862kp8Qj64Jo6az82+3Td9dZw=
gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
gopkg.in/yaml.v2 v2.2.7 h1:VUgggvou5XRW9mHwD/yXxIYSMtY0zoKQf/v226p2nyo=
gopkg.in/yaml.v2 v2.2.7/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
pond.go
+280
View File
@@ -0,0 +1,280 @@
package pond
import (
"fmt"
"runtime/debug"
"sync"
"sync/atomic"
"time"
)
const (
defaultIdleTimeout = 5 * time.Second
)
func defaultPanicHandler(panic interface{}) {
fmt.Printf("Worker exits from a panic: %v\nStack trace: %s\n", panic, string(debug.Stack()))
}
// Option represents an option that can be passed when building a worker pool to customize it
type Option func(*WorkerPool)
// IdleTimeout allows to change the idle timeout for a worker pool
func IdleTimeout(idleTimeout time.Duration) Option {
return func(pool *WorkerPool) {
pool.idleTimeout = idleTimeout
}
}
// PanicHandler allows to change the panic handler function for a worker pool
func PanicHandler(panicHandler func(interface{})) Option {
return func(pool *WorkerPool) {
pool.panicHandler = panicHandler
}
}
// WorkerPool models a pool of workers
type WorkerPool struct {
maxWorkers int
maxCapacity int
idleTimeout time.Duration
workerCount int32
tasks chan func()
dispatchedTasks chan func()
purgerQuit chan struct{}
stopOnce sync.Once
waitGroup sync.WaitGroup
panicHandler func(interface{})
}
// New creates a worker pool with that can scale up to the given number of workers and capacity
func New(maxWorkers, maxCapacity int, options ...Option) *WorkerPool {
pool := &WorkerPool{
maxWorkers: maxWorkers,
maxCapacity: maxCapacity,
idleTimeout: defaultIdleTimeout,
tasks: make(chan func(), maxCapacity),
dispatchedTasks: make(chan func(), maxWorkers),
purgerQuit: make(chan struct{}),
panicHandler: defaultPanicHandler,
}
// Apply all options
for _, opt := range options {
opt(pool)
}
// Start dispatcher goroutine
pool.waitGroup.Add(1)
go func() {
defer pool.waitGroup.Done()
pool.dispatch()
}()
// Start purger goroutine
pool.waitGroup.Add(1)
go func() {
defer pool.waitGroup.Done()
pool.purge()
}()
return pool
}
// Running returns the number of running workers
func (p *WorkerPool) Running() int {
return int(atomic.LoadInt32(&p.workerCount))
}
// Submit sends a task to this worker pool for execution. If the queue is full,
// it will wait until the task can be enqueued.
func (p *WorkerPool) Submit(task func()) {
if task == nil {
return
}
// Submit the task to the task channel
p.tasks <- task
}
// SubmitAndWait sends a task to this worker pool for execution and waits for it to complete
// before returning.
func (p *WorkerPool) SubmitAndWait(task func()) {
if task == nil {
return
}
done := make(chan struct{})
p.Submit(func() {
defer close(done)
task()
})
<-done
}
// Stop causes this pool to stop accepting tasks, without waiting for goroutines to exit
func (p *WorkerPool) Stop() {
p.stop(false)
}
// StopAndWait causes this pool to stop accepting tasks, waiting for all tasks in the queue to complete
func (p *WorkerPool) StopAndWait() {
p.stop(true)
}
// dispatch represents the work done by the dispatcher goroutine
func (p *WorkerPool) dispatch() {
for task := range p.tasks {
if task == nil {
// Received the signal to exit gracefully
break
}
select {
// Attempt to submit the task to a worker without blocking
case p.dispatchedTasks <- task:
if p.Running() == 0 {
p.startWorker()
}
default:
// Create a new worker if we haven't reached the limit yet
if p.Running() < p.maxWorkers {
p.startWorker()
}
// Block until a worker accepts this task
p.dispatchedTasks <- task
}
}
// Send signal to stop all workers
close(p.dispatchedTasks)
// Send signal to stop the purger
close(p.purgerQuit)
}
// purge represents the work done by the purger goroutine
func (p *WorkerPool) purge() {
ticker := time.NewTicker(p.idleTimeout)
defer ticker.Stop()
for {
select {
// Timed out waiting for any activity to happen, attempt to stop an idle worker
case <-ticker.C:
if p.Running() > 0 {
select {
case p.dispatchedTasks <- nil:
default:
// If dispatchedTasks channel is full, no need to kill the worker
}
}
// Received the signal to exit
case <-p.purgerQuit:
return
}
}
}
func (p *WorkerPool) startWorker() {
// Increment worker count
atomic.AddInt32(&p.workerCount, 1)
// Increment waiting group semaphore
p.waitGroup.Add(1)
worker(p.dispatchedTasks, func() {
// Decrement worker count
atomic.AddInt32(&p.workerCount, -1)
// Decrement waiting group semaphore
p.waitGroup.Done()
}, p.panicHandler)
}
// Stop causes this pool to stop accepting tasks, without waiting for goroutines to exit
func (p *WorkerPool) stop(wait bool) {
p.stopOnce.Do(func() {
if wait {
// Make sure all queued tasks complete before stopping the dispatcher
p.tasks <- nil
// Close the tasks channel to prevent receiving new tasks
close(p.tasks)
// Wait for all goroutines to exit
p.waitGroup.Wait()
} else {
// Close the tasks channel to prevent receiving new tasks
close(p.tasks)
}
})
}
// Group creates a new task group
func (p *WorkerPool) Group() *TaskGroup {
return &TaskGroup{
pool: p,
}
}
// worker launches a worker goroutine
func worker(tasks chan func(), exitHandler func(), panicHandler func(interface{})) {
go func() {
defer func() {
if panic := recover(); panic != nil {
// Handle panic
panicHandler(panic)
// Restart goroutine
worker(tasks, exitHandler, panicHandler)
} else {
// Handle exit
exitHandler()
}
}()
for task := range tasks {
if task == nil {
// We have received a signal to quit
return
}
// We have received a task, execute it
task()
}
}()
}
// TaskGroup represents a group of related tasks
type TaskGroup struct {
pool *WorkerPool
waitGroup sync.WaitGroup
}
// Submit adds a task to this group and sends it to the worker pool to be executed.
func (g *TaskGroup) Submit(task func()) {
g.waitGroup.Add(1)
g.pool.Submit(func() {
defer g.waitGroup.Done()
task()
})
}
// Wait waits until all the tasks in this group have completed. It returns
// a slice with all (non-nil) errors returned by tasks in this group.
func (g *TaskGroup) Wait() {
// Wait for all tasks to complete
g.waitGroup.Wait()
}
pond_benchmark_test.go
+113
View File
@@ -0,0 +1,113 @@
package pond_test
import (
"sync"
"testing"
"time"
"github.com/alitto/pond"
"github.com/gammazero/workerpool"
"github.com/panjf2000/ants/v2"
)
const (
taskCount = 1000000
taskDuration = 10 * time.Millisecond
workerCount = 200000
)
func BenchmarkPond(b *testing.B) {
var wg sync.WaitGroup
pool := pond.New(workerCount, taskCount)
defer pool.StopAndWait()
// Submit tasks
b.ResetTimer()
for i := 0; i < b.N; i++ {
wg.Add(taskCount)
for i := 0; i < taskCount; i++ {
pool.Submit(func() {
time.Sleep(taskDuration)
wg.Done()
})
}
wg.Wait()
}
b.StopTimer()
}
func BenchmarkPondGroup(b *testing.B) {
pool := pond.New(workerCount, taskCount)
defer pool.StopAndWait()
// Submit tasks
b.ResetTimer()
for i := 0; i < b.N; i++ {
group := pool.Group()
for i := 0; i < taskCount; i++ {
group.Submit(func() {
time.Sleep(taskDuration)
})
}
group.Wait()
}
b.StopTimer()
}
func BenchmarkGoroutines(b *testing.B) {
var wg sync.WaitGroup
// Submit tasks
b.ResetTimer()
for i := 0; i < b.N; i++ {
wg.Add(taskCount)
for i := 0; i < taskCount; i++ {
go func() {
time.Sleep(taskDuration)
wg.Done()
}()
}
wg.Wait()
}
b.StopTimer()
}
func BenchmarkGammazeroWorkerpool(b *testing.B) {
var wg sync.WaitGroup
wp := workerpool.New(workerCount)
defer wp.StopWait()
// Submit tasks
b.ResetTimer()
for i := 0; i < b.N; i++ {
wg.Add(taskCount)
for i := 0; i < taskCount; i++ {
wp.Submit(func() {
time.Sleep(taskDuration)
wg.Done()
})
}
wg.Wait()
}
b.StopTimer()
}
func BenchmarkAnts(b *testing.B) {
var wg sync.WaitGroup
p, _ := ants.NewPool(workerCount, ants.WithExpiryDuration(10*time.Second))
defer p.Release()
// Submit tasks
b.ResetTimer()
for i := 0; i < b.N; i++ {
wg.Add(taskCount)
for i := 0; i < taskCount; i++ {
_ = p.Submit(func() {
time.Sleep(taskDuration)
wg.Done()
})
}
wg.Wait()
}
b.StopTimer()
}
pond_test.go
+295
View File
@@ -0,0 +1,295 @@
package pond_test
import (
"fmt"
"sync/atomic"
"testing"
"time"
"github.com/alitto/pond"
"github.com/stretchr/testify/assert"
)
func TestSubmitAndStopWaiting(t *testing.T) {
assert := assert.New(t)
pool := pond.New(1, 5)
// Submit tasks
var doneCount int32
for i := 0; i < 17; i++ {
pool.Submit(func() {
time.Sleep(1 * time.Millisecond)
atomic.AddInt32(&doneCount, 1)
})
}
// Wait until all submitted tasks complete
pool.StopAndWait()
assert.Equal(int32(17), atomic.LoadInt32(&doneCount))
}
func TestSubmitAndStopWaitingWithMoreWorkersThanTasks(t *testing.T) {
assert := assert.New(t)
pool := pond.New(18, 5)
// Submit tasks
var doneCount int32
for i := 0; i < 17; i++ {
pool.Submit(func() {
time.Sleep(1 * time.Millisecond)
atomic.AddInt32(&doneCount, 1)
})
}
// Wait until all submitted tasks complete
pool.StopAndWait()
assert.Equal(int32(17), atomic.LoadInt32(&doneCount))
}
func TestSubmitAndStopWithoutWaiting(t *testing.T) {
assert := assert.New(t)
pool := pond.New(1, 5)
// Submit tasks
started := make(chan bool)
completed := make(chan bool)
var doneCount int32
for i := 0; i < 5; i++ {
pool.Submit(func() {
started <- true
time.Sleep(5 * time.Millisecond)
atomic.AddInt32(&doneCount, 1)
<-completed
})
}
// Make sure the first task started
<-started
// Stop without waiting for the rest of the tasks to start
pool.Stop()
// Let the first task complete now
completed <- true
// Only the first task should have been completed, the rest are discarded
assert.Equal(int32(1), atomic.LoadInt32(&doneCount))
// Make sure the exit lines in the worker pool are executed and covered
time.Sleep(6 * time.Millisecond)
}
func TestSubmitWithNilTask(t *testing.T) {
assert := assert.New(t)
pool := pond.New(2, 5)
// Submit nil task
pool.Submit(nil)
// Wait until all submitted tasks complete
pool.StopAndWait()
assert.Equal(0, pool.Running())
}
func TestSubmitAndWait(t *testing.T) {
assert := assert.New(t)
pool := pond.New(1, 5)
defer pool.StopAndWait()
// Submit a task and wait for it to complete
var doneCount int32
pool.SubmitAndWait(func() {
time.Sleep(5 * time.Millisecond)
atomic.AddInt32(&doneCount, 1)
})
assert.Equal(int32(1), atomic.LoadInt32(&doneCount))
}
func TestSubmitAndWaitWithNilTask(t *testing.T) {
assert := assert.New(t)
pool := pond.New(2, 5)
// Submit nil task
pool.SubmitAndWait(nil)
// Wait until all submitted tasks complete
pool.StopAndWait()
assert.Equal(0, pool.Running())
}
func TestRunning(t *testing.T) {
assert := assert.New(t)
workerCount := 5
taskCount := 10
pool := pond.New(workerCount, taskCount)
assert.Equal(0, pool.Running())
// Submit tasks
var started = make(chan struct{}, workerCount)
var completed = make(chan struct{}, workerCount)
for i := 0; i < taskCount; i++ {
pool.Submit(func() {
started <- struct{}{}
time.Sleep(1 * time.Millisecond)
<-completed
})
}
// Wait until half the tasks have started
for i := 0; i < taskCount/2; i++ {
<-started
}
assert.Equal(workerCount, pool.Running())
time.Sleep(1 * time.Millisecond)
// Make sure half the tasks tasks complete
for i := 0; i < taskCount/2; i++ {
completed <- struct{}{}
}
// Wait until the rest of the tasks have started
for i := 0; i < taskCount/2; i++ {
<-started
}
// Make sure all tasks complete
for i := 0; i < taskCount/2; i++ {
completed <- struct{}{}
}
pool.StopAndWait()
assert.Equal(0, pool.Running())
}
func TestSubmitWithPanic(t *testing.T) {
assert := assert.New(t)
pool := pond.New(1, 5)
assert.Equal(0, pool.Running())
// Submit a task that panics
var doneCount int32
pool.Submit(func() {
arr := make([]string, 0)
fmt.Printf("Out of range value: %s", arr[1])
atomic.AddInt32(&doneCount, 1)
})
// Submit a task that completes normally
pool.Submit(func() {
time.Sleep(2 * time.Millisecond)
atomic.AddInt32(&doneCount, 1)
})
pool.StopAndWait()
assert.Equal(0, pool.Running())
assert.Equal(int32(1), atomic.LoadInt32(&doneCount))
}
func TestSubmitWithIdleTimeout(t *testing.T) {
assert := assert.New(t)
pool := pond.New(1, 5, pond.IdleTimeout(2*time.Millisecond))
// Submit a task
started := make(chan bool)
completed := make(chan bool)
pool.Submit(func() {
<-started
time.Sleep(3 * time.Millisecond)
<-completed
})
// Make sure the first task has started
started <- true
// There should be 1 worker running
assert.Equal(1, pool.Running())
// Let the task complete
completed <- true
// Wait for idle timeout + 1ms
time.Sleep(3 * time.Millisecond)
// Worker should have been killed
assert.Equal(0, pool.Running())
pool.StopAndWait()
}
func TestSubmitWithPanicHandler(t *testing.T) {
assert := assert.New(t)
var capturedPanic interface{} = nil
panicHandler := func(panic interface{}) {
capturedPanic = panic
}
pool := pond.New(1, 5, pond.PanicHandler(panicHandler))
// Submit a task that panics
pool.Submit(func() {
panic("panic now!")
})
pool.StopAndWait()
// Panic should have been captured
assert.Equal("panic now!", capturedPanic)
}
func TestGroupSubmit(t *testing.T) {
assert := assert.New(t)
pool := pond.New(5, 5)
assert.Equal(0, pool.Running())
// Submit groups of tasks
var doneCount, taskCount int32
var groups []*pond.TaskGroup
for i := 0; i < 5; i++ {
group := pool.Group()
for j := 0; j < i+5; j++ {
group.Submit(func() {
time.Sleep(1 * time.Millisecond)
atomic.AddInt32(&doneCount, 1)
})
taskCount++
}
groups = append(groups, group)
}
// Wait for all groups to complete
for _, group := range groups {
group.Wait()
}
assert.Equal(int32(taskCount), atomic.LoadInt32(&doneCount))
}