Merge pull request #155 from sendgrid/rate_limit

Added optional rate limit handling.

sendgrid.go

@@ -2,14 +2,23 @@
 package sendgrid
 
 import (
+	"errors"
+	"net/http"
+	"strconv"
+	"time"
+
 	"github.com/sendgrid/rest" // depends on version 2.2.0
 	"github.com/sendgrid/sendgrid-go/helpers/mail"
 )
 
 // Version is this client library's current version
-const Version = "3.1.0"
+const (
+	Version        = "3.1.0"
+	rateLimitRetry = 5
+	rateLimitSleep = 1100
+)
 
-// Client ...
+// Client is the SendGrid Go client
 type Client struct {
 	// rest.Request
 	rest.Request
@@ -33,13 +42,13 @@ func GetRequest(key string, endpoint string, host string) rest.Request {
 	return request
 }
 
-//Send ...
+// Send sends an email through SendGrid
 func (cl *Client) Send(email *mail.SGMailV3) (*rest.Response, error) {
 	cl.Body = mail.GetRequestBody(email)
 	return API(cl.Request)
 }
 
-// NewSendClient ...
+// NewSendClient constructs a new SendGrid client given an API key
 func NewSendClient(key string) *Client {
 	request := GetRequest(key, "/v3/mail/send", "")
 	request.Method = "POST"
@@ -50,6 +59,69 @@ func NewSendClient(key string) *Client {
 var DefaultClient = rest.DefaultClient
 
 // API sets up the request to the SendGrid API, this is main interface.
+// This function is deprecated. Please use the MakeRequest or
+// MakeRequestAsync functions.
 func API(request rest.Request) (*rest.Response, error) {
 	return DefaultClient.API(request)
 }
+
+// MakeRequest attemps a SendGrid request synchronously.
+func MakeRequest(request rest.Request) (*rest.Response, error) {
+	return DefaultClient.API(request)
+}
+
+// MakeRequestRetry a synchronous request, but retry in the event of a rate
+// limited response.
+func MakeRequestRetry(request rest.Request) (*rest.Response, error) {
+	retry := 0
+	var response *rest.Response
+	var err error
+
+	for {
+		response, err = DefaultClient.API(request)
+		if err != nil {
+			return nil, err
+		}
+
+		if response.StatusCode != http.StatusTooManyRequests {
+			return response, nil
+		}
+
+		if retry > rateLimitRetry {
+			return nil, errors.New("Rate limit retry exceeded")
+		}
+		retry++
+
+		resetTime := time.Now().Add(rateLimitSleep * time.Millisecond)
+
+		reset, ok := response.Headers["X-RateLimit-Reset"]
+		if ok && len(reset) > 0 {
+			t, err := strconv.Atoi(reset[0])
+			if err == nil {
+				resetTime = time.Unix(int64(t), 0)
+			}
+		}
+		time.Sleep(resetTime.Sub(time.Now()))
+	}
+}
+
+// MakeRequestAsync attempts a request asynchronously in a new go
+// routine. This function returns two channels: responses
+// and errors. This function will retry in the case of a
+// rate limit.
+func MakeRequestAsync(request rest.Request) (chan *rest.Response, chan error) {
+	r := make(chan *rest.Response)
+	e := make(chan error)
+
+	go func() {
+		response, err := MakeRequestRetry(request)
+		if err != nil {
+			e <- err
+		}
+		if response != nil {
+			r <- response
+		}
+	}()
+
+	return r, e
+}

sendgrid_test.go

@@ -12,6 +12,7 @@ import (
 	"os/exec"
 	"path/filepath"
 	"runtime"
+	"strconv"
 	"strings"
 	"testing"
 	"time"
@@ -193,6 +194,104 @@ func TestCustomHTTPClient(t *testing.T) {
 	}
 }
 
+func TestRequestRetry_rateLimit(t *testing.T) {
+	fakeServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("X-RateLimit-Reset", strconv.Itoa(int(time.Now().Add(1*time.Second).Unix())))
+		w.WriteHeader(http.StatusTooManyRequests)
+	}))
+	defer fakeServer.Close()
+	apiKey := "SENDGRID_APIKEY"
+	host := fakeServer.URL
+	request := GetRequest(apiKey, "/v3/test_endpoint", host)
+	request.Method = "GET"
+	var custom rest.Client
+	custom.HTTPClient = &http.Client{Timeout: time.Millisecond * 10}
+	DefaultClient = &custom
+	_, err := MakeRequestRetry(request)
+	if err == nil {
+		t.Error("An error did not trigger")
+	}
+	if !strings.Contains(err.Error(), "Rate limit retry exceeded") {
+		t.Error("We did not receive the rate limit error")
+	}
+	DefaultClient = rest.DefaultClient
+}
+
+func TestRequestRetry_rateLimit_noHeader(t *testing.T) {
+	fakeServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusTooManyRequests)
+	}))
+	defer fakeServer.Close()
+	apiKey := "SENDGRID_APIKEY"
+	host := fakeServer.URL
+	request := GetRequest(apiKey, "/v3/test_endpoint", host)
+	request.Method = "GET"
+	var custom rest.Client
+	custom.HTTPClient = &http.Client{Timeout: time.Millisecond * 10}
+	DefaultClient = &custom
+	_, err := MakeRequestRetry(request)
+	if err == nil {
+		t.Error("An error did not trigger")
+	}
+	if !strings.Contains(err.Error(), "Rate limit retry exceeded") {
+		t.Error("We did not receive the rate limit error")
+	}
+	DefaultClient = rest.DefaultClient
+}
+
+func TestRequestAsync(t *testing.T) {
+	fakeServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer fakeServer.Close()
+	apiKey := "SENDGRID_APIKEY"
+	host := fakeServer.URL
+	request := GetRequest(apiKey, "/v3/test_endpoint", host)
+	request.Method = "GET"
+	var custom rest.Client
+	custom.HTTPClient = &http.Client{Timeout: time.Millisecond * 10}
+	DefaultClient = &custom
+	r, e := MakeRequestAsync(request)
+
+	select {
+	case <-r:
+	case err := <-e:
+		t.Errorf("Received an error,:%v", err)
+	case <-time.After(10 * time.Second):
+		t.Error("Timed out waiting for a response")
+	}
+	DefaultClient = rest.DefaultClient
+}
+
+func TestRequestAsync_rateLimit(t *testing.T) {
+	fakeServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("X-RateLimit-Reset", strconv.Itoa(int(time.Now().Add(1*time.Second).Unix())))
+		w.WriteHeader(http.StatusTooManyRequests)
+	}))
+	defer fakeServer.Close()
+	apiKey := "SENDGRID_APIKEY"
+	host := fakeServer.URL
+	request := GetRequest(apiKey, "/v3/test_endpoint", host)
+	request.Method = "GET"
+	var custom rest.Client
+	custom.HTTPClient = &http.Client{Timeout: time.Millisecond * 10}
+	DefaultClient = &custom
+	r, e := MakeRequestAsync(request)
+
+	select {
+	case <-r:
+		t.Error("Received a valid response")
+		return
+	case err := <-e:
+		if !strings.Contains(err.Error(), "Rate limit retry exceeded") {
+			t.Error("We did not receive the rate limit error")
+		}
+	case <-time.After(10 * time.Second):
+		t.Error("Timed out waiting for an error")
+	}
+	DefaultClient = rest.DefaultClient
+}
+
 func Test_test_access_settings_activity_get(t *testing.T) {
 	apiKey := "SENDGRID_APIKEY"
 	host := "http://localhost:4010"