satellite/audit,storagenode/gracefulexit: fixes to limiter

Ensure we don't rely on limiter to wait multiple times.

Change-Id: I75d48420236216d4c2fc6fa99293f51f80cd9c33

satellite/audit/worker.go

@@ -38,7 +38,7 @@ type Worker struct {
 	verifier    *Verifier
 	reporter    Reporter
 	Loop        *sync2.Cycle
-	limiter  *sync2.Limiter
+	concurrency int
 }
 
 // NewWorker instantiates Worker.
@@ -50,7 +50,7 @@ func NewWorker(log *zap.Logger, queues *Queues, verifier *Verifier, reporter Rep
 		verifier:    verifier,
 		reporter:    reporter,
 		Loop:        sync2.NewCycle(config.QueueInterval),
-		limiter:  sync2.NewLimiter(config.WorkerConcurrency),
+		concurrency: config.WorkerConcurrency,
 	}, nil
 }
 
@@ -58,9 +58,6 @@ func NewWorker(log *zap.Logger, queues *Queues, verifier *Verifier, reporter Rep
 func (worker *Worker) Run(ctx context.Context) (err error) {
 	defer mon.Task()(&ctx)(&err)
 
-	// Wait for all audits to run.
-	defer worker.limiter.Wait()
-
 	return worker.Loop.Run(ctx, func(ctx context.Context) (err error) {
 		defer mon.Task()(&ctx)(&err)
 		err = worker.process(ctx)
@@ -84,7 +81,9 @@ func (worker *Worker) process(ctx context.Context) (err error) {
 	// get the current queue
 	queue := worker.queues.Fetch()
 
-	worker.limiter.Wait()
+	limiter := sync2.NewLimiter(worker.concurrency)
+	defer limiter.Wait()
+
 	for {
 		segment, err := queue.Next()
 		if err != nil {
@@ -99,7 +98,7 @@ func (worker *Worker) process(ctx context.Context) (err error) {
 			return err
 		}
 
-		worker.limiter.Go(ctx, func() {
+		started := limiter.Go(ctx, func() {
 			err := worker.work(ctx, segment)
 			if err != nil {
 				worker.log.Error("error(s) during audit",
@@ -108,6 +107,9 @@ func (worker *Worker) process(ctx context.Context) (err error) {
 					zap.Error(err))
 			}
 		})
+		if !started {
+			return ctx.Err()
+		}
 	}
 }
 

storagenode/gracefulexit/chore.go

@@ -6,6 +6,7 @@ package gracefulexit
 import (
 	"context"
 	"sync"
+	"time"
 
 	"go.uber.org/zap"
 
@@ -46,8 +47,12 @@ func NewChore(log *zap.Logger, service Service, transferService piecetransfer.Se
 // Run starts the chore.
 func (chore *Chore) Run(ctx context.Context) (err error) {
 	defer mon.Task()(&ctx)(&err)
+	defer chore.limiter.Wait()
+	return chore.Loop.Run(ctx, chore.AddMissing)
+}
 
-	err = chore.Loop.Run(ctx, func(ctx context.Context) (err error) {
+// AddMissing starts any missing satellite chore.
+func (chore *Chore) AddMissing(ctx context.Context) (err error) {
 	defer mon.Task()(&ctx)(&err)
 
 	geSatellites, err := chore.service.ListPendingExits(ctx)
@@ -72,7 +77,7 @@ func (chore *Chore) Run(ctx context.Context) (err error) {
 			continue
 		}
 
-			chore.limiter.Go(ctx, func() {
+		started := chore.limiter.Go(ctx, func() {
 			err := worker.Run(ctx, func() {
 				chore.log.Debug("finished for satellite.", zap.Stringer("Satellite ID", satellite.SatelliteID))
 				chore.exitingMap.Delete(satellite.SatelliteID)
@@ -86,24 +91,36 @@ func (chore *Chore) Run(ctx context.Context) (err error) {
 				chore.log.Error("closing worker failed", zap.Error(err))
 			}
 		})
+		if !started {
+			return ctx.Err()
+		}
 	}
 
 	return nil
-	})
-
-	chore.limiter.Wait()
-
-	return err
 }
 
-// TestWaitForWorkers waits for any pending worker to finish.
-func (chore *Chore) TestWaitForWorkers() {
-	chore.limiter.Wait()
+// TestWaitForNoWorkers waits for any pending worker to finish.
+func (chore *Chore) TestWaitForNoWorkers(ctx context.Context) error {
+	for {
+		if !sync2.Sleep(ctx, 100*time.Millisecond) {
+			return ctx.Err()
+		}
+
+		found := false
+		chore.exitingMap.Range(func(key, value interface{}) bool {
+			found = true
+			return false
+		})
+		if !found {
+			return nil
+		}
+	}
 }
 
 // Close closes chore.
 func (chore *Chore) Close() error {
 	chore.Loop.Close()
+
 	chore.exitingMap.Range(func(key interface{}, value interface{}) bool {
 		worker := value.(*Worker)
 		err := worker.Close()

storagenode/gracefulexit/chore_test.go

@@ -107,7 +107,8 @@ func exitSatellite(ctx context.Context, t *testing.T, planet *testplanet.Planet,
 	// run the SN chore again to start processing transfers.
 	exitingNode.GracefulExit.Chore.Loop.TriggerWait()
 	// wait for workers to finish
-	exitingNode.GracefulExit.Chore.TestWaitForWorkers()
+	err = exitingNode.GracefulExit.Chore.TestWaitForNoWorkers(ctx)
+	require.NoError(t, err)
 
 	// check that there are no more items to process
 	queueItems, err = satellite1.DB.GracefulExit().GetIncomplete(ctx, exitStatus.NodeID, 10, 0)