net.(Listener).Close() 之後做了什麼？

回顧目前的 Server：

// file: ./internal/server/server.go
conn, err := s.listener.Accept()
if err != nil {
	slog.Error("failed to accept connection:", slog.Any("error", err))
    return err
}

其中的 s.listener.Accept() 會阻塞並等待新的連線，並且直到調用 s.listener.Close() 之後會回傳 net.ErrClosed。

處理「正在關閉」狀態

我們可以在 server.(*simpleSrv) 中加入 inShutdown：

// file: ./internal/server/server.go
type simpleSrv struct {
	handler  Handler
	restorer Restorer

	listener   net.Listener
+	inShutdown atomic.Bool
}

func (s *simpleSrv) Close() error {
+	s.inShutdown.Store(true)
	if s.listener == nil {
		return nil
	}

	return s.listener.Close()
}

使用 atomic.Bool 而非 bool 是因為我們可以預期到這個資料將會被多個 goroutine 一起存取，為了避免 race condition（競態條件），使用 atomic.Bool 會是更好的選擇。

當服務進入「正在關閉」狀態時，服務應該阻止新的連線被建立：

// file: ./internal/server/server.go
conn, err := s.listener.Accept()
if err != nil {
	if s.inShutdown.Load() {
    	return ErrServerClosed
	}

	// ...
}

註：在絕大多數的服務中，可以直接判斷 errors.Is(err, net.ErrClosed)，inShutdown 是來自標準庫 http.(*Server) 的實作。Ref

未完成的請求

截至目前，因為服務已經進入「正在關閉」狀態而不會再接受新的請求（會直接回傳 server.ErrServerClosed），但是已連接的連線還有些尚未完成的請求。

我們需要等待這些連線直到它們關閉，可以使用 sync.WaitGroup 來追蹤這些仍然活躍的連線…

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

type simpleSrv struct {
	handler  Handler
	restorer Restorer

	listener   net.Listener
	inShutdown atomic.Bool

+	wg sync.WaitGroup
}

將原本 go s.serve(conn) 改為 s.wg.Go()func() { s.serve(conn) } 即可，並且在 s.Close() 中用 s.wg.Wait() 阻塞尚未完成的請求。

註：sync.(*WaitGroup).Go() 是一個 Go 1.25 才被加入的函式，如果使用比較舊版的 Go 可以在 s.serve(conn) 前加上 s.wg.Add(1) 與 defer s.wg.Done()。

關閉閒置的連線

對於大多數的 Redis 客戶端（包括官方的 redis-cli），它會維持一條 TCP 連線以支援多組 Redis 指令（官方稱為 Redis Pipelining）。我們在 s.Close() 中使用了 s.wg.Wait()，它會被這些閒置的連線永久阻塞。

因此我們需要一個方法來追蹤這些連線，並且將它們一一關閉。

// file: ./internal/server/server.go
type simpleSrv struct {
	handler  Handler
	restorer Restorer

	listener   net.Listener
	inShutdown atomic.Bool

	wg    sync.WaitGroup
+	mu    sync.Mutex
+	conns map[net.Conn]struct{}
}

註：再次提醒，內建的 map 並不是並行安全的，務必使用 sync.Mutex 保護它。

當連線建立之後，將它放進 conns 中，並且當處理完成之後將它從 conns 移除：

// file: ./internal/server/server.go
func (s *simpleSrv) ListenAndServe() (err error) {
	if s.listener, err = net.Listen("tcp", ":16879"); err != nil {
		return
	}

	for {
		conn, err := s.listener.Accept()

		// ...

+		s.mu.Lock()
+		s.conns[conn] = struct{}{}
+		s.mu.Unlock()
		s.wg.Go(func() { s.serve(conn) })
	}
}

func (s *simpleSrv) serve(conn net.Conn) {
	defer conn.Close()
+	defer func() {
+		s.mu.Lock()
+		defer s.mu.Unlock()
+	
+		delete(s.conns, conn)
+	}()

	rd := resp.NewReader(conn)

	for {
		ret, err := s.handler.ServeRESP(context.Background(), rd)
		// ...

建立閒置連線驅逐機制：

// file: ./internal/server/server.go
func (s *simpleSrv) Close() 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)
	}

	s.wg.Wait()

	return errors.Join(errs...)
}

func (s *simpleSrv) closeConns() error {
	s.mu.Lock()
	conns := maps.Clone(s.conns)
	s.mu.Unlock()

	var errs []error
	for conn := range conns {
		errs = append(errs, conn.Close())
	}

	return errors.Join(errs...)
}

註：在 closeConns() 中我先在 Mutex Lock 中進行 maps.Clone() 再用迴圈一個個關閉，這是為了避免在遍歷 s.conn 時需要反覆操作 Mutex Lock。

現在，我們可以在 SET 命令中加入 time.Sleep()，並且確認在按下 Ctrl+C 之後 AOF 是否仍然會被寫入來驗證 Graceful Shutdown 是否生效。