fix(go-sdk): only retry on client credential requests that are 429 or…

… 5xx (#368)
openfga · May 30, 2024 · 346acb1 · 346acb1
2 parents e1b3cfd + 9ea9d77
commit 346acb1
Show file tree

Hide file tree

Showing 2 changed files with 140 additions and 43 deletions.
diff --git a/config/clients/go/template/oauth2/internal/token.go b/config/clients/go/template/oauth2/internal/token.go
@@ -288,18 +288,30 @@ func singleTokenRoundTrip(ctx context.Context, req *http.Request) (*Token, error
 }
 
 func doTokenRoundTrip(ctx context.Context, req *http.Request) (*Token, error) {
-    var token *Token
-    var err error
-
-    for i := 0; i < cMaxRetry; i++ {
-
-        token, err = singleTokenRoundTrip(ctx, req)
-        if err == nil {
-            return token, err
-        }
-        time.Sleep(time.Duration(internalutils.RandomTime(i, cMinWaitInMs)) * time.Millisecond)
-    }
-    return nil, err
+	var token *Token
+	var err error
+
+	for i := 0; i < cMaxRetry; i++ {
+
+		token, err = singleTokenRoundTrip(ctx, req)
+		if err == nil {
+			return token, err
+		}
+
+		// If this was a request error, then check if it was a 429 or 5xx error and retry.
+		// We do not want to retry any other error
+		if rErr, ok := err.(*RetrieveError); ok {
+			statusCode := rErr.Response.StatusCode
+			if statusCode == http.StatusTooManyRequests || (statusCode >= http.StatusInternalServerError && statusCode <= 599) {
+				time.Sleep(time.Duration(internalutils.RandomTime(i, cMinWaitInMs)) * time.Millisecond)
+				continue
+			}
+		}
+
+		return nil, err
+
+	}
+	return nil, err
 }
 
 type RetrieveError struct {

diff --git a/config/clients/go/template/oauth2/internal/token_test.go b/config/clients/go/template/oauth2/internal/token_test.go
@@ -8,13 +8,14 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"log"
 	"math"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
 	"testing"
 
-    "github.com/jarcoal/httpmock"
+	"github.com/jarcoal/httpmock"
 )
 
 func TestRetrieveToken_InParams(t *testing.T) {
@@ -72,24 +73,24 @@ func TestRetrieveTokenWithContextsRetry(t *testing.T) {
 	ResetAuthCache()
 	const clientID = "client-id"
 
-    type JSONResponse struct {
-        AccessToken string `json:"access_token"`
-        TokenType string `json:"token_type"`
-    }
-    expectedResponse := JSONResponse{
-        AccessToken: "ACCESS_TOKEN",
-        TokenType: "bearer",
-    }
-
-    httpmock.Activate()
-    defer httpmock.DeactivateAndReset()
-    firstMock := httpmock.NewStringResponder(429, "")
-    secondMock, _ := httpmock.NewJsonResponder(200, expectedResponse)
-
-    const testURL = "http://testserver.com"
-    httpmock.RegisterResponder("POST", testURL,
-        firstMock.Then(firstMock).Then(firstMock).Then(secondMock),
-    )
+	type JSONResponse struct {
+		AccessToken string `json:"access_token"`
+		TokenType   string `json:"token_type"`
+	}
+	expectedResponse := JSONResponse{
+		AccessToken: "ACCESS_TOKEN",
+		TokenType:   "bearer",
+	}
+
+	httpmock.Activate()
+	defer httpmock.DeactivateAndReset()
+	firstMock := httpmock.NewStringResponder(429, "")
+	secondMock, _ := httpmock.NewJsonResponder(200, expectedResponse)
+
+	const testURL = "http://testserver.com"
+	httpmock.RegisterResponder("POST", testURL,
+		firstMock.Then(firstMock).Then(firstMock).Then(secondMock),
+	)
 
 	_, err := RetrieveToken(context.Background(), clientID, "", testURL, url.Values{}, AuthStyleUnknown)
 	if err != nil {
@@ -101,23 +102,107 @@ func TestRetrieveTokenWithContextsFailure(t *testing.T) {
 	ResetAuthCache()
 	const clientID = "client-id"
 
-    type JSONResponse struct {
-        AccessToken string `json:"access_token"`
-        TokenType string `json:"token_type"`
-    }
+	type JSONResponse struct {
+		AccessToken string `json:"access_token"`
+		TokenType   string `json:"token_type"`
+	}
 
-    httpmock.Activate()
-    defer httpmock.DeactivateAndReset()
-    firstMock := httpmock.NewStringResponder(429, "")
+	httpmock.Activate()
+	defer httpmock.DeactivateAndReset()
+	firstMock := httpmock.NewStringResponder(429, "")
 
-    const testURL = "http://testserver.com"
-    httpmock.RegisterResponder("POST", testURL,
-        firstMock,
-    )
+	const testURL = "http://testserver.com"
+	httpmock.RegisterResponder("POST", testURL,
+		firstMock,
+	)
 
 	_, err := RetrieveToken(context.Background(), clientID, "", testURL, url.Values{}, AuthStyleUnknown)
 	if err == nil {
-        t.Errorf("Expect error to be returned when oauth server fails consistently")
+		t.Errorf("Expect error to be returned when oauth server fails consistently")
+	}
+}
+
+// Test that a 401 returns an error straight away
+func TestRetrieveTokenWithUnauthorizedErrors(t *testing.T) {
+	ResetAuthCache()
+	const clientID = "client-id"
+
+	httpmock.Activate()
+	defer httpmock.DeactivateAndReset()
+	firstMock := httpmock.NewStringResponder(401, "")
+
+	const testURL = "http://testserver.com"
+	httpmock.RegisterResponder("POST", testURL,
+		firstMock,
+	)
+
+	// Set AuthStyleInHeader to avoid making a discovery request
+	_, err := RetrieveToken(context.Background(), clientID, "", testURL, url.Values{}, AuthStyleInHeader)
+	if err == nil {
+		t.Errorf("Expect error to be returned when oauth server fails consistently")
+	}
+
+	log.Print(httpmock.GetTotalCallCount())
+
+	if httpmock.GetTotalCallCount() != 1 {
+		t.Errorf("Expected request to not be retried on a 401")
+	}
+}
+
+// Test that a 500 retries
+func TestRetrieveTokenWithServerErrorRetries(t *testing.T) {
+	ResetAuthCache()
+	const clientID = "client-id"
+
+	type JSONResponse struct {
+		AccessToken string `json:"access_token"`
+		TokenType   string `json:"token_type"`
+	}
+	expectedResponse := JSONResponse{
+		AccessToken: "ACCESS_TOKEN",
+		TokenType:   "bearer",
+	}
+
+	httpmock.Activate()
+	defer httpmock.DeactivateAndReset()
+	firstMock := httpmock.NewStringResponder(500, "")
+	secondMock, _ := httpmock.NewJsonResponder(200, expectedResponse)
+
+	const testURL = "http://testserver.com"
+	httpmock.RegisterResponder("POST", testURL,
+		firstMock.Then(firstMock).Then(firstMock).Then(secondMock),
+	)
+
+	_, err := RetrieveToken(context.Background(), clientID, "", testURL, url.Values{}, AuthStyleUnknown)
+	if err != nil {
+		t.Errorf("RetrieveToken (with background context) = %v; want no error", err)
+	}
+}
+
+// Test that constant 500s errors
+func TestRetrieveTokenWithServerErrorEventuallyErrors(t *testing.T) {
+	ResetAuthCache()
+	const clientID = "client-id"
+
+	httpmock.Activate()
+	defer httpmock.DeactivateAndReset()
+	firstMock := httpmock.NewStringResponder(500, "")
+
+	const testURL = "http://testserver.com"
+	httpmock.RegisterResponder("POST", testURL,
+		firstMock,
+	)
+
+	// Set AuthStyleInHeader to avoid making a discovery request
+	_, err := RetrieveToken(context.Background(), clientID, "", testURL, url.Values{}, AuthStyleInHeader)
+	if err == nil {
+		t.Errorf("Expect error to be returned when oauth server fails consistently")
+	}
+
+	log.Print(httpmock.GetTotalCallCount())
+
+	if httpmock.GetTotalCallCount() != cMaxRetry {
+		t.Errorf("Expected request to not be retried on a 401")
 	}
 }