Browse Source

Hardcoded reactions ratelimit (#398)

* add custom ratelimits

* check for nil ratelimiter

* Don't expose custom ratelimits to Session

* attempt to fix race conditions

* use defer instead

* Slightly improved ratelimiter

You shouldn't need to change the ratelimiters ratelimits while its
running so I removed the functions SetCustomRatelimit and
RemoveCustomRatelimit.
Necroforger 7 years ago
parent
commit
013faa1da4
2 changed files with 44 additions and 5 deletions
  1. 43 4
      ratelimit.go
  2. 1 1
      wsapi.go

+ 43 - 4
ratelimit.go

@@ -3,17 +3,26 @@ package discordgo
 import (
 	"net/http"
 	"strconv"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
 )
 
+// customRateLimit holds information for defining a custom rate limit
+type customRateLimit struct {
+	suffix   string
+	requests int
+	reset    time.Duration
+}
+
 // RateLimiter holds all ratelimit buckets
 type RateLimiter struct {
 	sync.Mutex
-	global          *int64
-	buckets         map[string]*Bucket
-	globalRateLimit time.Duration
+	global           *int64
+	buckets          map[string]*Bucket
+	globalRateLimit  time.Duration
+	customRateLimits []*customRateLimit
 }
 
 // NewRatelimiter returns a new RateLimiter
@@ -22,6 +31,13 @@ func NewRatelimiter() *RateLimiter {
 	return &RateLimiter{
 		buckets: make(map[string]*Bucket),
 		global:  new(int64),
+		customRateLimits: []*customRateLimit{
+			&customRateLimit{
+				suffix:   "//reactions//",
+				requests: 1,
+				reset:    200 * time.Millisecond,
+			},
+		},
 	}
 }
 
@@ -40,6 +56,14 @@ func (r *RateLimiter) getBucket(key string) *Bucket {
 		global:    r.global,
 	}
 
+	// Check if there is a custom ratelimit set for this bucket ID.
+	for _, rl := range r.customRateLimits {
+		if strings.HasSuffix(b.Key, rl.suffix) {
+			b.customRateLimit = rl
+			break
+		}
+	}
+
 	r.buckets[key] = b
 	return b
 }
@@ -76,13 +100,28 @@ type Bucket struct {
 	limit     int
 	reset     time.Time
 	global    *int64
+
+	lastReset       time.Time
+	customRateLimit *customRateLimit
 }
 
 // Release unlocks the bucket and reads the headers to update the buckets ratelimit info
 // and locks up the whole thing in case if there's a global ratelimit.
 func (b *Bucket) Release(headers http.Header) error {
-
 	defer b.Unlock()
+
+	// Check if the bucket uses a custom ratelimiter
+	if rl := b.customRateLimit; rl != nil {
+		if time.Now().Sub(b.lastReset) >= rl.reset {
+			b.remaining = rl.requests - 1
+			b.lastReset = time.Now()
+		}
+		if b.remaining < 1 {
+			b.reset = time.Now().Add(rl.reset)
+		}
+		return nil
+	}
+
 	if headers == nil {
 		return nil
 	}

+ 1 - 1
wsapi.go

@@ -199,7 +199,7 @@ type helloOp struct {
 	Trace             []string      `json:"_trace"`
 }
 
-// Number of heartbeat intervals to wait until forcing a connection restart.
+// FailedHeartbeatAcks is the Number of heartbeat intervals to wait until forcing a connection restart.
 const FailedHeartbeatAcks time.Duration = 5 * time.Millisecond
 
 // heartbeat sends regular heartbeats to Discord so it knows the client