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

     1  package augs
     2  
     3  import (
     4  	"strconv"
     5  	"strings"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/Rookout/GoSDK/pkg/com_ws"
    10  	"github.com/Rookout/GoSDK/pkg/config"
    11  	"github.com/Rookout/GoSDK/pkg/logger"
    12  	"github.com/Rookout/GoSDK/pkg/rookoutErrors"
    13  	"github.com/Rookout/GoSDK/pkg/types"
    14  	"github.com/Rookout/GoSDK/pkg/utils"
    15  )
    16  
    17  type LimitProvider struct {
    18  	GlobalRateLimiter *AugRateLimiter
    19  }
    20  
    21  var initOnce sync.Once
    22  var rookLimitProvider *LimitProvider
    23  
    24  func GetLimitProvider() *LimitProvider {
    25  	if rookLimitProvider == nil {
    26  		InitLimitProvider()
    27  	}
    28  
    29  	return rookLimitProvider
    30  }
    31  
    32  
    33  func InitLimitProvider() {
    34  	initOnce.Do(func() {
    35  		initializedSingleton := createLimitProvider()
    36  		rookLimitProvider = initializedSingleton
    37  	})
    38  }
    39  
    40  func createLimitProvider() *LimitProvider {
    41  	l := &LimitProvider{
    42  		GlobalRateLimiter: nil,
    43  	}
    44  	config.OnUpdate(l.updateConfig)
    45  	return l
    46  }
    47  
    48  func (l *LimitProvider) updateConfig() {
    49  	if config.RateLimiterConfig().GlobalRateLimit != "" {
    50  		l.tryToCreateGlobalRateLimiter()
    51  	} else {
    52  		l.DeleteGlobalRateLimiter()
    53  	}
    54  }
    55  
    56  func (l *LimitProvider) DeleteGlobalRateLimiter() {
    57  	l.GlobalRateLimiter = nil
    58  }
    59  
    60  func (l *LimitProvider) GetLimitManager(configuration types.AugConfiguration, augID types.AugID, output com_ws.Output) (LimitsManager, rookoutErrors.RookoutError) {
    61  	limitsManager := NewLimitsManager(augID, output)
    62  
    63  	rateLimiter, err := l.getRateLimiter(configuration)
    64  	if err != nil {
    65  		return nil, err
    66  	}
    67  
    68  	if rateLimiter != nil {
    69  		limitsManager.AddLimiter(rateLimiter)
    70  	}
    71  
    72  	limitsManager.AddLimiter(l.getAugTimeLimiter(configuration))
    73  
    74  	return limitsManager, nil
    75  }
    76  
    77  func (l *LimitProvider) getAugTimeLimiter(configuration types.AugConfiguration) *AugTimeLimiter {
    78  	maxAugTime := config.LocationsConfig().MaxAugTime
    79  	maxAugTimeStr, ok := configuration["maxAugTime"].(string)
    80  	if ok {
    81  		maxAugTimeMS, err := strconv.ParseInt(maxAugTimeStr, 10, 64)
    82  		if err != nil {
    83  			logger.Logger().WithError(err).Errorln("Failed to parse max aug time configuration")
    84  		} else {
    85  			maxAugTime = time.Duration(maxAugTimeMS) * time.Millisecond
    86  		}
    87  	}
    88  
    89  	maxAugTime = maxAugTime * time.Duration(config.LocationsConfig().MaxAugTimeMultiplier)
    90  
    91  	return NewAugTimeLimiter(maxAugTime)
    92  }
    93  
    94  func (l *LimitProvider) getRateLimiter(configuration types.AugConfiguration) (*AugRateLimiter, rookoutErrors.RookoutError) {
    95  	
    96  	if l.GlobalRateLimiter != nil {
    97  		return l.GlobalRateLimiter, nil
    98  	} else {
    99  		windowQuota := utils.MSToNS(200)
   100  		windowSize := utils.MSToNS(500)
   101  		rateLimitSpec, ok := configuration["rateLimit"].(string)
   102  
   103  		var err error
   104  		rateLimitModifier := 0
   105  		rateLimitModifierStr, ok := configuration["rateLimitModifier"].(string)
   106  		if ok {
   107  			rateLimitModifier, err = strconv.Atoi(rateLimitModifierStr)
   108  			if err != nil {
   109  				logger.Logger().WithError(err).Errorln("Failed to parse rate limit configuration")
   110  			}
   111  		}
   112  
   113  		rateLimiter, rookErr := l.createRateLimiter(rateLimitSpec, windowQuota, windowSize, rateLimitModifier)
   114  		if rookErr != nil {
   115  			return nil, rookErr
   116  		}
   117  
   118  		return rateLimiter, nil
   119  	}
   120  }
   121  
   122  func (l *LimitProvider) tryToCreateGlobalRateLimiter() {
   123  	globalRateLimit := config.RateLimiterConfig().GlobalRateLimit
   124  	globalRateLimiter, err := l.createRateLimiter(globalRateLimit, 0, 0, 0)
   125  	if globalRateLimiter == nil {
   126  		err = rookoutErrors.NewRookInvalidRateLimitConfiguration(globalRateLimit)
   127  	}
   128  
   129  	if err != nil {
   130  		logger.Logger().WithError(err).Warningln("Failed to create global rate limiter")
   131  		return
   132  	}
   133  
   134  	rookoutErrors.UsingGlobalRateLimiter = true
   135  	l.GlobalRateLimiter = globalRateLimiter
   136  }
   137  
   138  func (l *LimitProvider) createRateLimiter(limitsSpec string, defaultQuota time.Duration, defaultWindowSize time.Duration, activeLimit int) (*AugRateLimiter, rookoutErrors.RookoutError) {
   139  	quota := defaultQuota
   140  	windowSize := defaultWindowSize
   141  
   142  	if limitsSpec != "" {
   143  		limits := strings.Split(limitsSpec, "/")
   144  
   145  		if len(limits) == 2 {
   146  			var err error
   147  
   148  			quota, err = utils.StringMSToNS(limits[0])
   149  			if err != nil {
   150  				quota = defaultQuota
   151  			} else {
   152  				windowSize, err = utils.StringMSToNS(limits[1])
   153  				if err != nil {
   154  					quota = defaultQuota
   155  					windowSize = defaultWindowSize
   156  				}
   157  			}
   158  		}
   159  	}
   160  
   161  	if quota == 0 {
   162  		return nil, nil
   163  	}
   164  
   165  	if quota >= windowSize {
   166  		return nil, rookoutErrors.NewRookInvalidRateLimitConfiguration(limitsSpec)
   167  	}
   168  
   169  	return NewAugRateLimiter(quota, windowSize, activeLimit, config.RateLimiterConfig()), nil
   170  }