fix(api): wait for modules readiness before starting server (#752)

pkg/modules/api/api.go

@@ -17,6 +17,7 @@ import (
 	"go.uber.org/multierr"
 	"go.uber.org/zap"
 	"golang.org/x/net/http2"
+	"golang.org/x/sync/errgroup"
 
 	"github.com/gotenberg/gotenberg/v7/pkg/gotenberg"
 )
@@ -29,6 +30,7 @@ func init() {
 // middlewares or health checks.
 type Api struct {
 	port                      int
+	startTimeout              time.Duration
 	timeout                   time.Duration
 	rootPath                  string
 	traceHeader               string
@@ -37,6 +39,7 @@ type Api struct {
 	routes              []Route
 	externalMiddlewares []Middleware
 	healthChecks        []health.CheckerOption
+	readyFn             []func() error
 	fs                  *gotenberg.FileSystem
 	logger              *zap.Logger
 	srv                 *echo.Echo
@@ -145,6 +148,7 @@ type Middleware struct {
 // See https://github.com/alexliesenfeld/health for more details.
 type HealthChecker interface {
 	Checks() ([]health.CheckerOption, error)
+	Ready() error
 }
 
 // Descriptor returns an [Api]'s module descriptor.
@@ -155,6 +159,7 @@ func (a *Api) Descriptor() gotenberg.ModuleDescriptor {
 			fs := flag.NewFlagSet("api", flag.ExitOnError)
 			fs.Int("api-port", 3000, "Set the port on which the API should listen")
 			fs.String("api-port-from-env", "", "Set the environment variable with the port on which the API should listen - override the default port")
+			fs.Duration("api-start-timeout", time.Duration(30)*time.Second, "Set the time limit for the API to start")
 			fs.Duration("api-timeout", time.Duration(30)*time.Second, "Set the time limit for requests")
 			fs.String("api-root-path", "/", "Set the root path of the API - for service discovery via URL paths")
 			fs.String("api-trace-header", "Gotenberg-Trace", "Set the header name to use for identifying requests")
@@ -170,6 +175,7 @@ func (a *Api) Descriptor() gotenberg.ModuleDescriptor {
 func (a *Api) Provision(ctx *gotenberg.Context) error {
 	flags := ctx.ParsedFlags()
 	a.port = flags.MustInt("api-port")
+	a.startTimeout = flags.MustDuration("api-start-timeout")
 	a.timeout = flags.MustDuration("api-timeout")
 	a.rootPath = flags.MustString("api-root-path")
 	a.traceHeader = flags.MustString("api-trace-header")
@@ -259,6 +265,7 @@ func (a *Api) Provision(ctx *gotenberg.Context) error {
 		}
 
 		a.healthChecks = append(a.healthChecks, checks...)
+		a.readyFn = append(a.readyFn, healthChecker.Ready)
 	}
 
 	// Logger.
@@ -430,12 +437,25 @@ func (a *Api) Start() error {
 		func() echo.HandlerFunc {
 			checks := append(a.healthChecks, health.WithTimeout(a.timeout))
 			checker := health.NewChecker(checks...)
-
 			return echo.WrapHandler(health.NewHandler(checker))
 		}(),
 		hardTimeoutMiddleware(hardTimeout),
 	)
 
+	// Wait for all modules to be ready.
+	ctx, cancel := context.WithTimeout(context.Background(), a.startTimeout)
+	defer cancel()
+
+	eg, _ := errgroup.WithContext(ctx)
+	for _, f := range a.readyFn {
+		eg.Go(f)
+	}
+
+	err := eg.Wait()
+	if err != nil {
+		return fmt.Errorf("waiting for modules readiness: %w", err)
+	}
+
 	// As the following code is blocking, run it in a goroutine.
 	go func() {
 		server := &http2.Server{}

pkg/modules/api/api_test.go

@@ -10,6 +10,7 @@ import (
 	"os"
 	"reflect"
 	"testing"
+	"time"
 
 	"github.com/alexliesenfeld/health"
 	"github.com/labstack/echo/v4"
@@ -222,9 +223,6 @@ func TestApi_Provision(t *testing.T) {
 				mod.ValidateMock = func() error {
 					return errors.New("foo")
 				}
-				mod.ChecksMock = func() ([]health.CheckerOption, error) {
-					return nil, nil
-				}
 				return gotenberg.NewContext(
 					gotenberg.ParsedFlags{
 						FlagSet: new(Api).Descriptor().FlagSet,
@@ -347,6 +345,9 @@ func TestApi_Provision(t *testing.T) {
 				mod3.ChecksMock = func() ([]health.CheckerOption, error) {
 					return []health.CheckerOption{health.WithDisabledAutostart()}, nil
 				}
+				mod3.ReadyMock = func() error {
+					return nil
+				}
 
 				mod4 := &struct {
 					gotenberg.ModuleMock
@@ -643,141 +644,176 @@ func TestApi_Validate(t *testing.T) {
 }
 
 func TestApi_Start(t *testing.T) {
-	mod := new(Api)
-	mod.port = 3000
-	mod.rootPath = "/"
-	mod.disableHealthCheckLogging = true
-	mod.routes = []Route{
+	for _, tc := range []struct {
+		scenario    string
+		readyFn     []func() error
+		expectError bool
+	}{
 		{
-			Method:         http.MethodPost,
-			Path:           "/forms/foo",
-			IsMultipart:    true,
-			DisableLogging: true,
-			Handler: func(c echo.Context) error {
-				ctx := c.Get("context").(*Context)
-				ctx.outputPaths = []string{
-					"/tests/test/testdata/api/sample1.txt",
-				}
-
-				return nil
+			scenario: "at least one module not ready",
+			readyFn: []func() error{
+				func() error { return nil },
+				func() error { return errors.New("not ready") },
 			},
+			expectError: true,
 		},
 		{
-			Method:      http.MethodPost,
-			Path:        "/forms/bar",
-			IsMultipart: true,
-			Handler:     func(_ echo.Context) error { return errors.New("foo") },
+			scenario: "success",
+			readyFn: []func() error{
+				func() error { return nil },
+				func() error { return nil },
+			},
+			expectError: false,
 		},
-	}
-	mod.externalMiddlewares = []Middleware{
-		{
-			Stack: PreRouterStack,
-			Handler: func() echo.MiddlewareFunc {
-				return func(next echo.HandlerFunc) echo.HandlerFunc {
-					return func(c echo.Context) error {
-						return next(c)
-					}
-				}
-			}(),
-		},
-		{
-			Stack: MultipartStack,
-			Handler: func() echo.MiddlewareFunc {
-				return func(next echo.HandlerFunc) echo.HandlerFunc {
-					return func(c echo.Context) error {
-						return next(c)
-					}
-				}
-			}(),
-		},
-		{
-			Stack: DefaultStack,
-			Handler: func() echo.MiddlewareFunc {
-				return func(next echo.HandlerFunc) echo.HandlerFunc {
-					return func(c echo.Context) error {
-						return next(c)
-					}
-				}
-			}(),
-		},
-		{
-			Handler: func() echo.MiddlewareFunc {
-				return func(next echo.HandlerFunc) echo.HandlerFunc {
-					return func(c echo.Context) error {
-						return next(c)
-					}
-				}
-			}(),
-		},
-	}
-	mod.fs = gotenberg.NewFileSystem()
-	mod.logger = zap.NewNop()
+	} {
+		t.Run(tc.scenario, func(t *testing.T) {
+			mod := new(Api)
+			mod.port = 3000
+			mod.startTimeout = time.Duration(30) * time.Second
+			mod.rootPath = "/"
+			mod.disableHealthCheckLogging = true
+			mod.routes = []Route{
+				{
+					Method:         http.MethodPost,
+					Path:           "/forms/foo",
+					IsMultipart:    true,
+					DisableLogging: true,
+					Handler: func(c echo.Context) error {
+						ctx := c.Get("context").(*Context)
+						ctx.outputPaths = []string{
+							"/tests/test/testdata/api/sample1.txt",
+						}
 
-	err := mod.Start()
-	if err != nil {
-		t.Fatalf("expected no error but got: %v", err)
-	}
+						return nil
+					},
+				},
+				{
+					Method:      http.MethodPost,
+					Path:        "/forms/bar",
+					IsMultipart: true,
+					Handler:     func(_ echo.Context) error { return errors.New("foo") },
+				},
+			}
+			mod.externalMiddlewares = []Middleware{
+				{
+					Stack: PreRouterStack,
+					Handler: func() echo.MiddlewareFunc {
+						return func(next echo.HandlerFunc) echo.HandlerFunc {
+							return func(c echo.Context) error {
+								return next(c)
+							}
+						}
+					}(),
+				},
+				{
+					Stack: MultipartStack,
+					Handler: func() echo.MiddlewareFunc {
+						return func(next echo.HandlerFunc) echo.HandlerFunc {
+							return func(c echo.Context) error {
+								return next(c)
+							}
+						}
+					}(),
+				},
+				{
+					Stack: DefaultStack,
+					Handler: func() echo.MiddlewareFunc {
+						return func(next echo.HandlerFunc) echo.HandlerFunc {
+							return func(c echo.Context) error {
+								return next(c)
+							}
+						}
+					}(),
+				},
+				{
+					Handler: func() echo.MiddlewareFunc {
+						return func(next echo.HandlerFunc) echo.HandlerFunc {
+							return func(c echo.Context) error {
+								return next(c)
+							}
+						}
+					}(),
+				},
+			}
+			mod.readyFn = tc.readyFn
+			mod.fs = gotenberg.NewFileSystem()
+			mod.logger = zap.NewNop()
 
-	// health request.
-	recorder := httptest.NewRecorder()
-	healthRequest := httptest.NewRequest(http.MethodGet, "/health", nil)
-
-	mod.srv.ServeHTTP(recorder, healthRequest)
-	if recorder.Code != http.StatusOK {
-		t.Errorf("expected %d status code but got %d", http.StatusOK, recorder.Code)
-	}
-
-	// "multipart/form-data" request.
-	multipartRequest := func(url string) *http.Request {
-		body := &bytes.Buffer{}
-
-		writer := multipart.NewWriter(body)
-
-		defer func() {
-			err := writer.Close()
-			if err != nil {
+			err := mod.Start()
+			if !tc.expectError && err != nil {
 				t.Fatalf("expected no error but got: %v", err)
 			}
-		}()
 
-		err := writer.WriteField("foo", "foo")
-		if err != nil {
-			t.Fatalf("expected no error but got: %v", err)
-		}
+			if tc.expectError && err == nil {
+				t.Fatal("expected error but got none")
+			}
 
-		part, err := writer.CreateFormFile("foo.txt", "foo.txt")
-		if err != nil {
-			t.Fatalf("expected no error but got: %v", err)
-		}
+			if tc.expectError {
+				return
+			}
 
-		_, err = part.Write([]byte("foo"))
-		if err != nil {
-			t.Fatalf("expected no error but got: %v", err)
-		}
+			// health request.
+			recorder := httptest.NewRecorder()
+			healthRequest := httptest.NewRequest(http.MethodGet, "/health", nil)
 
-		req := httptest.NewRequest(http.MethodPost, url, body)
-		req.Header.Set(echo.HeaderContentType, writer.FormDataContentType())
+			mod.srv.ServeHTTP(recorder, healthRequest)
+			if recorder.Code != http.StatusOK {
+				t.Errorf("expected %d status code but got %d", http.StatusOK, recorder.Code)
+			}
 
-		return req
-	}
+			// "multipart/form-data" request.
+			multipartRequest := func(url string) *http.Request {
+				body := &bytes.Buffer{}
 
-	recorder = httptest.NewRecorder()
-	mod.srv.ServeHTTP(recorder, multipartRequest("/forms/foo"))
+				writer := multipart.NewWriter(body)
 
-	if recorder.Code != http.StatusOK {
-		t.Errorf("expected %d status code but got %d", http.StatusOK, recorder.Code)
-	}
+				defer func() {
+					err := writer.Close()
+					if err != nil {
+						t.Fatalf("expected no error but got: %v", err)
+					}
+				}()
 
-	recorder = httptest.NewRecorder()
-	mod.srv.ServeHTTP(recorder, multipartRequest("/forms/bar"))
+				err := writer.WriteField("foo", "foo")
+				if err != nil {
+					t.Fatalf("expected no error but got: %v", err)
+				}
 
-	if recorder.Code != http.StatusInternalServerError {
-		t.Errorf("expected %d status code but got %d", http.StatusInternalServerError, recorder.Code)
-	}
+				part, err := writer.CreateFormFile("foo.txt", "foo.txt")
+				if err != nil {
+					t.Fatalf("expected no error but got: %v", err)
+				}
 
-	err = mod.srv.Shutdown(context.TODO())
-	if err != nil {
-		t.Errorf("expected no error but got: %v", err)
+				_, err = part.Write([]byte("foo"))
+				if err != nil {
+					t.Fatalf("expected no error but got: %v", err)
+				}
+
+				req := httptest.NewRequest(http.MethodPost, url, body)
+				req.Header.Set(echo.HeaderContentType, writer.FormDataContentType())
+
+				return req
+			}
+
+			recorder = httptest.NewRecorder()
+			mod.srv.ServeHTTP(recorder, multipartRequest("/forms/foo"))
+
+			if recorder.Code != http.StatusOK {
+				t.Errorf("expected %d status code but got %d", http.StatusOK, recorder.Code)
+			}
+
+			recorder = httptest.NewRecorder()
+			mod.srv.ServeHTTP(recorder, multipartRequest("/forms/bar"))
+
+			if recorder.Code != http.StatusInternalServerError {
+				t.Errorf("expected %d status code but got %d", http.StatusInternalServerError, recorder.Code)
+			}
+
+			err = mod.srv.Shutdown(context.TODO())
+			if err != nil {
+				t.Errorf("expected no error but got: %v", err)
+			}
+		})
 	}
 }
 

pkg/modules/api/mocks.go

@@ -105,12 +105,17 @@ func (provider *MiddlewareProviderMock) Middlewares() ([]Middleware, error) {
 // HealthCheckerMock is mock for the [HealthChecker] interface.
 type HealthCheckerMock struct {
 	ChecksMock func() ([]health.CheckerOption, error)
+	ReadyMock  func() error
 }
 
 func (mod *HealthCheckerMock) Checks() ([]health.CheckerOption, error) {
 	return mod.ChecksMock()
 }
 
+func (mod *HealthCheckerMock) Ready() error {
+	return mod.ReadyMock()
+}
+
 // Interface guards.
 var (
 	_ Router             = (*RouterMock)(nil)

pkg/modules/api/mocks_test.go

@@ -148,10 +148,18 @@ func TestHealthCheckerMock(t *testing.T) {
 		ChecksMock: func() ([]health.CheckerOption, error) {
 			return nil, nil
 		},
+		ReadyMock: func() error {
+			return nil
+		},
 	}
 
 	_, err := mock.Checks()
 	if err != nil {
 		t.Errorf("expected no error from HealthCheckerMock.Checks, but got: %v", err)
 	}
+
+	err = mock.Ready()
+	if err != nil {
+		t.Errorf("expected no error from HealthCheckerMock.Ready, but got: %v", err)
+	}
 }