fix: chromium memory leaks (#705)

This commit is contained in:
Julien Neuhart
2023-10-23 17:52:30 +02:00
committed by GitHub
parent b5a59e4de0
commit 54daac329e
71 changed files with 4248 additions and 51761 deletions

View File

@@ -80,6 +80,28 @@ func (f *ParsedFlags) MustDeprecatedBool(deprecated string, newName string) bool
return f.MustBool(newName)
}
// MustInt64 returns the int64 value of a flag given by name.
// It panics if an error occurs.
func (f *ParsedFlags) MustInt64(name string) int64 {
val, err := f.GetInt64(name)
if err != nil {
panic(err)
}
return val
}
// MustDeprecatedInt64 returns the int64 value of a deprecated flag if it was
// explicitly set or the int64 value of the new flag.
// It panics if an error occurs.
func (f *ParsedFlags) MustDeprecatedInt64(deprecated string, newName string) int64 {
if f.Changed(deprecated) {
return f.MustInt64(deprecated)
}
return f.MustInt64(newName)
}
// MustInt returns the int value of a flag given by name.
// It panics if an error occurs.
func (f *ParsedFlags) MustInt(name string) int {

View File

@@ -252,6 +252,87 @@ func TestParsedFlags_MustDeprecatedBool(t *testing.T) {
}
}
func TestParsedFlags_MustInt64(t *testing.T) {
fs := flag.NewFlagSet("tests", flag.ContinueOnError)
fs.Int64("foo", 0, "")
err := fs.Parse([]string{"--foo=1"})
if err != nil {
t.Fatalf("expected no error but got: %v", err)
}
parsedFlags := ParsedFlags{FlagSet: fs}
for i, tc := range []struct {
name string
expectPanic bool
}{
{
name: "foo",
},
{
name: "bar",
expectPanic: true,
},
} {
func() {
if tc.expectPanic {
defer func() {
if r := recover(); r == nil {
t.Errorf("test %d: expected panic but got none", i)
}
}()
}
if !tc.expectPanic {
defer func() {
if r := recover(); r != nil {
t.Errorf("test %d: expected no panic but got: %v", i, r)
}
}()
}
parsedFlags.MustInt64(tc.name)
}()
}
}
func TestParsedFlags_MustDeprecatedInt64(t *testing.T) {
for i, tc := range []struct {
rawFlags []string
expectValue int64
}{
{
rawFlags: []string{"--foo=1"},
expectValue: 1,
},
{
rawFlags: []string{"--bar=2"},
expectValue: 2,
},
{
rawFlags: []string{"--foo=1", "--bar=2"},
expectValue: 1,
},
} {
fs := flag.NewFlagSet("tests", flag.ContinueOnError)
fs.Int64("foo", 0, "")
fs.Int64("bar", 0, "")
parsedFlags := ParsedFlags{FlagSet: fs}
err := parsedFlags.Parse(tc.rawFlags)
if err != nil {
t.Fatalf("test %d: expected no error but got: %v", i, err)
}
actual := parsedFlags.MustDeprecatedInt64("foo", "bar")
if actual != tc.expectValue {
t.Errorf("test %d: expected %d but got %d", i, tc.expectValue, actual)
}
}
}
func TestParsedFlags_MustInt(t *testing.T) {
fs := flag.NewFlagSet("tests", flag.ContinueOnError)
fs.Int("foo", 0, "")

View File

@@ -6,53 +6,106 @@ import (
"go.uber.org/zap"
)
// ModuleMock is a mock for the Module interface.
// ModuleMock is a mock for the [Module] interface.
type ModuleMock struct {
DescriptorMock func() ModuleDescriptor
}
func (mod ModuleMock) Descriptor() ModuleDescriptor {
func (mod *ModuleMock) Descriptor() ModuleDescriptor {
return mod.DescriptorMock()
}
// ValidatorMock is a mock for the Validator interface.
// ValidatorMock is a mock for the [Validator] interface.
type ValidatorMock struct {
ValidateMock func() error
}
func (mod ValidatorMock) Validate() error {
func (mod *ValidatorMock) Validate() error {
return mod.ValidateMock()
}
// PDFEngineMock is a mock for the PDFEngine interface.
// PDFEngineMock is a mock for the [PDFEngine] interface.
type PDFEngineMock struct {
MergeMock func(ctx context.Context, logger *zap.Logger, inputPaths []string, outputPath string) error
ConvertMock func(ctx context.Context, logger *zap.Logger, format, inputPath, outputPath string) error
}
func (engine PDFEngineMock) Merge(ctx context.Context, logger *zap.Logger, inputPaths []string, outputPath string) error {
func (engine *PDFEngineMock) Merge(ctx context.Context, logger *zap.Logger, inputPaths []string, outputPath string) error {
return engine.MergeMock(ctx, logger, inputPaths, outputPath)
}
func (engine PDFEngineMock) Convert(ctx context.Context, logger *zap.Logger, format, inputPath, outputPath string) error {
func (engine *PDFEngineMock) Convert(ctx context.Context, logger *zap.Logger, format, inputPath, outputPath string) error {
return engine.ConvertMock(ctx, logger, format, inputPath, outputPath)
}
// PDFEngineProviderMock is a mock for the PDFEngineProvider interface.
// PDFEngineProviderMock is a mock for the [PDFEngineProvider] interface.
type PDFEngineProviderMock struct {
PDFEngineMock func() (PDFEngine, error)
}
func (provider PDFEngineProviderMock) PDFEngine() (PDFEngine, error) {
func (provider *PDFEngineProviderMock) PDFEngine() (PDFEngine, error) {
return provider.PDFEngineMock()
}
// LoggerProviderMock is a mock for the LoggerProvider interface.
// ProcessMock is a mock for the [Process] interface.
type ProcessMock struct {
StartMock func(logger *zap.Logger) error
StopMock func(logger *zap.Logger) error
HealthyMock func(logger *zap.Logger) bool
}
func (p *ProcessMock) Start(logger *zap.Logger) error {
return p.StartMock(logger)
}
func (p *ProcessMock) Stop(logger *zap.Logger) error {
return p.StopMock(logger)
}
func (p *ProcessMock) Healthy(logger *zap.Logger) bool {
return p.HealthyMock(logger)
}
// ProcessSupervisorMock is a mock for the [ProcessSupervisor] interface.
type ProcessSupervisorMock struct {
LaunchMock func() error
ShutdownMock func() error
HealthyMock func() bool
RunMock func(ctx context.Context, logger *zap.Logger, task func() error) error
ReqQueueSizeMock func() int64
RestartsCountMock func() int64
}
func (s *ProcessSupervisorMock) Launch() error {
return s.LaunchMock()
}
func (s *ProcessSupervisorMock) Shutdown() error {
return s.ShutdownMock()
}
func (s *ProcessSupervisorMock) Healthy() bool {
return s.HealthyMock()
}
func (s *ProcessSupervisorMock) Run(ctx context.Context, logger *zap.Logger, task func() error) error {
return s.RunMock(ctx, logger, task)
}
func (s *ProcessSupervisorMock) ReqQueueSize() int64 {
return s.ReqQueueSizeMock()
}
func (s *ProcessSupervisorMock) RestartsCount() int64 {
return s.RestartsCountMock()
}
// LoggerProviderMock is a mock for the [LoggerProvider] interface.
type LoggerProviderMock struct {
LoggerMock func(mod Module) (*zap.Logger, error)
}
func (provider LoggerProviderMock) Logger(mod Module) (*zap.Logger, error) {
func (provider *LoggerProviderMock) Logger(mod Module) (*zap.Logger, error) {
return provider.LoggerMock(mod)
}
@@ -62,5 +115,7 @@ var (
_ Validator = (*ValidatorMock)(nil)
_ PDFEngine = (*PDFEngineMock)(nil)
_ PDFEngineProvider = (*PDFEngineProviderMock)(nil)
_ Process = (*ProcessMock)(nil)
_ ProcessSupervisor = (*ProcessSupervisorMock)(nil)
_ LoggerProvider = (*LoggerProviderMock)(nil)
)

View File

@@ -8,7 +8,7 @@ import (
)
func TestModuleMock(t *testing.T) {
mock := ModuleMock{
mock := &ModuleMock{
DescriptorMock: func() ModuleDescriptor {
return ModuleDescriptor{ID: "foo", New: func() Module {
return nil
@@ -17,12 +17,12 @@ func TestModuleMock(t *testing.T) {
}
if mock.Descriptor().ID != "foo" {
t.Errorf("expected ID '%s' from mock.Descriptor(), but got '%s'", "foo", mock.Descriptor().ID)
t.Errorf("expected ID '%s' from ModuleMock.Descriptor, but got '%s'", "foo", mock.Descriptor().ID)
}
}
func TestValidatorMock(t *testing.T) {
mock := ValidatorMock{
mock := &ValidatorMock{
ValidateMock: func() error {
return nil
},
@@ -30,12 +30,12 @@ func TestValidatorMock(t *testing.T) {
err := mock.Validate()
if err != nil {
t.Errorf("expected no error from mock.Validate(), but got: %v", err)
t.Errorf("expected no error from ValidatorMock.Validate, but got: %v", err)
}
}
func TestPDFEngineMock(t *testing.T) {
mock := PDFEngineMock{
mock := &PDFEngineMock{
MergeMock: func(ctx context.Context, logger *zap.Logger, inputPaths []string, outputPath string) error {
return nil
},
@@ -46,37 +46,119 @@ func TestPDFEngineMock(t *testing.T) {
err := mock.Merge(context.Background(), zap.NewNop(), nil, "")
if err != nil {
t.Errorf("expected no error from mock.Merge(), but got: %v", err)
t.Errorf("expected no error from PDFEngineMock.Merge, but got: %v", err)
}
err = mock.Convert(context.Background(), zap.NewNop(), "", "", "")
if err != nil {
t.Errorf("expected no error from mock.Convert(), but got: %v", err)
t.Errorf("expected no error from PDFEngineMock.Convert, but got: %v", err)
}
}
func TestPDFEngineProviderMock(t *testing.T) {
mock := PDFEngineProviderMock{
mock := &PDFEngineProviderMock{
PDFEngineMock: func() (PDFEngine, error) {
return PDFEngineMock{}, nil
return new(PDFEngineMock), nil
},
}
_, err := mock.PDFEngine()
if err != nil {
t.Errorf("expected no error from mock.PDFEngine(), but got: %v", err)
t.Errorf("expected no error from PDFEngineProviderMock.PDFEngine, but got: %v", err)
}
}
func TestProcessMock(t *testing.T) {
mock := &ProcessMock{
StartMock: func(logger *zap.Logger) error {
return nil
},
StopMock: func(logger *zap.Logger) error {
return nil
},
HealthyMock: func(logger *zap.Logger) bool {
return true
},
}
err := mock.Start(zap.NewNop())
if err != nil {
t.Errorf("expected no error from ProcessMock.Start, but got: %v", err)
}
err = mock.Stop(zap.NewNop())
if err != nil {
t.Errorf("expected no error from ProcessMock.Stop, but got: %v", err)
}
healthy := mock.Healthy(zap.NewNop())
if !healthy {
t.Error("expected true from ProcessMock.Healthy, but got false")
}
}
func TestProcessSupervisorMock(t *testing.T) {
mock := &ProcessSupervisorMock{
LaunchMock: func() error {
return nil
},
ShutdownMock: func() error {
return nil
},
HealthyMock: func() bool {
return true
},
RunMock: func(ctx context.Context, logger *zap.Logger, task func() error) error {
return nil
},
ReqQueueSizeMock: func() int64 {
return 0
},
RestartsCountMock: func() int64 {
return 0
},
}
err := mock.Launch()
if err != nil {
t.Errorf("expected no error from ProcessSupervisorMock.Launch, but got: %v", err)
}
err = mock.Shutdown()
if err != nil {
t.Errorf("expected no error from ProcessSupervisorMock.Shutdown, but got: %v", err)
}
healthy := mock.Healthy()
if !healthy {
t.Error("expected true from ProcessSupervisorMock.Healthy, but got false")
}
err = mock.Run(context.TODO(), zap.NewNop(), nil)
if err != nil {
t.Errorf("expected no error from ProcessSupervisorMock.Run, but got: %v", err)
}
size := mock.ReqQueueSize()
if size != 0 {
t.Errorf("expected 0 from ProcessSupervisorMock.ReqQueueSize, but got: %d", size)
}
restarts := mock.RestartsCount()
if restarts != 0 {
t.Errorf("expected 0 from ProcessSupervisorMock.RestartsCount, but got: %d", restarts)
}
}
func TestLoggerProviderMock(t *testing.T) {
mock := LoggerProviderMock{
mock := &LoggerProviderMock{
LoggerMock: func(mod Module) (*zap.Logger, error) {
return nil, nil
},
}
_, err := mock.Logger(ModuleMock{})
_, err := mock.Logger(new(ModuleMock))
if err != nil {
t.Errorf("expected no error from mock.Logger(), but got: %v", err)
t.Errorf("expected no error from LoggerProviderMock.Logger, but got: %v", err)
}
}

244
pkg/gotenberg/supervisor.go Normal file
View File

@@ -0,0 +1,244 @@
package gotenberg
import (
"context"
"fmt"
"sync/atomic"
"go.uber.org/zap"
)
// Process is an interface that represents an abstract process
// and provides methods for starting, stopping, and checking the health of the
// process.
//
// Implementations of this interface should handle the actual logic for
// starting, stopping, and ensuring the process's health.
type Process interface {
// Start initiates the process and returns an error if the process cannot
// be started.
Start(logger *zap.Logger) error
// Stop terminates the process and returns an error if the process cannot
// be stopped.
Stop(logger *zap.Logger) error
// Healthy checks the health of the process. It returns true if the process
// is healthy; otherwise, it returns false.
Healthy(logger *zap.Logger) bool
}
// ProcessSupervisor provides methods to manage a [Process], including
// starting, stopping, and ensuring its health.
//
// Additionally, it allows for the execution of tasks while managing the
// process's state and provides functionality for limiting the number of
// requests that can be handled by the process, as well as managing a request
// queue.
type ProcessSupervisor interface {
// Launch starts the managed [Process].
Launch() error
// Shutdown stops the managed [Process].
Shutdown() error
// Healthy checks and returns the health status of the managed [Process].
//
// If the process has not been started or is restarting, it is considered
// healthy and true is returned. Otherwise, it returns the health status of
// the actual process.
Healthy() bool
// Run executes a provided task while managing the state of the [Process].
//
// Run manages the request queue and may restart the process if it is not
// healthy or if the number of handled requests exceeds the maximum limit.
//
// It returns an error if the task cannot be run or if the process state
// cannot be managed properly.
Run(ctx context.Context, logger *zap.Logger, task func() error) error
// ReqQueueSize returns the current size of the request queue.
ReqQueueSize() int64
// RestartsCount returns the current number of restart.
RestartsCount() int64
}
type processSupervisor struct {
logger *zap.Logger
process Process
maxReqLimit int64
mutexChan chan struct{}
firstStart atomic.Bool
reqCounter atomic.Int64
reqQueueSize atomic.Int64
restartsCounter atomic.Int64
isRestarting atomic.Bool
}
// NewProcessSupervisor initializes a new [ProcessSupervisor].
func NewProcessSupervisor(logger *zap.Logger, process Process, maxReqLimit int64) ProcessSupervisor {
b := &processSupervisor{
logger: logger,
process: process,
mutexChan: make(chan struct{}, 1),
maxReqLimit: maxReqLimit,
}
b.reqCounter.Store(0)
b.reqQueueSize.Store(0)
b.restartsCounter.Store(0)
b.isRestarting.Store(false)
return b
}
func (s *processSupervisor) Launch() error {
s.logger.Debug("start process")
err := s.process.Start(s.logger)
if err != nil {
return fmt.Errorf("start process: %w", err)
}
s.firstStart.Store(true)
s.logger.Debug("process successfully started")
return nil
}
func (s *processSupervisor) Shutdown() error {
s.logger.Debug("shutdown process")
err := s.process.Stop(s.logger)
if err != nil {
return fmt.Errorf("shutdown process: %w", err)
}
s.logger.Debug("process successfully shutdown")
return nil
}
func (s *processSupervisor) restart() error {
if s.isRestarting.Load() {
s.logger.Debug("process already restarting, skip restart")
return nil
}
s.logger.Debug("restart process")
s.isRestarting.Store(true)
defer s.isRestarting.Store(false)
err := s.Shutdown()
if err != nil {
// No big deal? Chances are it's already stopped.
s.logger.Debug(fmt.Sprintf("stop process before restart: %s", err))
}
err = s.Launch()
if err != nil {
return fmt.Errorf("restart process: %w", err)
}
s.reqCounter.Store(0)
s.restartsCounter.Add(1)
s.logger.Debug("process successfully restarted")
return nil
}
func (s *processSupervisor) Healthy() bool {
if !s.firstStart.Load() {
// A non-started process is always healthy.
return true
}
if s.isRestarting.Load() {
// A restarting process is always healthy.
return true
}
return s.process.Healthy(s.logger)
}
func (s *processSupervisor) Run(ctx context.Context, logger *zap.Logger, task func() error) error {
s.reqQueueSize.Add(1)
select {
case s.mutexChan <- struct{}{}:
logger.Debug("process lock acquired")
s.reqQueueSize.Add(-1)
s.reqCounter.Add(1)
defer func() {
logger.Debug("process lock released")
<-s.mutexChan
}()
if !s.firstStart.Load() {
err := s.runWithDeadline(ctx, func() error {
return s.Launch()
})
if err != nil {
return fmt.Errorf("process first start: %w", err)
}
}
if !s.Healthy() {
s.logger.Debug("process is unhealthy, cannot handle task, restarting...")
err := s.runWithDeadline(ctx, func() error {
return s.restart()
})
if err != nil {
return fmt.Errorf("process restart before task: %w", err)
}
}
if s.maxReqLimit > 0 && s.reqCounter.Load() >= s.maxReqLimit {
s.logger.Debug("max request limit reached, restarting...")
err := s.runWithDeadline(ctx, func() error {
return s.restart()
})
if err != nil {
return fmt.Errorf("process restart before task: %w", err)
}
}
// FIXME: no error wrapping because it leaks on Chromium console exceptions output.
return s.runWithDeadline(ctx, task)
case <-ctx.Done():
logger.Debug("failed to acquire process lock before deadline")
s.reqQueueSize.Add(-1)
return fmt.Errorf("acquire process lock: %w", ctx.Err())
}
}
func (s *processSupervisor) runWithDeadline(ctx context.Context, task func() error) error {
runChan := make(chan error, 1)
go func() {
runChan <- task()
}()
for {
select {
case err := <-runChan:
return err
case <-ctx.Done():
return ctx.Err()
}
}
}
func (s *processSupervisor) ReqQueueSize() int64 {
return s.reqQueueSize.Load()
}
func (s *processSupervisor) RestartsCount() int64 {
return s.restartsCounter.Load()
}
// Interface guards.
var (
_ ProcessSupervisor = (*processSupervisor)(nil)
)

View File

@@ -0,0 +1,570 @@
package gotenberg
import (
"context"
"errors"
"sync"
"sync/atomic"
"testing"
"time"
"go.uber.org/zap"
)
func TestProcessSupervisor_Launch(t *testing.T) {
for _, tc := range []struct {
scenario string
startError error
expectError bool
firstStartSet bool
}{
{
scenario: "successful launch",
startError: nil,
expectError: false,
firstStartSet: true,
},
{
scenario: "failed launch",
startError: errors.New("start error"),
expectError: true,
firstStartSet: false,
},
{
scenario: "process already started",
startError: nil,
expectError: false,
firstStartSet: true,
},
} {
t.Run(tc.scenario, func(t *testing.T) {
logger := zap.NewNop()
process := &ProcessMock{
StartMock: func(logger *zap.Logger) error {
return tc.startError
},
}
ps := NewProcessSupervisor(logger, process, 5).(*processSupervisor)
if tc.firstStartSet {
ps.firstStart.Store(true)
}
err := ps.Launch()
if !tc.expectError && err != nil {
t.Fatalf("expected no error but got: %v", err)
}
if tc.expectError && err == nil {
t.Fatal("expected error but got none")
}
if tc.firstStartSet && !ps.firstStart.Load() {
t.Error("expected firstStart to be set but it was not")
}
})
}
}
func TestProcessSupervisor_Shutdown(t *testing.T) {
for _, tc := range []struct {
scenario string
stopError error
expectError bool
}{
{
scenario: "successful shutdown",
stopError: nil,
expectError: false,
},
{
scenario: "failed shutdown",
stopError: errors.New("stop error"),
expectError: true,
},
} {
t.Run(tc.scenario, func(t *testing.T) {
logger := zap.NewNop()
process := &ProcessMock{
StopMock: func(logger *zap.Logger) error {
return tc.stopError
},
}
ps := NewProcessSupervisor(logger, process, 5)
err := ps.Shutdown()
if !tc.expectError && err != nil {
t.Fatalf("expected no error but got: %v", err)
}
if tc.expectError && err == nil {
t.Fatal("expected error but got none")
}
})
}
}
func TestProcessSupervisor_restart(t *testing.T) {
for _, tc := range []struct {
scenario string
initiallyRestarting bool
startError error
stopError error
expectError bool
}{
{
scenario: "already restarting",
initiallyRestarting: true,
expectError: false,
},
{
scenario: "successful restart",
startError: nil,
stopError: nil,
expectError: false,
},
{
scenario: "failed to stop during restart",
startError: nil,
stopError: errors.New("stop error"),
expectError: false,
},
{
scenario: "failed to start during restart",
startError: errors.New("start error"),
stopError: nil,
expectError: true,
},
} {
t.Run(tc.scenario, func(t *testing.T) {
logger := zap.NewNop()
process := &ProcessMock{
StartMock: func(logger *zap.Logger) error {
return tc.startError
},
StopMock: func(logger *zap.Logger) error {
return tc.stopError
},
}
ps := NewProcessSupervisor(logger, process, 5).(*processSupervisor)
if tc.initiallyRestarting {
ps.isRestarting.Store(true)
}
err := ps.restart()
if !tc.expectError && err != nil {
t.Fatalf("expected no error but got: %v", err)
}
if tc.expectError && err == nil {
t.Fatal("expected error but got none")
}
})
}
}
func TestProcessSupervisor_Healthy(t *testing.T) {
for _, tc := range []struct {
scenario string
initiallyStarted bool
initiallyRestarting bool
processHealthy bool
expectHealthy bool
}{
{
scenario: "non-started process is always healthy",
initiallyStarted: false,
expectHealthy: true,
},
{
scenario: "restarting process is always healthy",
initiallyStarted: true,
initiallyRestarting: true,
expectHealthy: true,
},
{
scenario: "process reports as healthy",
initiallyStarted: true,
processHealthy: true,
expectHealthy: true,
},
{
scenario: "process reports as unhealthy",
initiallyStarted: true,
processHealthy: false,
expectHealthy: false,
},
} {
t.Run(tc.scenario, func(t *testing.T) {
logger := zap.NewNop()
process := &ProcessMock{
HealthyMock: func(logger *zap.Logger) bool {
return tc.processHealthy
},
}
ps := NewProcessSupervisor(logger, process, 5).(*processSupervisor)
if tc.initiallyStarted {
ps.firstStart.Store(true)
}
if tc.initiallyRestarting {
ps.isRestarting.Store(true)
}
healthy := ps.Healthy()
if healthy != tc.expectHealthy {
t.Fatalf("expected healthy to be %v but got %v", tc.expectHealthy, healthy)
}
})
}
}
func TestProcessSupervisor_Run(t *testing.T) {
for _, tc := range []struct {
scenario string
initiallyStarted bool
startError error
processHealthy bool
maxReqLimit int64
tasksToRun int
taskError error
expectError bool
expectedStartCalls int64
expectedHealthyCalls int64
expectedStopCalls int64
}{
{
scenario: "successfully run task on non-started process",
initiallyStarted: false,
processHealthy: true,
maxReqLimit: 2,
tasksToRun: 1,
expectError: false,
expectedStartCalls: 1,
expectedHealthyCalls: 1,
expectedStopCalls: 0,
},
{
scenario: "cannot launch non-started process",
initiallyStarted: false,
startError: errors.New("launch error"),
processHealthy: true,
maxReqLimit: 2,
tasksToRun: 1,
expectError: true,
expectedStartCalls: 1,
expectedHealthyCalls: 0,
expectedStopCalls: 0,
},
{
scenario: "run task with unhealthy process causing restart",
initiallyStarted: true,
processHealthy: false,
maxReqLimit: 2,
tasksToRun: 1,
expectError: false,
expectedStartCalls: 1,
expectedHealthyCalls: 1,
expectedStopCalls: 1,
},
{
scenario: "cannot restart unhealthy process",
startError: errors.New("start error"),
initiallyStarted: true,
processHealthy: false,
maxReqLimit: 2,
tasksToRun: 1,
expectError: true,
expectedStartCalls: 1,
expectedHealthyCalls: 1,
expectedStopCalls: 1,
},
{
scenario: "run tasks reaching max request limit causing restart",
initiallyStarted: true,
processHealthy: true,
maxReqLimit: 2,
tasksToRun: 3,
expectError: false,
expectedStartCalls: 1,
expectedHealthyCalls: 3,
expectedStopCalls: 1,
},
{
scenario: "cannot restart after reaching max request limit",
startError: errors.New("start error"),
initiallyStarted: true,
processHealthy: true,
maxReqLimit: 2,
tasksToRun: 2,
expectError: true,
expectedStartCalls: 1,
expectedHealthyCalls: 2,
expectedStopCalls: 1,
},
{
scenario: "task error",
initiallyStarted: true,
processHealthy: true,
maxReqLimit: 0,
tasksToRun: 1,
taskError: errors.New("task error"),
expectError: true,
expectedStartCalls: 0,
expectedHealthyCalls: 1,
expectedStopCalls: 0,
},
} {
t.Run(tc.scenario, func(t *testing.T) {
logger := zap.NewNop()
var startCalls, healthyCalls, stopCalls atomic.Int64
startCalls.Store(0)
healthyCalls.Store(0)
stopCalls.Store(0)
process := &ProcessMock{
StartMock: func(logger *zap.Logger) error {
startCalls.Add(1)
return tc.startError
},
StopMock: func(logger *zap.Logger) error {
stopCalls.Add(1)
return nil
},
HealthyMock: func(logger *zap.Logger) bool {
healthyCalls.Add(1)
return tc.processHealthy
},
}
ps := NewProcessSupervisor(logger, process, tc.maxReqLimit).(*processSupervisor)
if tc.initiallyStarted {
ps.firstStart.Store(true)
}
task := func() error {
return tc.taskError
}
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
var wg sync.WaitGroup
errorChan := make(chan error, tc.tasksToRun)
for i := 0; i < tc.tasksToRun; i++ {
wg.Add(1)
go func() {
defer wg.Done()
err := ps.Run(ctx, logger, task)
if err != nil {
errorChan <- err
}
}()
}
wg.Wait()
close(errorChan)
for err := range errorChan {
if tc.expectError && err == nil {
t.Fatal("expected an error but got none")
}
if !tc.expectError && err != nil {
t.Fatalf("expected no error but got: %v", err)
}
}
if startCalls.Load() != tc.expectedStartCalls {
t.Errorf("expected %d process.Start calls, got %d", tc.expectedStartCalls, startCalls.Load())
}
if healthyCalls.Load() != tc.expectedHealthyCalls {
t.Errorf("expected %d process.Healthy calls, got %d", tc.expectedHealthyCalls, healthyCalls.Load())
}
if stopCalls.Load() != tc.expectedStopCalls {
t.Errorf("expected %d process.Stop calls, got %d", tc.expectedStopCalls, stopCalls.Load())
}
})
}
}
func TestProcessSupervisor_runWithDeadline(t *testing.T) {
for _, tc := range []struct {
scenario string
ctxDone bool
expectError bool
}{
{
scenario: "task finished",
ctxDone: false,
expectError: false,
},
{
scenario: "context expired",
ctxDone: true,
expectError: true,
},
} {
t.Run(tc.scenario, func(t *testing.T) {
ps := NewProcessSupervisor(zap.NewNop(), new(ProcessMock), 0).(*processSupervisor)
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
defer cancel()
if tc.ctxDone {
cancel()
}
err := ps.runWithDeadline(ctx, func() error {
return nil
})
if tc.expectError && err == nil {
t.Fatal("expected an error but got none")
}
if !tc.expectError && err != nil {
t.Fatalf("expected no error but got: %v", err)
}
})
}
}
func TestProcessSupervisor_ReqQueueSize(t *testing.T) {
logger := zap.NewNop()
process := &ProcessMock{
StartMock: func(logger *zap.Logger) error {
return nil
},
HealthyMock: func(logger *zap.Logger) bool {
return true
},
}
ps := NewProcessSupervisor(logger, process, 0).(*processSupervisor)
// Simulating a lock.
ps.mutexChan <- struct{}{}
if ps.ReqQueueSize() != 0 {
t.Fatalf("expected queue size to be 0 but got %d", ps.ReqQueueSize())
}
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
var wg sync.WaitGroup
errorChan := make(chan error, 10)
for i := 0; i < 10; i++ {
wg.Add(1)
go func() {
defer wg.Done()
err := ps.Run(ctx, logger, func() error {
return nil
})
if err != nil {
errorChan <- err
}
}()
}
// We have to wait a little bit so that the request queue size may change.
time.Sleep(10 * time.Millisecond)
if ps.ReqQueueSize() != 10 {
t.Fatalf("expected queue size to be 10 but got %d", ps.ReqQueueSize())
}
wg.Wait()
close(errorChan)
for err := range errorChan {
if err == nil {
t.Error("expected a lock error but got none")
}
}
if ps.ReqQueueSize() != 0 {
t.Errorf("expected queue size to be 0 but got %d", ps.ReqQueueSize())
}
}
func TestProcessSupervisor_RestartsCount(t *testing.T) {
for _, tc := range []struct {
scenario string
initialRestartsCount int64
restartAttempts int
startError error
stopError error
expectedRestartsCount int64
}{
{
scenario: "no restarts, counter remains 0",
initialRestartsCount: 0,
restartAttempts: 0,
expectedRestartsCount: 0,
},
{
scenario: "successful restart increases counter",
initialRestartsCount: 0,
restartAttempts: 1,
startError: nil,
stopError: nil,
expectedRestartsCount: 1,
},
{
scenario: "failed to stop during restart, no impact",
initialRestartsCount: 0,
restartAttempts: 1,
startError: nil,
stopError: errors.New("stop error"),
expectedRestartsCount: 1,
},
{
scenario: "multiple successful restarts",
initialRestartsCount: 0,
restartAttempts: 3,
startError: nil,
stopError: nil,
expectedRestartsCount: 3,
},
} {
t.Run(tc.scenario, func(t *testing.T) {
logger := zap.NewNop()
process := &ProcessMock{
StartMock: func(logger *zap.Logger) error {
return tc.startError
},
StopMock: func(logger *zap.Logger) error {
return tc.stopError
},
}
ps := NewProcessSupervisor(logger, process, 0).(*processSupervisor)
ps.restartsCounter.Store(tc.initialRestartsCount)
for i := 0; i < tc.restartAttempts; i++ {
_ = ps.restart()
}
actualRestartsCount := ps.RestartsCount()
if actualRestartsCount != tc.expectedRestartsCount {
t.Fatalf("expected restarts count to be %d, but got %d", tc.expectedRestartsCount, actualRestartsCount)
}
})
}
}