我們不可能永遠等待 s.wg.Wait() 阻塞，一般服務都會定義強制關閉的超時機制。

註：絕大多數以虛擬機器為基礎的雲服務（AWS EC2, GCP Cloud Computing Engine 等）或是容器管理平台（如 K8s pods lifecycle）都會提供一個短暫的時間窗口給服務做 Graceful Shutdown。

使用 Shutdown() 取代 Close()

將 s.(*simpleSrv).Close() 重命名為 s.(*simpleSrv).Shutdown(context.Context)，因為我們要用 TimeoutContext 來控制超時。

// file: ./internal/server/server.go
type Server interface {
	ListenAndServe() error
	RestoreFromDisk() error

-	Close() error
+	Shutdown(context.Context) error
}

-func (s *simpleSrv) Close() error {
+func (s *simpleSrv) Shutdown(ctx context.Context) error {
	s.inShutdown.Store(true)
	if s.listener == nil {
		return nil
	}
	// ...
}

加入 TimeoutContext

// file: ./cmd/olivine/app.go
select {
case <-ctx.Done():
	slog.Info("closing olivine server")

	shutdownCtx, cancel := context.WithTimeout(context.Background(), time.Second*5)
	defer cancel()

	if err := app.srv.Shutdown(shutdownCtx); err != nil {
		return err
	}

	return <-errch
case err := <-errch:
	return err
}

TimeoutContext 是一個好用的 context.Context 變體，它會在給定的 time.Duration 前阻塞 <-ctx.Done()。可以從 Go Playground 來嘗試。

在 server.(*simpleSrv).Shutdown() 中使用

// file: ./internal/server/server.go

func (s *simpleSrv) Shutdown(ctx context.Context) error {
	s.inShutdown.Store(true)
	if s.listener == nil {
		return nil
	}

	var errs []error
	if err := s.listener.Close(); err != nil {
		errs = append(errs, err)
	}
	if err := s.closeConns(); err != nil {
		errs = append(errs, err)
	}

	done := make(chan struct{})
	go func() {
		s.wg.Wait()
		close(done)
	}()

	select {
	case <-done:
		return errors.Join(errs...)
	case <-ctx.Done():
		errs = append(errs, ctx.Err())
		return errors.Join(errs...)
	}
}

我們需要同時等待 <-ctx.Done() 與 s.wg.Wait()，因此使用 select 關鍵字。另一方面，因為 s.wg.Wait() 會阻塞，因此我們要將其放在一個獨立的 goroutine 中：

done := make(chan struct{})
go func() { 
	s.wg.Wait()
	close(done)
}()

select {
case <-done:
	// WaitGroup is done
case <-ctx.Done():
	// Timeout exceed
}

這是一個常見的 goroutine 範式，當配合一些會阻塞的操作（尤其是 sync.(*WaitGroup).Wait()）時非常好用。