Try 2: Add retry for oauth metadata endpoint (#4262)

* logic and commented out tests Signed-off-by: squiishyy <[email protected]> * make RetryOnSpecificErrorCodes, added functionality for calling sendhttprequestwithretry, need to work on config mappingn still, tests in Signed-off-by: squiishyy <[email protected]> * adding default config values Signed-off-by: squiishyy <[email protected]> * imports Signed-off-by: squiishyy <[email protected]> * updated tests to work Signed-off-by: squiishyy <[email protected]> * added test Signed-off-by: squiishyy <[email protected]> * spelling and info Signed-off-by: squiishyy <[email protected]> * working, removing comment Signed-off-by: squiishyy <[email protected]> * add back Signed-off-by: squiishyy <[email protected]> * code review Signed-off-by: squiishyy <[email protected]> * code review, refactoring to use retry.onerror Signed-off-by: squiishyy <[email protected]> * few code review comments, removing unused logic Signed-off-by: squiishyy <[email protected]> * retry on a wider range Signed-off-by: squiishyy <[email protected]> * comment Signed-off-by: squiishyy <[email protected]> * updated logs Signed-off-by: squiishyy <[email protected]> * lint Signed-off-by: squiishyy <[email protected]> * linting 2 Signed-off-by: squiishyy <[email protected]> --------- Signed-off-by: squiishyy <[email protected]>
flyteorg · Oct 19, 2023 · 9af6d84 · 9af6d84
1 parent 5c0b7ae
commit 9af6d84
Show file tree

Hide file tree

Showing 5 changed files with 167 additions and 2 deletions.
diff --git a/flyteadmin/auth/authzserver/metadata_provider.go b/flyteadmin/auth/authzserver/metadata_provider.go
@@ -2,14 +2,22 @@ package authzserver
 
 import (
 	"context"
+	"fmt"
 	"io/ioutil"
 	"net/http"
 	"net/url"
 	"strings"
+	"time"
+
+	"google.golang.org/grpc/codes"
+	"k8s.io/apimachinery/pkg/util/wait"
+	"k8s.io/client-go/util/retry"
 
 	"github.com/flyteorg/flyte/flyteadmin/auth"
 	authConfig "github.com/flyteorg/flyte/flyteadmin/auth/config"
+	flyteErrors "github.com/flyteorg/flyte/flyteadmin/pkg/errors"
 	"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/service"
+	"github.com/flyteorg/flyte/flytestdlib/logger"
 )
 
 type OAuth2MetadataProvider struct {
@@ -72,7 +80,7 @@ func (s OAuth2MetadataProvider) GetOAuth2Metadata(ctx context.Context, r *servic
 			httpClient.Transport = transport
 		}
 
-		response, err := httpClient.Get(externalMetadataURL.String())
+		response, err := sendAndRetryHTTPRequest(ctx, httpClient, externalMetadataURL.String(), s.cfg.AppAuth.ExternalAuthServer.RetryAttempts, s.cfg.AppAuth.ExternalAuthServer.RetryDelay.Duration)
 		if err != nil {
 			return nil, err
 		}
@@ -107,3 +115,41 @@ func NewService(config *authConfig.Config) OAuth2MetadataProvider {
 		cfg: config,
 	}
 }
+
+func sendAndRetryHTTPRequest(ctx context.Context, client *http.Client, url string, retryAttempts int, retryDelay time.Duration) (*http.Response, error) {
+	var response *http.Response
+	var err error
+	totalAttempts := retryAttempts + 1 // Add one for initial http request attempt
+
+	backoff := wait.Backoff{
+		Duration: retryDelay,
+		Steps:    totalAttempts,
+	}
+
+	retryableOauthMetadataError := flyteErrors.NewFlyteAdminError(codes.Internal, "Failed to get oauth metadata.")
+	err = retry.OnError(backoff,
+		func(err error) bool { // Determine if error is retryable
+			return err == retryableOauthMetadataError
+		}, func() error { // Send HTTP request
+			response, err = client.Get(url)
+			if err != nil {
+				logger.Errorf(ctx, "Failed to send oauth metadata HTTP request. Err: %v", err)
+				return err
+			}
+			if response != nil && response.StatusCode >= http.StatusUnauthorized && response.StatusCode <= http.StatusNetworkAuthenticationRequired {
+				logger.Errorf(ctx, "Failed to get oauth metadata, going to retry. StatusCode: %v Err: %v", response.StatusCode, err)
+				return retryableOauthMetadataError
+			}
+			return nil
+		})
+
+	if err != nil {
+		return nil, err
+	}
+
+	if response != nil && response.StatusCode != http.StatusOK {
+		return response, fmt.Errorf("failed to get oauth metadata with status code %v", response.StatusCode)
+	}
+
+	return response, nil
+}
diff --git a/flyteadmin/auth/authzserver/metadata_provider_test.go b/flyteadmin/auth/authzserver/metadata_provider_test.go
@@ -16,6 +16,8 @@ import (
 	config2 "github.com/flyteorg/flyte/flytestdlib/config"
 )
 
+var oauthMetadataFailureErrorMessage = "Failed to get oauth metadata."
+
 func TestOAuth2MetadataProvider_FlyteClient(t *testing.T) {
 	provider := NewService(&authConfig.Config{
 		AppAuth: authConfig.OAuth2Options{
@@ -111,3 +113,84 @@ func TestOAuth2MetadataProvider_OAuth2Metadata(t *testing.T) {
 		assert.Equal(t, "https://dev-14186422.okta.com", resp.Issuer)
 	})
 }
+
+func TestSendAndRetryHttpRequest(t *testing.T) {
+	t.Run("Retry into failure", func(t *testing.T) {
+		requestAttempts := 0
+		hf := func(w http.ResponseWriter, r *http.Request) {
+			switch strings.TrimSpace(r.URL.Path) {
+			case "/":
+				w.WriteHeader(500)
+				requestAttempts++
+			default:
+				http.NotFoundHandler().ServeHTTP(w, r)
+			}
+		}
+
+		server := httptest.NewServer(http.HandlerFunc(hf))
+		defer server.Close()
+		http.DefaultClient = server.Client()
+		retryAttempts := 5
+		totalAttempts := retryAttempts + 1 // 1 for the initial try
+
+		resp, err := sendAndRetryHttpRequest(context.Background(), server.Client(), server.URL, retryAttempts, 0 /* for testing */)
+		assert.Error(t, err)
+		assert.Equal(t, oauthMetadataFailureErrorMessage, err.Error())
+		assert.Nil(t, resp)
+		assert.Equal(t, totalAttempts, requestAttempts)
+	})
+
+	t.Run("Retry into success", func(t *testing.T) {
+		requestAttempts := 0
+		hf := func(w http.ResponseWriter, r *http.Request) {
+			switch strings.TrimSpace(r.URL.Path) {
+			case "/":
+				if requestAttempts > 2 {
+					w.WriteHeader(200)
+				} else {
+					requestAttempts++
+					w.WriteHeader(500)
+				}
+			default:
+				http.NotFoundHandler().ServeHTTP(w, r)
+			}
+		}
+
+		server := httptest.NewServer(http.HandlerFunc(hf))
+		defer server.Close()
+		http.DefaultClient = server.Client()
+		retryAttempts := 5
+		expectedRequestAttempts := 3
+
+		resp, err := sendAndRetryHttpRequest(context.Background(), server.Client(), server.URL, retryAttempts, 0 /* for testing */)
+		assert.NoError(t, err)
+		assert.NotNil(t, resp)
+		assert.Equal(t, 200, resp.StatusCode)
+		assert.Equal(t, expectedRequestAttempts, requestAttempts)
+	})
+
+	t.Run("Success", func(t *testing.T) {
+		requestAttempts := 0
+		hf := func(w http.ResponseWriter, r *http.Request) {
+			switch strings.TrimSpace(r.URL.Path) {
+			case "/":
+				w.WriteHeader(200)
+			default:
+				http.NotFoundHandler().ServeHTTP(w, r)
+			}
+		}
+
+		server := httptest.NewServer(http.HandlerFunc(hf))
+		defer server.Close()
+		http.DefaultClient = server.Client()
+		retryAttempts := 5
+		expectedRequestAttempts := 0
+
+		resp, err := sendAndRetryHttpRequest(context.Background(), server.Client(), server.URL, retryAttempts, 0 /* for testing */)
+		assert.NoError(t, err)
+		assert.NotNil(t, resp)
+		assert.Equal(t, 200, resp.StatusCode)
+		assert.Equal(t, expectedRequestAttempts, requestAttempts)
+	})
+
+}
diff --git a/flyteadmin/auth/config/config.go b/flyteadmin/auth/config/config.go
@@ -79,6 +79,10 @@ var (
 			},
 		},
 		AppAuth: OAuth2Options{
+			ExternalAuthServer: ExternalAuthorizationServer{
+				RetryAttempts: 5,
+				RetryDelay:    config.Duration{Duration: 1 * time.Second},
+			},
 			AuthServerType: AuthorizationServerTypeSelf,
 			ThirdParty: ThirdPartyConfigOptions{
 				FlyteClientConfig: FlyteClientConfig{
@@ -191,7 +195,9 @@ type ExternalAuthorizationServer struct {
 	AllowedAudience     []string   `json:"allowedAudience" pflag:",Optional: A list of allowed audiences. If not provided, the audience is expected to be the public Uri of the service."`
 	MetadataEndpointURL config.URL `json:"metadataUrl" pflag:",Optional: If the server doesn't support /.well-known/oauth-authorization-server, you can set a custom metadata url here.'"`
 	// HTTPProxyURL allows operators to access external OAuth2 servers using an external HTTP Proxy
-	HTTPProxyURL config.URL `json:"httpProxyURL" pflag:",OPTIONAL: HTTP Proxy to be used for OAuth requests."`
+	HTTPProxyURL  config.URL      `json:"httpProxyURL" pflag:",OPTIONAL: HTTP Proxy to be used for OAuth requests."`
+	RetryAttempts int             `json:"retryAttempts" pflag:", Optional: The number of attempted retries on a transient failure to get the OAuth metadata"`
+	RetryDelay    config.Duration `json:"retryDelay" pflag:", Optional, Duration to wait between retries"`
 }
 
 // OAuth2Options defines settings for app auth.

diff --git a/flyteadmin/auth/config/config_flags.go b/flyteadmin/auth/config/config_flags.go
diff --git a/flyteadmin/auth/config/config_flags_test.go b/flyteadmin/auth/config/config_flags_test.go