Golang使用channel實(shí)現(xiàn)一個優(yōu)雅退出功能

更新時間：2023年03月09日 10:56:07 作者：JustLorain

最近補(bǔ)?Golang?channel?方面八股的時候發(fā)現(xiàn)用?channel?實(shí)現(xiàn)一個優(yōu)雅退出功能好像不是很難，之前寫的?HTTP?框架剛好也不支持優(yōu)雅退出功能，于是就參考了?Hertz?優(yōu)雅退出方面的代碼，為我的?PIANO?補(bǔ)足了這個?feature

前言

最近補(bǔ) Golang channel 方面八股的時候發(fā)現(xiàn)用 channel 實(shí)現(xiàn)一個優(yōu)雅退出功能好像不是很難，之前寫的 HTTP 框架剛好也不支持優(yōu)雅退出功能，于是就參考了 Hertz 優(yōu)雅退出方面的代碼，為我的 PIANO 補(bǔ)足了這個 feature。

Hertz

字節(jié)跳動開源社區(qū) CloudWeGo 開源的一款高性能 HTTP 框架，具有高易用性、高性能、高擴(kuò)展性等特點(diǎn)。

PIANO

筆者自己實(shí)現(xiàn)的輕量級 HTTP 框架，具有中間件，三種不同的路由（靜態(tài)，通配，參數(shù)）方式，路由分組，優(yōu)雅退出等功能，迭代發(fā)展中。

實(shí)現(xiàn)思路

通過一個 os.Signal 類型的 chan 接收退出信號，收到信號后進(jìn)行對應(yīng)的退出收尾工作，利用 context.WithTimeout 或 time.After 等方式設(shè)置退出超時時間防止收尾等待時間過長。

讀源碼

由于 Hertz 的 Hook 功能中的 ShutdownHook 是 graceful shutdown 的一環(huán)，并且 Hook 功能的實(shí)現(xiàn)也不是很難所以這里就一起分析了，如果不想看直接跳到后面的章節(jié)即可 :)

Hook

Hook 函數(shù)是一個通用的概念，表示某事件觸發(fā)時所伴隨的操作，Hertz 提供了 StartHook 和 ShutdownHook 用于在服務(wù)觸發(fā)啟動后和退出前注入用戶自己的處理邏輯。

兩種 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
    
    ...
}

可以看到兩者都是函數(shù)數(shù)組的形式，并且是公開字段，所以可以直接 append，函數(shù)的簽名如下，OnShutdown 的函數(shù)不會返回 error 因?yàn)槎纪顺隽怂詻]法對錯誤進(jìn)行處理：

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

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

并且設(shè)置的 StartHook 會按照聲明順序依次調(diào)用，但是 ShutdownHook 會并發(fā)的進(jìn)行調(diào)用，這里的實(shí)現(xiàn)后面會講。

StartHook 的執(zhí)行時機(jī)

觸發(fā) Server 啟動后，框架會按函數(shù)聲明順序依次調(diào)用所有的 StartHook 函數(shù)，完成調(diào)用之后，才會正式開始端口監(jiān)聽，如果發(fā)生錯誤，則立刻終止服務(wù)。

上面是官方文檔中描述的 StartHook 的執(zhí)行時機(jī)，具體在源碼中就是下面的代碼：

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 的同學(xué)肯定知道 h.Spin() 方法調(diào)用后會正式啟動 Hertz 的 HTTP 服務(wù)，而上面的 engine.Run 方法則是被 h.Spin 異步調(diào)用的?？梢钥吹皆?engine.Run 方法里循環(huán)調(diào)用 engine.OnRun 數(shù)組中注冊的函數(shù)，最后執(zhí)行完成完成并且沒有 error 的情況下才會執(zhí)行 engine.listenAndServe() 正式開始端口監(jiān)聽，和官方文檔中說的一致，并且這里是通過 for 循環(huán)調(diào)用的所以也正如文檔所說框架會按函數(shù)聲明順序依次調(diào)用。

ShutdownHook 的執(zhí)行時機(jī)

Server 退出前，框架會并發(fā)地調(diào)用所有聲明的 ShutdownHook 函數(shù)，并且可以通過 server.WithExitWaitTime配置最大等待時長，默認(rèn)為5秒，如果超時，則立刻終止服務(wù)。

上面是官方文檔中描述的 ShutdownHook 的執(zhí)行時機(jī)，具體在源碼中就是下面的代碼：

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 函數(shù)都執(zhí)行完畢后給形參 ch 發(fā)送信號通知，注意這里每個 ShutdownHook 都起了一個協(xié)程，所以是并發(fā)執(zhí)行，這也是官方文檔所說的并發(fā)的進(jìn)行調(diào)用。

服務(wù)注冊與下線的執(zhí)行時機(jī)

服務(wù)注冊

Hertz 雖然是一個 HTTP 框架，但是 Hertz 的客戶端和服務(wù)端可以通過注冊中心進(jìn)行服務(wù)發(fā)現(xiàn)并進(jìn)行調(diào)用，并且 Hertz 也提供了大部分常用的注冊中心擴(kuò)展，在下面的 initOnRunHooks 方法中，通過注冊一個 StartHook 調(diào)用 Registry 接口的 Register 方法對服務(wù)進(jìn)行注冊。

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 方法中進(jìn)行調(diào)用 Deregister 取消注冊，可以看到剛剛提到的 executeOnShutdownHooks 的方法在開始異步執(zhí)行后就會進(jìn)行取消注冊操作。

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 字段以獲得更好的閱讀體驗(yàn)ww

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

如上所示，status 是一個 uint32 類型的內(nèi)部字段，用來表示 Hertz Engine 的狀態(tài)，具體具有四種狀態(tài)（Init 1, Running 2, Shutdown 3, Closed 4），由下面的常量定義。

const (
	_ uint32 = iota
	statusInitialized
	statusRunning
	statusShutdown
	statusClosed
)

下面列出了 Hertz Engine 狀態(tài)改變的時機(jī)：

函數(shù)	狀態(tài)改變前	狀態(tài)改變后
engine.Init	0	Init (1)
engine.Run	Init (1)	Running (2)
engine.Shutdown	Running (2)	Shutdown (3)
engine.Run defer	?	Closed (4)

對狀態(tài)的改變都是通過 atomic 包下的函數(shù)進(jìn)行更改的，保證了并發(fā)安全。

優(yōu)雅退出

Hertz Graceful Shutdown 功能的核心方法如下，signalToNotify 數(shù)組包含了所有會觸發(fā)退出的信號，觸發(fā)了的信號會傳向 signals 這個 channel，并且 Hertz 會根據(jù)收到信號類型決定進(jìn)行優(yōu)雅退出還是強(qiáng)制退出。

// 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 函數(shù)然后觸發(fā)立刻退出。前面也提到 h.Spin() 是以異步的方式調(diào)用 engine.Run，waitSignal 則由 h.Spin() 直接調(diào)用，所以運(yùn)行后 Hertz 會阻塞在 waitSignal 函數(shù)的 select 這里等待信號。

三個會觸發(fā) Shutdown 的信號區(qū)別如下：

syscall.SIGINT 表示中斷信號，通常由用戶在終端上按下 Ctrl+C 觸發(fā)，用于請求程序停止運(yùn)行；
syscall.SIGHUP 表示掛起信號，通常是由系統(tǒng)發(fā)送給進(jìn)程，用于通知進(jìn)程它的終端或控制臺已經(jīng)斷開連接或終止，進(jìn)程需要做一些清理工作；
syscall.SIGTERM 表示終止信號，通常也是由系統(tǒng)發(fā)送給進(jìn)程，用于請求進(jìn)程正常地終止運(yùn)行，進(jìn)程需要做一些清理工作；

如果 waitSignal 的返回值為 nil 則 h.Spin() 會進(jìn)行優(yōu)雅退出：

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 的方式設(shè)置了優(yōu)雅退出的超時時長，默認(rèn)為 5 秒，用戶可以通過 WithExitWaitTime 方法配置 server 的優(yōu)雅退出超時時長。將設(shè)置了超時時間的 ctx 傳入 Shutdown 方法，如果 ShutdownHook 先執(zhí)行完畢則 ch channel 收到信號后返回退出，否則 Context 超時收到信號強(qiáng)制返回退出。

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 優(yōu)雅退出部分的源碼分析，可以發(fā)現(xiàn) Hertz 多次利用了協(xié)程，通過 channel 傳遞信號進(jìn)行流程控制和信息傳遞，并通過 Context 的超時機(jī)制完成了整個優(yōu)雅退出流程。

自己實(shí)現(xiàn)

說是自己實(shí)現(xiàn)實(shí)際上也就是代碼搬運(yùn)工，把 Hertz 的 graceful shutdown 及其相關(guān)功能給 PIANO 進(jìn)行適配罷了ww

代碼實(shí)現(xiàn)都差不多，一些小細(xì)節(jié)根據(jù)我個人的習(xí)慣做了修改，完整修改參考這個 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)
	}
}