fix: graceful shutdown with asynchronous processes - fixes #1022

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 {
-	return a.srv.Shutdown(ctx)
+	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.