propagate context

Author: shijiesheng

internal/internal_worker_base.go

@@ -252,7 +252,7 @@ func (bw *baseWorker) runPoller() {
 		select {
 		case <-bw.shutdownCh:
 			return
-		case <-bw.dynamic.TaskPermit.AcquireChan(): // don't poll unless there is a task permit
+		case <-bw.dynamic.TaskPermit.AcquireChan(bw.limiterContext): // don't poll unless there is a task permit
 			// 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/dynamic_params.go

@@ -38,7 +38,7 @@ type DynamicParams struct {
 // Permit is an adaptive
 type Permit interface {
 	Acquire(context.Context, int) error
-	AcquireChan() <-chan struct{}
+	AcquireChan(context.Context) <-chan struct{}
 	Quota() int
 	SetQuota(int)
 	Count() int
@@ -63,10 +63,10 @@ func (p *permit) Acquire(ctx context.Context, count int) error {
 }
 
 // AcquireChan returns a permit ready channel. It's closed then permit is acquired
-func (p *permit) AcquireChan() <-chan struct{} {
+func (p *permit) AcquireChan(ctx context.Context) <-chan struct{} {
 	ch := make(chan struct{})
 	go func() {
-		if err := p.sem.Acquire(nil, 1); err != nil { // nil context indicates no need to exit on context done
+		if err := p.sem.Acquire(ctx, 1); err != nil {
 			close(ch)
 			return
 		}