add ConcurrencyLimit to worker to enable dynamic tuning of concurrencies (#1410)

What changed?

[High Risk]

replaced buffered channel with resizable semaphore to control task concurrency
[Low Risk]

added worker package for modularity
added ConcurrencyLimit entity to worker
removed unused methods of autoscaler interface
Why?

needed as first step to enable dynamic tuning of poller and task concurrencies

How did you test it?

Unit Test

Author: shijiesheng

internal/common/autoscaler/autoscaler.go

@@ -24,12 +24,6 @@ package autoscaler
 type (
 	AutoScaler interface {
 		Estimator
-		// Acquire X ResourceUnit of resource
-		Acquire(ResourceUnit) error
-		// Release X ResourceUnit of resource
-		Release(ResourceUnit)
-		// GetCurrent ResourceUnit of resource
-		GetCurrent() ResourceUnit
 		// Start starts the autoscaler go routine that scales the ResourceUnit according to Estimator
 		Start()
 		// Stop stops the autoscaler if started or do nothing if not yet started

internal/internal_poller_autoscaler.go

@@ -26,11 +26,11 @@ import (
 	"sync"
 	"time"
 
-	"github.com/marusama/semaphore/v2"
 	"go.uber.org/atomic"
 	"go.uber.org/zap"
 
 	"go.uber.org/cadence/internal/common/autoscaler"
+	"go.uber.org/cadence/internal/worker"
 )
 
 // defaultPollerScalerCooldownInSeconds
@@ -53,7 +53,7 @@ type (
 		isDryRun     bool
 		cooldownTime time.Duration
 		logger       *zap.Logger
-		sem          semaphore.Semaphore // resizable semaphore to control number of concurrent pollers
+		permit       worker.Permit
 		ctx          context.Context
 		cancel       context.CancelFunc
 		wg           *sync.WaitGroup // graceful stop
@@ -82,6 +82,7 @@ type (
 func newPollerScaler(
 	options pollerAutoScalerOptions,
 	logger *zap.Logger,
+	permit worker.Permit,
 	hooks ...func()) *pollerAutoScaler {
 	if !options.Enabled {
 		return nil
@@ -91,7 +92,7 @@ func newPollerScaler(
 		isDryRun:             options.DryRun,
 		cooldownTime:         options.Cooldown,
 		logger:               logger,
-		sem:                  semaphore.New(options.InitCount),
+		permit:               permit,
 		wg:                   &sync.WaitGroup{},
 		ctx:                  ctx,
 		cancel:               cancel,
@@ -107,21 +108,6 @@ func newPollerScaler(
 	}
 }
 
-// Acquire concurrent poll quota
-func (p *pollerAutoScaler) Acquire(resource autoscaler.ResourceUnit) error {
-	return p.sem.Acquire(p.ctx, int(resource))
-}
-
-// Release concurrent poll quota
-func (p *pollerAutoScaler) Release(resource autoscaler.ResourceUnit) {
-	p.sem.Release(int(resource))
-}
-
-// GetCurrent poll quota
-func (p *pollerAutoScaler) GetCurrent() autoscaler.ResourceUnit {
-	return autoscaler.ResourceUnit(p.sem.GetLimit())
-}
-
 // Start an auto-scaler go routine and returns a done to stop it
 func (p *pollerAutoScaler) Start() {
 	logger := p.logger.Sugar()
@@ -133,7 +119,7 @@ func (p *pollerAutoScaler) Start() {
 			case <-p.ctx.Done():
 				return
 			case <-time.After(p.cooldownTime):
-				currentResource := autoscaler.ResourceUnit(p.sem.GetLimit())
+				currentResource := autoscaler.ResourceUnit(p.permit.Quota())
 				currentUsages, err := p.pollerUsageEstimator.Estimate()
 				if err != nil {
 					logger.Warnw("poller autoscaler skip due to estimator error", "error", err)
@@ -146,7 +132,7 @@ func (p *pollerAutoScaler) Start() {
 					"recommend", uint64(proposedResource),
 					"isDryRun", p.isDryRun)
 				if !p.isDryRun {
-					p.sem.SetLimit(int(proposedResource))
+					p.permit.SetQuota(int(proposedResource))
 				}
 				p.pollerUsageEstimator.Reset()
 

internal/internal_poller_autoscaler_test.go

@@ -21,12 +21,14 @@
 package internal
 
 import (
+	"context"
 	"math/rand"
 	"sync"
 	"testing"
 	"time"
 
 	"go.uber.org/cadence/internal/common/testlogger"
+	"go.uber.org/cadence/internal/worker"
 
 	"github.com/stretchr/testify/assert"
 	"go.uber.org/atomic"
@@ -171,6 +173,7 @@ func Test_pollerAutoscaler(t *testing.T) {
 					TargetUtilization: float64(tt.args.targetMilliUsage) / 1000,
 				},
 				testlogger.NewZap(t),
+				worker.NewResizablePermit(tt.args.initialPollerCount),
 				// hook function that collects number of iterations
 				func() {
 					autoscalerEpoch.Add(1)
@@ -190,18 +193,19 @@ func Test_pollerAutoscaler(t *testing.T) {
 				go func() {
 					defer wg.Done()
 					for pollResult := range pollChan {
-						pollerScaler.Acquire(1)
+						err := pollerScaler.permit.Acquire(context.Background())
+						assert.NoError(t, err)
 						pollerScaler.CollectUsage(pollResult)
-						pollerScaler.Release(1)
+						pollerScaler.permit.Release()
 					}
 				}()
 			}
 
 			assert.Eventually(t, func() bool {
 				return autoscalerEpoch.Load() == uint64(tt.args.autoScalerEpoch)
-			}, tt.args.cooldownTime+20*time.Millisecond, 10*time.Millisecond)
+			}, tt.args.cooldownTime+100*time.Millisecond, 10*time.Millisecond)
 			pollerScaler.Stop()
-			res := pollerScaler.GetCurrent()
+			res := pollerScaler.permit.Quota() - pollerScaler.permit.Count()
 			assert.Equal(t, tt.want, int(res))
 		})
 	}

internal/internal_worker_base.go

@@ -33,6 +33,7 @@ import (
 	"time"
 
 	"go.uber.org/cadence/internal/common/debug"
+	"go.uber.org/cadence/internal/worker"
 
 	"github.com/uber-go/tally"
 	"go.uber.org/zap"
@@ -141,7 +142,7 @@ type (
 		logger               *zap.Logger
 		metricsScope         tally.Scope
 
-		pollerRequestCh    chan struct{}
+		concurrency        *worker.ConcurrencyLimit
 		pollerAutoScaler   *pollerAutoScaler
 		taskQueueCh        chan interface{}
 		sessionTokenBucket *sessionTokenBucket
@@ -167,11 +168,17 @@ func createPollRetryPolicy() backoff.RetryPolicy {
 func newBaseWorker(options baseWorkerOptions, logger *zap.Logger, metricsScope tally.Scope, sessionTokenBucket *sessionTokenBucket) *baseWorker {
 	ctx, cancel := context.WithCancel(context.Background())
 
+	concurrency := &worker.ConcurrencyLimit{
+		PollerPermit: worker.NewResizablePermit(options.pollerCount),
+		TaskPermit:   worker.NewResizablePermit(options.maxConcurrentTask),
+	}
+
 	var pollerAS *pollerAutoScaler
 	if pollerOptions := options.pollerAutoScaler; pollerOptions.Enabled {
 		pollerAS = newPollerScaler(
 			pollerOptions,
 			logger,
+			concurrency.PollerPermit,
 		)
 	}
 
@@ -182,7 +189,7 @@ func newBaseWorker(options baseWorkerOptions, logger *zap.Logger, metricsScope t
 		retrier:              backoff.NewConcurrentRetrier(pollOperationRetryPolicy),
 		logger:               logger.With(zapcore.Field{Key: tagWorkerType, Type: zapcore.StringType, String: options.workerType}),
 		metricsScope:         tagScope(metricsScope, tagWorkerType, options.workerType),
-		pollerRequestCh:      make(chan struct{}, options.maxConcurrentTask),
+		concurrency:          concurrency,
 		pollerAutoScaler:     pollerAS,
 		taskQueueCh:          make(chan interface{}), // no buffer, so poller only able to poll new task after previous is dispatched.
 		limiterContext:       ctx,
@@ -241,14 +248,19 @@ func (bw *baseWorker) runPoller() {
 	bw.metricsScope.Counter(metrics.PollerStartCounter).Inc(1)
 
 	for {
+		permitChannel, channelDone := bw.concurrency.TaskPermit.AcquireChan(bw.limiterContext)
 		select {
 		case <-bw.shutdownCh:
+			channelDone()
 			return
-		case <-bw.pollerRequestCh:
-			bw.metricsScope.Gauge(metrics.ConcurrentTaskQuota).Update(float64(cap(bw.pollerRequestCh)))
-			// This metric is used to monitor how many poll requests have been allocated
-			// and can be used to approximate number of concurrent task running (not pinpoint accurate)
-			bw.metricsScope.Gauge(metrics.PollerRequestBufferUsage).Update(float64(cap(bw.pollerRequestCh) - len(bw.pollerRequestCh)))
+		case <-permitChannel: // don't poll unless there is a task permit
+			channelDone()
+			// TODO move to a centralized place inside the worker
+			// emit metrics on concurrent task permit quota and current task permit count
+			// NOTE task permit doesn't mean there is a task running, it still needs to poll until it gets a task to process
+			// thus the metrics is only an estimated value of how many tasks are running concurrently
+			bw.metricsScope.Gauge(metrics.ConcurrentTaskQuota).Update(float64(bw.concurrency.TaskPermit.Quota()))
+			bw.metricsScope.Gauge(metrics.PollerRequestBufferUsage).Update(float64(bw.concurrency.TaskPermit.Count()))
 			if bw.sessionTokenBucket != nil {
 				bw.sessionTokenBucket.waitForAvailableToken()
 			}
@@ -260,10 +272,6 @@ func (bw *baseWorker) runPoller() {
 func (bw *baseWorker) runTaskDispatcher() {
 	defer bw.shutdownWG.Done()
 
-	for i := 0; i < bw.options.maxConcurrentTask; i++ {
-		bw.pollerRequestCh <- struct{}{}
-	}
-
 	for {
 		// wait for new task or shutdown
 		select {
@@ -294,10 +302,10 @@ func (bw *baseWorker) pollTask() {
 	var task interface{}
 
 	if bw.pollerAutoScaler != nil {
-		if pErr := bw.pollerAutoScaler.Acquire(1); pErr == nil {
-			defer bw.pollerAutoScaler.Release(1)
+		if pErr := bw.concurrency.PollerPermit.Acquire(bw.limiterContext); pErr == nil {
+			defer bw.concurrency.PollerPermit.Release()
 		} else {
-			bw.logger.Warn("poller auto scaler acquire error", zap.Error(pErr))
+			bw.logger.Warn("poller permit acquire error", zap.Error(pErr))
 		}
 	}
 
@@ -333,7 +341,7 @@ func (bw *baseWorker) pollTask() {
 		case <-bw.shutdownCh:
 		}
 	} else {
-		bw.pollerRequestCh <- struct{}{} // poll failed, trigger a new poll
+		bw.concurrency.TaskPermit.Release() // poll failed, trigger a new poll by returning a task permit
 	}
 }
 
@@ -368,7 +376,7 @@ func (bw *baseWorker) processTask(task interface{}) {
 		}
 
 		if isPolledTask {
-			bw.pollerRequestCh <- struct{}{}
+			bw.concurrency.TaskPermit.Release() // task processed, trigger a new poll by returning a task permit
 		}
 	}()
 	err := bw.options.taskWorker.ProcessTask(task)

internal/worker/concurrency.go

@@ -0,0 +1,41 @@
+// Copyright (c) 2017-2021 Uber Technologies Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package worker
+
+import "context"
+
+var _ Permit = (*resizablePermit)(nil)
+
+// ConcurrencyLimit contains synchronization primitives for dynamically controlling the concurrencies in workers
+type ConcurrencyLimit struct {
+	PollerPermit Permit // controls concurrency of pollers
+	TaskPermit   Permit // controls concurrency of task processing
+}
+
+// Permit is an adaptive permit issuer to control concurrency
+type Permit interface {
+	Acquire(context.Context) error
+	AcquireChan(context.Context) (channel <-chan struct{}, done func())
+	Count() int
+	Quota() int
+	Release()
+	SetQuota(int)
+}