[Day 9, 10] Graceful Shutdown

Close() in Server

Add Close() error in Server{} interface:

type Server interface {
	ListenAndServe() error
	RestoreFromDisk() error

	Close() error
}

Handle signals

Use signal.NotifyContext() to handle shutdown signals:

func (app *App) Run() error {
	ctx, stop := signal.NotifyContext(context.TODO(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	slog.Info("restoring data from disk")
	if err := app.srv.RestoreFromDisk(); err != nil {
		return
	}

	go func() {
		slog.Info("starting olivine server")
		if err := app.srv.ListenAndServe(); err != nil {
			return
		}
	}()

	<-ctx.Done()

	slog.Info("closing olivine server")
	return app.srv.Close()
}

I moved app.srv.ListenAndServe() into a goroutine, because it includes a infinite loop to handle connections, it blocks main process and we cannot handle signal by <-ctx.Done() directly.

Error handling with goroutine

I ignored app.srv.ListenAndServe() error in goroutine, it’s not a good practice.

Use channel to handle errors in goroutine and select statement to work with <-ctx.Done()

func (app *App) Run() error {
	ctx, stop := signal.NotifyContext(context.TODO(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	slog.Info("restoring data from disk")
	if err := app.srv.RestoreFromDisk(); err != nil {
		return err
	}

	errch := make(chan error, 1)
	go func() {
		slog.Info("starting olivine server")
		errch <- app.srv.ListenAndServe()
	}()

	select {
	case <-ctx.Done():
		slog.Info("closing olivine server")
		if err := app.srv.Close(); err != nil {
			return err
		}

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

Note: Why don’t put app.srv.RestoreFromDisk() into goroutine?

If user send signal when storing from disk, it closes listener doesn’t exist. Process ListenAndServe after restored and block with infinite for loop.

Normal Flow: Start -> RestoreFromDisk -> Restored -> ListenAndServe -> Ctrl+C -> Close -> Exit Unexpected Flow: Start -> RestoreFromDisk -> Ctrl+C -> Close Unexists Listener -> Restored -> ListenAndServe -> Do not exit

What happen after listener Close()?

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

s.listener.Accept() is a blocking operation, it waits for new connections or error.

When s.listener.Close() is called in s.Close(), it returns net.ErrClosed, and returns to errch <- app.srv.ListenAndServe().

	go func() {
		slog.Info("starting olivine server")
		errch <- app.srv.ListenAndServe()
	}()

It’s an unexpected behavior: user send signals to close server, but shows error message failed to run app error="accept tcp [::]:16879: use of closed network connection".

	if err := app.Run(); err != nil {
		slog.Error("failed to run app", slog.Any("error", err))
		return
	}

Handle “shutdowning” status

Add inShutdown into simpleSrv{} as a flag to check if it is “shutdowning”.

// 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()
}

Check Accept() error

	for {
		conn, err := s.listener.Accept()
		if err != nil {
			if s.inShutdown.Load() {
				return ErrServerClosed
			}

			slog.Error("failed to accept connection:", slog.Any("error", err))
			return err
		}

		go s.serve(conn)
	}
}

When s.listener.Accept() returns err, check if it is “shutdowning” status, and return ErrServerClosed

	go func() {
		slog.Info("starting olivine server")
		if err := app.srv.ListenAndServe(); err != nil && !errors.Is(err, server.ErrServerClosed) {
			errch <- err
		}
		errch <- nil
	}()

Check if app.srv.ListenAndServe() returns server.ErrServerClosed, and send err into errch.

Note: it’s ok to check if net.ErrClosed in most of services, inShutdown is a workaround of http.Server in built-in library implementation. https://github.com/golang/go/blob/go1.26.4/src/net/http/server.go#L3434-L3438

Incomplete Connections

It doesn’t accept new connection after s.listener.Close() (returns net.ErrClosed immediately), but connected connections may not complete.

We need to wait active connections are closed, use sync.WaitGroup to track active connections.

type simpleSrv struct {
	handler  Handler
	restorer Restorer

	listener   net.Listener
	inShutdown atomic.Bool

+	wg sync.WaitGroup
}

Use s.wg.Go(func() { s.serve(conn) }) to handle connections by sync.WaitGroup. (Line 62)

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

	for {
		conn, err := s.listener.Accept()
		if err != nil {
			if s.inShutdown.Load() {
				return ErrServerClosed
			}

			slog.Error("failed to accept connection:", slog.Any("error", err))
			return err
		}

		s.wg.Go(func() { s.serve(conn) })
	}
}

And block it in Close(). (Line 96)

func (s *simpleSrv) Close() error {
	s.inShutdown.Store(true)
	if s.listener == nil {
		return nil
	}
	if err := s.listener.Close(); err != nil {
		return err
	}

	s.wg.Wait()

	return nil
}

Note: s.wg.Go() is added in Go 1.25, use s.wg.Add(1) before s.serve(conn) and defer s.wg.Done() in s.serve().

Close Idle Connections

Most of redis clients, including redis-cli, handle a TCP connection and send multiple commands. That’s why using infinite loop to handle s.handler.ServeRESP:

	for {
		ret, err := s.handler.ServeRESP(context.Background(), rd)
		if err != nil && errors.Is(err, ErrServer) {
			return
		}
		if ret == nil {
			ret = resp.NewNullBulkString()
		}

		if _, err := conn.Write(ret.Marshal()); err != nil {
			slog.Error("failed to write to conn:", slog.Any("error", err))
			return
		}
	}

With s.wg.Wait() in s.Close(), it could block forever if idle connections alive. We need to close active connections when s.Close().

Note: close an active connection make conn.Read() and conn.Write() returns net.ErrClosed immediately, but it doesn’t interrupt processing data (e.g. AOF writing).

Add conns map[net.Conn]struct{} and mu sync.Mutex in simpleSrv{}

type simpleSrv struct {
	handler  Handler
	restorer Restorer

	listener   net.Listener
	inShutdown atomic.Bool

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

Note: built-in map is not concurrency safe, it requires sync.Mutex to protect when reading and writing.

Tracking connections in conns when accepted:

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) })
	}
}

Delete connections from map when done:

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)
		// ...

Add closeConns() for simpleSrv{}

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...)
}

Note: To avoid mutex lock racing, clone the map in lock and close outside.

Finally, close connection between s.listener.Close() and s.wg.Wait().

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...)
}

Note: Is it safe to close connections after listener close too early?

There’s a little data racing problem if a connection accepted but haven’t put into conns map, the connection cannot be closed and block s.wg.Wait(), but we ignored at this moment.

Now, we can add a time.Sleep() into SET command, check AOF will be written even if in shutting down status.

Timeout

We can not wait sync.WaitGroup forever if it is blocked. Services defines a timeout to force terminate process even if it doesn’t shutdown gracefully.

Note: most of VM based cloud services (AWS EC2, GCP Cloud Computing, …) or container management platform (K8s pod lifecycle) provides a little gime for graceful shutdown.

Shutdown() instead of Close()

Rename Close() to Shutdown(context.Context), because we’ll use TimeoutContext to handle timeout.

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
	}
	// ...
}

Add Timeout Context

	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
	}

Timeout context is an useful context variant, it blocks <-ctx.Done() by given time.Duration.

Try in playground: https://go.dev/play/p/bnO8OIyXhS7

Handle Timeout in Shutdown()

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...)
	}
}

We need to wait <-ctx.Done() and s.wg.Wait() parallelly, use select statement, but the s.wg.Wait() blocks process, we need to put it into goroutine.

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

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