diff --git a/pkg/aksmachine/client_armapi.go b/pkg/aksmachine/client_armapi.go index 8748060c..4132c4c8 100644 --- a/pkg/aksmachine/client_armapi.go +++ b/pkg/aksmachine/client_armapi.go @@ -7,9 +7,11 @@ import ( "log/slog" "net/http" "strings" + "time" "github.com/Azure/azure-sdk-for-go/sdk/azcore" "github.com/Azure/azure-sdk-for-go/sdk/azcore/arm" + "github.com/Azure/azure-sdk-for-go/sdk/azcore/policy" "github.com/Azure/azure-sdk-for-go/sdk/azidentity" "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerservice/armcontainerservice/v8" @@ -17,6 +19,16 @@ import ( "github.com/Azure/AKSFlexNode/pkg/config" ) +const ( + // ARM can throttle concurrent node bootstrap requests for several minutes. + // Keep the backoff bounded while allowing a realistic subscription burst to + // clear without requiring the whole bootstrap workflow to restart. + armMachineMaxRetries = 30 + armMachineTryTimeout = 2 * time.Minute + armMachineRetryDelay = 5 * time.Second + armMachineMaxRetryDelay = time.Minute +) + type armMachineClient struct { machineID *arm.ResourceID client *armcontainerservice.MachinesClient @@ -34,7 +46,7 @@ func newARMClient(cfg *config.Config, logger *slog.Logger) (MachineClient, error if err != nil { return nil, fmt.Errorf("resolve ARM credential: %w", err) } - armOpts := &arm.ClientOptions{ClientOptions: clientOpts} + armOpts := armMachineClientOptions(clientOpts) client, err := armcontainerservice.NewMachinesClient(machineID.SubscriptionID, cred, armOpts) if err != nil { return nil, fmt.Errorf("create machines client: %w", err) @@ -159,6 +171,18 @@ func azureClientOptionsFromConfig(cfg *config.Config) azcore.ClientOptions { return azclient.ClientOptionsFromConfig(cfg) } +func armMachineClientOptions(clientOpts azcore.ClientOptions) *arm.ClientOptions { + // The SDK uses ARM's Retry-After, Retry-After-Ms, and x-ms-retry-after-ms + // response headers before falling back to this jittered exponential delay. + clientOpts.Retry = policy.RetryOptions{ + MaxRetries: armMachineMaxRetries, + TryTimeout: armMachineTryTimeout, + RetryDelay: armMachineRetryDelay, + MaxRetryDelay: armMachineMaxRetryDelay, + } + return &arm.ClientOptions{ClientOptions: clientOpts} +} + func getCredential(cfg *config.Config, logger *slog.Logger, clientOpts azcore.ClientOptions) (azcore.TokenCredential, error) { switch { case cfg.IsSPConfigured(): diff --git a/pkg/aksmachine/client_armapi_proxy.go b/pkg/aksmachine/client_armapi_proxy.go index 9974a33f..94c17d8b 100644 --- a/pkg/aksmachine/client_armapi_proxy.go +++ b/pkg/aksmachine/client_armapi_proxy.go @@ -10,7 +10,6 @@ import ( "time" "github.com/Azure/azure-sdk-for-go/sdk/azcore" - "github.com/Azure/azure-sdk-for-go/sdk/azcore/arm" "github.com/Azure/azure-sdk-for-go/sdk/azcore/policy" "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerservice/armcontainerservice/v8" @@ -37,11 +36,8 @@ func newARMProxyClient(cfg *config.Config, logger *slog.Logger) (MachineClient, client, err := armcontainerservice.NewMachinesClient( machineID.SubscriptionID, staticARMProxyCredential{}, - &arm.ClientOptions{ - ClientOptions: policy.ClientOptions{ - Transport: transport, - }, - }) + armMachineClientOptions(policy.ClientOptions{Transport: transport}), + ) if err != nil { return nil, fmt.Errorf("create proxied machines client: %w", err) } diff --git a/pkg/aksmachine/client_armapi_test.go b/pkg/aksmachine/client_armapi_test.go index 0644171d..aca4c216 100644 --- a/pkg/aksmachine/client_armapi_test.go +++ b/pkg/aksmachine/client_armapi_test.go @@ -1,14 +1,22 @@ package aksmachine import ( + "context" + "encoding/json" + "errors" "io" "log/slog" "math" + "net/http" "path/filepath" "strings" + "sync/atomic" "testing" + "time" + "github.com/Azure/azure-sdk-for-go/sdk/azcore" "github.com/Azure/azure-sdk-for-go/sdk/azcore/cloud" + "github.com/Azure/azure-sdk-for-go/sdk/azcore/policy" "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerservice/armcontainerservice/v8" "github.com/Azure/AKSFlexNode/pkg/config" @@ -165,6 +173,136 @@ func TestAzureClientOptionsFromConfig(t *testing.T) { } } +func TestARMMachineClientOptions(t *testing.T) { + t.Parallel() + + transport := roundTripFunc(func(*http.Request) (*http.Response, error) { + return nil, errors.New("not implemented") + }) + opts := armMachineClientOptions(azcore.ClientOptions{Transport: transport}) + + if opts.Transport == nil { + t.Fatal("Transport was not preserved") + } + if opts.Retry.MaxRetries != armMachineMaxRetries { + t.Fatalf("MaxRetries = %d, want %d", opts.Retry.MaxRetries, armMachineMaxRetries) + } + if opts.Retry.TryTimeout != armMachineTryTimeout { + t.Fatalf("TryTimeout = %s, want %s", opts.Retry.TryTimeout, armMachineTryTimeout) + } + if opts.Retry.RetryDelay != armMachineRetryDelay { + t.Fatalf("RetryDelay = %s, want %s", opts.Retry.RetryDelay, armMachineRetryDelay) + } + if opts.Retry.MaxRetryDelay != armMachineMaxRetryDelay { + t.Fatalf("MaxRetryDelay = %s, want %s", opts.Retry.MaxRetryDelay, armMachineMaxRetryDelay) + } + if opts.Retry.StatusCodes != nil { + t.Fatalf("StatusCodes = %v, want nil to preserve Azure SDK transient status defaults", opts.Retry.StatusCodes) + } +} + +func TestARMMachineClientRetriesThrottledRequests(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + call func(context.Context, *armMachineClient) error + }{ + { + name: "get", + call: func(ctx context.Context, client *armMachineClient) error { + _, err := client.Get(ctx) + return err + }, + }, + { + name: "create or update", + call: func(ctx context.Context, client *armMachineClient) error { + _, err := client.Create(ctx, GoalState{KubernetesVersion: "1.34.0"}) + return err + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + var attempts atomic.Int32 + client := newTestARMMachineClient(t, roundTripFunc(func(request *http.Request) (*http.Response, error) { + if attempts.Add(1) == 1 { + return throttledMachineResponse(request, "1"), nil + } + return successfulMachineResponse(t, request), nil + })) + + if err := tt.call(t.Context(), client); err != nil { + t.Fatalf("machine request error = %v", err) + } + if got := attempts.Load(); got != 2 { + t.Fatalf("request attempts = %d, want 2", got) + } + }) + } +} + +func TestARMMachineClientHonorsRetryAfter(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + header string + headerValue string + }{ + {name: "seconds", header: "Retry-After", headerValue: "5"}, + {name: "milliseconds", header: "Retry-After-Ms", headerValue: "5000"}, + {name: "ARM milliseconds", header: "x-ms-retry-after-ms", headerValue: "5000"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + var attempts atomic.Int32 + client := newTestARMMachineClient(t, roundTripFunc(func(request *http.Request) (*http.Response, error) { + attempts.Add(1) + response := throttledMachineResponse(request, "") + response.Header.Set(tt.header, tt.headerValue) + return response, nil + })) + ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond) + defer cancel() + + _, err := client.Get(ctx) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("Get() error = %v, want context deadline exceeded", err) + } + if got := attempts.Load(); got != 1 { + t.Fatalf("request attempts before %s elapsed = %d, want 1", tt.header, got) + } + }) + } +} + +func TestARMMachineClientStopsAfterRetryBudget(t *testing.T) { + t.Parallel() + + var attempts atomic.Int32 + client := newTestARMMachineClient(t, roundTripFunc(func(request *http.Request) (*http.Response, error) { + attempts.Add(1) + return throttledMachineResponse(request, "1"), nil + })) + + _, err := client.Get(t.Context()) + var responseErr *azcore.ResponseError + if !errors.As(err, &responseErr) || responseErr.StatusCode != http.StatusTooManyRequests { + t.Fatalf("Get() error = %v, want final HTTP 429 response error", err) + } + if got, want := attempts.Load(), int32(armMachineMaxRetries+1); got != want { + t.Fatalf("request attempts = %d, want %d", got, want) + } +} + func TestClientCertificateCredentialOptionsSendCertificateChain(t *testing.T) { t.Parallel() @@ -441,6 +579,67 @@ func TestMachineFromARMBackfillsIdentity(t *testing.T) { } } +func newTestARMMachineClient(t *testing.T, transport policy.Transporter) *armMachineClient { + t.Helper() + + machineID, err := machineResourceIDFromConfig(testARMConfig(testClusterResourceID, "flex-node-1", "1.34.0")) + if err != nil { + t.Fatalf("machineResourceIDFromConfig() error = %v", err) + } + sdkClient, err := armcontainerservice.NewMachinesClient( + machineID.SubscriptionID, + staticARMProxyCredential{}, + armMachineClientOptions(azcore.ClientOptions{Transport: transport}), + ) + if err != nil { + t.Fatalf("NewMachinesClient() error = %v", err) + } + return &armMachineClient{ + machineID: machineID, + client: sdkClient, + logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + } +} + +func throttledMachineResponse(request *http.Request, retryAfterMS string) *http.Response { + header := make(http.Header) + header.Set("Content-Type", "application/json") + if retryAfterMS != "" { + header.Set("Retry-After-Ms", retryAfterMS) + } + return &http.Response{ + StatusCode: http.StatusTooManyRequests, + Header: header, + Body: io.NopCloser(strings.NewReader(`{"error":{"code":"SubscriptionRequestsThrottled","message":"try again later"}}`)), + Request: request, + } +} + +func successfulMachineResponse(t *testing.T, request *http.Request) *http.Response { + t.Helper() + + body, err := json.Marshal(armcontainerservice.Machine{ + ID: ptr(testClusterResourceID + "/agentPools/aksflexnodes/machines/flex-node-1"), + Name: ptr("flex-node-1"), + Properties: &armcontainerservice.MachineProperties{ + ETag: ptr("settings-1"), + Kubernetes: &armcontainerservice.MachineKubernetesProfile{ + OrchestratorVersion: ptr("1.34.0"), + }, + ProvisioningState: ptr("Succeeded"), + }, + }) + if err != nil { + t.Fatalf("Marshal() error = %v", err) + } + return &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(string(body))), + Request: request, + } +} + func testARMConfig(clusterResourceID, nodeName, kubernetesVersion string) *config.Config { return &config.Config{ Azure: config.AzureConfig{