add PermitChannel

Author: shijiesheng

internal/internal_worker_base.go

@@ -170,12 +170,11 @@ func newBaseWorker(options baseWorkerOptions, logger *zap.Logger, metricsScope t
 
 	concurrency := &worker.ConcurrencyLimit{
 		PollerPermit: worker.NewResizablePermit(options.pollerCount),
-		TaskPermit:   worker.NewChannelPermit(options.maxConcurrentTask),
+		TaskPermit:   worker.NewResizablePermit(options.maxConcurrentTask),
 	}
 
 	var pollerAS *pollerAutoScaler
 	if pollerOptions := options.pollerAutoScaler; pollerOptions.Enabled {
-		concurrency.PollerPermit = worker.NewResizablePermit(pollerOptions.InitCount)
 		pollerAS = newPollerScaler(
 			pollerOptions,
 			logger,
@@ -249,10 +248,14 @@ func (bw *baseWorker) runPoller() {
 	bw.metricsScope.Counter(metrics.PollerStartCounter).Inc(1)
 
 	for {
+		// permitChannel can be blocking without passing context because shutdownCh is used
+		permitChannel := bw.concurrency.PollerPermit.AcquireChan(context.Background())
 		select {
 		case <-bw.shutdownCh:
+			permitChannel.Close()
 			return
-		case <-bw.concurrency.TaskPermit.GetChan(): // don't poll unless there is a task permit
+		case <-permitChannel.C(): // don't poll unless there is a task permit
+			permitChannel.Close()
 			// 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

internal/worker/channel_permit.go

@@ -23,19 +23,26 @@ package worker
 import "context"
 
 var _ Permit = (*resizablePermit)(nil)
-var _ ChannelPermit = (*channelPermit)(nil)
 
 // ConcurrencyLimit contains synchronization primitives for dynamically controlling the concurrencies in workers
 type ConcurrencyLimit struct {
-	PollerPermit Permit        // controls concurrency of pollers
-	TaskPermit   ChannelPermit // controls concurrency of task processing
+	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) PermitChannel
 	Count() int
 	Quota() int
 	Release()
 	SetQuota(int)
 }
+
+// PermitChannel is a channel that can be used to wait for a permit to be available
+// Remember to call Close() to avoid goroutine leak
+type PermitChannel interface {
+	C() <-chan struct{}
+	Close()
+}

internal/worker/concurrency.go

@@ -23,6 +23,7 @@ package worker
 import (
 	"context"
 	"fmt"
+	"sync"
 
 	"github.com/marusama/semaphore/v2"
 )
@@ -45,14 +46,17 @@ func (p *resizablePermit) Acquire(ctx context.Context) error {
 	return nil
 }
 
+// Release release one permit
 func (p *resizablePermit) Release() {
 	p.sem.Release(1)
 }
 
+// Quota returns the maximum number of permits that can be acquired
 func (p *resizablePermit) Quota() int {
 	return p.sem.GetLimit()
 }
 
+// SetQuota sets the maximum number of permits that can be acquired
 func (p *resizablePermit) SetQuota(c int) {
 	p.sem.SetLimit(c)
 }
@@ -61,3 +65,53 @@ func (p *resizablePermit) SetQuota(c int) {
 func (p *resizablePermit) Count() int {
 	return p.sem.GetCount()
 }
+
+// AcquireChan creates a PermitChannel that can be used to wait for a permit
+// After usage:
+// 1. avoid goroutine leak by calling permitChannel.Close()
+// 2. release permit by calling permit.Release()
+func (p *resizablePermit) AcquireChan(ctx context.Context) PermitChannel {
+	ctx, cancel := context.WithCancel(ctx)
+	pc := &permitChannel{
+		p:      p,
+		c:      make(chan struct{}),
+		ctx:    ctx,
+		cancel: cancel,
+		wg:     &sync.WaitGroup{},
+	}
+	pc.start()
+	return pc
+}
+
+type permitChannel struct {
+	p      Permit
+	c      chan struct{}
+	ctx    context.Context
+	cancel context.CancelFunc
+	wg     *sync.WaitGroup
+}
+
+func (ch *permitChannel) C() <-chan struct{} {
+	return ch.c
+}
+
+func (ch *permitChannel) start() {
+	ch.wg.Add(1)
+	go func() {
+		defer ch.wg.Done()
+		if err := ch.p.Acquire(ch.ctx); err != nil {
+			return
+		}
+		// avoid blocking on sending to the channel
+		select {
+		case ch.c <- struct{}{}:
+		case <-ch.ctx.Done(): // release if acquire is successful but fail to send to the channel
+			ch.p.Release()
+		}
+	}()
+}
+
+func (ch *permitChannel) Close() {
+	ch.cancel()
+	ch.wg.Wait()
+}

internal/worker/concurrency_auto_scaler.go

@@ -99,7 +99,9 @@ func TestPermit_Simulation(t *testing.T) {
 			failures := atomic.NewInt32(0)
 			ctx, cancel := context.WithTimeout(context.Background(), tt.maxTestDuration)
 			defer cancel()
-			for i := 0; i < tt.goroutines; i++ {
+
+			aquireChan := tt.goroutines / 2
+			for i := 0; i < tt.goroutines-aquireChan; i++ {
 				wg.Add(1)
 				go func() {
 					defer wg.Done()
@@ -111,10 +113,25 @@ func TestPermit_Simulation(t *testing.T) {
 					permit.Release()
 				}()
 			}
+			for i := 0; i < aquireChan; i++ {
+				wg.Add(1)
+				go func() {
+					defer wg.Done()
+					permitChan := permit.AcquireChan(ctx)
+					select {
+					case <-permitChan.C():
+						time.Sleep(time.Duration(100+rand.Intn(50)) * time.Millisecond)
+						permit.Release()
+					case <-ctx.Done():
+						failures.Inc()
+					}
+					permitChan.Close()
+				}()
+			}
 
 			wg.Wait()
 			// sanity check
-			assert.Equal(t, 0, permit.Count())
+			assert.Equal(t, 0, permit.Count(), "all permit should be released")
 			assert.Equal(t, tt.capacity[len(tt.capacity)-1], permit.Quota())
 
 			// expect failures in range