github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/dm/master/agent_pool.go (about)

     1  // Copyright 2019 PingCAP, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package master
    15  
    16  import (
    17  	"context"
    18  	"math"
    19  
    20  	"github.com/pingcap/tiflow/dm/pkg/log"
    21  	"go.uber.org/zap"
    22  	"golang.org/x/time/rate"
    23  )
    24  
    25  // rate limit related constant value.
    26  const (
    27  	DefaultRate      = float64(10)
    28  	DefaultBurst     = 40
    29  	ErrorNoEmitToken = "fail to get emit opportunity for %s"
    30  )
    31  
    32  type emitFunc func(args ...interface{})
    33  
    34  // AgentPool is a pool to control communication with dm-workers
    35  // It provides rate limit control for agent acquire, including dispatch rate r
    36  // and permits bursts of at most b tokens.
    37  // caller shouldn't to hold agent to avoid deadlock.
    38  type AgentPool struct {
    39  	requests chan int
    40  	agents   chan *Agent
    41  	cfg      *RateLimitConfig
    42  	limiter  *rate.Limiter
    43  }
    44  
    45  // RateLimitConfig holds rate limit config.
    46  type RateLimitConfig struct {
    47  	rate  float64 // dispatch rate
    48  	burst int     // max permits bursts
    49  }
    50  
    51  // Agent communicate with dm-workers.
    52  type Agent struct {
    53  	ID int
    54  }
    55  
    56  // NewAgentPool returns a agent pool.
    57  func NewAgentPool(cfg *RateLimitConfig) *AgentPool {
    58  	requests := make(chan int, int(math.Ceil(1/cfg.rate))+cfg.burst)
    59  	agents := make(chan *Agent, cfg.burst)
    60  	limiter := rate.NewLimiter(rate.Limit(cfg.rate), cfg.burst)
    61  
    62  	return &AgentPool{
    63  		requests: requests,
    64  		agents:   agents,
    65  		cfg:      cfg,
    66  		limiter:  limiter,
    67  	}
    68  }
    69  
    70  // Apply applies for a agent
    71  // if ctx is canceled before we get an agent, returns nil.
    72  func (ap *AgentPool) Apply(ctx context.Context, id int) *Agent {
    73  	select {
    74  	case <-ctx.Done():
    75  		return nil
    76  	case ap.requests <- id:
    77  	}
    78  
    79  	select {
    80  	case <-ctx.Done():
    81  		return nil
    82  	case agent := <-ap.agents:
    83  		return agent
    84  	}
    85  }
    86  
    87  // Start starts AgentPool background dispatcher.
    88  func (ap *AgentPool) Start(ctx context.Context) {
    89  	for {
    90  		select {
    91  		case <-ctx.Done():
    92  			return
    93  		case id := <-ap.requests:
    94  			err := ap.limiter.Wait(ctx)
    95  			if err != nil {
    96  				if err != context.Canceled {
    97  					log.L().Fatal("agent limiter wait meets unexpected error", zap.Error(err))
    98  				}
    99  				return
   100  			}
   101  			select {
   102  			case <-ctx.Done():
   103  				return
   104  			case ap.agents <- &Agent{ID: id}:
   105  			}
   106  		}
   107  	}
   108  }
   109  
   110  // Emit applies for an agent to communicates with dm-worker.
   111  func (ap *AgentPool) Emit(ctx context.Context, id int, fn emitFunc, errFn emitFunc, args ...interface{}) {
   112  	if agent := ap.Apply(ctx, id); agent == nil {
   113  		errFn(args...)
   114  	} else {
   115  		fn(args...)
   116  	}
   117  }