Advertisement
Guest User

Untitled

a guest
Aug 20th, 2018
80
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Go 2.70 KB | None | 0 0
  1. func (r *RateLimit) HandleResponse(key string, res *http.Response, content []byte) {
  2.     var err error
  3.     var global bool
  4.     var limit uint64
  5.     var remaining uint64
  6.     var reset int64
  7.     var body *ratelimitBody
  8.     var noBody bool
  9.  
  10.     // read body as well
  11.     if len(content) == 0 {
  12.         noBody = true
  13.     } else {
  14.         err = json.Unmarshal(content, body)
  15.         if err != nil {
  16.             return
  17.         }
  18.     }
  19.  
  20.     // global?
  21.     if res.Header.Get(XRateLimitGlobal) == "true" || (!noBody && body.Global) {
  22.         global = true
  23.     }
  24.  
  25.     // max number of request before reset
  26.     if res.Header.Get(XRateLimitLimit) != "" || (!noBody && body.Global) {
  27.         limit, err = strconv.ParseUint(res.Header.Get(XRateLimitLimit), 10, 64)
  28.         if err != nil {
  29.             // TODO: logging
  30.         }
  31.     }
  32.  
  33.     // remaining requests before reset
  34.     remainingStr := res.Header.Get(XRateLimitRemaining)
  35.     if remainingStr != "" {
  36.         remaining, err = strconv.ParseUint(remainingStr, 10, 64)
  37.         if err != nil {
  38.             // TODO: logging
  39.         }
  40.     }
  41.  
  42.     // reset unix timestamp
  43.     resetStr := res.Header.Get(XRateLimitReset)
  44.     if resetStr != "" {
  45.         // here we get a unix timestamp in seconds, which we convert to milliseconds
  46.         reset, err = strconv.ParseInt(remainingStr, 10, 64)
  47.         if err == nil {
  48.             reset *= 1000 // => milliseconds
  49.         } else {
  50.             // TODO: logging
  51.         }
  52.     } else if res.Header.Get(RateLimitRetryAfter) != "" || (!noBody && body.RetryAfter > 0) {
  53.         // here we are given a delay in millisecond, which we need to convert into a timestamp
  54.         if res.Header.Get(RateLimitRetryAfter) != "" {
  55.             reset, err = strconv.ParseInt(res.Header.Get(RateLimitRetryAfter), 10, 64)
  56.             if err != nil {
  57.                 reset = 0
  58.             }
  59.         } else if !noBody && body.RetryAfter > 0 {
  60.             reset = body.RetryAfter
  61.         }
  62.  
  63.         // convert diff to timestamp
  64.         reset += time.Now().UnixNano() / 1000
  65.     }
  66.  
  67.     if global {
  68.         r.global.mu.Lock()
  69.         defer r.global.mu.Unlock()
  70.  
  71.         if limit != 0 {
  72.             r.global.limit = limit
  73.         }
  74.         if remaining != 0 {
  75.             r.global.remaining = remaining
  76.         }
  77.         if reset != 0 {
  78.             r.global.reset = reset
  79.         }
  80.     } else {
  81.         bucket := r.Bucket(key)
  82.         bucket.mu.Lock()
  83.         defer bucket.mu.Unlock()
  84.  
  85.         if limit != 0 {
  86.             bucket.limit = limit
  87.         }
  88.         if remaining != 0 {
  89.             bucket.remaining = remaining
  90.         }
  91.         if reset != 0 {
  92.             bucket.reset = reset
  93.         }
  94.     }
  95. }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement