diff --git a/pkg/gotenberg/supervisor.go b/pkg/gotenberg/supervisor.go index b55f47b4..1df1e554 100644 --- a/pkg/gotenberg/supervisor.go +++ b/pkg/gotenberg/supervisor.go @@ -13,6 +13,10 @@ import ( // to restart an already restarting [Process]. var ErrProcessAlreadyRestarting = errors.New("process already restarting") +// ErrMaximumQueueSizeExceeded happens if Run() is called but the maximum queue +// size is already used. +var ErrMaximumQueueSizeExceeded = errors.New("maximum queue size exceeded") + // Process is an interface that represents an abstract process // and provides methods for starting, stopping, and checking the health of the // process. @@ -74,6 +78,7 @@ type processSupervisor struct { logger *zap.Logger process Process maxReqLimit int64 + maxQueueSize int64 mutexChan chan struct{} firstStart atomic.Bool reqCounter atomic.Int64 @@ -83,12 +88,13 @@ type processSupervisor struct { } // NewProcessSupervisor initializes a new [ProcessSupervisor]. -func NewProcessSupervisor(logger *zap.Logger, process Process, maxReqLimit int64) ProcessSupervisor { +func NewProcessSupervisor(logger *zap.Logger, process Process, maxReqLimit int64, maxQueueSize int64) ProcessSupervisor { b := &processSupervisor{ - logger: logger, - process: process, - mutexChan: make(chan struct{}, 1), - maxReqLimit: maxReqLimit, + logger: logger, + process: process, + mutexChan: make(chan struct{}, 1), + maxReqLimit: maxReqLimit, + maxQueueSize: maxQueueSize, } b.reqCounter.Store(0) b.reqQueueSize.Store(0) @@ -167,6 +173,11 @@ func (s *processSupervisor) Healthy() bool { } func (s *processSupervisor) Run(ctx context.Context, logger *zap.Logger, task func() error) error { + currentQueueSize := s.reqQueueSize.Load() + if s.maxQueueSize > 0 && currentQueueSize >= s.maxQueueSize { + return ErrMaximumQueueSizeExceeded + } + s.reqQueueSize.Add(1) for { diff --git a/pkg/gotenberg/supervisor_test.go b/pkg/gotenberg/supervisor_test.go index aa046d8c..ef7eb4b6 100644 --- a/pkg/gotenberg/supervisor_test.go +++ b/pkg/gotenberg/supervisor_test.go @@ -46,7 +46,7 @@ func TestProcessSupervisor_Launch(t *testing.T) { }, } - ps := NewProcessSupervisor(logger, process, 5).(*processSupervisor) + ps := NewProcessSupervisor(logger, process, 5, 0).(*processSupervisor) if tc.firstStartSet { ps.firstStart.Store(true) } @@ -94,7 +94,7 @@ func TestProcessSupervisor_Shutdown(t *testing.T) { }, } - ps := NewProcessSupervisor(logger, process, 5) + ps := NewProcessSupervisor(logger, process, 5, 0) err := ps.Shutdown() if !tc.expectError && err != nil { @@ -154,7 +154,7 @@ func TestProcessSupervisor_restart(t *testing.T) { }, } - ps := NewProcessSupervisor(logger, process, 5).(*processSupervisor) + ps := NewProcessSupervisor(logger, process, 5, 0).(*processSupervisor) if tc.initiallyRestarting { ps.isRestarting.Store(true) } @@ -217,7 +217,7 @@ func TestProcessSupervisor_Healthy(t *testing.T) { }, } - ps := NewProcessSupervisor(logger, process, 5).(*processSupervisor) + ps := NewProcessSupervisor(logger, process, 5, 0).(*processSupervisor) if tc.initiallyStarted { ps.firstStart.Store(true) } @@ -249,6 +249,8 @@ func TestProcessSupervisor_Run(t *testing.T) { expectedStartCalls int64 expectedHealthyCalls int64 expectedStopCalls int64 + currentQueueSize int64 + maxQueueSize int64 }{ { scenario: "successfully run task on non-started process", @@ -348,6 +350,34 @@ func TestProcessSupervisor_Run(t *testing.T) { expectedHealthyCalls: 1, expectedStopCalls: 0, }, + { + scenario: "queue size exceeded", + initiallyStarted: false, + isRestarting: false, + processHealthy: true, + maxReqLimit: 2, + tasksToRun: 1, + expectError: true, + expectedStartCalls: 0, + expectedHealthyCalls: 0, + expectedStopCalls: 0, + currentQueueSize: 1, + maxQueueSize: 1, + }, + { + scenario: "queue size not exceeded", + initiallyStarted: false, + isRestarting: false, + processHealthy: true, + maxReqLimit: 2, + tasksToRun: 1, + expectError: true, + expectedStartCalls: 1, + expectedHealthyCalls: 1, + expectedStopCalls: 0, + currentQueueSize: 1, + maxQueueSize: 2, + }, } { t.Run(tc.scenario, func(t *testing.T) { logger := zap.NewNop() @@ -372,13 +402,16 @@ func TestProcessSupervisor_Run(t *testing.T) { }, } - ps := NewProcessSupervisor(logger, process, tc.maxReqLimit).(*processSupervisor) + ps := NewProcessSupervisor(logger, process, tc.maxReqLimit, tc.maxQueueSize).(*processSupervisor) if tc.initiallyStarted { ps.firstStart.Store(true) } if tc.isRestarting { ps.isRestarting.Store(true) } + if tc.currentQueueSize > 0 { + ps.reqQueueSize.Store(tc.currentQueueSize) + } task := func() error { return tc.taskError @@ -451,7 +484,7 @@ func TestProcessSupervisor_runWithDeadline(t *testing.T) { }, } { t.Run(tc.scenario, func(t *testing.T) { - ps := NewProcessSupervisor(zap.NewNop(), new(ProcessMock), 0).(*processSupervisor) + ps := NewProcessSupervisor(zap.NewNop(), new(ProcessMock), 0, 0).(*processSupervisor) ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) defer cancel() @@ -485,7 +518,7 @@ func TestProcessSupervisor_ReqQueueSize(t *testing.T) { return true }, } - ps := NewProcessSupervisor(logger, process, 0).(*processSupervisor) + ps := NewProcessSupervisor(logger, process, 0, 0).(*processSupervisor) // Simulating a lock. ps.mutexChan <- struct{}{} @@ -586,7 +619,7 @@ func TestProcessSupervisor_RestartsCount(t *testing.T) { }, } - ps := NewProcessSupervisor(logger, process, 0).(*processSupervisor) + ps := NewProcessSupervisor(logger, process, 0, 0).(*processSupervisor) ps.restartsCounter.Store(tc.initialRestartsCount) for i := 0; i < tc.restartAttempts; i++ { diff --git a/pkg/modules/chromium/chromium.go b/pkg/modules/chromium/chromium.go index 9a085534..beb37aa8 100644 --- a/pkg/modules/chromium/chromium.go +++ b/pkg/modules/chromium/chromium.go @@ -295,6 +295,7 @@ func (mod *Chromium) Descriptor() gotenberg.ModuleDescriptor { fs.Bool("chromium-clear-cookies", false, "Clear Chromium cookies between each conversion") fs.Bool("chromium-disable-javascript", false, "Disable JavaScript") fs.Bool("chromium-disable-routes", false, "Disable the routes") + fs.Int64("chromium-max-queue-size", 0, "Maximum request queue size for chromium. Set to 0 to disable this feature") return fs }(), @@ -344,7 +345,7 @@ func (mod *Chromium) Provision(ctx *gotenberg.Context) error { // Process. mod.browser = newChromiumBrowser(mod.args) - mod.supervisor = gotenberg.NewProcessSupervisor(mod.logger, mod.browser, flags.MustInt64("chromium-restart-after")) + mod.supervisor = gotenberg.NewProcessSupervisor(mod.logger, mod.browser, flags.MustInt64("chromium-restart-after"), flags.MustInt64("chromium-max-queue-size")) // PDF Engine. provider, err := ctx.Module(new(gotenberg.PdfEngineProvider)) diff --git a/pkg/modules/chromium/routes.go b/pkg/modules/chromium/routes.go index 97b01691..5f1ab92d 100644 --- a/pkg/modules/chromium/routes.go +++ b/pkg/modules/chromium/routes.go @@ -540,6 +540,16 @@ func convertUrl(ctx *api.Context, chromium Api, engine gotenberg.PdfEngine, url ) } + if errors.Is(err, gotenberg.ErrMaximumQueueSizeExceeded) { + return api.WrapError( + fmt.Errorf("convert to PDF: %w", err), + api.NewSentinelHttpError( + http.StatusTooManyRequests, + "The maximum queue size has been reached", + ), + ) + } + return fmt.Errorf("convert to PDF: %w", err) } diff --git a/pkg/modules/libreoffice/api/api.go b/pkg/modules/libreoffice/api/api.go index d3b82128..27a3a05c 100644 --- a/pkg/modules/libreoffice/api/api.go +++ b/pkg/modules/libreoffice/api/api.go @@ -80,6 +80,7 @@ func (a *Api) Descriptor() gotenberg.ModuleDescriptor { FlagSet: func() *flag.FlagSet { fs := flag.NewFlagSet("api", flag.ExitOnError) fs.Int64("libreoffice-restart-after", 10, "Number of conversions after which LibreOffice will automatically restart. Set to 0 to disable this feature") + fs.Int64("libreoffice-max-queue-size", 0, "Maximum request queue size for libreoffice. Set to 0 to disable this feature") fs.Bool("libreoffice-auto-start", false, "Automatically launch LibreOffice upon initialization if set to true; otherwise, LibreOffice will start at the time of the first conversion") fs.Duration("libreoffice-start-timeout", time.Duration(20)*time.Second, "Maximum duration to wait for LibreOffice to start or restart") @@ -123,7 +124,7 @@ func (a *Api) Provision(ctx *gotenberg.Context) error { // Process. a.libreOffice = newLibreOfficeProcess(a.args) - a.supervisor = gotenberg.NewProcessSupervisor(a.logger, a.libreOffice, flags.MustInt64("libreoffice-restart-after")) + a.supervisor = gotenberg.NewProcessSupervisor(a.logger, a.libreOffice, flags.MustInt64("libreoffice-restart-after"), flags.MustInt64("libreoffice-max-queue-size")) return nil } diff --git a/pkg/modules/libreoffice/routes.go b/pkg/modules/libreoffice/routes.go index a6bee965..95e352e7 100644 --- a/pkg/modules/libreoffice/routes.go +++ b/pkg/modules/libreoffice/routes.go @@ -68,6 +68,16 @@ func convertRoute(libreOffice libreofficeapi.Uno, engine gotenberg.PdfEngine) ap err = libreOffice.Pdf(ctx, ctx.Log(), inputPath, outputPaths[i], options) if err != nil { + if errors.Is(err, gotenberg.ErrMaximumQueueSizeExceeded) { + return api.WrapError( + fmt.Errorf("convert to PDF: %w", err), + api.NewSentinelHttpError( + http.StatusTooManyRequests, + "The maximum queue size has been reached", + ), + ) + } + if errors.Is(err, libreofficeapi.ErrInvalidPdfFormats) { return api.WrapError( fmt.Errorf("convert to PDF: %w", err),