feat: handle multiple http servers with one graceful service

graceful/CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## To be released
 
+* feat: allow multiple http servers to be started with a single graceful service
+
 ## v1.1.3
 
 * build(deps): remove patch version number from go version 1.22.4 => 1.22

graceful/README.md

@@ -34,3 +34,17 @@ s := graceful.NewService(
 
 err := s.ListenAndServe(ctx, "tcp", ":9000", handler)
 ```
+
+### Configuration for multiple servers
+
+```
+s := graceful.NewService(
+	graceful.WithWaitDuration(30 * time.Second),
+	graceful.WithReloadWaitDuration(time.Hour),
+	graceful.WithPIDFile("/var/run/service.pid"),
+	graceful.WithNumServers(2),
+)
+
+err := s.ListenAndServe(ctx, "tcp", ":9000", handler)
+err := s.ListenAndServe(ctx, "tcp", ":9001", handler2)
+```

graceful/service.go

@@ -3,23 +3,22 @@ package graceful
 import (
 	"context"
 	"crypto/tls"
-	"fmt"
 	"net/http"
 	"os"
 	"sync"
 	"time"
 
-	"github.com/Scalingo/go-utils/errors/v2"
-
 	"github.com/cloudflare/tableflip"
 
+	"github.com/Scalingo/go-utils/errors/v2"
 	"github.com/Scalingo/go-utils/logger"
 )
 
 type Service struct {
-	httpServer *http.Server
-	graceful   *tableflip.Upgrader
-	wg         *sync.WaitGroup
+	httpServers []*http.Server
+	upg         *tableflip.Upgrader
+	mx          sync.Mutex
+	wg          *sync.WaitGroup
 	// waitDuration is the duration which is waited for all connections to stop
 	// in order to graceful shutdown the server. If some connections are still up
 	// after this timer they'll be cut aggressively.
@@ -28,6 +27,8 @@ type Service struct {
 	// connection to close when a graceful restart has been ordered. The new
 	// process is already working as expecting.
 	reloadWaitDuration time.Duration
+	// numServers is the number of servers to register before being ready
+	numServers int
 	// pidFile tracks the pid of the last child among the chain of graceful restart
 	// Required for daemon manager to track the service
 	pidFile string
@@ -37,9 +38,11 @@ type Option func(*Service)
 
 func NewService(opts ...Option) *Service {
 	s := &Service{
+		httpServers:        make([]*http.Server, 0),
 		wg:                 &sync.WaitGroup{},
 		waitDuration:       time.Minute,
 		reloadWaitDuration: 30 * time.Minute,
+		numServers:         1,
 	}
 	for _, opt := range opts {
 		opt(s)
@@ -65,6 +68,28 @@ func WithPIDFile(path string) Option {
 	})
 }
 
+func WithNumServers(n int) Option {
+	return Option(func(s *Service) {
+		s.numServers = n
+	})
+}
+
+func (s *Service) initTableflipUpgrader(ctx context.Context) error {
+	var err error
+	s.mx.Lock()
+	if s.upg == nil {
+		s.upg, err = tableflip.New(tableflip.Options{
+			UpgradeTimeout: s.reloadWaitDuration,
+			PIDFile:        s.pidFile,
+		})
+		if err != nil {
+			return errors.Wrap(ctx, err, "create tableflip upgrader")
+		}
+	}
+	s.mx.Unlock()
+	return nil
+}
+
 func (s *Service) ListenAndServeTLS(ctx context.Context, proto string, addr string, handler http.Handler, tlsConfig *tls.Config) error {
 	httpServer := &http.Server{
 		Addr:      addr,
@@ -83,32 +108,25 @@ func (s *Service) ListenAndServe(ctx context.Context, proto string, addr string,
 }
 
 func (s *Service) listenAndServe(ctx context.Context, _ string, addr string, server *http.Server) error {
+	err := s.initTableflipUpgrader(ctx)
+	if err != nil {
+		return errors.Wrap(ctx, err, "init tableflip upgrader")
+	}
+
 	log := logger.Get(ctx)
 
-	if s.pidFile != "" {
-		pid := os.Getpid()
-		err := os.WriteFile(s.pidFile, []byte(fmt.Sprintf("%d\n", pid)), 0600)
+	if len(s.httpServers) == 0 {
+		err := s.prepare(ctx)
 		if err != nil {
-			return errors.Wrap(ctx, err, "fail to write PID file")
+			// purposefully do not wrap error here, as it is wrapped in prepare
+			return err
 		}
 	}
 
-	// Use tableflip to handle graceful restart requests
-	upg, err := tableflip.New(tableflip.Options{
-		UpgradeTimeout: s.reloadWaitDuration,
-		PIDFile:        s.pidFile,
-	})
-	if err != nil {
-		return errors.Wrap(ctx, err, "creating tableflip upgrader")
-	}
-	s.graceful = upg
-	defer upg.Stop()
-
-	// setup the signal handling
-	go s.setupSignals(ctx)
+	s.httpServers = append(s.httpServers, server)
 
 	// Listen must be called before Ready
-	ln, err := upg.Listen("tcp", addr)
+	ln, err := s.upg.Listen("tcp", addr)
 	if err != nil {
 		return errors.Wrap(ctx, err, "upgrader listen")
 	}
@@ -117,36 +135,84 @@ func (s *Service) listenAndServe(ctx context.Context, _ string, addr string, ser
 		ln = tls.NewListener(ln, server.TLSConfig)
 	}
 
-	s.httpServer = server
-
 	go func() {
 		err := server.Serve(ln)
 		if !errors.Is(err, http.ErrServerClosed) {
-			log.WithError(err).Error("http server serve")
+			log.WithError(err).Error("Http server serve")
 		}
 	}()
 
-	log.Info("ready")
-	if err := upg.Ready(); err != nil {
+	if len(s.httpServers) == s.numServers {
+		err := s.finalize(ctx)
+		if err != nil {
+			// purposefully do not wrap error here, as it is wrapped in finalize
+			return err
+		}
+	}
+
+	return nil
+}
+
+// prepare is called before the first server is started.
+func (s *Service) prepare(ctx context.Context) error {
+	if s.pidFile != "" {
+		err := os.Remove(s.pidFile)
+		if err != nil && !os.IsNotExist(err) {
+			return errors.Wrap(ctx, err, "remove PID file")
+		}
+	}
+
+	// setup the signal handling
+	go s.setupSignals(ctx)
+
+	return nil
+}
+
+// finalize is called when all servers are started.
+func (s *Service) finalize(ctx context.Context) error {
+	log := logger.Get(ctx)
+
+	defer s.upg.Stop()
+
+	log.Info("Ready")
+	if err := s.upg.Ready(); err != nil {
 		return errors.Wrapf(ctx, err, "upgrader notify ready")
 	}
-	<-upg.Exit()
+
+	// Once the service has started, it will be blocked here until a signal is received.
+	<-s.upg.Exit()
+	log.Info("Upgrader finished")
 
 	// Normally the server should be always gracefully stopped and entering the
 	// above condition when server is closed If by any mean the serve stops
 	// without error, we're stopping the server ourselves here.  This code is a
 	// security to free resource but should be unreachable
 	ctx, cancel := context.WithTimeout(ctx, s.waitDuration)
 	defer cancel()
-	err = s.shutdown(ctx)
+	err := s.shutdown(ctx)
 	if err != nil {
 		return errors.Wrapf(ctx, err, "fail to shutdown service")
 	}
 
 	// Wait for connections to drain.
-	err = server.Shutdown(ctx)
-	if err != nil {
-		return errors.Wrap(ctx, err, "server shutdown")
+	errChan := make(chan error, len(s.httpServers))
+	for i, httpServer := range s.httpServers {
+		err = httpServer.Shutdown(ctx)
+		if err != nil {
+			errChan <- errors.Wrapf(ctx, err, "server shutdown %d", i)
+		}
+	}
+	close(errChan)
+	var shutdownErr error
+	for err := range errChan {
+		if shutdownErr == nil {
+			shutdownErr = err
+		} else {
+			shutdownErr = errors.Wrap(ctx, shutdownErr, err.Error())
+		}
+	}
+	if shutdownErr != nil {
+		return shutdownErr
 	}
 
 	return nil
@@ -175,19 +241,51 @@ func (s *Service) DecConnCount(ctx context.Context) {
 // developer.
 func (s *Service) shutdown(ctx context.Context) error {
 	log := logger.Get(ctx)
-	log.Info("shutting down http server")
-	err := s.httpServer.Shutdown(ctx)
-	if err != nil {
-		return errors.Wrapf(ctx, err, "fail to shutdown http server")
+
+	errChan := make(chan error, len(s.httpServers))
+	var wg sync.WaitGroup
+
+	for i, httpServer := range s.httpServers {
+		wg.Add(1)
+		go func(i int, httpServer *http.Server) {
+			defer wg.Done()
+			log := logger.Get(ctx)
+			if len(s.httpServers) > 1 {
+				log = log.WithField("index", i)
+			}
+			log.Info("Shutting down http server")
+			err := httpServer.Shutdown(ctx)
+			if err != nil {
+				log.WithError(err).Error("fail to shutdown http server")
+				errChan <- errors.Wrapf(ctx, err, "shutdown http server %d", i)
+			} else {
+				log.Info("Http server is stopped")
+			}
+		}(i, httpServer)
+	}
+
+	wg.Wait()
+	close(errChan)
+
+	var shutdownErr error
+	for err := range errChan {
+		if shutdownErr == nil {
+			shutdownErr = err
+		} else {
+			shutdownErr = errors.Wrap(ctx, shutdownErr, err.Error())
+		}
+	}
+
+	if shutdownErr != nil {
+		return shutdownErr
 	}
-	log.Info("http server is stopped")
 
-	log.Info("wait hijacked connections")
-	err = s.waitHijackedConnections(ctx)
+	log.Info("Wait hijacked connections")
+	err := s.waitHijackedConnections(ctx)
 	if err != nil {
 		return errors.Wrapf(ctx, err, "fail to wait hijacked connections")
 	}
-	log.Info("no more connection running")
+	log.Info("No more connection running")
 
 	return nil
 }

graceful/service_signals.go

@@ -21,13 +21,13 @@ func (s *Service) setupSignals(ctx context.Context) {
 		sig := <-ch
 		switch sig {
 		case syscall.SIGINT, syscall.SIGTERM:
-			s.graceful.Stop()
+			s.upg.Stop()
 			return
 		case syscall.SIGHUP:
-			log.Info("request graceful restart")
-			err := s.graceful.Upgrade()
+			log.Info("Request graceful restart")
+			err := s.upg.Upgrade()
 			if err != nil {
-				log.WithError(err).Error("fail to start new service")
+				log.WithError(err).Error("Fail to start new service")
 			}
 		}
 	}