Fix possibly failing to reserve slots due to stale issued number (#1870)

Author: Sushisource

internal/tuning.go

@@ -78,7 +78,8 @@ type SlotReservationInfo interface {
 	WorkerBuildId() string
 	// WorkerBuildId returns the build ID of the worker that is reserving the slot.
 	WorkerIdentity() string
-	// NumIssuedSlots returns the number of slots that have already been issued by the supplier.
+	// NumIssuedSlots returns the current number of slots that have already been issued by the
+	// supplier. This value may change over the course of the reservation.
 	NumIssuedSlots() int
 	// Logger returns an appropriately tagged logger.
 	Logger() log.Logger
@@ -328,7 +329,7 @@ type slotReserveInfoImpl struct {
 	taskQueue      string
 	workerBuildId  string
 	workerIdentity string
-	issuedSlots    int
+	issuedSlots    *atomic.Int32
 	logger         log.Logger
 	metrics        metrics.Handler
 }
@@ -346,7 +347,7 @@ func (s slotReserveInfoImpl) WorkerIdentity() string {
 }
 
 func (s slotReserveInfoImpl) NumIssuedSlots() int {
-	return s.issuedSlots
+	return int(s.issuedSlots.Load())
 }
 
 func (s slotReserveInfoImpl) Logger() log.Logger {
@@ -442,7 +443,7 @@ func (t *trackingSlotSupplier) ReserveSlot(
 		taskQueue:      data.taskQueue,
 		workerBuildId:  t.workerBuildId,
 		workerIdentity: t.workerIdentity,
-		issuedSlots:    int(t.issuedSlotsAtomic.Load()),
+		issuedSlots:    &t.issuedSlotsAtomic,
 		logger:         t.logger,
 		metrics:        t.metrics,
 	})
@@ -465,7 +466,7 @@ func (t *trackingSlotSupplier) TryReserveSlot(data *slotReservationData) *SlotPe
 		taskQueue:      data.taskQueue,
 		workerBuildId:  t.workerBuildId,
 		workerIdentity: t.workerIdentity,
-		issuedSlots:    int(t.issuedSlotsAtomic.Load()),
+		issuedSlots:    &t.issuedSlotsAtomic,
 		logger:         t.logger,
 		metrics:        t.metrics,
 	})

test/worker_tuner_test.go

@@ -100,6 +100,35 @@ func (ts *WorkerTunerTestSuite) TestCompositeWorkerTuner() {
 	ts.runTheWorkflow(tuner, ctx)
 }
 
+func (ts *WorkerTunerTestSuite) TestResourceBasedSmallSlots() {
+	ctx, cancel := context.WithTimeout(context.Background(), ctxTimeout)
+	defer cancel()
+
+	wfSS, err := worker.NewFixedSizeSlotSupplier(10)
+	ts.NoError(err)
+	controllerOpts := resourcetuner.DefaultResourceControllerOptions()
+	controllerOpts.MemTargetPercent = 0.8
+	controllerOpts.CpuTargetPercent = 0.9
+	controller := resourcetuner.NewResourceController(controllerOpts)
+	actSS, err := resourcetuner.NewResourceBasedSlotSupplier(controller,
+		resourcetuner.ResourceBasedSlotSupplierOptions{
+			MinSlots:     1,
+			MaxSlots:     4,
+			RampThrottle: 0,
+		})
+	ts.NoError(err)
+	laCss, err := worker.NewFixedSizeSlotSupplier(5)
+	ts.NoError(err)
+	tuner, err := worker.NewCompositeTuner(worker.CompositeTunerOptions{
+		WorkflowSlotSupplier: wfSS, ActivitySlotSupplier: actSS, LocalActivitySlotSupplier: laCss})
+	ts.NoError(err)
+
+	// The bug this is verifying was triggered by a race, so run this a bunch to verify it's not hit
+	for i := 0; i < 10; i++ {
+		ts.runTheWorkflow(tuner, ctx)
+	}
+}
+
 func (ts *WorkerTunerTestSuite) runTheWorkflow(tuner worker.WorkerTuner, ctx context.Context) {
 	workerOptions := worker.Options{Tuner: tuner}
 	myWorker := worker.New(ts.client, ts.taskQueueName, workerOptions)