Skip to content

Commit

Permalink
Run gofmt
Browse files Browse the repository at this point in the history 
Signed-off-by: Daniel Hu <tao.hu@merico.dev>
  • Loading branch information
@daniel-hutao
daniel-hutao committed Jul 26, 2023
1 parent 460a9a5 commit 7dc1fce
Show file tree
Hide file tree
Showing 4 changed files with 211 additions and 212 deletions.
226 changes: 113 additions & 113 deletions gopool.go
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
package gopool

import (
"sync"
"time"
"sync"
"time"
)

// task represents a function that will be executed by a worker.
Expand All @@ -11,151 +11,151 @@ type task func() (interface{}, error)

// goPool represents a pool of workers.
type goPool struct {
workers []*worker
workerStack []int
maxWorkers int
// Set by WithMinWorkers(), used to adjust the number of workers. Default equals to maxWorkers.
minWorkers int
// tasks are added to this channel first, then dispatched to workers. Default buffer size is 1 million.
taskQueue chan task
// Set by WithRetryCount(), used to retry a task when it fails. Default is 0.
retryCount int
lock sync.Locker
cond *sync.Cond
// Set by WithTimeout(), used to set a timeout for a task. Default is 0, which means no timeout.
timeout time.Duration
// Set by WithResultCallback(), used to handle the result of a task. Default is nil.
resultCallback func(interface{})
// Set by WithErrorCallback(), used to handle the error of a task. Default is nil.
errorCallback func(error)
// adjustInterval is the interval to adjust the number of workers. Default is 1 second.
adjustInterval time.Duration
workers []*worker
workerStack []int
maxWorkers int
// Set by WithMinWorkers(), used to adjust the number of workers. Default equals to maxWorkers.
minWorkers int
// tasks are added to this channel first, then dispatched to workers. Default buffer size is 1 million.
taskQueue chan task
// Set by WithRetryCount(), used to retry a task when it fails. Default is 0.
retryCount int
lock sync.Locker
cond *sync.Cond
// Set by WithTimeout(), used to set a timeout for a task. Default is 0, which means no timeout.
timeout time.Duration
// Set by WithResultCallback(), used to handle the result of a task. Default is nil.
resultCallback func(interface{})
// Set by WithErrorCallback(), used to handle the error of a task. Default is nil.
errorCallback func(error)
// adjustInterval is the interval to adjust the number of workers. Default is 1 second.
adjustInterval time.Duration
}

// NewGoPool creates a new pool of workers.
func NewGoPool(maxWorkers int, opts ...Option) *goPool {
pool := &goPool{
maxWorkers: maxWorkers,
// Set minWorkers to maxWorkers by default
minWorkers: maxWorkers,
workers: make([]*worker, maxWorkers),
workerStack: make([]int, maxWorkers),
taskQueue: make(chan task, 1e6),
retryCount: 0,
lock: new(sync.Mutex),
timeout: 0,
adjustInterval: 1 * time.Second,
}
// Apply options
for _, opt := range opts {
opt(pool)
}
if pool.cond == nil {
pool.cond = sync.NewCond(pool.lock)
}
// Create workers with the minimum number. Don't use pushWorker() here.
for i := 0; i < pool.minWorkers; i++ {
worker := newWorker()
pool.workers[i] = worker
pool.workerStack[i] = i
worker.start(pool, i)
}
go pool.adjustWorkers()
go pool.dispatch()
return pool
pool := &goPool{
maxWorkers: maxWorkers,
// Set minWorkers to maxWorkers by default
minWorkers: maxWorkers,
workers: make([]*worker, maxWorkers),
workerStack: make([]int, maxWorkers),
taskQueue: make(chan task, 1e6),
retryCount: 0,
lock: new(sync.Mutex),
timeout: 0,
adjustInterval: 1 * time.Second,
}
// Apply options
for _, opt := range opts {
opt(pool)
}
if pool.cond == nil {
pool.cond = sync.NewCond(pool.lock)
}
// Create workers with the minimum number. Don't use pushWorker() here.
for i := 0; i < pool.minWorkers; i++ {
worker := newWorker()
pool.workers[i] = worker
pool.workerStack[i] = i
worker.start(pool, i)
}
go pool.adjustWorkers()
go pool.dispatch()
return pool
}

// AddTask adds a task to the pool.
func (p *goPool) AddTask(t task) {
p.taskQueue <- t
p.taskQueue <- t
}

// Wait waits for all tasks to be dispatched.
func (p *goPool) Wait() {
for len(p.taskQueue) > 0 {
time.Sleep(100 * time.Millisecond)
}
for len(p.taskQueue) > 0 {
time.Sleep(100 * time.Millisecond)
}
}

// Release stops all workers and releases resources.
func (p *goPool) Release() {
close(p.taskQueue)
p.cond.L.Lock()
for len(p.workerStack) != p.minWorkers {
p.cond.Wait()
}
p.cond.L.Unlock()
for _, worker := range p.workers {
close(worker.taskQueue)
}
p.workers = nil
p.workerStack = nil
func (p *goPool) Release() {
close(p.taskQueue)
p.cond.L.Lock()
for len(p.workerStack) != p.minWorkers {
p.cond.Wait()
}
p.cond.L.Unlock()
for _, worker := range p.workers {
close(worker.taskQueue)
}
p.workers = nil
p.workerStack = nil
}

func (p *goPool) popWorker() int {
p.lock.Lock()
workerIndex := p.workerStack[len(p.workerStack)-1]
p.workerStack = p.workerStack[:len(p.workerStack)-1]
p.lock.Unlock()
return workerIndex
p.lock.Lock()
workerIndex := p.workerStack[len(p.workerStack)-1]
p.workerStack = p.workerStack[:len(p.workerStack)-1]
p.lock.Unlock()
return workerIndex
}

func (p *goPool) pushWorker(workerIndex int) {
p.lock.Lock()
p.workerStack = append(p.workerStack, workerIndex)
p.lock.Unlock()
p.cond.Signal()
p.lock.Lock()
p.workerStack = append(p.workerStack, workerIndex)
p.lock.Unlock()
p.cond.Signal()
}

// adjustWorkers adjusts the number of workers according to the number of tasks in the queue.
func (p *goPool) adjustWorkers() {
ticker := time.NewTicker(p.adjustInterval)
defer ticker.Stop()
ticker := time.NewTicker(p.adjustInterval)
defer ticker.Stop()

for range ticker.C {
p.cond.L.Lock()
if len(p.taskQueue) > len(p.workerStack)*3/4 && len(p.workerStack) < p.maxWorkers {
// Double the number of workers until it reaches the maximum
newWorkers := min(len(p.workerStack)*2, p.maxWorkers) - len(p.workerStack)
for i := 0; i < newWorkers; i++ {
worker := newWorker()
p.workers = append(p.workers, worker)
p.workerStack = append(p.workerStack, len(p.workers)-1)
worker.start(p, len(p.workers)-1)
}
} else if len(p.taskQueue) == 0 && len(p.workerStack) > p.minWorkers {
// Halve the number of workers until it reaches the minimum
removeWorkers := max((len(p.workerStack)-p.minWorkers)/2, p.minWorkers)
p.workers = p.workers[:len(p.workers)-removeWorkers]
p.workerStack = p.workerStack[:len(p.workerStack)-removeWorkers]
}
p.cond.L.Unlock()
}
for range ticker.C {
p.cond.L.Lock()
if len(p.taskQueue) > len(p.workerStack)*3/4 && len(p.workerStack) < p.maxWorkers {
// Double the number of workers until it reaches the maximum
newWorkers := min(len(p.workerStack)*2, p.maxWorkers) - len(p.workerStack)
for i := 0; i < newWorkers; i++ {
worker := newWorker()
p.workers = append(p.workers, worker)
p.workerStack = append(p.workerStack, len(p.workers)-1)
worker.start(p, len(p.workers)-1)
}
} else if len(p.taskQueue) == 0 && len(p.workerStack) > p.minWorkers {
// Halve the number of workers until it reaches the minimum
removeWorkers := max((len(p.workerStack)-p.minWorkers)/2, p.minWorkers)
p.workers = p.workers[:len(p.workers)-removeWorkers]
p.workerStack = p.workerStack[:len(p.workerStack)-removeWorkers]
}
p.cond.L.Unlock()
}
}

// dispatch dispatches tasks to workers.
func (p *goPool) dispatch() {
for t := range p.taskQueue {
p.cond.L.Lock()
for len(p.workerStack) == 0 {
p.cond.Wait()
}
p.cond.L.Unlock()
workerIndex := p.popWorker()
p.workers[workerIndex].taskQueue <- t
}
for t := range p.taskQueue {
p.cond.L.Lock()
for len(p.workerStack) == 0 {
p.cond.Wait()
}
p.cond.L.Unlock()
workerIndex := p.popWorker()
p.workers[workerIndex].taskQueue <- t
}
}

func min(a, b int) int {
if a < b {
return a
}
return b
if a < b {
return a
}
return b
}

func max(a, b int) int {
if a > b {
return a
}
return b
if a > b {
return a
}
return b
}
50 changes: 25 additions & 25 deletions gopool_test.go
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
package gopool

import (
"errors"
"sync"
"testing"
"time"
"errors"

"github.com/daniel-hutao/spinlock"
)
Expand Down Expand Up @@ -93,35 +93,35 @@ func BenchmarkGoroutines(b *testing.B) {
}

func TestGoPoolWithError(t *testing.T) {
var errTaskError = errors.New("task error")
pool := NewGoPool(100, WithErrorCallback(func(err error) {
if err != errTaskError {
t.Errorf("Expected error %v, but got %v", errTaskError, err)
}
}))
var errTaskError = errors.New("task error")
pool := NewGoPool(100, WithErrorCallback(func(err error) {
if err != errTaskError {
t.Errorf("Expected error %v, but got %v", errTaskError, err)
}
}))
defer pool.Release()

for i := 0; i< 1000; i++ {
pool.AddTask(func() (interface{}, error) {
return nil, errTaskError
})
}
pool.Wait()
for i := 0; i < 1000; i++ {
pool.AddTask(func() (interface{}, error) {
return nil, errTaskError
})
}
pool.Wait()
}

func TestGoPoolWithResult(t *testing.T) {
var expectedResult = "task result"
pool := NewGoPool(100, WithResultCallback(func(result interface{}) {
if result != expectedResult {
t.Errorf("Expected result %v, but got %v", expectedResult, result)
}
}))
var expectedResult = "task result"
pool := NewGoPool(100, WithResultCallback(func(result interface{}) {
if result != expectedResult {
t.Errorf("Expected result %v, but got %v", expectedResult, result)
}
}))
defer pool.Release()

for i := 0; i< 1000; i++ {
pool.AddTask(func() (interface{}, error) {
return expectedResult, nil
})
}
pool.Wait()
for i := 0; i < 1000; i++ {
pool.AddTask(func() (interface{}, error) {
return expectedResult, nil
})
}
pool.Wait()
}
Loading

0 comments on commit 7dc1fce

Please sign in to comment.