fix: graceful shutdown with asynchronous processes - fixes #1022

cmd/gotenberg.go

@@ -2,6 +2,7 @@ package gotenbergcmd
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"os"
 	"os/signal"
@@ -173,7 +174,9 @@ func Run() {
 				id := app.(gotenberg.Module).Descriptor().ID
 
 				err = app.Stop(gracefulShutdownCtx)
-				if err != nil {
+				if errors.Is(err, gotenberg.ErrCancelGracefulShutdownContext) {
+					cancel()
+				} else if err != nil {
 					return fmt.Errorf("stopping %s: %w", id, err)
 				}
 

pkg/gotenberg/shutdown.go

@@ -0,0 +1,8 @@
+package gotenberg
+
+import "errors"
+
+// ErrCancelGracefulShutdownContext tells that a module wants to abort a
+// graceful shutdown and stops Gotenberg right away as there are no more
+// ongoing processes.
+var ErrCancelGracefulShutdownContext = errors.New("cancel graceful shutdown's context")

pkg/modules/api/api.go

@@ -48,6 +48,7 @@ type Api struct {
 	externalMiddlewares []Middleware
 	healthChecks        []health.CheckerOption
 	readyFn             []func() error
+	asyncCounters       []AsynchronousCounter
 	fs                  *gotenberg.FileSystem
 	logger              *zap.Logger
 	srv                 *echo.Echo
@@ -166,6 +167,14 @@ type HealthChecker interface {
 	Ready() error
 }
 
+// AsynchronousCounter is a module interface that returns the number of active
+// asynchronous requests.
+//
+// See https://github.com/gotenberg/gotenberg/issues/1022.
+type AsynchronousCounter interface {
+	AsyncCount() int64
+}
+
 // Descriptor returns an [Api]'s module descriptor.
 func (a *Api) Descriptor() gotenberg.ModuleDescriptor {
 	return gotenberg.ModuleDescriptor{
@@ -307,6 +316,17 @@ func (a *Api) Provision(ctx *gotenberg.Context) error {
 		a.readyFn = append(a.readyFn, healthChecker.Ready)
 	}
 
+	// Get asynchronous counters.
+	mods, err = ctx.Modules(new(AsynchronousCounter))
+	if err != nil {
+		return fmt.Errorf("get asynchronous counters: %w", err)
+	}
+
+	a.asyncCounters = make([]AsynchronousCounter, len(mods))
+	for i, asyncCounter := range mods {
+		a.asyncCounters[i] = asyncCounter.(AsynchronousCounter)
+	}
+
 	// Logger.
 	loggerProvider, err := ctx.Module(new(gotenberg.LoggerProvider))
 	if err != nil {
@@ -597,7 +617,28 @@ func (a *Api) StartupMessage() string {
 
 // Stop stops the HTTP server.
 func (a *Api) Stop(ctx context.Context) error {
+	for {
+		count := int64(0)
+		for _, asyncCounter := range a.asyncCounters {
+			count += asyncCounter.AsyncCount()
+		}
+		select {
+		case <-ctx.Done():
 			return a.srv.Shutdown(ctx)
+		default:
+			a.logger.Debug(fmt.Sprintf("%d asynchronous requests", count))
+			if count > 0 {
+				time.Sleep(1 * time.Second)
+				continue
+			}
+			a.logger.Debug("no more asynchronous requests, continue with shutdown")
+			err := a.srv.Shutdown(ctx)
+			if err != nil {
+				return fmt.Errorf("shutdown: %w", err)
+			}
+			return gotenberg.ErrCancelGracefulShutdownContext
+		}
+	}
 }
 
 // Interface guards.

pkg/modules/webhook/middleware.go

@@ -173,10 +173,13 @@ func webhookMiddleware(w *Webhook) api.Middleware {
 						}
 					}
 
+					w.asyncCount.Add(1)
+
 					// As a webhook URL has been given, we handle the request in a
 					// goroutine and return immediately.
 					go func() {
 						defer cancel()
+						defer w.asyncCount.Add(-1)
 
 						// Call the next middleware in the chain.
 						err := next(c)

pkg/modules/webhook/webhook.go

@@ -1,6 +1,7 @@
 package webhook
 
 import (
+	"sync/atomic"
 	"time"
 
 	"github.com/dlclark/regexp2"
@@ -25,6 +26,7 @@ type Webhook struct {
 	retryMinWait   time.Duration
 	retryMaxWait   time.Duration
 	clientTimeout  time.Duration
+	asyncCount     atomic.Int64
 	disable        bool
 }
 
@@ -62,6 +64,7 @@ func (w *Webhook) Provision(ctx *gotenberg.Context) error {
 	w.retryMaxWait = flags.MustDuration("webhook-retry-max-wait")
 	w.clientTimeout = flags.MustDuration("webhook-client-timeout")
 	w.disable = flags.MustBool("webhook-disable")
+	w.asyncCount.Store(0)
 
 	return nil
 }
@@ -77,9 +80,15 @@ func (w *Webhook) Middlewares() ([]api.Middleware, error) {
 	}, nil
 }
 
+// AsyncCount returns the number of asynchronous requests.
+func (w *Webhook) AsyncCount() int64 {
+	return w.asyncCount.Load()
+}
+
 // Interface guards.
 var (
 	_ gotenberg.Module        = (*Webhook)(nil)
 	_ gotenberg.Provisioner   = (*Webhook)(nil)
 	_ api.MiddlewareProvider  = (*Webhook)(nil)
+	_ api.AsynchronousCounter = (*Webhook)(nil)
 )