防止粗鲁结束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实现，后续继续补充。

本文参考链接：https://blog.csdn.net/neweastsun/article/details/108244101