audit contexts

.tool-versions

@@ -1,2 +1,2 @@
 golang 1.20
-mockery 2.12.1
+mockery 2.43.2

pkg/v2/coordinator/coordinator.go

@@ -32,6 +32,7 @@ import (
 
 	"github.com/smartcontractkit/chainlink-automation/pkg/util"
 	ocr2keepers "github.com/smartcontractkit/chainlink-automation/pkg/v2"
+	"github.com/smartcontractkit/chainlink-common/pkg/services"
 )
 
 var (
@@ -74,7 +75,7 @@ type reportCoordinator struct {
 
 	// run state data
 	running atomic.Bool
-	chStop  chan struct{}
+	chStop  services.StopChan
 }
 
 // NewReportCoordinator provides a new report coordinator. The coordinator
@@ -228,7 +229,7 @@ func (rc *reportCoordinator) checkLogs(ctx context.Context) error {
 		}
 	}
 
-	staleReportLogs, err = rc.logs.StaleReportLogs(context.Background())
+	staleReportLogs, err = rc.logs.StaleReportLogs(ctx)
 	// It can happen that in between the time the report is generated and it gets
 	// confirmed on chain something changes and it becomes stale. Current scenarios are:
 	//    - Another report for the upkeep is transmitted making this report stale
@@ -386,13 +387,20 @@ func (rc *reportCoordinator) Close() error {
 func (rc *reportCoordinator) run() {
 	cadence := time.Second
 	timer := time.NewTimer(cadence)
+	defer timer.Stop()
+
+	ctx, cancel := rc.chStop.NewCtx()
+	defer cancel()
 
 	for {
 		select {
 		case <-timer.C:
 			startTime := time.Now()
 
-			if err := rc.checkLogs(context.Background()); err != nil {
+			if err := rc.checkLogs(ctx); err != nil {
+				if ctx.Err() != nil {
+					return
+				}
 				rc.logger.Printf("failed to check perform and stale report logs: %s", err)
 			}
 
@@ -407,8 +415,7 @@ func (rc *reportCoordinator) run() {
 				// wait the difference between the cadence and the time taken
 				timer.Reset(cadence - diff)
 			}
-		case <-rc.chStop:
-			timer.Stop()
+		case <-ctx.Done():
 			return
 		}
 	}

pkg/v3/coordinator/coordinator.go

@@ -7,9 +7,9 @@ import (
 	"log"
 	"time"
 
+	"github.com/smartcontractkit/chainlink-common/pkg/services"
 	common "github.com/smartcontractkit/chainlink-common/pkg/types/automation"
 
-	internalutil "github.com/smartcontractkit/chainlink-automation/internal/util"
 	"github.com/smartcontractkit/chainlink-automation/pkg/util"
 	"github.com/smartcontractkit/chainlink-automation/pkg/v3/config"
 	"github.com/smartcontractkit/chainlink-automation/pkg/v3/types"
@@ -21,7 +21,9 @@ const (
 )
 
 type coordinator struct {
-	closer internalutil.Closer
+	services.StateMachine
+	stopCh services.StopChan
+	done   chan struct{}
 	logger *log.Logger
 
 	eventsProvider   types.TransmitEventProvider
@@ -46,6 +48,8 @@ type record struct {
 func NewCoordinator(transmitEventProvider types.TransmitEventProvider, upkeepTypeGetter types.UpkeepTypeGetter, conf config.OffchainConfig, logger *log.Logger) *coordinator {
 	performLockoutWindow := time.Duration(conf.PerformLockoutWindow) * time.Millisecond
 	return &coordinator{
+		stopCh:               make(chan struct{}),
+		done:                 make(chan struct{}),
 		logger:               logger,
 		eventsProvider:       transmitEventProvider,
 		upkeepTypeGetter:     upkeepTypeGetter,
@@ -208,16 +212,24 @@ func (c *coordinator) checkEvents(ctx context.Context) error {
 	return nil
 }
 
-func (c *coordinator) run(ctx context.Context) {
+func (c *coordinator) run() {
+	defer close(c.done)
+
 	timer := time.NewTimer(cadence)
 	defer timer.Stop()
 
+	ctx, cancel := c.stopCh.NewCtx()
+	defer cancel()
+
 	for {
 		select {
 		case <-timer.C:
 			startTime := time.Now()
 
-			if err := c.checkEvents(context.Background()); err != nil {
+			if err := c.checkEvents(ctx); err != nil {
+				if ctx.Err() != nil {
+					return
+				}
 				c.logger.Printf("failed to check for transmit events: %s", err)
 			}
 
@@ -239,32 +251,29 @@ func (c *coordinator) run(ctx context.Context) {
 }
 
 // Start starts all subprocesses
-func (c *coordinator) Start(pctx context.Context) error {
-	ctx, cancel := context.WithCancel(pctx)
-	defer cancel()
-
-	if !c.closer.Store(cancel) {
-		return fmt.Errorf("process already running")
+func (c *coordinator) Start(_ context.Context) error {
+	if err := c.StateMachine.StartOnce("Coordinator", func() error { return nil }); err != nil {
+		return err
 	}
 
 	go c.cache.Start(defaultCacheClean)
 	go c.visited.Start(defaultCacheClean)
 
-	c.run(ctx)
-
+	c.run()
 	return nil
 }
 
 // Close terminates all subprocesses
 func (c *coordinator) Close() error {
-	if !c.closer.Close() {
-		return fmt.Errorf("process not running")
-	}
+	return c.StateMachine.StopOnce("Coordinator", func() error {
+		close(c.stopCh)
+		<-c.done
 
-	c.cache.Stop()
-	c.visited.Stop()
+		c.cache.Stop()
+		c.visited.Stop()
 
-	return nil
+		return nil
+	})
 }
 
 func (c *coordinator) visitedID(e common.TransmitEvent) string {

pkg/v3/coordinator/coordinator_test.go

@@ -145,19 +145,6 @@ func TestNewCoordinator(t *testing.T) {
 		wg.Wait()
 		assert.True(t, strings.Contains(memLog.String(), "check database indexes and other performance improvements"))
 	})
-
-	t.Run("starting an already started coordinator returns an error", func(t *testing.T) {
-		c := NewCoordinator(nil, nil, config.OffchainConfig{PerformLockoutWindow: 3600 * 1000, MinConfirmations: 2}, nil)
-		c.closer.Store(context.CancelFunc(func() {}))
-		err := c.Start(context.Background())
-		assert.Error(t, err)
-	})
-
-	t.Run("closing an already closed coordinator returns an error", func(t *testing.T) {
-		c := NewCoordinator(nil, nil, config.OffchainConfig{PerformLockoutWindow: 3600 * 1000, MinConfirmations: 2}, nil)
-		err := c.Close()
-		assert.Error(t, err)
-	})
 }
 
 func TestNewCoordinator_checkEvents(t *testing.T) {