Graceful shutdown concurrent Go program with WaitGroup and Context
One of the well known advantage of Go is its support of concurrency. Thanks for the goroutine and channel, it makes writing high performance concurrent code become much easier. It is also fun to implement different concurrent patterns. I personally use this pattern a lot in some crawler and downloading resource concurrently, hope it helps!
Let start with a simple go program:
There is a task() function just sleeping for 1 second to simulate a time consuming task. And we want it to run concurrently so add a go keyword in front of the function call to start a goroutine.
➜ go run main.go
main exiting...As expected, the program will exit immediately because the main function doesn’t wait the goroutine to finish.
To fix it, we can simply add a channel to block the main function:
We create a non-buffered channel with empty struct{} type (since we just use the channel for signalling, the type doesn’t matters), after starting a goroutine, we immediately get the data from the channel by <-ch , it will block the main() until we can get something from the channel. When the task() finish, it will send an empty struct data to the ch , at the point, the main() can finally get something from the ch and continue to run.
➜ go run main.go
task finished!
main exiting...read and write to a non-buffered channel are blocking operations, it can be used to synchronize and communicate with different goroutines. Whereas buffered channel doesn’t block unless the buffer is full.
Beside using a channel, we can also use a WaitGroup to let the main function waits, it may also be more handy if there are multiple goroutines:
What we need to do is quite simple, just declare a sync.WaitGroup variable. When starting a concurrent job, call the wg.Add(1) to increment the counter, when to job is done, call the wg.Done() to decrease the counter. And at the end of main() we need to call wg.Wait() , it will block until the counter become zero.
➜ go run main.go
waiting...
task 0 finished!
task 1 finished!
task 2 finished!
main exiting...
The things becomes more interesting when we implement a worker pool pattern:
First we declare a buffered int channel pool , then we create a goroutine which will keep fetching the data from the pool and execute the task , this is the worker, if we want multiple workers, we can simply copy the goroutine code multiple times or wrap it with a loop. The worker is ready and blocking because nothing is inside the pool, now we need to feed some jobs to the pool. We can do that by a simple for loop and sending the loop index as a task id to the pool. As soon as the worker can get something from the pool , it will start working. Finally don’t forget to close() the channel if the sender(main) finish their work, otherwise the receiver(worker) will block there forever waiting for the new data comes in the channel and produce a deadlock.
➜ go run main.go
task 1 added!
task 2 added!
task 3 added!
task 4 added!
task 5 added!
waiting...
task 1 finished!
task 2 finished!
task 3 finished!
task 4 finished!
task 5 finished!
main exiting...In real world situation, it is more likely that we don’t know how many jobs we need to do or we just want to keep feeding the jobs unless we stop it.
To achieve that, we can remove the loop condition to make it infinite loop and wrap it inside a goroutine to make it non blocking. Then we also need a channel with os.Signal type to block the main() . The program will capture SIGTERM, SIGINT and send to the channel by signal.Notify()
➜ go run main.go
task 1 added!
task 2 added!
task 1 finished!
task 3 added!
task 4 added!
^Ctask 2 finished!
panic: send on closed channelgoroutine 34 [running]:
main.main.func2()
/Users/yk/Project/test/main.go:77 +0x59
created by main.main
/Users/yk/Project/test/main.go:74 +0x185
exit status 2
What!? panic…It’s because we close the pool channel after we received the quit signal but the producer goroutine still trying to send the job to the pool channel and panic happens. We also need a way to stop the producer goroutine.
Worker pool graceful shutdown with WaitGroup and Context:
Base on the previous version, we create a context ctx using the signal.NotifyContext() function, it will close the Done channel of the context when the corresponding SIGTERM, SIGINT arrives. In the producer goroutine, instead of a simple for loop, we also need to add a select{} statement. If we receive SIGTERM, SIGINT , it will notify to close the Done channel of the context and enter the case <- ctx.Done(): and exit the goroutine. Otherwise, it will just run the default case to feed jobs to the pool.
➜ go run main.go
task 1 added!
task 2 added!
task 3 added!
task 1 finished!
^Cstop filling the pool!
task 2 finished!
task 3 finished!
main exiting...Now when we send the SIGTERM, SIGINT to the program, it will first stop feeding more jobs to the pool and exit from the producer goroutine, then it waits for the worker goroutine to finsish all the existing task, then exit the main program.
If you like this blog, please give me a clap and follow, your precious feedback is my motivation!
