防止粗鲁结束Golang程序
防止粗鲁结束Golang程序
本文介绍如何防止用户粗鲁结束应用程序，正在运行的程序可能一些任务正执行一半，导致数据不一致或资源未回收。
1. 问题描述
首先我们通过一个程序模拟问题场景。
package main
import (
"fmt"
"time"
)
type Task struct {
ticker *time.Ticker
}
// 每隔一段时间重复执行
func (t *Task) Run() {
for {
select {
case <-t.ticker.C:
handle()
}
}
}
// 模拟执行的任务
func handle() {
for i := 0; i < 10; i++ {
fmt.Print("#")
time.Sleep(time.Millisecond * 200)
}
fmt.Println()
}
func main() {
task := &Task{
ticker: time.NewTicker(time.Second * 2),
}
task.Run()
}
在2秒时间间隔内运行 handle任务，仅打印10个#号字符，每200ms打印一次。如果通过 ctrl+c终止程序，则可能部分任务没有执行完毕（每个任务打印10个#号）：
##########
#####
我们希望能够捕获终止信号，在 handle 任务执行完成之后才结束应用，实际应用中可能执行一些清理任务。下面我们先解决捕获终止信号问题。
2. 捕获 ctrl+c信号
go的实现依赖通道，因此首先要定义os.Signal类型通道，并带有一个缓冲区空间，不想丢失任何信号；接着告诉系统希望捕获 os.Interrupt信号至我们创建的通道；最后在协程中等待信号到来。
根据上面的思路，我们仅修改 main函数，代码如下：
func main() {
task := &Task{
ticker: time.NewTicker(time.Second * 2),
}
c := make(chan os.Signal, 1)
signal.Notify(c, os.Interrupt) // 还可以增加多个信号，如：signal.Notify(killSignal, os.Interrupt)
go func() {
select {
case sig := <-c:
fmt.Printf("Got %s signal. Aborting...\n", sig)
os.Exit(1)
}
}()
task.Run()
}
现在，如果终止 handle任务，运行结果：
##########
#######Got interrupt signal. Aborting...
很好，除了看到捕获的信号输出，其他都没有变。下面我们实现最后完整程序。
3. 防止粗鲁结束程序
我们利用通道实现优雅地结束程序模式：
type Task struct {
closed chan struct{}
ticker *time.Ticker
}
通道用于通知所有感兴趣的伙伴，有终止信号想停止执行的任务。因此我们命令通道为 closed，当然这没有强制规定。这种类型的通道并不做具体事宜，因此通常使用 struct{}类型，重要的是从通道中能够接收到值。
所有希望优雅地结束长时间运行任务，除了执行自身实际任务外，还需从该通道侦听值，如果存在值则终止执行。因此我们修改 run函数：
func (t *Task) Run() {
for {
select {
case <-t.closed:
return
case <-t.ticker.C:
handle()
}
}
}
如果从 closed通道中接收到值则通过 return结束 run，即不再继续执行新的任务。
为了表示终止任务意图，我们需要向通道发送值。但我们可以做得更好，因为从已关闭通道接收会立即返回零值，所以我们可以直接关闭通道。
func (t *Task) Stop() {
close(t.closed)
}
收到中断信号我们调用该函数。因此需先创建关闭通道：
func main() {
task := &Task{
closed: make(chan struct{}),
ticker: time.NewTicker(time.Second * 2),
}
c := make(chan os.Signal, 1)
signal.Notify(c, os.Interrupt)
go func() {
select {
case sig := <-c:
fmt.Printf("Got %s signal. Aborting...\n", sig)
task.Stop()
}
}()
task.Run()
}
完整版本代码：
package main
import (
"fmt"
"os"
"os/signal"
"time"
)
type Task struct {
closed chan struct{}
ticker *time.Ticker
}
func (t *Task) Run() {
for {
select {
case <-t.closed:
return
case <-t.ticker.C:
handle()
}
}
}
func (t *Task) Stop() {
close(t.closed)
}
func handle() {
for i := 0; i < 10; i++ {
fmt.Print("#")
time.Sleep(time.Millisecond * 200)
}
fmt.Println()
}
func main() {
task := &Task{
closed: make(chan struct{}),
ticker: time.NewTicker(time.Second * 2),
}
c := make(chan os.Signal, 1)
signal.Notify(c, os.Interrupt)
go func() {
select {
case sig := <-c:
fmt.Printf("Got %s signal. Aborting...\n", sig)
task.Stop()
}
}()
task.Run()
}
现在如果在 handle执行一半时中断应用，输出结果：
######Got interrupt signal. Aborting...
####
很好，尽管收到中断信号，当前任务仍正常打印完成。
4. 完善——等待协程结束
上面程序已经可以工作，但有个问题。task.Run()在主协程中，而处理中断信号工作是在另一个协程中。当捕获到中断信号后，调用task.Stop()后该协程生命周期结束。而此时主协程继续执行 select中的 Run方法，从 t.closed通道中接收值并返回。在这过程中程序不再接收任何中断信号。
那么不在主协程中执行 task.Run会怎样？
func main() {
// previous code...
go task.Run()
select {
case sig := <-c:
fmt.Printf("Got %s signal. Aborting...\n", sig)
task.Stop()
}
}
如果现在中断执行，当前运行的 handle将不会正常完成，因为捕获到终止信号后主协程结束，导致其他协程立刻终止了。因此需引入 sync.WaitGroup 解决该问题。首先再Task中加入同步等待组：
type Task struct {
closed chan struct{}
wg sync.WaitGroup
ticker *time.Ticker
}
我们让同步等待组区等待后台正在执行的任务完成——即 task.Run
func main() {
// previous code...
task.wg.Add(1)
go func() { defer task.wg.Done(); task.Run() }()
// other code...
}
最终我们需要实际等待task.Run 完成，在 Stop中加入：
func (t *Task) Stop() {
close(t.closed)
t.wg.Wait()
}
完整代码如下：
package main
import (
"fmt"
"os"
"os/signal"
"sync"
"time"
)
type Task struct {
closed chan struct{}
wg sync.WaitGroup
ticker *time.Ticker
}
func (t *Task) Run() {
for {
select {
case <-t.closed:
return
case <-t.ticker.C:
handle()
}
}
}
func (t *Task) Stop() {
close(t.closed)
t.wg.Wait()
}
func handle() {
for i := 0; i < 10; i++ {
fmt.Print("#")
time.Sleep(time.Millisecond * 200)
}
fmt.Println()
}
func main() {
task := &Task{
closed: make(chan struct{}),
ticker: time.NewTicker(time.Second * 2),
}
c := make(chan os.Signal, 1)
signal.Notify(c, os.Interrupt)
task.wg.Add(1)
go func() {
defer task.wg.Done()
task.Run()
}()
select {
case sig := <-c:
fmt.Printf("Got %s signal. Aborting...\n", sig)
task.Stop()
}
}
5. 总结
本文介绍了如何捕获 ctrl+c终止程序信号，优雅地结束程序。类似功能也可以通过 context实现，后续继续补充。