From 0a8b24339b6c720980228eeeacba543d91fe532f Mon Sep 17 00:00:00 2001
From: Maximilian von Lindern <48887425+mavolin@users.noreply.github.com>
Date: Wed, 25 Nov 2020 21:08:42 +0100
Subject: [PATCH] API: Added timeout if deadline is after rate limit (#173)

* Rate: don't sleep if sleep exceeds context deadline

* Httputil: add Client.Timeout

* Bot: set default API timeout to 5 minutes

* Rate: reduce calls to time.Now in Acquire

* API: Optimize to use deadline instead of recalculating

Co-authored-by: diamondburned <datutbrus@gmail.com>
---
 api/rate/rate.go         | 28 ++++++++++++++++------------
 bot/ctx.go               |  4 ++++
 utils/httputil/client.go | 17 ++++++++++++++++-
 3 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/api/rate/rate.go b/api/rate/rate.go
index 9dd1dd9..a379b10 100644
--- a/api/rate/rate.go
+++ b/api/rate/rate.go
@@ -18,6 +18,10 @@ import (
 // RE: Those who want others to fix it for them: release the source code then.
 const ExtraDelay = 250 * time.Millisecond
 
+// ErrTimedOutEarly is the error returned by Limiter.Acquire, if a rate limit
+// exceeds the deadline of the context.Context.
+var ErrTimedOutEarly = errors.New("rate: rate limit exceeds context deadline")
+
 // This makes me suicidal.
 // https://github.com/bwmarrin/discordgo/blob/master/ratelimit.go
 
@@ -94,28 +98,28 @@ func (l *Limiter) Acquire(ctx context.Context, path string) error {
 		return err
 	}
 
-	// Time to sleep
-	var sleep time.Duration
+	// Deadline until the limiter is released.
+	until := time.Time{}
+	now := time.Now()
 
-	if b.remaining == 0 && b.reset.After(time.Now()) {
+	if b.remaining == 0 && b.reset.After(now) {
 		// out of turns, gotta wait
-		sleep = time.Until(b.reset)
+		until = b.reset
 	} else {
 		// maybe global rate limit has it
-		now := time.Now()
-		until := time.Unix(0, atomic.LoadInt64(l.global))
-
-		if until.After(now) {
-			sleep = until.Sub(now)
-		}
+		until = time.Unix(0, atomic.LoadInt64(l.global))
 	}
 
-	if sleep > 0 {
+	if until.After(now) {
+		if deadline, ok := ctx.Deadline(); ok && until.After(deadline) {
+			return ErrTimedOutEarly
+		}
+
 		select {
 		case <-ctx.Done():
 			b.lock.Unlock()
 			return ctx.Err()
-		case <-time.After(sleep):
+		case <-time.After(until.Sub(now)):
 		}
 	}
 
diff --git a/bot/ctx.go b/bot/ctx.go
index c278820..acfe6c6 100644
--- a/bot/ctx.go
+++ b/bot/ctx.go
@@ -7,6 +7,7 @@ import (
 	"os/signal"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/pkg/errors"
 
@@ -151,6 +152,9 @@ func Start(
 		return nil, errors.Wrap(err, "failed to create a dgo session")
 	}
 
+	// fail api request if they (will) take up more than 5 minutes
+	s.Client.Client.Timeout = 5 * time.Minute
+
 	c, err := New(s, cmd)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to create rfrouter")
diff --git a/utils/httputil/client.go b/utils/httputil/client.go
index f7ab85c..46c849b 100644
--- a/utils/httputil/client.go
+++ b/utils/httputil/client.go
@@ -7,6 +7,7 @@ import (
 	"context"
 	"io"
 	"mime/multipart"
+	"time"
 
 	"github.com/pkg/errors"
 
@@ -32,6 +33,11 @@ type Client struct {
 	// errors out. The error returned will override Do's if it's not nil.
 	OnResponse []ResponseFunc
 
+	// Timeout is the maximum amount of time the client will wait for a request
+	// to finish. If this is 0 or smaller the Client won't time out. Otherwise,
+	// the timeout will be used as deadline for context of every request.
+	Timeout time.Duration
+
 	// Default to the global Retries variable (5).
 	Retries uint
 
@@ -143,10 +149,19 @@ func (c *Client) Request(method, url string, opts ...RequestOption) (httpdriver.
 	var r httpdriver.Response
 	var status int
 
+	ctx := c.context
+
+	if c.Timeout > 0 {
+		var cancel func()
+
+		ctx, cancel = context.WithTimeout(ctx, c.Timeout)
+		defer cancel()
+	}
+
 	// The c.Retries < 1 check ensures that we retry forever if that field is
 	// less than 1.
 	for i := uint(0); c.Retries < 1 || i < c.Retries; i++ {
-		q, err := c.Client.NewRequest(c.context, method, url)
+		q, err := c.Client.NewRequest(ctx, method, url)
 		if err != nil {
 			return nil, RequestError{err}
 		}