github.com/unionj-cloud/go-doudou@v1.3.8-0.20221011095552-0088008e5b31/framework/ratelimit/limit.go (about)

     1  package ratelimit
     2  
     3  import (
     4  	"github.com/pkg/errors"
     5  	"strconv"
     6  	"strings"
     7  	"time"
     8  )
     9  
    10  type Limit struct {
    11  	Rate   float64
    12  	Burst  int
    13  	Period time.Duration
    14  }
    15  
    16  type LimiterOption func(Limiter)
    17  
    18  func PerSecond(rate float64) Limit {
    19  	return Limit{
    20  		Rate:   rate,
    21  		Period: time.Second,
    22  		Burst:  1,
    23  	}
    24  }
    25  
    26  func PerMinute(rate float64) Limit {
    27  	return Limit{
    28  		Rate:   rate,
    29  		Period: time.Minute,
    30  		Burst:  1,
    31  	}
    32  }
    33  
    34  func PerHour(rate float64) Limit {
    35  	return Limit{
    36  		Rate:   rate,
    37  		Period: time.Hour,
    38  		Burst:  1,
    39  	}
    40  }
    41  
    42  func PerDay(rate float64) Limit {
    43  	return Limit{
    44  		Rate:   rate,
    45  		Period: time.Hour * 24,
    46  		Burst:  1,
    47  	}
    48  }
    49  
    50  func PerSecondBurst(rate float64, burst int) Limit {
    51  	return Limit{
    52  		Rate:   rate,
    53  		Period: time.Second,
    54  		Burst:  burst,
    55  	}
    56  }
    57  
    58  func PerMinuteBurst(rate float64, burst int) Limit {
    59  	return Limit{
    60  		Rate:   rate,
    61  		Period: time.Minute,
    62  		Burst:  burst,
    63  	}
    64  }
    65  
    66  func PerHourBurst(rate float64, burst int) Limit {
    67  	return Limit{
    68  		Rate:   rate,
    69  		Period: time.Hour,
    70  		Burst:  burst,
    71  	}
    72  }
    73  
    74  func PerDayBurst(rate float64, burst int) Limit {
    75  	return Limit{
    76  		Rate:   rate,
    77  		Period: time.Hour * 24,
    78  		Burst:  burst,
    79  	}
    80  }
    81  
    82  // Parse is borrowed from https://github.com/ulule/limiter Copyright (c) 2015-2018 Ulule
    83  // You can use the simplified format "<limit>-<period>"(burst is 1 by default) or "<limit>-<period>-<burst>", with the given
    84  // periods:
    85  //
    86  // * "S": second
    87  // * "M": minute
    88  // * "H": hour
    89  // * "D": day
    90  //
    91  // Examples for "<limit>-<period>" format:
    92  //
    93  // * 5.5 reqs/second: "5.5-S"
    94  // * 10 reqs/minute: "10-M"
    95  // * 1000 reqs/hour: "1000-H"
    96  // * 2000 reqs/day: "2000-D"
    97  //
    98  // Examples:
    99  //
   100  // * 0.0055 reqs/second with burst 20: "0.0055-S-20" https://github.com/go-redis/redis_rate/issues/63
   101  // * 10 reqs/minute with burst 200: "10-M-200"
   102  // * 1000 reqs/hour with burst 500: "1000-H-500"
   103  // * 2000 reqs/day with burst 1000: "2000-D-1000"
   104  //
   105  func Parse(value string) (Limit, error) {
   106  	var l Limit
   107  
   108  	splits := strings.Split(value, "-")
   109  	if len(splits) != 2 && len(splits) != 3 {
   110  		return l, errors.Errorf("incorrect format '%s'", value)
   111  	}
   112  
   113  	periods := map[string]time.Duration{
   114  		"S": time.Second,    // Second
   115  		"M": time.Minute,    // Minute
   116  		"H": time.Hour,      // Hour
   117  		"D": time.Hour * 24, // Day
   118  	}
   119  
   120  	if len(splits) == 2 {
   121  		r, period := splits[0], strings.ToUpper(splits[1])
   122  
   123  		p, ok := periods[period]
   124  		if !ok {
   125  			return l, errors.Errorf("incorrect period '%s'", period)
   126  		}
   127  
   128  		rate, err := strconv.ParseFloat(r, 64)
   129  		if err != nil {
   130  			return l, errors.Errorf("incorrect rate '%s'", r)
   131  		}
   132  
   133  		l = Limit{
   134  			Rate:   rate,
   135  			Burst:  1,
   136  			Period: p,
   137  		}
   138  
   139  		return l, nil
   140  	}
   141  
   142  	r, period, b := splits[0], strings.ToUpper(splits[1]), splits[2]
   143  
   144  	p, ok := periods[period]
   145  	if !ok {
   146  		return l, errors.Errorf("incorrect period '%s'", period)
   147  	}
   148  
   149  	rate, err := strconv.ParseFloat(r, 64)
   150  	if err != nil {
   151  		return l, errors.Errorf("incorrect rate '%s'", r)
   152  	}
   153  
   154  	burst, err := strconv.Atoi(b)
   155  	if err != nil {
   156  		return l, errors.Errorf("incorrect burst '%s'", b)
   157  	}
   158  
   159  	l = Limit{
   160  		Rate:   rate,
   161  		Burst:  burst,
   162  		Period: p,
   163  	}
   164  
   165  	return l, nil
   166  }