github.com/Rookout/GoSDK@v0.1.48/pkg/augs/aug_rate_limiter.go (about)

     1  package augs
     2  
     3  import (
     4  	"sync"
     5  	"time"
     6  
     7  	"github.com/Rookout/GoSDK/pkg/config"
     8  	"github.com/Rookout/GoSDK/pkg/rookoutErrors"
     9  	"github.com/Rookout/GoSDK/pkg/types"
    10  )
    11  
    12  type AugRateLimiter struct {
    13  	config config.RateLimiterConfiguration
    14  
    15  	WindowSize   time.Duration
    16  	Quota        time.Duration
    17  	activeWeight int64
    18  
    19  	windowUsages map[int64]time.Duration
    20  
    21  	dataLock    sync.Mutex
    22  	currentData map[string]executionData
    23  }
    24  
    25  type executionData struct {
    26  	windowIndex int64
    27  	startTime   time.Time
    28  }
    29  
    30  func NewAugRateLimiter(quota time.Duration, windowSize time.Duration, activeLimit int, config config.RateLimiterConfiguration) *AugRateLimiter {
    31  	if windowSize <= 0 {
    32  		windowSize = 1
    33  	}
    34  
    35  	var activeWeight int64
    36  	if activeLimit != 0 {
    37  		activeWeight = int64(quota) / int64(activeLimit)
    38  	} else {
    39  		activeWeight = 0
    40  	}
    41  
    42  	return &AugRateLimiter{
    43  		config:       config,
    44  		WindowSize:   windowSize,
    45  		Quota:        quota,
    46  		activeWeight: activeWeight,
    47  		windowUsages: make(map[int64]time.Duration),
    48  		currentData:  make(map[string]executionData),
    49  	}
    50  }
    51  
    52  func (r *AugRateLimiter) activeCount() int64 {
    53  	return int64(len(r.currentData))
    54  }
    55  
    56  func (r *AugRateLimiter) currentUsage(startTime time.Time, currentWindowIndex int64) int64 {
    57  	
    58  	currentWindowUsage, ok := r.windowUsages[currentWindowIndex]
    59  	if !ok {
    60  		
    61  		r.windowUsages[currentWindowIndex] = 0
    62  		currentWindowUsage = 0
    63  	}
    64  
    65  	
    66  	prevWindowUsage, ok := r.windowUsages[currentWindowIndex-1]
    67  	if !ok {
    68  		prevWindowUsage = 0
    69  	}
    70  	
    71  	timeInWindow := float64(startTime.UnixNano() % r.WindowSize.Nanoseconds())
    72  	
    73  	
    74  	prevWeight := 1 - (timeInWindow / float64(r.WindowSize))
    75  
    76  	
    77  	prevCount := int64(float64(prevWindowUsage.Nanoseconds()) * prevWeight)
    78  	activeCount := r.activeCount() * r.activeWeight
    79  	return prevCount + activeCount + currentWindowUsage.Nanoseconds()
    80  }
    81  
    82  func (r *AugRateLimiter) cleanup(currentWindowIndex int64) {
    83  	if len(r.windowUsages) > 10 {
    84  		for windowIndex := range r.windowUsages {
    85  			if windowIndex < currentWindowIndex-5 {
    86  				delete(r.windowUsages, windowIndex)
    87  			}
    88  		}
    89  	}
    90  }
    91  
    92  func (r *AugRateLimiter) BeforeRun(executionId string) (types.AugStatus, rookoutErrors.RookoutError) {
    93  	r.dataLock.Lock()
    94  	defer r.dataLock.Unlock()
    95  
    96  	startTime := time.Now()
    97  	currentWindowIndex := startTime.UnixNano() / r.WindowSize.Nanoseconds() 
    98  	r.currentData[executionId] = executionData{
    99  		startTime:   startTime,
   100  		windowIndex: currentWindowIndex,
   101  	}
   102  
   103  	r.cleanup(currentWindowIndex)
   104  
   105  	if r.Quota > 0 && r.currentUsage(startTime, currentWindowIndex) > int64(r.Quota) {
   106  		return types.Warning, rookoutErrors.NewRookRuleRateLimited()
   107  	}
   108  
   109  	return types.Active, nil
   110  }
   111  
   112  func (r *AugRateLimiter) CancelRun(executionId string) {
   113  	r.dataLock.Lock()
   114  	defer r.dataLock.Unlock()
   115  
   116  	delete(r.currentData, executionId)
   117  }
   118  
   119  func (r *AugRateLimiter) AfterRun(executionId string) (types.AugStatus, rookoutErrors.RookoutError) {
   120  	r.dataLock.Lock()
   121  	defer r.dataLock.Unlock()
   122  
   123  	data := r.currentData[executionId]
   124  	duration := time.Since(data.startTime)
   125  	if duration < r.config.MinRateLimitValue {
   126  		duration = r.config.MinRateLimitValue
   127  	}
   128  	if totalUsage, ok := r.windowUsages[data.windowIndex]; ok {
   129  		r.windowUsages[data.windowIndex] = totalUsage + duration
   130  	}
   131  
   132  	delete(r.currentData, executionId)
   133  
   134  	return types.Active, nil
   135  }