sigs.k8s.io/prow@v0.0.0-20240503223140-c5e374dc7eb1/pkg/throttle/throttle.go (about)

     1  /*
     2  Copyright 2023 The Kubernetes Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package throttle
    18  
    19  import (
    20  	"context"
    21  	"fmt"
    22  	"sync"
    23  	"sync/atomic"
    24  	"time"
    25  
    26  	"github.com/sirupsen/logrus"
    27  )
    28  
    29  const throttlerGlobalKey = "*"
    30  
    31  type Throttler struct {
    32  	ticker   map[string]*time.Ticker
    33  	throttle map[string]chan time.Time
    34  	slow     map[string]*int32 // Helps log once when requests start/stop being throttled
    35  	lock     sync.RWMutex
    36  }
    37  
    38  func (t *Throttler) Wait(ctx context.Context, org string) error {
    39  	start := time.Now()
    40  	log := logrus.WithFields(logrus.Fields{"throttled": true})
    41  	defer func() {
    42  		waitTime := time.Since(start)
    43  		switch {
    44  		case waitTime > 15*time.Minute:
    45  			log.WithField("throttle-duration", waitTime.String()).Warn("Throttled clientside for more than 15 minutes")
    46  		case waitTime > time.Minute:
    47  			log.WithField("throttle-duration", waitTime.String()).Debug("Throttled clientside for more than a minute")
    48  		}
    49  	}()
    50  	t.lock.RLock()
    51  	defer t.lock.RUnlock()
    52  	if _, found := t.ticker[org]; !found {
    53  		org = throttlerGlobalKey
    54  	}
    55  	if _, hasThrottler := t.ticker[org]; !hasThrottler {
    56  		return nil
    57  	}
    58  
    59  	var more bool
    60  	select {
    61  	case _, more = <-t.throttle[org]:
    62  		// If we were throttled and the channel is now somewhat (25%+) full, note this
    63  		if len(t.throttle[org]) > cap(t.throttle[org])/4 && atomic.CompareAndSwapInt32(t.slow[org], 1, 0) {
    64  			log.Debug("Unthrottled")
    65  		}
    66  		if !more {
    67  			log.Debug("Throttle channel closed")
    68  		}
    69  		return nil
    70  	default: // Do not wait if nothing is available right now
    71  	}
    72  	// If this is the first time we are waiting, note this
    73  	if slow := atomic.SwapInt32(t.slow[org], 1); slow == 0 {
    74  		log.Debug("Throttled")
    75  	}
    76  
    77  	select {
    78  	case _, more = <-t.throttle[org]:
    79  		if !more {
    80  			log.Debug("Throttle channel closed")
    81  		}
    82  	case <-ctx.Done():
    83  		return ctx.Err()
    84  	}
    85  
    86  	return nil
    87  }
    88  
    89  func (t *Throttler) Refund(org string) {
    90  	t.lock.RLock()
    91  	defer t.lock.RUnlock()
    92  	if _, found := t.ticker[org]; !found {
    93  		org = throttlerGlobalKey
    94  	}
    95  	if _, hasThrottler := t.ticker[org]; !hasThrottler {
    96  		return
    97  	}
    98  	select {
    99  	case t.throttle[org] <- time.Now():
   100  	default:
   101  	}
   102  }
   103  
   104  // Throttle client to a rate of at most hourlyTokens requests per hour,
   105  // allowing burst tokens.
   106  func (t *Throttler) Throttle(hourlyTokens, burst int, orgs ...string) error {
   107  	org := "*"
   108  	if len(orgs) > 0 {
   109  		if len(orgs) > 1 {
   110  			return fmt.Errorf("may only pass one org for throttling, got %d", len(orgs))
   111  		}
   112  		org = orgs[0]
   113  	}
   114  	t.lock.Lock()
   115  	defer t.lock.Unlock()
   116  	if hourlyTokens <= 0 || burst <= 0 { // Disable throttle
   117  		if t.throttle[org] != nil {
   118  			delete(t.throttle, org)
   119  			delete(t.slow, org)
   120  			t.ticker[org].Stop()
   121  			delete(t.ticker, org)
   122  		}
   123  		return nil
   124  	}
   125  	period := time.Hour / time.Duration(hourlyTokens) // Duration between token refills
   126  	ticker := time.NewTicker(period)
   127  	throttle := make(chan time.Time, burst)
   128  	for i := 0; i < burst; i++ { // Fill up the channel
   129  		throttle <- time.Now()
   130  	}
   131  	go func() {
   132  		// Before refilling, wait the amount of time it would have taken to refill the burst channel.
   133  		// This prevents granting too many tokens in the first hour due to the initial burst.
   134  		for i := 0; i < burst; i++ {
   135  			<-ticker.C
   136  		}
   137  		// Refill the channel
   138  		for t := range ticker.C {
   139  			select {
   140  			case throttle <- t:
   141  			default:
   142  			}
   143  		}
   144  	}()
   145  
   146  	if t.ticker == nil {
   147  		t.ticker = map[string]*time.Ticker{}
   148  	}
   149  	t.ticker[org] = ticker
   150  
   151  	if t.throttle == nil {
   152  		t.throttle = map[string]chan time.Time{}
   153  	}
   154  	t.throttle[org] = throttle
   155  
   156  	if t.slow == nil {
   157  		t.slow = map[string]*int32{}
   158  	}
   159  	var i int32
   160  	t.slow[org] = &i
   161  
   162  	return nil
   163  }