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