使用 channel 实现退出功能

前言

最近补 Golang channel 方面八股的时候发现用 channel 实现一个优雅退出功能好像不是很难，之前写的 HTTP 框架刚好也不支持优雅退出功能，于是就参考了 Hertz 优雅退出方面的代码，为我的 PIANO 补足了这个 feature。

• Hertz

  字节跳动开源社区 CloudWeGo 开源的一款高性能 HTTP 框架，具有高易用性、高性能、高扩展性等特点。

• PIANO

  笔者自己实现的轻量级 HTTP 框架，具有中间件，三种不同的路由（静态，通配，参数）方式，路由分组，优雅退出等功能，迭代发展中。

实现思路

通过一个 os.Signal 类型的 chan 接收退出信号，收到信号后进行对应的退出收尾工作，利用 context.WithTimeout 或 time.After 等方式设置退出超时时间防止收尾等待时间过长。

读源码

由于 Hertz 的 Hook 功能中的 ShutdownHook 是 graceful shutdown 的一环，并且 Hook 功能的实现也不是很难所以这里就一起分析了，如果不想看直接跳到后面的章节即可 :)

Hook

Hook 函数是一个通用的概念，表示某事件触发时所伴随的操作，Hertz 提供了 StartHook 和 ShutdownHook 用于在服务触发启动后和退出前注入用户自己的处理逻辑。

两种 Hook 具体是作为两种不同类型的 Hertz Engine 字段，用户可以直接以 append 的方式添加自己的 Hooks，下面是作为 Hertz Engine 字段的代码：

type Engine struct {
    ...
    
    // Hook functions get triggered sequentially when engine start
	OnRun []CtxErrCallback

	// Hook functions get triggered simultaneously when engine shutdown
	OnShutdown []CtxCallback
    
    ...
}

可以看到两者都是函数数组的形式，并且是公开字段，所以可以直接 append，函数的签名如下，OnShutdown 的函数不会返回 error 因为都退出了所以没法对错误进行处理：

// OnRun
type CtxCallback func(ctx context.Context)

// OnShutdown
type CtxErrCallback func(ctx context.Context) error

并且设置的 StartHook 会按照声明顺序依次调用，但是 ShutdownHook 会并发的进行调用，这里的实现后面会讲。

StartHook 的执行时机

触发 Server 启动后，框架会按函数声明顺序依次调用所有的 StartHook 函数，完成调用之后，才会正式开始端口监听，如果发生错误，则立刻终止服务。

上面是官方文档中描述的 StartHook 的执行时机，具体在源码中就是下面的代码：

func (engine *Engine) Run() (err error) {
	...

	// trigger hooks if any
	ctx := context.Background()
	for i := range engine.OnRun {
		if err = engine.OnRun[i](ctx); err != nil {
			return err
		}
	}

	return engine.listenAndServe()
}

熟悉或使用过 Hertz 的同学肯定知道 h.Spin() 方法调用后会正式启动 Hertz 的 HTTP 服务，而上面的 engine.Run 方法则是被 h.Spin 异步调用的。可以看到在 engine.Run 方法里循环调用 engine.OnRun 数组中注册的函数，最后执行完成完成并且没有 error 的情况下才会执行 engine.listenAndServe() 正式开始端口监听，和官方文档中说的一致，并且这里是通过 for 循环调用的所以也正如文档所说框架会按函数声明顺序依次调用。

ShutdownHook 的执行时机

Server 退出前，框架会并发地调用所有声明的 ShutdownHook 函数，并且可以通过 server.WithExitWaitTime配置最大等待时长，默认为5秒，如果超时，则立刻终止服务。

上面是官方文档中描述的 ShutdownHook 的执行时机，具体在源码中就是下面的代码：

func (engine *Engine) executeOnShutdownHooks(ctx context.Context, ch chan struct{}) {
	wg := sync.WaitGroup{}
	for i := range engine.OnShutdown {
		wg.Add(1)
		go func(index int) {
			defer wg.Done()
			engine.OnShutdown[index](ctx)
		}(i)
	}
	wg.Wait()
	ch <- struct{}{}
}

通过 sync.WaitGroup 保证每个 ShutdownHook 函数都执行完毕后给形参 ch 发送信号通知，注意这里每个 ShutdownHook 都起了一个协程，所以是并发执行，这也是官方文档所说的并发的进行调用。

服务注册与下线的执行时机

• 服务注册

  Hertz 虽然是一个 HTTP 框架，但是 Hertz 的客户端和服务端可以通过注册中心进行服务发现并进行调用，并且 Hertz 也提供了大部分常用的注册中心扩展，在下面的 initOnRunHooks 方法中，通过注册一个 StartHook 调用 Registry 接口的 Register 方法对服务进行注册。

  func (h *Hertz) initOnRunHooks(errChan chan error) {
  	// add register func to runHooks
  	opt := h.GetOptions()
  	h.OnRun = append(h.OnRun, func(ctx context.Context) error {
  		go func() {
  			// delay register 1s
  			time.Sleep(1 * time.Second)
  			if err := opt.Registry.Register(opt.RegistryInfo); err != nil {
  				hlog.SystemLogger().Errorf("Register error=%v", err)
  				// pass err to errChan
  				errChan <- err
  			}
  		}()
  		return nil
  	})
  }
  

• 取消注册

  在 Shutdown 方法中进行调用 Deregister 取消注册，可以看到刚刚提到的 executeOnShutdownHooks 的方法在开始异步执行后就会进行取消注册操作。

  func (engine *Engine) Shutdown(ctx context.Context) (err error) {
  	...
  
  	ch := make(chan struct{})
  	// trigger hooks if any
  	go engine.executeOnShutdownHooks(ctx, ch)
  
  	defer func() {
  		// ensure that the hook is executed until wait timeout or finish
  		select {
  		case <-ctx.Done():
  			hlog.SystemLogger().Infof("Execute OnShutdownHooks timeout: error=%v", ctx.Err())
  			return
  		case <-ch:
  			hlog.SystemLogger().Info("Execute OnShutdownHooks finish")
  			return
  		}
  	}()
  
  	if opt := engine.options; opt != nil && opt.Registry != nil {
  		if err = opt.Registry.Deregister(opt.RegistryInfo); err != nil {
  			hlog.SystemLogger().Errorf("Deregister error=%v", err)
  			return err
  		}
  	}
  
  	...
  }
  

Engine Status

讲 graceful shutdown 之前最好了解一下 Hertz Engine 的 status 字段以获得更好的阅读体验ww

type Engine struct {
    ...
    
    // Indicates the engine status (Init/Running/Shutdown/Closed).
    status uint32
    
    ...
}

如上所示，status 是一个 uint32 类型的内部字段，用来表示 Hertz Engine 的状态，具体具有四种状态（Init 1, Running 2, Shutdown 3, Closed 4），由下面的常量定义。

const (
	_ uint32 = iota
	statusInitialized
	statusRunning
	statusShutdown
	statusClosed
)

下面列出了 Hertz Engine 状态改变的时机：

函数	状态改变前	状态改变后
engine.Init	0	Init (1)
engine.Run	Init (1)	Running (2)
engine.Shutdown	Running (2)	Shutdown (3)
engine.Run defer	?	Closed (4)

对状态的改变都是通过 atomic 包下的函数进行更改的，保证了并发安全。

优雅退出

Hertz Graceful Shutdown 功能的核心方法如下，signalToNotify 数组包含了所有会触发退出的信号，触发了的信号会传向 signals 这个 channel，并且 Hertz 会根据收到信号类型决定进行优雅退出还是强制退出。

// Default implementation for signal waiter.
// SIGTERM triggers immediately close.
// SIGHUP|SIGINT triggers graceful shutdown.
func waitSignal(errCh chan error) error {
	signalToNotify := []os.Signal{syscall.SIGINT, syscall.SIGHUP, syscall.SIGTERM}
	if signal.Ignored(syscall.SIGHUP) {
		signalToNotify = []os.Signal{syscall.SIGINT, syscall.SIGTERM}
	}

	signals := make(chan os.Signal, 1)
	signal.Notify(signals, signalToNotify...)

	select {
	case sig := <-signals:
		switch sig {
		case syscall.SIGTERM:
			// force exit
			return errors.New(sig.String()) // nolint
		case syscall.SIGHUP, syscall.SIGINT:
			hlog.SystemLogger().Infof("Received signal: %s\n", sig)
			// graceful shutdown
			return nil
		}
	case err := <-errCh:
		// error occurs, exit immediately
		return err
	}

	return nil
}

如果 engine.Run 方法返回了一个错误则会通过 errCh 传入 waitSignal 函数然后触发立刻退出。前面也提到 h.Spin() 是以异步的方式调用 engine.Run，waitSignal 则由 h.Spin() 直接调用，所以运行后 Hertz 会阻塞在 waitSignal 函数的 select 这里等待信号。

三个会触发 Shutdown 的信号区别如下：

• syscall.SIGINT 表示中断信号，通常由用户在终端上按下 Ctrl C 触发，用于请求程序停止运行；
• syscall.SIGHUP 表示挂起信号，通常是由系统发送给进程，用于通知进程它的终端或控制台已经断开连接或终止，进程需要做一些清理工作；
• syscall.SIGTERM 表示终止信号，通常也是由系统发送给进程，用于请求进程正常地终止运行，进程需要做一些清理工作；

如果 waitSignal 的返回值为 nil 则 h.Spin() 会进行优雅退出：

func (h *Hertz) Spin() {
	errCh := make(chan error)
	h.initOnRunHooks(errCh)
	go func() {
		errCh <- h.Run()
	}()

	signalWaiter := waitSignal
	if h.signalWaiter != nil {
		signalWaiter = h.signalWaiter
	}

	if err := signalWaiter(errCh); err != nil {
		hlog.SystemLogger().Errorf("Receive close signal: error=%v", err)
		if err := h.Engine.Close(); err != nil {
			hlog.SystemLogger().Errorf("Close error=%v", err)
		}
		return
	}

	hlog.SystemLogger().Infof("Begin graceful shutdown, wait at most num=%d seconds...", h.GetOptions().ExitWaitTimeout/time.Second)

	ctx, cancel := context.WithTimeout(context.Background(), h.GetOptions().ExitWaitTimeout)
	defer cancel()

	if err := h.Shutdown(ctx); err != nil {
		hlog.SystemLogger().Errorf("Shutdown error=%v", err)
	}
}

并且 Hertz 通过 context.WithTimeout 的方式设置了优雅退出的超时时长，默认为 5 秒，用户可以通过 WithExitWaitTime 方法配置 server 的优雅退出超时时长。将设置了超时时间的 ctx 传入 Shutdown 方法，如果 ShutdownHook 先执行完毕则 ch channel 收到信号后返回退出，否则 Context 超时收到信号强制返回退出。

func (engine *Engine) Shutdown(ctx context.Context) (err error) {
	...

	ch := make(chan struct{})
	// trigger hooks if any
	go engine.executeOnShutdownHooks(ctx, ch)

	defer func() {
		// ensure that the hook is executed until wait timeout or finish
		select {
		case <-ctx.Done():
			hlog.SystemLogger().Infof("Execute OnShutdownHooks timeout: error=%v", ctx.Err())
			return
		case <-ch:
			hlog.SystemLogger().Info("Execute OnShutdownHooks finish")
			return
		}
	}()

	...
	return
}

自己实现

说是自己实现实际上也就是代码搬运工，把 Hertz 的 graceful shutdown 及其相关功能给 PIANO 进行适配罢了ww

代码实现都差不多，一些小细节根据我个人的习惯做了修改，完整修改参考这个 commit，对 PIANO 感兴趣的话欢迎 Star !

适配 Hook

type Engine struct {
    ...

	// hook
	OnRun      []HookFuncWithErr
	OnShutdown []HookFunc

	...
}

type (
	HookFunc        func(ctx context.Context)
	HookFuncWithErr func(ctx context.Context) error
)

func (e *Engine) executeOnRunHooks(ctx context.Context) error {
	for _, h := range e.OnRun {
		if err := h(ctx); err != nil {
			return err
		}
	}
	return nil
}

func (e *Engine) executeOnShutdownHooks(ctx context.Context, ch chan struct{}) {
	wg := sync.WaitGroup{}
	for _, h := range e.OnShutdown {
		wg.Add(1)
		go func(hook HookFunc) {
			defer wg.Done()
			hook(ctx)
		}(h)
	}
	wg.Wait()
	ch <- struct{}{}
}

适配 Engine Status

type Engine struct {
	...
    
	// initialized | running | shutdown | closed
	status uint32

    ...
}

const (
	_ uint32 = iota
	statusInitialized
	statusRunning
	statusShutdown
	statusClosed
)

适配 Graceful Shutdown

// Play the PIANO now
func (p *Piano) Play() {
	errCh := make(chan error)
	go func() {
		errCh <- p.Run()
	}()
	waitSignal := func(errCh chan error) error {
		signalToNotify := []os.Signal{syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM}
		if signal.Ignored(syscall.SIGHUP) {
			signalToNotify = signalToNotify[1:]
		}
		signalCh := make(chan os.Signal, 1)
		signal.Notify(signalCh, signalToNotify...)
		select {
		case sig := <-signalCh:
			switch sig {
			case syscall.SIGTERM:
				// force exit
				return errors.New(sig.String())
			case syscall.SIGHUP, syscall.SIGINT:
				// graceful shutdown
				log.Infof("---PIANO--- Receive signal: %v", sig)
				return nil
			}
		case err := <-errCh:
			return err
		}
		return nil
	}
	if err := waitSignal(errCh); err != nil {
		log.Errorf("---PIANO--- Receive close signal error: %v", err)
		return
	}
	log.Infof("---PIANO--- Begin graceful shutdown, wait up to %d seconds", p.Options().ShutdownTimeout/time.Second)
	ctx, cancel := context.WithTimeout(context.Background(), p.Options().ShutdownTimeout)
	defer cancel()
	if err := p.Shutdown(ctx); err != nil {
		log.Errorf("---PIANO--- Shutdown err: %v", err)
	}
}

总结

本文通过对 Hertz 优雅退出功能的实现做了源码分析并对自己的 HTTP 框架进行了适配，希望可以帮助读者利用 channel 实现一个优雅退出功能提供参考和思路，如果哪里有问题或者错误欢迎评论或者私信，以上。

参考列表

• github.com/B1NARY-GR0U…

• www.cloudwego.io/zh/docs/her…

• github.com/cloudwego/h…