diff --git a/cmd/gotenberg.go b/cmd/gotenberg.go index 5e456a5f..8e558064 100644 --- a/cmd/gotenberg.go +++ b/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) } diff --git a/pkg/gotenberg/shutdown.go b/pkg/gotenberg/shutdown.go new file mode 100644 index 00000000..053ed5d1 --- /dev/null +++ b/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") diff --git a/pkg/modules/api/api.go b/pkg/modules/api/api.go index 7662fe38..72a831ce 100644 --- a/pkg/modules/api/api.go +++ b/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. diff --git a/pkg/modules/webhook/middleware.go b/pkg/modules/webhook/middleware.go index ab8363d7..ec2904f2 100644 --- a/pkg/modules/webhook/middleware.go +++ b/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) diff --git a/pkg/modules/webhook/webhook.go b/pkg/modules/webhook/webhook.go index 54661a20..0ff6f3f7 100644 --- a/pkg/modules/webhook/webhook.go +++ b/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) + _ gotenberg.Module = (*Webhook)(nil) + _ gotenberg.Provisioner = (*Webhook)(nil) + _ api.MiddlewareProvider = (*Webhook)(nil) + _ api.AsynchronousCounter = (*Webhook)(nil) )