fix: cancel error group properly

Author: jsiebens

internal/auth/server.go

@@ -1,6 +1,7 @@
 package auth
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"github.com/hashicorp/go-bexpr"
@@ -22,7 +23,7 @@ import (
 
 const authCachePrefix = "a_"
 
-func StartServer(config *config.Config) error {
+func StartServer(ctx context.Context, config *config.Config) error {
 	v, r := version.GetReleaseInfo()
 	logrus.Infof("Starting brink auth server. Version %s - %s", v, r)
 
@@ -44,7 +45,7 @@ func StartServer(config *config.Config) error {
 	}
 	authServer.RegisterRoutes(e, true)
 
-	return server.Start(config, e)
+	return server.Start(ctx, config, e)
 }
 
 func NewServer(config config.Auth, cache cache.Cache) (*Server, error) {

internal/cmd/server.go

@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"context"
 	"github.com/jsiebens/brink/internal/auth"
 	"github.com/jsiebens/brink/internal/config"
 	"github.com/jsiebens/brink/internal/proxy"
@@ -27,7 +28,7 @@ func serverProxyCommand() *coral.Command {
 	return createServerCommand("proxy", "Start a proxy server with a configuration file.", proxy.StartServer)
 }
 
-func createServerCommand(use, short string, start func(*config.Config) error) *coral.Command {
+func createServerCommand(use, short string, start func(context.Context, *config.Config) error) *coral.Command {
 	command := &coral.Command{
 		Use:          use,
 		Short:        short,
@@ -43,7 +44,7 @@ func createServerCommand(use, short string, start func(*config.Config) error) *c
 		if err != nil {
 			return err
 		}
-		return start(c)
+		return start(command.Context(), c)
 	}
 
 	return command

internal/proxy/server.go

@@ -1,6 +1,7 @@
 package proxy
 
 import (
+	"context"
 	"fmt"
 	"github.com/hashicorp/yamux"
 	"github.com/jsiebens/brink/internal/api"
@@ -25,7 +26,7 @@ import (
 const authCachePrefix = "pa_"
 const proxyCachePrefix = "pp_"
 
-func StartServer(config *config.Config) error {
+func StartServer(ctx context.Context, config *config.Config) error {
 	v, r := version.GetReleaseInfo()
 	logrus.Infof("Starting brink proxy server. Version %s - %s", v, r)
 
@@ -70,7 +71,7 @@ func StartServer(config *config.Config) error {
 	}
 	proxyServer.RegisterRoutes(e)
 
-	return server.Start(config, e)
+	return server.Start(ctx, config, e)
 }
 
 func NewServer(config config.Proxy, cache cache.Cache, registry auth.SessionRegistry) (*Server, error) {

internal/server/server.go

@@ -15,31 +15,26 @@ import (
 	"time"
 )
 
-func Start(config *config.Config, e *echo.Echo) error {
+func Start(ctx context.Context, config *config.Config, e *echo.Echo) error {
 	m := mon.Echo()
 
 	registerDefaultRoutes(e)
 	e.Use(mon.Middleware())
 
-	done := make(chan os.Signal, 1)
-	signal.Notify(done, os.Interrupt, syscall.SIGTERM)
+	return serve(contextWithSigterm(ctx), e, m, config)
+}
+
+func serve(ctx context.Context, e *echo.Echo, p *echo.Echo, config *config.Config) error {
+	g, gCtx := errgroup.WithContext(ctx)
 
 	go func() {
-		<-done
-		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		<-gCtx.Done()
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 		defer cancel()
-		_ = e.Shutdown(ctx)
-		_ = m.Shutdown(ctx)
+		_ = e.Shutdown(shutdownCtx)
+		_ = p.Shutdown(shutdownCtx)
 	}()
 
-	logrus.Infof("server listening on %s", config.ListenAddr)
-	logrus.Infof("metrics listening on %s", config.Metrics.ListenAddr)
-
-	return serve(e, m, config)
-}
-
-func serve(e *echo.Echo, p *echo.Echo, config *config.Config) error {
-	g := new(errgroup.Group)
 	g.Go(func() error {
 		if config.Tls.Disable {
 			if err := e.Start(config.ListenAddr); err != nil && !errors.Is(err, http.ErrServerClosed) {
@@ -52,12 +47,17 @@ func serve(e *echo.Echo, p *echo.Echo, config *config.Config) error {
 		}
 		return nil
 	})
+
 	g.Go(func() error {
 		if err := p.Start(config.Metrics.ListenAddr); err != nil && !errors.Is(err, http.ErrServerClosed) {
 			return err
 		}
 		return nil
 	})
+
+	logrus.Infof("server listening on %s", config.ListenAddr)
+	logrus.Infof("metrics listening on %s", config.Metrics.ListenAddr)
+
 	return g.Wait()
 }
 
@@ -66,3 +66,20 @@ func registerDefaultRoutes(e *echo.Echo) {
 		return c.Render(http.StatusOK, "index.html", nil)
 	})
 }
+
+func contextWithSigterm(ctx context.Context) context.Context {
+	ctxWithCancel, cancel := context.WithCancel(ctx)
+	go func() {
+		defer cancel()
+
+		signalCh := make(chan os.Signal, 1)
+		signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM)
+
+		select {
+		case <-signalCh:
+		case <-ctx.Done():
+		}
+	}()
+
+	return ctxWithCancel
+}