github.com/pf-qiu/concourse/v6@v6.7.3-0.20201207032516-1f455d73275f/atc/db/resource_check_rate_limiter.go (about)

     1  package db
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"sync"
     7  	"time"
     8  
     9  	"code.cloudfoundry.org/clock"
    10  	"golang.org/x/time/rate"
    11  )
    12  
    13  type ResourceCheckRateLimiter struct {
    14  	checkLimiter *rate.Limiter
    15  
    16  	refreshConn    Conn
    17  	checkInterval  time.Duration
    18  	refreshLimiter *rate.Limiter
    19  
    20  	clock clock.Clock
    21  	mut   *sync.Mutex
    22  }
    23  
    24  func NewResourceCheckRateLimiter(
    25  	checksPerSecond rate.Limit,
    26  	checkInterval time.Duration,
    27  	refreshConn Conn,
    28  	refreshInterval time.Duration,
    29  	clock clock.Clock,
    30  ) *ResourceCheckRateLimiter {
    31  	limiter := &ResourceCheckRateLimiter{
    32  		clock: clock,
    33  		mut:   new(sync.Mutex),
    34  	}
    35  
    36  	if checksPerSecond < 0 {
    37  		checksPerSecond = rate.Inf
    38  	}
    39  
    40  	if checksPerSecond != 0 {
    41  		limiter.checkLimiter = rate.NewLimiter(checksPerSecond, 1)
    42  	} else {
    43  		limiter.checkInterval = checkInterval
    44  		limiter.refreshConn = refreshConn
    45  		limiter.refreshLimiter = rate.NewLimiter(rate.Every(refreshInterval), 1)
    46  	}
    47  
    48  	return limiter
    49  }
    50  
    51  func (limiter *ResourceCheckRateLimiter) Wait(ctx context.Context) error {
    52  	limiter.mut.Lock()
    53  	defer limiter.mut.Unlock()
    54  
    55  	if limiter.refreshLimiter != nil && limiter.refreshLimiter.AllowN(limiter.clock.Now(), 1) {
    56  		err := limiter.refreshCheckLimiter()
    57  		if err != nil {
    58  			return fmt.Errorf("refresh: %w", err)
    59  		}
    60  	}
    61  
    62  	reservation := limiter.checkLimiter.ReserveN(limiter.clock.Now(), 1)
    63  
    64  	delay := reservation.DelayFrom(limiter.clock.Now())
    65  	if delay == 0 {
    66  		return nil
    67  	}
    68  
    69  	timer := limiter.clock.NewTimer(delay)
    70  	defer timer.Stop()
    71  
    72  	select {
    73  	case <-timer.C():
    74  		return nil
    75  	case <-ctx.Done():
    76  		reservation.Cancel()
    77  		return ctx.Err()
    78  	}
    79  }
    80  
    81  func (limiter *ResourceCheckRateLimiter) Limit() rate.Limit {
    82  	limiter.mut.Lock()
    83  	defer limiter.mut.Unlock()
    84  
    85  	return limiter.checkLimiter.Limit()
    86  }
    87  
    88  func (limiter *ResourceCheckRateLimiter) refreshCheckLimiter() error {
    89  	var count int
    90  	err := psql.Select("COUNT(id)").
    91  		From("resource_config_scopes").
    92  		RunWith(limiter.refreshConn).
    93  		QueryRow().
    94  		Scan(&count)
    95  	if err != nil {
    96  		return err
    97  	}
    98  
    99  	limit := rate.Limit(float64(count) / limiter.checkInterval.Seconds())
   100  	if count == 0 {
   101  		// don't bother waiting if there aren't any checkables
   102  		limit = rate.Inf
   103  	}
   104  
   105  	if limiter.checkLimiter == nil {
   106  		limiter.checkLimiter = rate.NewLimiter(limit, 1)
   107  	} else if limit != limiter.checkLimiter.Limit() {
   108  		limiter.checkLimiter.SetLimit(limit)
   109  	}
   110  
   111  	return nil
   112  }