├── .github ├── ISSUE_TEMPLATE │ ├── bug_report.md │ └── feature_request.md └── workflows │ ├── docs.yml │ └── go.yml ├── .gitignore ├── LICENSE ├── README.md ├── examples ├── deadline_worker │ └── deadlineworker.go ├── multiple_workers │ └── multipleworkers.go ├── passing_fields │ └── passingfields.go ├── quickstart │ └── quickstart.go └── timeout_worker │ └── timeoutworker.go ├── go.mod ├── go.sum ├── workers.go └── workers_test.go /.github/ISSUE_TEMPLATE/bug_report.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Bug report 3 | about: Create a report to help us improve 4 | title: '' 5 | labels: '' 6 | assignees: '' 7 | 8 | --- 9 | 10 | **Describe the bug** 11 | A clear and concise description of what the bug is. 12 | 13 | **To Reproduce** 14 | Steps to reproduce the behavior: 15 | 1. Go to '...' 16 | 2. Click on '....' 17 | 3. Scroll down to '....' 18 | 4. See error 19 | 20 | **Expected behavior** 21 | A clear and concise description of what you expected to happen. 22 | 23 | **Screenshots** 24 | If applicable, add screenshots to help explain your problem. 25 | 26 | **Desktop (please complete the following information):** 27 | - OS: [e.g. iOS] 28 | - Browser [e.g. chrome, safari] 29 | - Version [e.g. 22] 30 | 31 | **Smartphone (please complete the following information):** 32 | - Device: [e.g. iPhone6] 33 | - OS: [e.g. iOS8.1] 34 | - Browser [e.g. stock browser, safari] 35 | - Version [e.g. 22] 36 | 37 | **Additional context** 38 | Add any other context about the problem here. 39 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/feature_request.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Feature request 3 | about: Suggest an idea for this project 4 | title: '' 5 | labels: '' 6 | assignees: '' 7 | 8 | --- 9 | 10 | **Is your feature request related to a problem? Please describe.** 11 | A clear and concise description of what the problem is. Ex. I'm always frustrated when [...] 12 | 13 | **Describe the solution you'd like** 14 | A clear and concise description of what you want to happen. 15 | 16 | **Describe alternatives you've considered** 17 | A clear and concise description of any alternative solutions or features you've considered. 18 | 19 | **Additional context** 20 | Add any other context or screenshots about the feature request here. 21 | -------------------------------------------------------------------------------- /.github/workflows/docs.yml: -------------------------------------------------------------------------------- 1 | on: 2 | release: 3 | types: 4 | - created 5 | tags: 6 | - 'v[0-9]+.[0-9]+.[0-9]+' 7 | - '**/v[0-9]+.[0-9]+.[0-9]+' 8 | 9 | jobs: 10 | build: 11 | name: Renew documentation 12 | runs-on: ubuntu-latest 13 | steps: 14 | - name: Pull new module version 15 | uses: andrewslotin/go-proxy-pull-action@master 16 | - name: actions-goveralls 17 | uses: shogo82148/actions-goveralls@v1.4.2 18 | 19 | -------------------------------------------------------------------------------- /.github/workflows/go.yml: -------------------------------------------------------------------------------- 1 | name: Go 2 | 3 | on: 4 | push: 5 | branches: [ master ] 6 | pull_request: 7 | branches: [ master ] 8 | 9 | jobs: 10 | 11 | build: 12 | runs-on: ubuntu-latest 13 | steps: 14 | - uses: actions/checkout@v2 15 | 16 | - name: Set up Go 17 | uses: actions/setup-go@v2 18 | with: 19 | go-version: 1.15 20 | 21 | - name: Build 22 | run: go build -v ./... 23 | 24 | - name: Test 25 | run: go test -v ./... 26 | 27 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /go-workers.iml 2 | /.idea/ 3 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 Robert Catmull 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ![go workers](https://raw.githubusercontent.com/catmullet/go-workers/assets/constworker_header_anim.gif) 2 | 3 | [![Mentioned in Awesome Go](https://awesome.re/mentioned-badge-flat.svg)](https://github.com/avelino/awesome-go#goroutines) 4 | [![Maintainability](https://api.codeclimate.com/v1/badges/402fee86fbd1e24defb2/maintainability)](https://codeclimate.com/github/catmullet/go-workers/maintainability) 5 | [![GoCover](http://gocover.io/_badge/github.com/catmullet/go-workers)](http://gocover.io/github.com/catmullet/go-workers) 6 | [![Go Reference](https://pkg.go.dev/badge/github.com/catmullet/go-workers.svg)](https://pkg.go.dev/github.com/catmullet/go-workers) 7 | 8 | # Examples 9 | * [Quickstart](https://github.com/catmullet/go-workers/blob/master/examples/quickstart/quickstart.go) 10 | * [Multiple Go Workers](https://github.com/catmullet/go-workers/blob/master/examples/multiple_workers/multipleworkers.go) 11 | * [Passing Fields](https://github.com/catmullet/go-workers/blob/master/examples/passing_fields/passingfields.go) 12 | # Getting Started 13 | ### Pull in the dependency 14 | ```zsh 15 | go get github.com/catmullet/go-workers 16 | ``` 17 | 18 | ### Add the import to your project 19 | giving an alias helps since go-workers doesn't exactly follow conventions. 20 | _(If you're using a JetBrains IDE it should automatically give it an alias)_ 21 | ```go 22 | import ( 23 | workers "github.com/catmullet/go-workers" 24 | ) 25 | ``` 26 | ### Create a new worker worker 27 | The NewWorker factory method returns a new worker. 28 | _(Method chaining can be performed on this method like calling .Work() immediately after.)_ 29 | ```go 30 | type MyWorker struct {} 31 | 32 | func NewMyWorker() Worker { 33 | return &MyWorker{} 34 | } 35 | 36 | func (my *MyWorker) Work(in interface{}, out chan<- interface{}) error { 37 | // work iteration here 38 | } 39 | 40 | runner := workers.NewRunner(ctx, NewMyWorker(), numberOfWorkers) 41 | ``` 42 | ### Send work to worker 43 | Send accepts an interface. So send it anything you want. 44 | ```go 45 | runner.Send("Hello World") 46 | ``` 47 | ### Wait for the worker to finish and handle errors 48 | Any error that bubbles up from your worker functions will return here. 49 | ```go 50 | if err := runner.Wait(); err != nil { 51 | //Handle error 52 | } 53 | ``` 54 | 55 | ## Working With Multiple Workers 56 | ### Passing work form one worker to the next 57 | 58 | By using the InFrom method you can tell `workerTwo` to accept output from `workerOne` 59 | ```go 60 | runnerOne := workers.NewRunner(ctx, NewMyWorker(), 100).Work() 61 | runnerTwo := workers.NewRunner(ctx, NewMyWorkerTwo(), 100).InFrom(workerOne).Work() 62 | ``` 63 | ### Accepting output from multiple workers 64 | It is possible to accept output from more than one worker but it is up to you to determine what is coming from which worker. (They will send on the same channel.) 65 | ```go 66 | runnerOne := workers.NewRunner(ctx, NewMyWorker(), 100).Work() 67 | runnerTwo := workers.NewRunner(ctx, NewMyWorkerTwo(), 100).Work() 68 | runnerThree := workers.NewRunner(ctx, NewMyWorkerThree(), 100).InFrom(workerOne, workerTwo).Work() 69 | ``` 70 | 71 | ## Passing Fields To Workers 72 | ### Adding Values 73 | Fields can be passed via the workers object. Be sure as with any concurrency in Golang that your variables are concurrent safe. Most often the golang documentation will state the package or parts of it are concurrent safe. If it does not state so there is a good chance it isn't. Use the sync package to lock and unlock for writes on unsafe variables. (It is good practice NOT to defer in the work function.) 74 | 75 | worker **ONLY** use the `Send()` method to get data into your worker. It is not shared memory unlike the worker objects values. 76 | 77 | ```go 78 | type MyWorker struct { 79 | message string 80 | } 81 | 82 | func NewMyWorker(message string) Worker { 83 | return &MyWorker{message} 84 | } 85 | 86 | func (my *MyWorker) Work(in interface{}, out chan<- interface{}) error { 87 | fmt.Println(my.message) 88 | } 89 | 90 | runner := workers.NewRunner(ctx, NewMyWorker(), 100).Work() 91 | ``` 92 | 93 | ### Setting Timeouts or Deadlines 94 | If your workers needs to stop at a deadline or you just need to have a timeout use the SetTimeout or SetDeadline methods. (These must be in place before setting the workers off to work.) 95 | ```go 96 | // Setting a timeout of 2 seconds 97 | timeoutRunner.SetTimeout(2 * time.Second) 98 | 99 | // Setting a deadline of 4 hours from now 100 | deadlineRunner.SetDeadline(time.Now().Add(4 * time.Hour)) 101 | 102 | func workerFunction(in interface{}, out chan<- interface{} error { 103 | fmt.Println(in) 104 | time.Sleep(1 * time.Second) 105 | } 106 | ``` 107 | 108 | 109 | ## Performance Hints 110 | ### Buffered Writer 111 | If you want to write out to a file or just stdout you can use SetWriterOut(writer io.Writer). The worker will have the following methods available 112 | ```go 113 | runner.Println() 114 | runner.Printf() 115 | runner.Print() 116 | ``` 117 | The workers use a buffered writer for output and can be up to 3 times faster than the fmt package. Just be mindful it won't write out to the console as quickly as an unbuffered writer. It will sync and eventually flush everything at the end, making it ideal for writing out to a file. 118 | 119 | ### Using GOGC env variable 120 | If your application is based solely around using workers, consider upping the percentage of when the scheduler will garbage collect. (ex. GOGC=200) 200% -> 300% is a good starting point. Make sure your machine has some good memory behind it. 121 | By upping the percentage your application will interupt the workers less, meaning they get more work done. However, be aware of the rest of your applications needs when modifying this variable. 122 | 123 | ### Using GOMAXPROCS env variable 124 | For workers that run quick bursts of lots of simple data consider lowering the GOMAXPROCS. Be carfeful though, this can affect your entire applicaitons performance. Profile your application and benchmark it. See where your application runs best. 125 | -------------------------------------------------------------------------------- /examples/deadline_worker/deadlineworker.go: -------------------------------------------------------------------------------- 1 | // +build ignore 2 | 3 | package main 4 | 5 | import ( 6 | "context" 7 | "fmt" 8 | "github.com/catmullet/go-workers" 9 | "time" 10 | ) 11 | 12 | func main() { 13 | ctx := context.Background() 14 | t := time.Now() 15 | 16 | deadlineWorker := workers.NewRunner(ctx, NewDeadlineWorker(), 100). 17 | SetDeadline(t.Add(200 * time.Millisecond)).Start() 18 | 19 | for i := 0; i < 1000000; i++ { 20 | deadlineWorker.Send("hello") 21 | } 22 | 23 | err := deadlineWorker.Wait() 24 | if err != nil { 25 | fmt.Println(err) 26 | } 27 | fmt.Println("finished") 28 | } 29 | 30 | type DeadlineWorker struct{} 31 | 32 | func NewDeadlineWorker() workers.Worker { 33 | return &DeadlineWorker{} 34 | } 35 | 36 | func (dlw *DeadlineWorker) Work(in interface{}, out chan<- interface{}) error { 37 | fmt.Println(in) 38 | time.Sleep(1 * time.Second) 39 | return nil 40 | } 41 | -------------------------------------------------------------------------------- /examples/multiple_workers/multipleworkers.go: -------------------------------------------------------------------------------- 1 | // +build ignore 2 | 3 | package main 4 | 5 | import ( 6 | "context" 7 | "fmt" 8 | "github.com/catmullet/go-workers" 9 | "math/rand" 10 | "sync" 11 | ) 12 | 13 | var ( 14 | count = make(map[string]int) 15 | mut = sync.RWMutex{} 16 | ) 17 | 18 | func main() { 19 | ctx := context.Background() 20 | 21 | workerOne := workers.NewRunner(ctx, NewWorkerOne(), 1000).Start() 22 | workerTwo := workers.NewRunner(ctx, NewWorkerTwo(), 1000).InFrom(workerOne).Start() 23 | 24 | go func() { 25 | for i := 0; i < 100000; i++ { 26 | workerOne.Send(rand.Intn(100)) 27 | } 28 | if err := workerOne.Wait(); err != nil { 29 | fmt.Println(err) 30 | } 31 | }() 32 | 33 | if err := workerTwo.Wait(); err != nil { 34 | fmt.Println(err) 35 | } 36 | 37 | fmt.Println("worker_one", count["worker_one"]) 38 | fmt.Println("worker_two", count["worker_two"]) 39 | fmt.Println("finished") 40 | } 41 | 42 | type WorkerOne struct { 43 | } 44 | type WorkerTwo struct { 45 | } 46 | 47 | func NewWorkerOne() workers.Worker { 48 | return &WorkerOne{} 49 | } 50 | 51 | func NewWorkerTwo() workers.Worker { 52 | return &WorkerTwo{} 53 | } 54 | 55 | func (wo *WorkerOne) Work(in interface{}, out chan<- interface{}) error { 56 | var workerOne = "worker_one" 57 | mut.Lock() 58 | if val, ok := count[workerOne]; ok { 59 | count[workerOne] = val + 1 60 | } else { 61 | count[workerOne] = 1 62 | } 63 | mut.Unlock() 64 | 65 | total := in.(int) * 2 66 | fmt.Println("worker1", fmt.Sprintf("%d * 2 = %d", in.(int), total)) 67 | out <- total 68 | return nil 69 | } 70 | 71 | func (wt *WorkerTwo) Work(in interface{}, out chan<- interface{}) error { 72 | var workerTwo = "worker_two" 73 | mut.Lock() 74 | if val, ok := count[workerTwo]; ok { 75 | count[workerTwo] = val + 1 76 | } else { 77 | count[workerTwo] = 1 78 | } 79 | mut.Unlock() 80 | 81 | totalFromWorkerOne := in.(int) 82 | fmt.Println("worker2", fmt.Sprintf("%d * 4 = %d", totalFromWorkerOne, totalFromWorkerOne*4)) 83 | return nil 84 | } 85 | -------------------------------------------------------------------------------- /examples/passing_fields/passingfields.go: -------------------------------------------------------------------------------- 1 | // +build ignore 2 | 3 | package main 4 | 5 | import ( 6 | "context" 7 | "fmt" 8 | "github.com/catmullet/go-workers" 9 | "math/rand" 10 | ) 11 | 12 | func main() { 13 | ctx := context.Background() 14 | workerOne := workers.NewRunner(ctx, NewWorkerOne(2), 100).Start() 15 | workerTwo := workers.NewRunner(ctx, NewWorkerTwo(4), 100).InFrom(workerOne).Start() 16 | 17 | for i := 0; i < 15; i++ { 18 | workerOne.Send(rand.Intn(100)) 19 | } 20 | 21 | if err := workerOne.Wait(); err != nil { 22 | fmt.Println(err) 23 | } 24 | 25 | if err := workerTwo.Wait(); err != nil { 26 | fmt.Println(err) 27 | } 28 | } 29 | 30 | type WorkerOne struct { 31 | amountToMultiply int 32 | } 33 | type WorkerTwo struct { 34 | amountToMultiply int 35 | } 36 | 37 | func NewWorkerOne(amountToMultiply int) workers.Worker { 38 | return &WorkerOne{ 39 | amountToMultiply: amountToMultiply, 40 | } 41 | } 42 | 43 | func NewWorkerTwo(amountToMultiply int) workers.Worker { 44 | return &WorkerTwo{ 45 | amountToMultiply, 46 | } 47 | } 48 | 49 | func (wo *WorkerOne) Work(in interface{}, out chan<- interface{}) error { 50 | total := in.(int) * wo.amountToMultiply 51 | fmt.Println("worker1", fmt.Sprintf("%d * %d = %d", in.(int), wo.amountToMultiply, total)) 52 | out <- total 53 | return nil 54 | } 55 | 56 | func (wt *WorkerTwo) Work(in interface{}, out chan<- interface{}) error { 57 | totalFromWorkerOne := in.(int) 58 | fmt.Println("worker2", fmt.Sprintf("%d * %d = %d", totalFromWorkerOne, wt.amountToMultiply, totalFromWorkerOne*wt.amountToMultiply)) 59 | return nil 60 | } 61 | -------------------------------------------------------------------------------- /examples/quickstart/quickstart.go: -------------------------------------------------------------------------------- 1 | // +build ignore 2 | 3 | package main 4 | 5 | import ( 6 | "context" 7 | "fmt" 8 | "github.com/catmullet/go-workers" 9 | "math/rand" 10 | "time" 11 | ) 12 | 13 | func main() { 14 | ctx := context.Background() 15 | t := time.Now() 16 | rnr := workers.NewRunner(ctx, NewWorker(), 100).Start() 17 | 18 | for i := 0; i < 1000000; i++ { 19 | rnr.Send(rand.Intn(100)) 20 | } 21 | 22 | if err := rnr.Wait(); err != nil { 23 | fmt.Println(err) 24 | } 25 | 26 | totalTime := time.Since(t).Milliseconds() 27 | fmt.Printf("total time %dms\n", totalTime) 28 | } 29 | 30 | type WorkerOne struct { 31 | } 32 | 33 | func NewWorker() workers.Worker { 34 | return &WorkerOne{} 35 | } 36 | 37 | func (wo *WorkerOne) Work(in interface{}, out chan<- interface{}) error { 38 | total := in.(int) * 2 39 | fmt.Println(fmt.Sprintf("%d * 2 = %d", in.(int), total)) 40 | return nil 41 | } 42 | -------------------------------------------------------------------------------- /examples/timeout_worker/timeoutworker.go: -------------------------------------------------------------------------------- 1 | // +build ignore 2 | 3 | package main 4 | 5 | import ( 6 | "context" 7 | "fmt" 8 | "github.com/catmullet/go-workers" 9 | "time" 10 | ) 11 | 12 | func main() { 13 | ctx := context.Background() 14 | 15 | timeoutWorker := workers.NewRunner(ctx, NewTimeoutWorker(), 10).SetTimeout(100 * time.Millisecond).Start() 16 | 17 | for i := 0; i < 1000000; i++ { 18 | timeoutWorker.Send("hello") 19 | } 20 | 21 | err := timeoutWorker.Wait() 22 | if err != nil { 23 | fmt.Println(err) 24 | } 25 | } 26 | 27 | type TimeoutWorker struct{} 28 | 29 | func NewTimeoutWorker() workers.Worker { 30 | return &TimeoutWorker{} 31 | } 32 | 33 | func (tw *TimeoutWorker) Work(in interface{}, out chan<- interface{}) error { 34 | fmt.Println(in) 35 | time.Sleep(1 * time.Second) 36 | return nil 37 | } 38 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/catmullet/go-workers 2 | 3 | go 1.15 4 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/catmullet/go-workers/0a8e61ee4bd9b2ac1ed5ae44b3000f41950b15a9/go.sum -------------------------------------------------------------------------------- /workers.go: -------------------------------------------------------------------------------- 1 | package workers 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "os" 7 | "os/signal" 8 | "sync" 9 | "syscall" 10 | "time" 11 | ) 12 | 13 | var defaultWatchSignals = []os.Signal{syscall.SIGINT, syscall.SIGTERM, syscall.SIGKILL} 14 | 15 | // Worker Contains the work function. Allows an input and output to a channel or another worker for pipeline work. 16 | // Return nil if you want the Runner to continue otherwise any error will cause the Runner to shutdown and return the 17 | // error. 18 | type Worker interface { 19 | Work(in interface{}, out chan<- interface{}) error 20 | } 21 | 22 | // Runner Handles the running the Worker logic. 23 | type Runner interface { 24 | BeforeFunc(func(ctx context.Context) error) Runner 25 | AfterFunc(func(ctx context.Context, err error) error) Runner 26 | SetDeadline(t time.Time) Runner 27 | SetTimeout(duration time.Duration) Runner 28 | SetFollower() 29 | Send(in interface{}) 30 | InFrom(w ...Runner) Runner 31 | SetOut(chan interface{}) 32 | Start() Runner 33 | Stop() chan error 34 | Wait() error 35 | } 36 | 37 | type runner struct { 38 | ctx context.Context 39 | cancel context.CancelFunc 40 | inChan chan interface{} 41 | outChan chan interface{} 42 | errChan chan error 43 | signalChan chan os.Signal 44 | limiter chan struct{} 45 | 46 | afterFunc func(ctx context.Context, err error) error 47 | workFunc func(in interface{}, out chan<- interface{}) error 48 | beforeFunc func(ctx context.Context) error 49 | 50 | timeout time.Duration 51 | deadline time.Duration 52 | 53 | isLeader bool 54 | stopCalled bool 55 | 56 | numWorkers int64 57 | lock *sync.RWMutex 58 | wg *sync.WaitGroup 59 | done *sync.Once 60 | once *sync.Once 61 | } 62 | 63 | // NewRunner Factory function for a new Runner. The Runner will handle running the workers logic. 64 | func NewRunner(ctx context.Context, w Worker, numWorkers int64) Runner { 65 | var runnerCtx, runnerCancel = context.WithCancel(ctx) 66 | var runner = &runner{ 67 | ctx: runnerCtx, 68 | cancel: runnerCancel, 69 | inChan: make(chan interface{}, numWorkers), 70 | outChan: nil, 71 | errChan: make(chan error, 1), 72 | signalChan: make(chan os.Signal, 1), 73 | limiter: make(chan struct{}, numWorkers), 74 | afterFunc: func(ctx context.Context, err error) error { return err }, 75 | workFunc: w.Work, 76 | beforeFunc: func(ctx context.Context) error { return nil }, 77 | numWorkers: numWorkers, 78 | isLeader: true, 79 | lock: new(sync.RWMutex), 80 | wg: new(sync.WaitGroup), 81 | once: new(sync.Once), 82 | done: new(sync.Once), 83 | } 84 | runner.waitForSignal(defaultWatchSignals...) 85 | return runner 86 | } 87 | 88 | // Send Send an object to the worker for processing. 89 | func (r *runner) Send(in interface{}) { 90 | select { 91 | case <-r.ctx.Done(): 92 | return 93 | case r.inChan <- in: 94 | } 95 | } 96 | 97 | // InFrom Set a worker to accept output from another worker(s). 98 | func (r *runner) InFrom(w ...Runner) Runner { 99 | r.SetFollower() 100 | for _, wr := range w { 101 | wr.SetOut(r.inChan) 102 | } 103 | return r 104 | } 105 | 106 | // SetFollower Sets the worker as a follower and does not need to close it's in channel. 107 | func (r *runner) SetFollower() { 108 | r.lock.Lock() 109 | r.isLeader = false 110 | r.lock.Unlock() 111 | } 112 | 113 | // Start Starts the worker on processing. 114 | func (r *runner) Start() Runner { 115 | r.startWork() 116 | return r 117 | } 118 | 119 | // BeforeFunc Function to be run before worker starts processing. 120 | func (r *runner) BeforeFunc(f func(ctx context.Context) error) Runner { 121 | r.beforeFunc = f 122 | return r 123 | } 124 | 125 | // AfterFunc Function to be run after worker has stopped. 126 | func (r *runner) AfterFunc(f func(ctx context.Context, err error) error) Runner { 127 | r.afterFunc = f 128 | return r 129 | } 130 | 131 | // SetOut Allows the setting of a workers out channel, if not already set. 132 | func (r *runner) SetOut(c chan interface{}) { 133 | if r.outChan != nil { 134 | return 135 | } 136 | r.outChan = c 137 | } 138 | 139 | // SetDeadline allows a time to be set when the workers should stop. 140 | // Deadline needs to be handled by the IsDone method. 141 | func (r *runner) SetDeadline(t time.Time) Runner { 142 | r.lock.Lock() 143 | defer r.lock.Unlock() 144 | r.ctx, r.cancel = context.WithDeadline(r.ctx, t) 145 | return r 146 | } 147 | 148 | // SetTimeout allows a time duration to be set when the workers should stop. 149 | // Timeout needs to be handled by the IsDone method. 150 | func (r *runner) SetTimeout(duration time.Duration) Runner { 151 | r.lock.Lock() 152 | defer r.lock.Unlock() 153 | r.timeout = duration 154 | return r 155 | } 156 | 157 | // Wait calls stop on workers and waits for the channel to drain. 158 | // !!Should only be called when certain nothing will send to worker. 159 | func (r *runner) Wait() error { 160 | r.waitForDrain() 161 | if err := <-r.Stop(); err != nil && !errors.Is(err, context.Canceled) { 162 | return err 163 | } 164 | return nil 165 | } 166 | 167 | // Stop Stops the processing of a worker and closes it's channel in. 168 | // Returns a blocking channel with type error. 169 | // !!Should only be called when certain nothing will send to worker. 170 | func (r *runner) Stop() chan error { 171 | r.done.Do(func() { 172 | if r.inChan != nil && r.isLeader { 173 | close(r.inChan) 174 | } 175 | }) 176 | return r.errChan 177 | } 178 | 179 | // IsDone returns a channel signaling the workers context has been canceled. 180 | func (r *runner) IsDone() <-chan struct{} { 181 | return r.ctx.Done() 182 | } 183 | 184 | // waitForSignal make sure we wait for a term signal and shutdown correctly 185 | func (r *runner) waitForSignal(signals ...os.Signal) { 186 | go func() { 187 | signal.Notify(r.signalChan, signals...) 188 | <-r.signalChan 189 | if r.cancel != nil { 190 | r.cancel() 191 | } 192 | }() 193 | } 194 | 195 | // waitForDrain Waits for the limiter to be zeroed out and the in channel to be empty. 196 | func (r *runner) waitForDrain() { 197 | for len(r.limiter) > 0 || len(r.inChan) > 0 { 198 | // Wait for the drain. 199 | } 200 | } 201 | 202 | // startWork Runs the before function and starts processing until one of three things happen. 203 | // 1. A term signal is received or cancellation of context. 204 | // 2. Stop function is called. 205 | // 3. Worker returns an error. 206 | func (r *runner) startWork() { 207 | var err error 208 | if err = r.beforeFunc(r.ctx); err != nil { 209 | r.errChan <- err 210 | return 211 | } 212 | if r.timeout > 0 { 213 | r.ctx, r.cancel = context.WithTimeout(r.ctx, r.timeout) 214 | } 215 | r.wg.Add(1) 216 | go func() { 217 | var workerWG = new(sync.WaitGroup) 218 | var closeOnce = new(sync.Once) 219 | 220 | // write out error if not nil on exit. 221 | defer func() { 222 | workerWG.Wait() 223 | r.errChan <- err 224 | closeOnce.Do(func() { 225 | if r.outChan != nil { 226 | close(r.outChan) 227 | } 228 | }) 229 | r.wg.Done() 230 | }() 231 | for in := range r.inChan { 232 | input := in 233 | r.limiter <- struct{}{} 234 | workerWG.Add(1) 235 | go func() { 236 | defer func() { 237 | <-r.limiter 238 | workerWG.Done() 239 | }() 240 | if err := r.workFunc(input, r.outChan); err != nil { 241 | r.once.Do(func() { 242 | r.errChan <- err 243 | r.cancel() 244 | return 245 | }) 246 | } 247 | }() 248 | } 249 | }() 250 | } 251 | -------------------------------------------------------------------------------- /workers_test.go: -------------------------------------------------------------------------------- 1 | package workers 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "fmt" 7 | "math/rand" 8 | "os" 9 | "runtime" 10 | "runtime/debug" 11 | "sync" 12 | "testing" 13 | "time" 14 | ) 15 | 16 | const ( 17 | workerCount = 100000 18 | workerTimeout = time.Millisecond * 300 19 | runTimes = 100000 20 | ) 21 | 22 | type WorkerOne struct { 23 | Count int 24 | sync.Mutex 25 | } 26 | type WorkerTwo struct { 27 | Count int 28 | sync.Mutex 29 | } 30 | 31 | func NewWorkerOne() *WorkerOne { 32 | return &WorkerOne{} 33 | } 34 | 35 | func NewWorkerTwo() *WorkerTwo { 36 | return &WorkerTwo{} 37 | } 38 | 39 | func (wo *WorkerOne) CurrentCount() int { 40 | wo.Lock() 41 | defer wo.Unlock() 42 | return wo.Count 43 | } 44 | 45 | func (wo *WorkerOne) Work(in interface{}, out chan<- interface{}) error { 46 | mut.Lock() 47 | wo.Count = wo.Count + 1 48 | mut.Unlock() 49 | 50 | total := in.(int) * 2 51 | out <- total 52 | return nil 53 | } 54 | 55 | func (wt *WorkerTwo) CurrentCount() int { 56 | wt.Lock() 57 | defer wt.Unlock() 58 | return wt.Count 59 | } 60 | 61 | func (wt *WorkerTwo) Work(in interface{}, out chan<- interface{}) error { 62 | mut.Lock() 63 | wt.Count = wt.Count + 1 64 | mut.Unlock() 65 | return nil 66 | } 67 | 68 | var ( 69 | mut = sync.RWMutex{} 70 | err = errors.New("test error") 71 | deadline = func() time.Time { return time.Now().Add(workerTimeout) } 72 | workerTestScenarios = []workerTest{ 73 | { 74 | name: "work basic", 75 | worker: NewTestWorkerObject(workBasic()), 76 | numWorkers: workerCount, 77 | }, 78 | { 79 | name: "work basic with timeout", 80 | timeout: workerTimeout, 81 | worker: NewTestWorkerObject(workBasic()), 82 | numWorkers: workerCount, 83 | }, 84 | { 85 | name: "work basic with deadline", 86 | deadline: deadline, 87 | worker: NewTestWorkerObject(workBasic()), 88 | numWorkers: workerCount, 89 | }, 90 | { 91 | name: "work with return of error", 92 | worker: NewTestWorkerObject(workWithError(err)), 93 | errExpected: true, 94 | numWorkers: workerCount, 95 | }, 96 | { 97 | name: "work with return of error with timeout", 98 | timeout: workerTimeout, 99 | worker: NewTestWorkerObject(workWithError(err)), 100 | errExpected: true, 101 | numWorkers: workerCount, 102 | }, 103 | { 104 | name: "work with return of error with deadline", 105 | deadline: deadline, 106 | worker: NewTestWorkerObject(workWithError(err)), 107 | errExpected: true, 108 | numWorkers: workerCount, 109 | }, 110 | } 111 | 112 | getWorker = func(ctx context.Context, wt workerTest) Runner { 113 | worker := NewRunner(ctx, wt.worker, wt.numWorkers) 114 | if wt.timeout > 0 { 115 | worker.SetTimeout(wt.timeout) 116 | } 117 | if wt.deadline != nil { 118 | worker.SetDeadline(wt.deadline()) 119 | } 120 | return worker 121 | } 122 | ) 123 | 124 | type workerTest struct { 125 | name string 126 | timeout time.Duration 127 | deadline func() time.Time 128 | worker Worker 129 | numWorkers int64 130 | testSignal bool 131 | errExpected bool 132 | } 133 | 134 | type TestWorkerObject struct { 135 | workFunc func(in interface{}, out chan<- interface{}) error 136 | } 137 | 138 | func NewTestWorkerObject(wf func(in interface{}, out chan<- interface{}) error) Worker { 139 | return &TestWorkerObject{wf} 140 | } 141 | 142 | func (tw *TestWorkerObject) Work(in interface{}, out chan<- interface{}) error { 143 | return tw.workFunc(in, out) 144 | } 145 | 146 | func workBasicNoOut() func(in interface{}, out chan<- interface{}) error { 147 | return func(in interface{}, out chan<- interface{}) error { 148 | _ = in.(int) 149 | return nil 150 | } 151 | } 152 | 153 | func workBasic() func(in interface{}, out chan<- interface{}) error { 154 | return func(in interface{}, out chan<- interface{}) error { 155 | i := in.(int) 156 | out <- i 157 | return nil 158 | } 159 | } 160 | 161 | func workWithError(err error) func(in interface{}, out chan<- interface{}) error { 162 | return func(in interface{}, out chan<- interface{}) error { 163 | i := in.(int) 164 | total := i * rand.Intn(1000) 165 | if i == 100 { 166 | return err 167 | } 168 | out <- total 169 | return nil 170 | } 171 | } 172 | 173 | func TestMain(m *testing.M) { 174 | debug.SetGCPercent(500) 175 | runtime.GOMAXPROCS(2) 176 | code := m.Run() 177 | os.Exit(code) 178 | } 179 | 180 | func TestWorkers(t *testing.T) { 181 | for _, tt := range workerTestScenarios { 182 | t.Run(tt.name, func(t *testing.T) { 183 | ctx := context.Background() 184 | workerOne := getWorker(ctx, tt).Start() 185 | // always need a consumer for the out tests so using basic here. 186 | workerTwo := NewRunner(ctx, NewTestWorkerObject(workBasicNoOut()), workerCount).InFrom(workerOne).Start() 187 | 188 | for i := 0; i < runTimes; i++ { 189 | workerOne.Send(i) 190 | } 191 | 192 | if err := workerOne.Wait(); err != nil && (!tt.errExpected) { 193 | t.Error(err) 194 | } 195 | if err := workerTwo.Wait(); err != nil && !tt.errExpected { 196 | t.Error(err) 197 | } 198 | }) 199 | } 200 | } 201 | 202 | func TestWorkersFinish100(t *testing.T) { 203 | const workCount = 100 204 | ctx := context.Background() 205 | w1 := NewWorkerOne() 206 | w2 := NewWorkerTwo() 207 | workerOne := NewRunner(ctx, w1, 1000).Start() 208 | workerTwo := NewRunner(ctx, w2, 1000).InFrom(workerOne).Start() 209 | 210 | for i := 0; i < workCount; i++ { 211 | workerOne.Send(rand.Intn(100)) 212 | } 213 | 214 | if err := workerOne.Wait(); err != nil { 215 | fmt.Println(err) 216 | } 217 | 218 | if err := workerTwo.Wait(); err != nil { 219 | fmt.Println(err) 220 | } 221 | 222 | if w1.CurrentCount() != workCount { 223 | t.Log("worker one failed to finish,", "worker_one count", w1.CurrentCount(), "/ 100000") 224 | t.Fail() 225 | } 226 | if w2.CurrentCount() != workCount { 227 | t.Log("worker two failed to finish,", "worker_two count", w2.CurrentCount(), "/ 100000") 228 | t.Fail() 229 | } 230 | 231 | t.Logf("worker_one count: %d, worker_two count: %d", w1.CurrentCount(), w2.CurrentCount()) 232 | } 233 | 234 | func TestWorkersFinish100000(t *testing.T) { 235 | const workCount = 100000 236 | ctx := context.Background() 237 | w1 := NewWorkerOne() 238 | w2 := NewWorkerTwo() 239 | workerOne := NewRunner(ctx, w1, 1000).Start() 240 | workerTwo := NewRunner(ctx, w2, 1000).InFrom(workerOne).Start() 241 | 242 | for i := 0; i < workCount; i++ { 243 | workerOne.Send(rand.Intn(100)) 244 | } 245 | 246 | if err := workerOne.Wait(); err != nil { 247 | fmt.Println(err) 248 | } 249 | 250 | if err := workerTwo.Wait(); err != nil { 251 | fmt.Println(err) 252 | } 253 | 254 | if w1.CurrentCount() != workCount { 255 | t.Log("worker one failed to finish,", "worker_one count", w1.CurrentCount(), "/ 100000") 256 | t.Fail() 257 | } 258 | if w2.CurrentCount() != workCount { 259 | t.Log("worker two failed to finish,", "worker_two count", w2.CurrentCount(), "/ 100000") 260 | t.Fail() 261 | } 262 | 263 | t.Logf("worker_one count: %d, worker_two count: %d", w1.CurrentCount(), w2.CurrentCount()) 264 | } 265 | 266 | func TestWorkersFinish1000000(t *testing.T) { 267 | const workCount = 1000000 268 | ctx := context.Background() 269 | w1 := NewWorkerOne() 270 | w2 := NewWorkerTwo() 271 | workerOne := NewRunner(ctx, w1, 1000).Start() 272 | workerTwo := NewRunner(ctx, w2, 1000).InFrom(workerOne).Start() 273 | 274 | for i := 0; i < workCount; i++ { 275 | workerOne.Send(rand.Intn(100)) 276 | } 277 | 278 | if err := workerOne.Wait(); err != nil { 279 | fmt.Println(err) 280 | } 281 | 282 | if err := workerTwo.Wait(); err != nil { 283 | fmt.Println(err) 284 | } 285 | 286 | if w1.CurrentCount() != workCount { 287 | t.Log("worker one failed to finish,", "worker_one count", w1.CurrentCount(), "/ 100000") 288 | t.Fail() 289 | } 290 | if w2.CurrentCount() != workCount { 291 | t.Log("worker two failed to finish,", "worker_two count", w2.CurrentCount(), "/ 100000") 292 | t.Fail() 293 | } 294 | 295 | t.Logf("worker_one count: %d, worker_two count: %d", w1.CurrentCount(), w2.CurrentCount()) 296 | } 297 | 298 | func BenchmarkGoWorkers1to1(b *testing.B) { 299 | worker := NewRunner(context.Background(), NewTestWorkerObject(workBasicNoOut()), 1000).Start() 300 | 301 | b.ResetTimer() 302 | for i := 0; i < b.N; i++ { 303 | for j := 0; j < 1000; j++ { 304 | worker.Send(j) 305 | } 306 | } 307 | b.StopTimer() 308 | 309 | if err := worker.Wait(); err != nil { 310 | b.Error(err) 311 | } 312 | } 313 | 314 | func Benchmark100GoWorkers(b *testing.B) { 315 | b.ReportAllocs() 316 | worker := NewRunner(context.Background(), NewTestWorkerObject(workBasicNoOut()), 100).Start() 317 | 318 | b.ResetTimer() 319 | for i := 0; i < b.N; i++ { 320 | worker.Send(i) 321 | } 322 | 323 | if err := worker.Wait(); err != nil { 324 | b.Error(err) 325 | } 326 | } 327 | 328 | func Benchmark1000GoWorkers(b *testing.B) { 329 | b.ReportAllocs() 330 | worker := NewRunner(context.Background(), NewTestWorkerObject(workBasicNoOut()), 1000).Start() 331 | 332 | b.ResetTimer() 333 | for i := 0; i < b.N; i++ { 334 | worker.Send(i) 335 | } 336 | 337 | if err := worker.Wait(); err != nil { 338 | b.Error(err) 339 | } 340 | } 341 | 342 | func Benchmark10000GoWorkers(b *testing.B) { 343 | b.ReportAllocs() 344 | worker := NewRunner(context.Background(), NewTestWorkerObject(workBasicNoOut()), 10000).Start() 345 | 346 | b.ResetTimer() 347 | for i := 0; i < b.N; i++ { 348 | worker.Send(i) 349 | } 350 | 351 | if err := worker.Wait(); err != nil { 352 | b.Error(err) 353 | } 354 | } 355 | --------------------------------------------------------------------------------