stackrox · msugakov · Jan 2, 2023 · Nov 29, 2022 · Nov 30, 2022 · Nov 29, 2022
diff --git a/central/cve/suppress/reprocessor.go b/central/cve/suppress/reprocessor.go
@@ -68,8 +68,7 @@ func newLoopWithDuration(tickerDuration time.Duration, cveStores ...vulnsStore)
 		cveStores:                 cveStores,
 		cveSuppressTickerDuration: tickerDuration,
 
-		stopChan: concurrency.NewSignal(),
-		stopped:  concurrency.NewSignal(),
+		stopper: concurrency.NewStopper(),
 	}
 }
 
@@ -79,8 +78,7 @@ type cveUnsuppressLoopImpl struct {
 
 	cveStores []vulnsStore
 
-	stopChan concurrency.Signal
-	stopped  concurrency.Signal
+	stopper concurrency.Stopper
 }
 
 // Start starts the CVE unsuppress loop.
@@ -91,13 +89,15 @@ func (l *cveUnsuppressLoopImpl) Start() {
 
 // Stop stops the CVE unsuppress loop.
 func (l *cveUnsuppressLoopImpl) Stop() {
-	l.stopChan.Signal()
-	l.stopped.Wait()
+	l.stopper.Client().Stop()
+	_ = l.stopper.Client().Stopped().Wait()
 }
 
 func (l *cveUnsuppressLoopImpl) unsuppressCVEsWithExpiredSuppressState() {
-	if l.stopped.IsDone() {
+	select {
+	case <-l.stopper.Flow().StopRequested():
 		return
+	default:
 	}
 
 	totalUnsuppressedCVEs := 0
@@ -142,13 +142,13 @@ func getCVEsWithExpiredSuppressState(cveStore vulnsStore) ([]string, error) {
 }
 
 func (l *cveUnsuppressLoopImpl) loop() {
-	defer l.stopped.Signal()
+	defer l.stopper.Flow().ReportStopped()
 	defer l.cveSuppressTicker.Stop()
 
 	go l.unsuppressCVEsWithExpiredSuppressState()
 	for {
 		select {
-		case <-l.stopChan.Done():
+		case <-l.stopper.Flow().StopRequested():
 			return
 		case <-l.cveSuppressTicker.C:
 			l.unsuppressCVEsWithExpiredSuppressState()

diff --git a/central/processindicator/datastore/datastore.go b/central/processindicator/datastore/datastore.go
@@ -36,7 +36,7 @@ type DataStore interface {
 	WalkAll(ctx context.Context, fn func(pi *storage.ProcessIndicator) error) error
 
 	// Stop signals all goroutines associated with this object to terminate.
-	Stop() bool
+	Stop()
 	// Wait waits until all goroutines associated with this object have terminated, or cancelWhen gets triggered.
 	// A return value of false indicates that cancelWhen was triggered.
 	Wait(cancelWhen concurrency.Waitable) bool
@@ -50,8 +50,7 @@ func New(store store.Store, indexer index.Indexer, searcher search.Searcher, pru
 		searcher:              searcher,
 		prunerFactory:         prunerFactory,
 		prunedArgsLengthCache: make(map[processindicator.ProcessWithContainerInfo]int),
-		stopSig:               concurrency.NewSignal(),
-		stoppedSig:            concurrency.NewSignal(),
+		stopper:               concurrency.NewStopper(),
 	}
 	ctx := sac.WithAllAccess(context.Background())
 	if err := d.buildIndex(ctx); err != nil {

diff --git a/central/processindicator/datastore/datastore_impl.go b/central/processindicator/datastore/datastore_impl.go
@@ -39,7 +39,7 @@ type datastoreImpl struct {
 	prunerFactory         pruner.Factory
 	prunedArgsLengthCache map[processindicator.ProcessWithContainerInfo]int
 
-	stopSig, stoppedSig concurrency.Signal
+	stopper concurrency.Stopper
 }
 
 func checkReadAccess(ctx context.Context, indicator *storage.ProcessIndicator) (bool, error) {
@@ -170,19 +170,19 @@ func (ds *datastoreImpl) RemoveProcessIndicatorsByPod(ctx context.Context, id st
 }
 
 func (ds *datastoreImpl) prunePeriodically(ctx context.Context) {
-	defer ds.stoppedSig.Signal()
+	defer ds.stopper.Flow().ReportStopped()
 
 	if ds.prunerFactory == nil {
 		return
 	}
 
 	t := time.NewTicker(ds.prunerFactory.Period())
 	defer t.Stop()
-	for !ds.stopSig.IsDone() {
+	for {
 		select {
 		case <-t.C:
 			ds.prune(ctx)
-		case <-ds.stopSig.Done():
+		case <-ds.stopper.Flow().StopRequested():
 			return
 		}
 	}
@@ -248,12 +248,12 @@ func (ds *datastoreImpl) prune(ctx context.Context) {
 	}
 }
 
-func (ds *datastoreImpl) Stop() bool {
-	return ds.stopSig.Signal()
+func (ds *datastoreImpl) Stop() {
+	ds.stopper.Client().Stop()
 }
 
 func (ds *datastoreImpl) Wait(cancelWhen concurrency.Waitable) bool {
-	return concurrency.WaitInContext(&ds.stoppedSig, cancelWhen)
+	return concurrency.WaitInContext(ds.stopper.Client().Stopped(), cancelWhen)
 }
 
 func (ds *datastoreImpl) fullReindex(ctx context.Context) error {

diff --git a/central/processindicator/datastore/mocks/datastore.go b/central/processindicator/datastore/mocks/datastore.go
diff --git a/central/pruning/pruning.go b/central/pruning/pruning.go
@@ -97,8 +97,7 @@ func newGarbageCollector(alerts alertDatastore.DataStore,
 		k8sRoles:        k8sRoles,
 		k8sRoleBindings: k8sRoleBindings,
 		logimbueStore:   logimbueStore,
-		stopSig:         concurrency.NewSignal(),
-		stoppedSig:      concurrency.NewSignal(),
+		stopper:         concurrency.NewStopper(),
 	}
 }
 
@@ -120,8 +119,7 @@ type garbageCollectorImpl struct {
 	k8sRoles        k8sRoleDataStore.DataStore
 	k8sRoleBindings roleBindingDataStore.DataStore
 	logimbueStore   logimbueDataStore.Store
-	stopSig         concurrency.Signal
-	stoppedSig      concurrency.Signal
+	stopper         concurrency.Stopper
 }
 
 func (g *garbageCollectorImpl) Start() {
@@ -159,6 +157,8 @@ func (g *garbageCollectorImpl) pruneBasedOnConfig() {
 }
 
 func (g *garbageCollectorImpl) runGC() {
+	defer g.stopper.Flow().ReportStopped()
+
 	lastClusterPruneTime = time.Now().Add(-24 * time.Hour)
 	lastLogImbuePruneTime = time.Now().Add(-24 * time.Hour)
 	g.pruneBasedOnConfig()
@@ -168,8 +168,7 @@ func (g *garbageCollectorImpl) runGC() {
 		select {
 		case <-t.C:
 			g.pruneBasedOnConfig()
-		case <-g.stopSig.Done():
-			g.stoppedSig.Signal()
+		case <-g.stopper.Flow().StopRequested():
 			return
 		}
 	}
@@ -876,6 +875,6 @@ func (g *garbageCollectorImpl) pruneLogImbues() {
 }
 
 func (g *garbageCollectorImpl) Stop() {
-	g.stopSig.Signal()
-	<-g.stoppedSig.Done()
+	g.stopper.Client().Stop()
+	_ = g.stopper.Client().Stopped().Wait()
 }
@@ -136,8 +136,7 @@ type scheduler struct {
 
 	reportsToRun chan *ReportRequest
 
-	stoppedSig concurrency.Signal
-	stopped    concurrency.Signal
+	stopper concurrency.Stopper
 
 	Schema *graphql.Schema
 }
@@ -178,8 +177,7 @@ func New(reportConfigDS reportConfigDS.DataStore, notifierDS notifierDataStore.D
 		reportsToRun:           make(chan *ReportRequest, 100),
 		Schema:                 ourSchema,
 
-		stoppedSig: concurrency.NewSignal(),
-		stopped:    concurrency.NewSignal(),
+		stopper: concurrency.NewStopper(),
 	}
 	return s
 }
@@ -230,10 +228,10 @@ func (s *scheduler) SubmitReport(reportRequest *ReportRequest) {
 }
 
 func (s *scheduler) runReports() {
-	defer s.stopped.Signal()
-	for !s.stoppedSig.IsDone() {
+	defer s.stopper.Flow().ReportStopped()
+	for {
 		select {
-		case <-s.stoppedSig.Done():
+		case <-s.stopper.Flow().StopRequested():
 			return
 		case req := <-s.reportsToRun:
 			log.Infof("Executing report '%s' at %v", req.ReportConfig.GetName(), time.Now().Format(time.RFC822))
@@ -416,6 +414,6 @@ func (s *scheduler) Start() {
 }
 
 func (s *scheduler) Stop() {
-	s.stoppedSig.Signal()
-	s.stopped.Wait()
+	s.stopper.Client().Stop()
+	_ = s.stopper.Client().Stopped().Wait()
 }
diff --git a/central/vulnerabilityrequest/manager/requestmgr/manager_impl.go b/central/vulnerabilityrequest/manager/requestmgr/manager_impl.go
@@ -58,8 +58,7 @@ type managerImpl struct {
 	reObserveTimedDeferralsTickerDuration     time.Duration
 	reObserveWhenFixedDeferralsTickerDuration time.Duration
 
-	stopSig concurrency.Signal
-	stopped concurrency.Signal
+	stopper concurrency.Stopper
 }
 
 func (m *managerImpl) Start() {
@@ -70,8 +69,8 @@ func (m *managerImpl) Start() {
 }
 
 func (m *managerImpl) Stop() {
-	m.stopSig.Signal()
-	m.stopped.Wait()
+	m.stopper.Client().Stop()
+	_ = m.stopper.Client().Stopped().Wait()
 }
 
 func (m *managerImpl) Create(ctx context.Context, req *storage.VulnerabilityRequest) error {
@@ -372,7 +371,7 @@ func (m *managerImpl) getExpiredDeferrals() ([]*storage.VulnerabilityRequest, er
 }
 
 func (m *managerImpl) reObserveExpiredDeferrals() {
-	if m.stopped.IsDone() {
+	if m.stopper.Client().Stopped().IsDone() {
 		return
 	}
 
@@ -438,7 +437,7 @@ func (m *managerImpl) getFixableDeferrals() ([]*storage.VulnerabilityRequest, er
 }
 
 func (m *managerImpl) reObserveFixableDeferrals() {
-	if m.stopped.IsDone() {
+	if m.stopper.Client().Stopped().IsDone() {
 		return
 	}
 
@@ -459,7 +458,7 @@ func (m *managerImpl) reObserveFixableDeferrals() {
 }
 
 func (m *managerImpl) runExpiredDeferralsProcessor() {
-	defer m.stopped.Signal()
+	defer m.stopper.Flow().ReportStopped()
 	reObserveTimedDeferralsTicker := time.NewTicker(m.reObserveTimedDeferralsTickerDuration)
 	defer reObserveTimedDeferralsTicker.Stop()
 	reObserveWhenFixedDeferralsTicker := time.NewTicker(m.reObserveWhenFixedDeferralsTickerDuration)
@@ -471,7 +470,7 @@ func (m *managerImpl) runExpiredDeferralsProcessor() {
 
 	for {
 		select {
-		case <-m.stopSig.Done():
+		case <-m.stopper.Flow().StopRequested():
 			return
 		case <-reObserveTimedDeferralsTicker.C:
 			go m.reObserveExpiredDeferrals()

diff --git a/central/vulnerabilityrequest/manager/requestmgr/manager_impl_test.go b/central/vulnerabilityrequest/manager/requestmgr/manager_impl_test.go
@@ -171,8 +171,7 @@ func (s *VulnRequestManagerTestSuite) SetupTest() {
 		activeReqCache:                        s.activeReqCache,
 		reObserveTimedDeferralsTickerDuration: expiryLoopDurationForTest,
 		reObserveWhenFixedDeferralsTickerDuration: expiryLoopDurationForTest,
-		stopSig: concurrency.NewSignal(),
-		stopped: concurrency.NewSignal(),
+		stopper: concurrency.NewStopper(),
 	}
 }
 

diff --git a/central/vulnerabilityrequest/manager/requestmgr/singleton.go b/central/vulnerabilityrequest/manager/requestmgr/singleton.go
@@ -63,7 +63,6 @@ func New(
 
 		reObserveTimedDeferralsTickerDuration:     env.VulnDeferralTimedReObserveInterval.DurationSetting(),
 		reObserveWhenFixedDeferralsTickerDuration: env.VulnDeferralFixableReObserveInterval.DurationSetting(),
-		stopSig: concurrency.NewSignal(),
-		stopped: concurrency.NewSignal(),
+		stopper: concurrency.NewStopper(),
 	}
 }
diff --git a/compliance/collection/auditlog/auditlog.go b/compliance/collection/auditlog/auditlog.go
@@ -18,15 +18,15 @@ type Reader interface {
 	// StartReader will start the audit log reader process which will continuously read and send events until stopped.
 	// Returns true if the reader can be started (log exists and can be read). Log file missing is not considered an error.
 	StartReader(ctx context.Context) (bool, error)
-	// StopReader will stop the reader if it's started. Will return false if it was already stopped.
-	StopReader() bool
+	// StopReader will stop the reader if it's started.
+	StopReader()
 }
 
 // NewReader returns a new instance of Reader
 func NewReader(client sensor.ComplianceService_CommunicateClient, nodeName string, clusterID string, startState *storage.AuditLogFileState) Reader {
 	return &auditLogReaderImpl{
 		logPath:    defaultLogPath,
-		stopC:      concurrency.NewSignal(),
+		stopper:    concurrency.NewStopper(),
 		sender:     newAuditLogSender(client, nodeName, clusterID),
 		startState: startState,
 	}