fix: run maxNumberOfPollers poller routines instead of initialNumberOfPollers (#2105)

* Update internal_worker_base.go

* scalableTaskPoller concurrency scaling tests

Author: jjaegerman

internal/internal_worker_base.go

@@ -885,7 +885,7 @@ func newScalableTaskPoller(
 	}
 	switch p := pollerBehavior.(type) {
 	case *pollerBehaviorAutoscaling:
-		tw.pollerCount = p.initialNumberOfPollers
+		tw.pollerCount = p.maximumNumberOfPollers
 		tw.pollerSemaphore = newPollerSemaphore(p.initialNumberOfPollers)
 		tw.pollerAutoscalerReportHandle = newPollScalerReportHandle(pollScalerReportHandleOptions{
 			initialPollerCount: p.initialNumberOfPollers,

internal/internal_worker_base_test.go

@@ -1,12 +1,18 @@
 package internal
 
 import (
+	"context"
+	"sync/atomic"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
 	enumspb "go.temporal.io/api/enums/v1"
 	"go.temporal.io/api/serviceerror"
+	"go.temporal.io/sdk/internal/common/metrics"
+	ilog "go.temporal.io/sdk/internal/log"
 )
 
 type (
@@ -19,6 +25,16 @@ func TestPollScalerReportHandleSuite(t *testing.T) {
 	suite.Run(t, new(PollScalerReportHandleSuite))
 }
 
+type (
+	ScalableTaskPollerSuite struct {
+		suite.Suite
+	}
+)
+
+func TestScalableTaskPollerSuite(t *testing.T) {
+	suite.Run(t, new(ScalableTaskPollerSuite))
+}
+
 type testTask struct {
 	psd pollerScaleDecision
 }
@@ -115,3 +131,196 @@ func (s *PollScalerReportHandleSuite) TestScaleUpOnDelay() {
 	assert.Equal(s.T(), 10, targetSuggestion)
 
 }
+
+func (s *ScalableTaskPollerSuite) TestAutoscalingConcurrencyScalesUpToMaximum() {
+	behavior := &pollerBehaviorAutoscaling{
+		initialNumberOfPollers: 2,
+		maximumNumberOfPollers: 3,
+		minimumNumberOfPollers: 1,
+	}
+
+	blockingPoller := newSemaphoreProbeTaskPoller()
+	poller := newScalableTaskPoller(blockingPoller, ilog.NewNopLogger(), behavior)
+	poller.taskPollerType = "test"
+
+	bw := newBaseWorker(baseWorkerOptions{
+		slotSupplier:     &testSlotSupplier{},
+		maxTaskPerSecond: 1000,
+		taskPollers:      []scalableTaskPoller{poller},
+		taskProcessor:    noopTaskProcessor{},
+		workerType:       "AutoscalingTest",
+		logger:           ilog.NewNopLogger(),
+		stopTimeout:      time.Second,
+		metricsHandler:   metrics.NopHandler,
+	})
+
+	bw.Start()
+	defer func() {
+		allowBlockedPollers(blockingPoller, poller.pollerSemaphore)
+		blockingPoller.Close()
+		bw.Stop()
+	}()
+
+	eventuallySemaphoreState(s.T(), blockingPoller, poller.pollerSemaphore, 2, 2, "expected initial poller to start")
+
+	require.Never(s.T(), func() bool {
+		allowBlockedPollers(blockingPoller, poller.pollerSemaphore)
+		permits, _ := readSemaphoreState(poller.pollerSemaphore)
+		return permits > 2
+	}, 200*time.Millisecond, 10*time.Millisecond, "should not exceed initial concurrency")
+
+	poller.pollerAutoscalerReportHandle.updateTarget(func(int64) int64 { return 3 })
+
+	eventuallySemaphoreState(s.T(), blockingPoller, poller.pollerSemaphore, 3, 3, "expected concurrency to scale up to maximum")
+
+	require.Never(s.T(), func() bool {
+		allowBlockedPollers(blockingPoller, poller.pollerSemaphore)
+		permits, _ := readSemaphoreState(poller.pollerSemaphore)
+		return permits > 3
+	}, 200*time.Millisecond, 10*time.Millisecond, "should not exceed maximum concurrency")
+}
+
+func (s *ScalableTaskPollerSuite) TestAutoscalingScalesDownToMinimum() {
+	behavior := &pollerBehaviorAutoscaling{
+		initialNumberOfPollers: 2,
+		maximumNumberOfPollers: 3,
+		minimumNumberOfPollers: 1,
+	}
+
+	blockingPoller := newSemaphoreProbeTaskPoller()
+	poller := newScalableTaskPoller(blockingPoller, ilog.NewNopLogger(), behavior)
+	poller.taskPollerType = "test"
+
+	bw := newBaseWorker(baseWorkerOptions{
+		slotSupplier:     &testSlotSupplier{},
+		maxTaskPerSecond: 1000,
+		taskPollers:      []scalableTaskPoller{poller},
+		taskProcessor:    noopTaskProcessor{},
+		workerType:       "AutoscalingTest",
+		logger:           ilog.NewNopLogger(),
+		stopTimeout:      time.Second,
+		metricsHandler:   metrics.NopHandler,
+	})
+
+	bw.Start()
+	defer func() {
+		allowBlockedPollers(blockingPoller, poller.pollerSemaphore)
+		blockingPoller.Close()
+		bw.Stop()
+	}()
+
+	eventuallySemaphoreState(s.T(), blockingPoller, poller.pollerSemaphore, 2, 2, "expected initial concurrency")
+
+	poller.pollerAutoscalerReportHandle.updateTarget(func(target int64) int64 { return 1 })
+
+	eventuallySemaphoreState(s.T(), blockingPoller, poller.pollerSemaphore, 1, 1, "expected concurrency to reduce to minimum")
+
+	require.Never(s.T(), func() bool {
+		allowBlockedPollers(blockingPoller, poller.pollerSemaphore)
+		permits, _ := readSemaphoreState(poller.pollerSemaphore)
+		return permits == 0
+	}, 200*time.Millisecond, 10*time.Millisecond, "should not scale below minimum")
+}
+
+type semaphoreProbeTaskPoller struct {
+	signals chan struct{}
+	closed  atomic.Bool
+}
+
+func newSemaphoreProbeTaskPoller() *semaphoreProbeTaskPoller {
+	return &semaphoreProbeTaskPoller{
+		signals: make(chan struct{}, 32),
+	}
+}
+
+// PollTask implements taskPoller and blocks until a signal is provided so the semaphore permits stay acquired.
+func (p *semaphoreProbeTaskPoller) PollTask() (taskForWorker, error) {
+	_, ok := <-p.signals
+	if !ok {
+		return nil, nil
+	}
+	return nil, nil
+}
+
+// Cleanup implements taskPoller.
+func (p *semaphoreProbeTaskPoller) Cleanup() error {
+	p.Close()
+	return nil
+}
+
+func (p *semaphoreProbeTaskPoller) Allow(n int) {
+	for range n {
+		for {
+			if p.closed.Load() {
+				return
+			}
+			select {
+			case p.signals <- struct{}{}:
+				goto next
+			default:
+				time.Sleep(1 * time.Millisecond)
+			}
+		}
+	next:
+	}
+}
+
+func (p *semaphoreProbeTaskPoller) Close() {
+	if p.closed.CompareAndSwap(false, true) {
+		close(p.signals)
+	}
+}
+
+func allowBlockedPollers(p *semaphoreProbeTaskPoller, sem *pollerSemaphore) {
+	if p == nil || sem == nil {
+		return
+	}
+	permits, _ := readSemaphoreState(sem)
+	if permits > 0 {
+		p.Allow(permits)
+	}
+}
+
+func eventuallySemaphoreState(t *testing.T, blockingPoller *semaphoreProbeTaskPoller, sem *pollerSemaphore, expectedPermits, expectedMax int, msg string) {
+	require.Eventually(t, func() bool {
+		allowBlockedPollers(blockingPoller, sem)
+		permits, max := readSemaphoreState(sem)
+		return permits == expectedPermits && max == expectedMax
+	}, time.Second, 10*time.Millisecond, msg)
+}
+
+func readSemaphoreState(ps *pollerSemaphore) (permits int, max int) {
+	if ps == nil {
+		return 0, 0
+	}
+	barrier := <-ps.bs
+	permits = ps.permits
+	max = ps.maxPermits
+	ps.bs <- barrier
+	return
+}
+
+type testSlotSupplier struct{}
+
+func (s *testSlotSupplier) ReserveSlot(ctx context.Context, info SlotReservationInfo) (*SlotPermit, error) {
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	default:
+	}
+	return &SlotPermit{}, nil
+}
+
+func (s *testSlotSupplier) TryReserveSlot(SlotReservationInfo) *SlotPermit {
+	return &SlotPermit{}
+}
+
+func (s *testSlotSupplier) MarkSlotUsed(SlotMarkUsedInfo) {}
+
+func (s *testSlotSupplier) ReleaseSlot(SlotReleaseInfo) {}
+
+func (s *testSlotSupplier) MaxSlots() int { return 0 }
+
+type noopTaskProcessor struct{}
+
+func (noopTaskProcessor) ProcessTask(any) error { return nil }