fix: use errors instead of fatal for serve cmd

Author: aeneasr

cmd/daemon/serve.go

@@ -2,10 +2,9 @@ package daemon
 
 import (
 	"crypto/tls"
-	"net/http"
-	"sync"
-
 	"github.com/ory/kratos/schema"
+	"golang.org/x/sync/errgroup"
+	"net/http"
 
 	"github.com/ory/kratos/selfservice/flow/recovery"
 
@@ -50,7 +49,7 @@ type options struct {
 	ctx stdctx.Context
 }
 
-func newOptions(ctx stdctx.Context, opts []Option) *options {
+func NewOptions(ctx stdctx.Context, opts []Option) *options {
 	o := new(options)
 	o.ctx = ctx
 	for _, f := range opts {
@@ -73,9 +72,8 @@ func WithContext(ctx stdctx.Context) Option {
 	}
 }
 
-func ServePublic(r driver.Registry, wg *sync.WaitGroup, cmd *cobra.Command, args []string, opts ...Option) {
-	defer wg.Done()
-	modifiers := newOptions(cmd.Context(), opts)
+func ServePublic(r driver.Registry, cmd *cobra.Command, args []string, opts ...Option) error {
+	modifiers := NewOptions(cmd.Context(), opts)
 	ctx := modifiers.ctx
 
 	c := r.Config(cmd.Context())
@@ -144,14 +142,15 @@ func ServePublic(r driver.Registry, wg *sync.WaitGroup, cmd *cobra.Command, args
 		}
 		return server.ServeTLS(listener, "", "")
 	}, server.Shutdown); err != nil {
-		l.Fatalf("Failed to gracefully shutdown public httpd: %s", err)
+		l.Errorf("Failed to gracefully shutdown public httpd: %s", err)
+		return err
 	}
 	l.Println("Public httpd was shutdown gracefully")
+	return nil
 }
 
-func ServeAdmin(r driver.Registry, wg *sync.WaitGroup, cmd *cobra.Command, args []string, opts ...Option) {
-	defer wg.Done()
-	modifiers := newOptions(cmd.Context(), opts)
+func ServeAdmin(r driver.Registry, cmd *cobra.Command, args []string, opts ...Option) error {
+	modifiers := NewOptions(cmd.Context(), opts)
 	ctx := modifiers.ctx
 
 	c := r.Config(ctx)
@@ -206,9 +205,11 @@ func ServeAdmin(r driver.Registry, wg *sync.WaitGroup, cmd *cobra.Command, args
 		}
 		return server.ServeTLS(listener, "", "")
 	}, server.Shutdown); err != nil {
-		l.Fatalf("Failed to gracefully shutdown admin httpd: %s", err)
+		l.Errorf("Failed to gracefully shutdown admin httpd: %s", err)
+		return err
 	}
 	l.Println("Admin httpd was shutdown gracefully")
+	return nil
 }
 
 func sqa(ctx stdctx.Context, cmd *cobra.Command, d driver.Registry) *metricsx.Service {
@@ -280,23 +281,34 @@ func sqa(ctx stdctx.Context, cmd *cobra.Command, d driver.Registry) *metricsx.Se
 	)
 }
 
-func bgTasks(d driver.Registry, wg *sync.WaitGroup, cmd *cobra.Command, args []string, opts ...Option) {
-	defer wg.Done()
-	modifiers := newOptions(cmd.Context(), opts)
+func bgTasks(d driver.Registry, cmd *cobra.Command, args []string, opts ...Option) error {
+	modifiers := NewOptions(cmd.Context(), opts)
 	ctx := modifiers.ctx
 
 	if d.Config(ctx).IsBackgroundCourierEnabled() {
 		go courier.Watch(ctx, d)
 	}
+	return nil
 }
 
-func ServeAll(d driver.Registry, opts ...Option) func(cmd *cobra.Command, args []string) {
-	return func(cmd *cobra.Command, args []string) {
-		var wg sync.WaitGroup
-		wg.Add(3)
-		go ServePublic(d, &wg, cmd, args, opts...)
-		go ServeAdmin(d, &wg, cmd, args, opts...)
-		go bgTasks(d, &wg, cmd, args, opts...)
-		wg.Wait()
+func ServeAll(d driver.Registry, opts ...Option) func(cmd *cobra.Command, args []string) error {
+	return func(cmd *cobra.Command, args []string) error {
+		mods := NewOptions(cmd.Context(), opts)
+		ctx := mods.ctx
+
+		g, ctx := errgroup.WithContext(ctx)
+		cmd.SetContext(ctx)
+		opts = append(opts, WithContext(ctx))
+
+		g.Go(func() error {
+			return ServePublic(d, cmd, args, opts...)
+		})
+		g.Go(func() error {
+			return ServeAdmin(d, cmd, args, opts...)
+		})
+		g.Go(func() error {
+			return bgTasks(d, cmd, args, opts...)
+		})
+		return g.Wait()
 	}
 }

cmd/serve/root.go

@@ -29,7 +29,7 @@ func NewServeCmd() (serveCmd *cobra.Command) {
 	serveCmd = &cobra.Command{
 		Use:   "serve",
 		Short: "Run the Ory Kratos server",
-		Run: func(cmd *cobra.Command, args []string) {
+		RunE: func(cmd *cobra.Command, args []string) error {
 			d := driver.New(cmd.Context(), cmd.ErrOrStderr(), configx.WithFlags(cmd.Flags()))
 
 			if d.Config(cmd.Context()).IsInsecureDevMode() {
@@ -50,7 +50,7 @@ DON'T DO THIS IN PRODUCTION!
 				d.Logger().Warnf("Config version is '%s' but kratos runs on version '%s'", configVersion, config.Version)
 			}
 
-			daemon.ServeAll(d)(cmd, args)
+			return daemon.ServeAll(d)(cmd, args)
 		},
 	}
 	configx.RegisterFlags(serveCmd.PersistentFlags())