From beb54504d122d41d258eb9fb9815193b044f5dbc Mon Sep 17 00:00:00 2001 From: Yingjie Date: Mon, 26 Oct 2026 18:29:18 +0800 Subject: [PATCH 1/2] fix: honor Options.Parallelism when fetching built-in policy definitions --- alzlib.go | 327 +++++++++++++------------ builtin_fetch.go | 145 +++++++++++ builtin_fetch_test.go | 328 +++++++++++++++++++++++++ deployment/hierarchy.go | 108 +++++++- deployment/hierarchy_fetch_test.go | 178 ++++++++++++++ go.mod | 1 + go.sum | 2 + internal/testutil/fakearm/fakearm.go | 354 +++++++++++++++++++++++++++ 8 files changed, 1281 insertions(+), 162 deletions(-) create mode 100644 builtin_fetch.go create mode 100644 builtin_fetch_test.go create mode 100644 deployment/hierarchy_fetch_test.go create mode 100644 internal/testutil/fakearm/fakearm.go diff --git a/alzlib.go b/alzlib.go index 5ff5a23..286147a 100644 --- a/alzlib.go +++ b/alzlib.go @@ -52,6 +52,9 @@ type AlzLib struct { cache BuiltInCache clients *azureClients mu sync.RWMutex // mu is a mutex to concurrency protect the AlzLib maps + // fetchMu serialises concurrent GetDefinitionsFromAzure callers sharing this AlzLib, so that they + // do not fetch and add the same definitions at the same time. + fetchMu sync.Mutex } type azureClients struct { @@ -65,6 +68,8 @@ type Options struct { AllowOverwrite bool // Parallelism is the number of parallel requests to make to Azure APIs when getting policy definitions // and policy set definitions. + // Requests for the same definition name are always made serially. + // Set to 1 to fetch all definitions serially. Parallelism int // UniqueRoleDefinitions indicates whether to update the role definitions to be unique per management group. // If this is not set, you may end up with conflicting role definition names. @@ -726,7 +731,11 @@ func (az *AlzLib) Init(ctx context.Context, libs ...LibraryReference) error { // before falling back to Azure API calls. // The cache is retained for the lifetime of AlzLib; callers can explicitly clear it by calling // AddCache(nil) when they no longer need it. +// Definitions are fetched concurrently, bounded by [Options.Parallelism]. func (az *AlzLib) GetDefinitionsFromAzure(ctx context.Context, reqs []BuiltInRequest) error { + az.fetchMu.Lock() + defer az.fetchMu.Unlock() + policyDefsToGet := make([]BuiltInRequest, 0, len(reqs)) policySetDefsToGet := make([]BuiltInRequest, 0, len(reqs)) @@ -922,145 +931,177 @@ func (az *AlzLib) policySetDefinitionFromCache(name string, version *string) (*a // getBuiltInPolicies retrieves the built-in policy definitions with the given names // and adds them to the AlzLib struct. +// Requests are grouped by definition name and the groups are fetched concurrently. func (az *AlzLib) getBuiltInPolicies(ctx context.Context, reqs []BuiltInRequest) error { var ( versionedClient *armpolicy.DefinitionVersionsClient client *armpolicy.DefinitionsClient ) - for _, req := range reqs { - if az.PolicyDefinitionExists(req.ResourceID.Name, req.Version) { - continue - } - - // Check cache before calling Azure. - pdv, cacheErr := az.policyDefinitionFromCache(req.ResourceID.Name, req.Version) - if cacheErr != nil { - return fmt.Errorf( - "Alzlib.getBuiltInPolicies: looking up cached built-in policy definition %s: %w", - JoinNameAndVersion(req.ResourceID.Name, req.Version), - cacheErr, - ) - } - - if pdv != nil { - if err := az.AddPolicyDefinitions(pdv); err != nil { - return fmt.Errorf( - "Alzlib.getBuiltInPolicies: adding cached built-in policy definition %s: %w", - JoinNameAndVersion(req.ResourceID.Name, req.Version), - err, - ) - } + // Azure SDK clients are safe for concurrent use, so they are created once and shared. + if policyClient := az.policyClientFactory(); policyClient != nil { + versionedClient = policyClient.NewDefinitionVersionsClient() + client = policyClient.NewDefinitionsClient() + } - continue - } + return forEachGroupConcurrently( + ctx, + az.effectiveParallelism(), + groupByName(reqs, builtInRequestName), + func(ctx context.Context, req BuiltInRequest) error { + return az.getBuiltInPolicy(ctx, client, versionedClient, req) + }, + ) +} - // Fall back to Azure. - az.mu.RLock() - policyClient := az.clients.policyClient - az.mu.RUnlock() +// getBuiltInPolicy retrieves a single built-in policy definition, from the cache if possible, +// and adds it to the AlzLib struct. +func (az *AlzLib) getBuiltInPolicy( + ctx context.Context, + client *armpolicy.DefinitionsClient, + versionedClient *armpolicy.DefinitionVersionsClient, + req BuiltInRequest, +) error { + if az.PolicyDefinitionExists(req.ResourceID.Name, req.Version) { + return nil + } - if policyClient == nil { - return errors.New("Alzlib.getBuiltInPolicies: policy client not set") - } + // Check cache before calling Azure. + pdv, cacheErr := az.policyDefinitionFromCache(req.ResourceID.Name, req.Version) + if cacheErr != nil { + return fmt.Errorf( + "Alzlib.getBuiltInPolicies: looking up cached built-in policy definition %s: %w", + JoinNameAndVersion(req.ResourceID.Name, req.Version), + cacheErr, + ) + } - // Lazily initialise clients on first Azure call. - if client == nil { - versionedClient = policyClient.NewDefinitionVersionsClient() - client = policyClient.NewDefinitionsClient() + if pdv != nil { + if err := az.AddPolicyDefinitions(pdv); err != nil { + return fmt.Errorf( + "Alzlib.getBuiltInPolicies: adding cached built-in policy definition %s: %w", + JoinNameAndVersion(req.ResourceID.Name, req.Version), + err, + ) } - var err error + return nil + } - switch req.Version { - case nil: - err = az.addLatestPolicyDefinition(ctx, client, req) - default: - err = az.addVersionedPolicyDefinition(ctx, versionedClient, req) - } + // Fall back to Azure. + if client == nil || versionedClient == nil { + return errors.New("Alzlib.getBuiltInPolicies: policy client not set") + } - if err != nil { - return err - } + if req.Version == nil { + return az.addLatestPolicyDefinition(ctx, client, req) } - return nil + return az.addVersionedPolicyDefinition(ctx, versionedClient, req) } // getBuiltInPolicySets retrieves the built-in policy set definitions with the given names // and adds them to the AlzLib struct. +// Requests are grouped by definition name and the groups are fetched concurrently. +// Once all policy set definitions have been retrieved, the policy definitions they reference +// are retrieved. func (az *AlzLib) getBuiltInPolicySets(ctx context.Context, reqs []BuiltInRequest) error { - processedRequests := make([]*assets.PolicySetDefinition, 0, len(reqs)) - var ( versionedClient *armpolicy.SetDefinitionVersionsClient client *armpolicy.SetDefinitionsClient ) - for _, req := range reqs { - if az.PolicySetDefinitionExists(req.ResourceID.Name, req.Version) { - continue - } + // Azure SDK clients are safe for concurrent use, so they are created once and shared. + if policyClient := az.policyClientFactory(); policyClient != nil { + versionedClient = policyClient.NewSetDefinitionVersionsClient() + client = policyClient.NewSetDefinitionsClient() + } - // Check cache before calling Azure. - psdv, cacheErr := az.policySetDefinitionFromCache(req.ResourceID.Name, req.Version) - if cacheErr != nil { - return fmt.Errorf( - "Alzlib.getBuiltInPolicySets: looking up cached built-in policy set definition %s: %w", - req.String(), - cacheErr, - ) - } + // Each request writes only to its own slot, so the processed policy set definitions keep the + // request order regardless of the order in which they are fetched. + processed := make([]*assets.PolicySetDefinition, len(reqs)) + indices := make([]int, len(reqs)) - if psdv != nil { - if err := az.AddPolicySetDefinitions(psdv); err != nil { - return fmt.Errorf( - "Alzlib.getBuiltInPolicySets: adding cached built-in policy set definition %s: %w", - req.String(), - err, - ) - } + for i := range indices { + indices[i] = i + } - processedRequests = append(processedRequests, psdv) + err := forEachGroupConcurrently( + ctx, + az.effectiveParallelism(), + groupByName(indices, func(i int) string { return builtInRequestName(reqs[i]) }), + func(ctx context.Context, i int) error { + psd, err := az.getBuiltInPolicySet(ctx, client, versionedClient, reqs[i]) + processed[i] = psd - continue - } + return err + }, + ) + if err != nil { + return err + } - // Fall back to Azure. - az.mu.RLock() - policyClient := az.clients.policyClient - az.mu.RUnlock() + processedRequests := make([]*assets.PolicySetDefinition, 0, len(reqs)) - if policyClient == nil { - return errors.New("Alzlib.getBuiltInPolicySets: policy client not set") + for _, psd := range processed { + if psd != nil { + processedRequests = append(processedRequests, psd) } + } - // Lazily initialise clients on first Azure call. - if client == nil { - versionedClient = policyClient.NewSetDefinitionVersionsClient() - client = policyClient.NewSetDefinitionsClient() - } + return az.ensureReferencedPolicyDefinitions(ctx, processedRequests) +} - var ( - psd *assets.PolicySetDefinition - err error +// getBuiltInPolicySet retrieves a single built-in policy set definition, from the cache if possible, +// and adds it to the AlzLib struct. +// It returns the added policy set definition, or nil if it already existed. +func (az *AlzLib) getBuiltInPolicySet( + ctx context.Context, + client *armpolicy.SetDefinitionsClient, + versionedClient *armpolicy.SetDefinitionVersionsClient, + req BuiltInRequest, +) (*assets.PolicySetDefinition, error) { + if az.PolicySetDefinitionExists(req.ResourceID.Name, req.Version) { + return nil, nil + } + + // Check cache before calling Azure. + psdv, cacheErr := az.policySetDefinitionFromCache(req.ResourceID.Name, req.Version) + if cacheErr != nil { + return nil, fmt.Errorf( + "Alzlib.getBuiltInPolicySets: looking up cached built-in policy set definition %s: %w", + req.String(), + cacheErr, ) + } - switch req.Version { - case nil: - psd, err = az.addLatestPolicySetDefinition(ctx, client, req) - default: - psd, err = az.addVersionedPolicySetDefinition(ctx, versionedClient, req) + if psdv != nil { + if err := az.AddPolicySetDefinitions(psdv); err != nil { + return nil, fmt.Errorf( + "Alzlib.getBuiltInPolicySets: adding cached built-in policy set definition %s: %w", + req.String(), + err, + ) } - if err != nil { - return err - } + return psdv, nil + } - processedRequests = append(processedRequests, psd) + // Fall back to Azure. + if client == nil || versionedClient == nil { + return nil, errors.New("Alzlib.getBuiltInPolicySets: policy client not set") } - return az.ensureReferencedPolicyDefinitions(ctx, processedRequests) + if req.Version == nil { + return az.addLatestPolicySetDefinition(ctx, client, req) + } + + return az.addVersionedPolicySetDefinition(ctx, versionedClient, req) +} + +// builtInRequestName returns the definition name of a built-in request. +func builtInRequestName(req BuiltInRequest) string { + return req.ResourceID.Name } func (az *AlzLib) addLatestPolicyDefinition( @@ -1271,58 +1312,56 @@ func (az *AlzLib) ensureReferencedPolicyDefinitions( return nil } - // Lazily initialised Azure clients, shared across all references within this invocation. - var ( - definitionsClient *armpolicy.DefinitionsClient - definitionsVersionedClient *armpolicy.DefinitionVersionsClient - ) + // Validate every reference before making any Azure calls, so that malformed references fail early + // and deterministically. + refs := make([]definitionReference, 0, len(processedRequests)) for _, def := range processedRequests { for _, ref := range def.Properties.PolicyDefinitions { - err := az.ensureReferencedPolicyDefinition( - ctx, - def, - ref, - &definitionsClient, - &definitionsVersionedClient, - ) + defRef, err := parseDefinitionReference(def, ref) if err != nil { return err } + + if defRef != nil { + refs = append(refs, *defRef) + } } } - return nil -} - -func (az *AlzLib) ensureReferencedPolicyDefinition( - ctx context.Context, - setDefinition *assets.PolicySetDefinition, - ref *armpolicy.DefinitionReference, - definitionsClient **armpolicy.DefinitionsClient, - definitionsVersionedClient **armpolicy.DefinitionVersionsClient, -) error { - if ref == nil { + if len(refs) == 0 { return nil } - if ref.PolicyDefinitionID == nil { - return fmt.Errorf( - "Alzlib.getBuiltInPolicySets: error getting policy definition references for policy set definition `%s`: "+ - "policy definition ID is nil", - JoinNameAndVersion(*setDefinition.Name, setDefinition.Properties.Version), - ) - } + var ( + definitionsClient *armpolicy.DefinitionsClient + definitionsVersionedClient *armpolicy.DefinitionVersionsClient + ) - resID, err := arm.ParseResourceID(*ref.PolicyDefinitionID) - if err != nil { - return fmt.Errorf( - "Alzlib.getBuiltInPolicySets: error parsing resource id `%s` referenced in policy set `%s`", - *ref.PolicyDefinitionID, - JoinNameAndVersion(*setDefinition.Name, setDefinition.Properties.Version), - ) + // Azure SDK clients are safe for concurrent use, so they are created once and shared. + if policyClient := az.policyClientFactory(); policyClient != nil { + definitionsClient = policyClient.NewDefinitionsClient() + definitionsVersionedClient = policyClient.NewDefinitionVersionsClient() } + return forEachGroupConcurrently( + ctx, + az.effectiveParallelism(), + groupByName(refs, func(r definitionReference) string { return r.resID.Name }), + func(ctx context.Context, r definitionReference) error { + return az.ensureReferencedPolicyDefinition(ctx, r, definitionsClient, definitionsVersionedClient) + }, + ) +} + +func (az *AlzLib) ensureReferencedPolicyDefinition( + ctx context.Context, + defRef definitionReference, + definitionsClient *armpolicy.DefinitionsClient, + definitionsVersionedClient *armpolicy.DefinitionVersionsClient, +) error { + setDefinition, ref, resID := defRef.setDefinition, defRef.ref, defRef.resID + // Check AlzLib first to avoid redundant work. if az.PolicyDefinitionExists(resID.Name, ref.DefinitionVersion) { return nil @@ -1355,28 +1394,16 @@ func (az *AlzLib) ensureReferencedPolicyDefinition( } // Fall back to Azure. - az.mu.RLock() - policyClient := az.clients.policyClient - az.mu.RUnlock() - - if policyClient == nil { + if definitionsClient == nil || definitionsVersionedClient == nil { return errors.New("Alzlib.ensureReferencedPolicyDefinition: policy client not set") } if ref.DefinitionVersion != nil { - if *definitionsVersionedClient == nil { - *definitionsVersionedClient = policyClient.NewDefinitionVersionsClient() - } - - return az.fetchReferencedPolicyDefinitionVersions(ctx, *definitionsVersionedClient, resID.Name, setDefinition, ref) - } - - if *definitionsClient == nil { - *definitionsClient = policyClient.NewDefinitionsClient() + return az.fetchReferencedPolicyDefinitionVersions(ctx, definitionsVersionedClient, resID.Name, setDefinition, ref) } return az.fetchLatestReferencedPolicyDefinition( - ctx, *definitionsClient, resID.Name, setDefinition, ref.DefinitionVersion, + ctx, definitionsClient, resID.Name, setDefinition, ref.DefinitionVersion, ) } @@ -1424,7 +1451,7 @@ func (az *AlzLib) fetchReferencedPolicyDefinitionVersions( *ref.DefinitionVersion, definitionName, JoinNameAndVersion(*setDefinition.Name, setDefinition.Properties.Version), - az.policyDefinitions[definitionName].Versions(), + az.policyDefinitionVersionsString(definitionName), ) } diff --git a/builtin_fetch.go b/builtin_fetch.go new file mode 100644 index 0000000..f0d5fa3 --- /dev/null +++ b/builtin_fetch.go @@ -0,0 +1,145 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package alzlib + +import ( + "context" + "fmt" + "strings" + + "github.com/Azure/alzlib/assets" + "github.com/Azure/azure-sdk-for-go/sdk/azcore/arm" + "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armpolicy" + "golang.org/x/sync/errgroup" +) + +// definitionReference is a validated policy definition reference from a built-in policy set definition. +type definitionReference struct { + setDefinition *assets.PolicySetDefinition + ref *armpolicy.DefinitionReference + resID *arm.ResourceID +} + +// effectiveParallelism returns the maximum number of concurrent Azure API calls used when fetching +// built-in definitions. Values below 1 fall back to serial fetching. +func (az *AlzLib) effectiveParallelism() int { + if az.Options == nil || az.Options.Parallelism < 1 { + return 1 + } + + return az.Options.Parallelism +} + +// policyClientFactory returns the policy client factory, or nil if it has not been set. +func (az *AlzLib) policyClientFactory() *armpolicy.ClientFactory { + az.mu.RLock() + defer az.mu.RUnlock() + + return az.clients.policyClient +} + +// policyDefinitionVersionsString returns the stored versions of the named policy definition for use in +// error messages. +func (az *AlzLib) policyDefinitionVersionsString(name string) string { + az.mu.RLock() + defer az.mu.RUnlock() + + pdvs, ok := az.policyDefinitions[name] + if !ok || pdvs == nil { + return "[]" + } + + return fmt.Sprint(pdvs.Versions()) +} + +// parseDefinitionReference validates a policy definition reference of the supplied policy set definition. +// It returns nil if the reference is nil, as there is nothing to fetch. +func parseDefinitionReference( + setDefinition *assets.PolicySetDefinition, + ref *armpolicy.DefinitionReference, +) (*definitionReference, error) { + if ref == nil { + return nil, nil + } + + if ref.PolicyDefinitionID == nil { + return nil, fmt.Errorf( + "Alzlib.getBuiltInPolicySets: error getting policy definition references for policy set definition `%s`: "+ + "policy definition ID is nil", + JoinNameAndVersion(*setDefinition.Name, setDefinition.Properties.Version), + ) + } + + resID, err := arm.ParseResourceID(*ref.PolicyDefinitionID) + if err != nil { + return nil, fmt.Errorf( + "Alzlib.getBuiltInPolicySets: error parsing resource id `%s` referenced in policy set `%s`", + *ref.PolicyDefinitionID, + JoinNameAndVersion(*setDefinition.Name, setDefinition.Properties.Version), + ) + } + + return &definitionReference{ + setDefinition: setDefinition, + ref: ref, + resID: resID, + }, nil +} + +// groupByName groups items by case-insensitive definition name, preserving first-seen order. +// Items in one group are processed serially, so two requests for the same definition can never race +// to add it, while different definitions are fetched concurrently. +func groupByName[T any](items []T, name func(T) string) [][]T { + groups := make([][]T, 0, len(items)) + index := make(map[string]int, len(items)) + + for _, item := range items { + key := strings.ToLower(name(item)) + + if i, ok := index[key]; ok { + groups[i] = append(groups[i], item) + continue + } + + index[key] = len(groups) + groups = append(groups, []T{item}) + } + + return groups +} + +// forEachGroupConcurrently calls fn for every item, running at most limit groups at once. +// Items within a group are processed in order. +// The first error cancels the remaining work and is returned. +func forEachGroupConcurrently[T any]( + ctx context.Context, + limit int, + groups [][]T, + fn func(context.Context, T) error, +) error { + if len(groups) == 0 { + return nil + } + + grp, grpCtx := errgroup.WithContext(ctx) + grp.SetLimit(max(limit, 1)) + + for _, group := range groups { + grp.Go(func() error { + for _, item := range group { + if err := grpCtx.Err(); err != nil { + return err + } + + if err := fn(grpCtx, item); err != nil { + return err + } + } + + return nil + }) + } + + return grp.Wait() +} diff --git a/builtin_fetch_test.go b/builtin_fetch_test.go new file mode 100644 index 0000000..0df8e46 --- /dev/null +++ b/builtin_fetch_test.go @@ -0,0 +1,328 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package alzlib + +import ( + "context" + "fmt" + "net/http" + "strings" + "sync" + "testing" + "time" + + "github.com/Azure/alzlib/internal/testutil/fakearm" + "github.com/Azure/alzlib/to" + "github.com/Azure/azure-sdk-for-go/sdk/azcore/arm" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const fakeLatency = 20 * time.Millisecond + +// newFakeAlzLib returns an AlzLib whose policy client talks to a new fake ARM server. +func newFakeAlzLib(tb testing.TB, parallelism int) (*AlzLib, *fakearm.Server) { + tb.Helper() + + srv := fakearm.New() + cf, err := srv.NewClientFactory() + require.NoError(tb, err) + + az := NewAlzLib(&Options{ + Parallelism: parallelism, + AllowOverwrite: false, + UniqueRoleDefinitions: true, + }) + az.AddPolicyClient(cf) + + return az, srv +} + +func policyDefinitionRequest(tb testing.TB, name string, version *string) BuiltInRequest { + tb.Helper() + + resID, err := arm.ParseResourceID(fmt.Sprintf("/providers/Microsoft.Authorization/policyDefinitions/%s", name)) + require.NoError(tb, err) + + return BuiltInRequest{ResourceID: resID, Version: version} +} + +func policySetDefinitionRequest(tb testing.TB, name string) BuiltInRequest { + tb.Helper() + + resID, err := arm.ParseResourceID(fmt.Sprintf("/providers/Microsoft.Authorization/policySetDefinitions/%s", name)) + require.NoError(tb, err) + + return BuiltInRequest{ResourceID: resID} +} + +// addFakePolicyDefinitions registers count policy definitions named pd-NNN and returns versionless +// requests for them. +func addFakePolicyDefinitions(tb testing.TB, srv *fakearm.Server, count int) []BuiltInRequest { + tb.Helper() + + reqs := make([]BuiltInRequest, 0, count) + + for i := range count { + name := fmt.Sprintf("pd-%03d", i) + srv.AddPolicyDefinition(name, "1.0.0") + reqs = append(reqs, policyDefinitionRequest(tb, name, nil)) + } + + return reqs +} + +func TestGetDefinitionsFromAzureRespectsParallelism(t *testing.T) { + t.Parallel() + + for _, p := range []int{1, 4, 10} { + t.Run(fmt.Sprintf("parallelism-%d", p), func(t *testing.T) { + t.Parallel() + + az, srv := newFakeAlzLib(t, p) + srv.Latency = fakeLatency + reqs := addFakePolicyDefinitions(t, srv, 40) + + require.NoError(t, az.GetDefinitionsFromAzure(t.Context(), reqs)) + + for _, req := range reqs { + assert.True(t, az.PolicyDefinitionExists(req.ResourceID.Name, nil), req.ResourceID.Name) + assert.Equal(t, 1, srv.Calls(req.ResourceID.Name), req.ResourceID.Name) + } + + assert.LessOrEqual(t, srv.PeakInFlight(), p) + + if p == 1 { + assert.Equal(t, 1, srv.PeakInFlight()) + } else { + assert.Greater(t, srv.PeakInFlight(), 1) + } + }) + } +} + +func TestGetDefinitionsFromAzureReferencedDefinitionsFetchedConcurrently(t *testing.T) { + t.Parallel() + + const parallelism = 8 + + az, srv := newFakeAlzLib(t, parallelism) + srv.Latency = fakeLatency + + refs := make([]fakearm.Reference, 0, 30) + + for i := range 30 { + name := fmt.Sprintf("ref-pd-%03d", i) + srv.AddPolicyDefinition(name, "1.0.0", "1.1.0") + + var version *string + if i%2 == 0 { + version = to.Ptr("1.*.*") + } + + refs = append(refs, fakearm.Reference{Name: name, Version: version}) + } + + srv.AddPolicySetDefinition("psd", []string{"1.0.0"}, refs...) + + reqs := []BuiltInRequest{policySetDefinitionRequest(t, "psd")} + require.NoError(t, az.GetDefinitionsFromAzure(t.Context(), reqs)) + + require.True(t, az.PolicySetDefinitionExists("psd", nil)) + + for _, ref := range refs { + assert.True(t, az.PolicyDefinitionExists(ref.Name, ref.Version), ref.Name) + assert.Equal(t, 1, srv.Calls(ref.Name), ref.Name) + } + + assert.Greater(t, srv.PeakInFlight(), 1) + assert.LessOrEqual(t, srv.PeakInFlight(), parallelism) +} + +func TestGetDefinitionsFromAzureSharedDefinitionsFetchedOnce(t *testing.T) { + t.Parallel() + + az, srv := newFakeAlzLib(t, 10) + srv.Latency = fakeLatency + + shared := []string{"shared-a", "shared-b", "shared-c"} + for _, name := range shared { + srv.AddPolicyDefinition(name, "1.0.0") + } + + srv.AddPolicyDefinition("direct-dup", "1.0.0") + srv.AddPolicySetDefinition("psd-1", []string{"1.0.0"}, + fakearm.Reference{Name: "shared-a"}, + fakearm.Reference{Name: "shared-b", Version: to.Ptr("1.*.*")}, + fakearm.Reference{Name: "shared-c"}, + ) + srv.AddPolicySetDefinition("psd-2", []string{"1.0.0"}, + fakearm.Reference{Name: "shared-a", Version: to.Ptr("1.*.*")}, + fakearm.Reference{Name: "shared-b"}, + fakearm.Reference{Name: "shared-c", Version: to.Ptr("1.0.*")}, + ) + + reqs := []BuiltInRequest{ + policyDefinitionRequest(t, "direct-dup", nil), + policyDefinitionRequest(t, "direct-dup", nil), + policyDefinitionRequest(t, "direct-dup", to.Ptr("1.*.*")), + policySetDefinitionRequest(t, "psd-1"), + policySetDefinitionRequest(t, "psd-2"), + policySetDefinitionRequest(t, "psd-1"), + } + + require.NoError(t, az.GetDefinitionsFromAzure(t.Context(), reqs)) + + for _, name := range append(shared, "direct-dup", "psd-1", "psd-2") { + assert.Equal(t, 1, srv.Calls(name), name) + } +} + +func TestGetDefinitionsFromAzureConcurrentCallers(t *testing.T) { + t.Parallel() + + const callers = 4 + + az, srv := newFakeAlzLib(t, 10) + srv.Latency = fakeLatency + reqs := addFakePolicyDefinitions(t, srv, 20) + + var wg sync.WaitGroup + + errs := make([]error, callers) + + for i := range callers { + wg.Go(func() { + errs[i] = az.GetDefinitionsFromAzure(t.Context(), reqs) + }) + } + + wg.Wait() + + for _, err := range errs { + require.NoError(t, err) + } + + for _, req := range reqs { + assert.Equal(t, 1, srv.Calls(req.ResourceID.Name), req.ResourceID.Name) + } +} + +func TestGetDefinitionsFromAzureFirstErrorCancelsRemainingWork(t *testing.T) { + t.Parallel() + + const count = 50 + + az, srv := newFakeAlzLib(t, 4) + srv.Latency = fakeLatency + reqs := addFakePolicyDefinitions(t, srv, count) + srv.FailDefinition(reqs[0].ResourceID.Name, http.StatusNotFound) + + err := az.GetDefinitionsFromAzure(t.Context(), reqs) + require.Error(t, err) + assert.Contains(t, err.Error(), reqs[0].ResourceID.Name) + assert.Less(t, srv.TotalCalls(), count) +} + +func TestGetDefinitionsFromAzureHonoursCallerCancellation(t *testing.T) { + t.Parallel() + + az, srv := newFakeAlzLib(t, 4) + srv.Latency = time.Second + reqs := addFakePolicyDefinitions(t, srv, 20) + + ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond) + defer cancel() + + start := time.Now() + err := az.GetDefinitionsFromAzure(ctx, reqs) + + require.Error(t, err) + assert.Less(t, time.Since(start), 2*time.Second) +} + +func TestGetDefinitionsFromAzureMissingReferencedVersionReportsAvailableVersions(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + versions []string + want string + }{ + {name: "other versions available", versions: []string{"1.0.0"}, want: "Versions available: [1.0.0]"}, + {name: "no versions available", versions: nil, want: "Versions available: []"}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + az, srv := newFakeAlzLib(t, 10) + srv.AddPolicyDefinition("pd-missing-version", tc.versions...) + srv.AddPolicySetDefinition("psd", []string{"1.0.0"}, + fakearm.Reference{Name: "pd-missing-version", Version: to.Ptr("2.*.*")}, + ) + + err := az.GetDefinitionsFromAzure(t.Context(), []BuiltInRequest{policySetDefinitionRequest(t, "psd")}) + require.Error(t, err) + assert.Contains(t, err.Error(), tc.want) + }) + } +} + +func TestEffectiveParallelism(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + opts *Options + want int + }{ + {name: "nil options", opts: nil, want: 1}, + {name: "zero", opts: &Options{Parallelism: 0}, want: 1}, + {name: "negative", opts: &Options{Parallelism: -3}, want: 1}, + {name: "positive", opts: &Options{Parallelism: 7}, want: 7}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + az := &AlzLib{Options: tc.opts} + assert.Equal(t, tc.want, az.effectiveParallelism()) + }) + } +} + +func TestGroupByNameIsCaseInsensitiveAndOrdered(t *testing.T) { + t.Parallel() + + items := []string{"b", "A", "c", "a", "B", "b"} + groups := groupByName(items, strings.Clone) + + assert.Equal(t, [][]string{{"b", "B", "b"}, {"A", "a"}, {"c"}}, groups) + assert.Empty(t, groupByName([]string{}, strings.Clone)) +} + +// BenchmarkGetDefinitionsFromAzureSimulatedLatency measures fetching 100 policy definitions with +// 10ms of simulated latency per request, at different levels of parallelism. +func BenchmarkGetDefinitionsFromAzureSimulatedLatency(b *testing.B) { + for _, p := range []int{1, 10} { + b.Run(fmt.Sprintf("parallelism-%d", p), func(b *testing.B) { + for b.Loop() { + b.StopTimer() + + az, srv := newFakeAlzLib(b, p) + srv.Latency = 10 * time.Millisecond + reqs := addFakePolicyDefinitions(b, srv, 100) + + b.StartTimer() + + if err := az.GetDefinitionsFromAzure(b.Context(), reqs); err != nil { + b.Fatal(err) + } + } + }) + } +} diff --git a/deployment/hierarchy.go b/deployment/hierarchy.go index 72e8c54..5f03c01 100644 --- a/deployment/hierarchy.go +++ b/deployment/hierarchy.go @@ -117,6 +117,17 @@ func (h *Hierarchy) FromArchitecture( if architecture == nil { return fmt.Errorf("Hierarchy.FromArchitecture: getting architecture `%s`", arch) } + + // Fetch the built-in definitions for the whole architecture in one batch, so that they are spread + // across the configured parallelism rather than requested one management group at a time. + if err := h.alzlib.GetDefinitionsFromAzure(ctx, h.architectureBuiltInRequests(architecture)); err != nil { + return fmt.Errorf( + "Hierarchy.FromArchitecture: getting policy definitions from Azure for architecture `%s`: %w", + arch, + err, + ) + } + // Get the architecture root management groups. for _, a := range architecture.RootMgs() { if err := recurseAddManagementGroup(ctx, h, a, externalParentID, location, true, 0); err != nil { @@ -252,6 +263,85 @@ func recurseAddManagementGroup( return nil } +// architectureBuiltInRequests returns the requests for the definitions referenced by the policy +// assignments of every management group in the architecture. +// Policy assignments that cannot be resolved are skipped, so that addManagementGroup reports them +// with its usual error when the management group is added. +func (h *Hierarchy) architectureBuiltInRequests(architecture *alzlib.Architecture) []alzlib.BuiltInRequest { + reqs := make([]alzlib.BuiltInRequest, 0, builtInRequestSliceCapacity) + seen := mapset.NewThreadUnsafeSet[string]() + + var walk func(mgs []*alzlib.ArchitectureManagementGroup) + + walk = func(mgs []*alzlib.ArchitectureManagementGroup) { + slices.SortFunc(mgs, func(a, b *alzlib.ArchitectureManagementGroup) int { + return strings.Compare(a.ID(), b.ID()) + }) + + for _, mg := range mgs { + for _, pa := range sortedSet(archetypesPolicyAssignments(mg.Archetypes())) { + if !seen.Add(pa) { + continue + } + + req, err := h.builtInRequestForPolicyAssignment(pa) + if err != nil { + continue + } + + reqs = append(reqs, req) + } + + walk(mg.Children()) + } + } + + walk(architecture.RootMgs()) + + return reqs +} + +// archetypesPolicyAssignments combines the policy assignments of the supplied archetypes into a single set. +func archetypesPolicyAssignments(archetypes []*alzlib.Archetype) mapset.Set[string] { + res := mapset.NewThreadUnsafeSet[string]() + for _, archetype := range archetypes { + res = res.Union(archetype.PolicyAssignments) + } + + return res +} + +// errPolicyAssignmentNotFound is returned by builtInRequestForPolicyAssignment when the policy assignment +// does not exist in the library. +var errPolicyAssignmentNotFound = errors.New("policy assignment does not exist in the library") + +// builtInRequestForPolicyAssignment returns the request for the definition referenced by the named +// policy assignment. +func (h *Hierarchy) builtInRequestForPolicyAssignment(name string) (alzlib.BuiltInRequest, error) { + polAssign := h.alzlib.PolicyAssignment(name) + if polAssign == nil { + return alzlib.BuiltInRequest{}, errPolicyAssignmentNotFound + } + + referencedResourceID, definitionVersion, err := polAssign.ReferencedPolicyDefinitionResourceIDAndVersion() + if err != nil { + return alzlib.BuiltInRequest{}, err + } + + return alzlib.BuiltInRequest{ + ResourceID: referencedResourceID, + Version: definitionVersion, + }, nil +} + +// sortedSet returns the members of the set as a sorted slice. +func sortedSet(s mapset.Set[string]) []string { + res := s.ToSlice() + slices.Sort(res) + + return res +} + // addManagementGroup adds a management group to the hierarchy, with a parent if specified. // If the parent is not specified, the management group is considered the root of the hierarchy. // The archetype should have been obtained using the `AlzLib.CopyArchetype` method. @@ -314,14 +404,11 @@ func (h *Hierarchy) addManagementGroup( assignedPolicyDefinitions := make([]alzlib.BuiltInRequest, 0, builtInRequestSliceCapacity) // Combine all assignments from all supplied archetypes into a single set - allPolicyAssignments := mapset.NewThreadUnsafeSet[string]() - for _, archetype := range req.archetypes { - allPolicyAssignments = allPolicyAssignments.Union(archetype.PolicyAssignments) - } + allPolicyAssignments := archetypesPolicyAssignments(req.archetypes) for pa := range allPolicyAssignments.Iter() { - polAssign := h.alzlib.PolicyAssignment(pa) - if polAssign == nil { + builtInReq, err := h.builtInRequestForPolicyAssignment(pa) + if errors.Is(err, errPolicyAssignmentNotFound) { return nil, fmt.Errorf( "Hierarchy.AddManagementGroup: policy assignment `%s` referenced in management group `%s` "+ "does not exist in the library", @@ -330,7 +417,6 @@ func (h *Hierarchy) addManagementGroup( ) } - referencedResourceID, definitionVersion, err := polAssign.ReferencedPolicyDefinitionResourceIDAndVersion() if err != nil { return nil, fmt.Errorf( "Hierarchy.AddManagementGroup: getting referenced policy definition resource ID "+ @@ -341,13 +427,11 @@ func (h *Hierarchy) addManagementGroup( ) } - assignedPolicyDefinitions = append(assignedPolicyDefinitions, alzlib.BuiltInRequest{ - ResourceID: referencedResourceID, - Version: definitionVersion, - }) + assignedPolicyDefinitions = append(assignedPolicyDefinitions, builtInReq) } - // Definitions are retrieved only at the time of adding the management group to the hierarchy. + // Definitions are normally prefetched for the whole architecture by FromArchitecture, in which case + // this is a no-op. It is kept so that management groups added any other way still get their definitions. if err := h.alzlib.GetDefinitionsFromAzure(ctx, assignedPolicyDefinitions); err != nil { return nil, fmt.Errorf( "Hierarchy.AddManagementGroup: adding mg `%s` error getting policy definitions from Azure: %w", diff --git a/deployment/hierarchy_fetch_test.go b/deployment/hierarchy_fetch_test.go new file mode 100644 index 0000000..4147ff8 --- /dev/null +++ b/deployment/hierarchy_fetch_test.go @@ -0,0 +1,178 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package deployment + +import ( + "fmt" + "net/http" + "testing" + "testing/fstest" + "time" + + "github.com/Azure/alzlib" + "github.com/Azure/alzlib/internal/testutil/fakearm" + "github.com/Azure/alzlib/to" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const ( + fetchTestArchitecture = "fetchtest" + fetchTestPdIDFmt = "/providers/Microsoft.Authorization/policyDefinitions/%s" + fetchTestPsdIDFmt = "/providers/Microsoft.Authorization/policySetDefinitions/%s" +) + +// fetchTestArchitectureYAML has four management groups whose archetypes reference overlapping +// built-in definitions. +const fetchTestArchitectureYAML = `name: fetchtest +management_groups: + - id: root + display_name: root + archetypes: [a] + exists: false + parent_id: null + - id: child1 + display_name: child1 + archetypes: [b] + exists: false + parent_id: root + - id: child2 + display_name: child2 + archetypes: [c] + exists: false + parent_id: root + - id: grandchild + display_name: grandchild + archetypes: [a, c] + exists: false + parent_id: child1 +` + +func fetchTestArchetypeYAML(name string, assignments ...string) string { + res := fmt.Sprintf("name: %s\npolicy_assignments:\n", name) + for _, pa := range assignments { + res += fmt.Sprintf(" - %s\n", pa) + } + + return res + "policy_definitions: []\npolicy_set_definitions: []\nrole_definitions: []\n" +} + +func fetchTestAssignmentJSON(name, definitionID string) string { + return fmt.Sprintf(`{ + "name": %q, + "type": "Microsoft.Authorization/policyAssignments", + "apiVersion": "2022-06-01", + "location": "${default_location}", + "identity": {"type": "None"}, + "properties": { + "displayName": %q, + "description": %q, + "enforcementMode": "Default", + "notScopes": [], + "parameters": {}, + "policyDefinitionId": %q, + "scope": "/providers/Microsoft.Management/managementGroups/PLACEHOLDER" + } +}`, name, name, name+" description", definitionID) +} + +// newFetchTestHierarchy returns a hierarchy whose AlzLib is initialised from an in-memory library and +// fetches built-in definitions from a fake ARM server. +func newFetchTestHierarchy(t *testing.T, archetypeC []string, assignments map[string]string) (*Hierarchy, *fakearm.Server) { + t.Helper() + + lib := fstest.MapFS{ + "fetchtest.alz_architecture_definition.yaml": {Data: []byte(fetchTestArchitectureYAML)}, + "a.alz_archetype_definition.yaml": {Data: []byte(fetchTestArchetypeYAML("a", "pa-pd-1", "pa-psd-1"))}, + "b.alz_archetype_definition.yaml": {Data: []byte(fetchTestArchetypeYAML("b", "pa-pd-2", "pa-psd-1"))}, + "c.alz_archetype_definition.yaml": {Data: []byte(fetchTestArchetypeYAML("c", archetypeC...))}, + } + + for name, definitionID := range assignments { + lib[name+".alz_policy_assignment.json"] = &fstest.MapFile{Data: []byte(fetchTestAssignmentJSON(name, definitionID))} + } + + srv := fakearm.New() + srv.Latency = 20 * time.Millisecond + + for _, pd := range []string{"pd-1", "pd-2", "pd-3", "pd-4", "pd-5"} { + srv.AddPolicyDefinition(pd, "1.0.0") + } + + srv.AddPolicySetDefinition("psd-1", []string{"1.0.0"}, + fakearm.Reference{Name: "pd-1", Version: to.Ptr("1.*.*")}, + fakearm.Reference{Name: "pd-4"}, + fakearm.Reference{Name: "pd-5"}, + ) + + cf, err := srv.NewClientFactory() + require.NoError(t, err) + + az := alzlib.NewAlzLib(nil) + az.AddPolicyClient(cf) + require.NoError(t, az.Init(t.Context(), alzlib.NewCustomLibraryReferenceFromFS("fetchtest", lib))) + + return NewHierarchy(az), srv +} + +func fetchTestAssignments() map[string]string { + return map[string]string{ + "pa-pd-1": fmt.Sprintf(fetchTestPdIDFmt, "pd-1"), + "pa-pd-2": fmt.Sprintf(fetchTestPdIDFmt, "pd-2"), + "pa-pd-3": fmt.Sprintf(fetchTestPdIDFmt, "pd-3"), + "pa-psd-1": fmt.Sprintf(fetchTestPsdIDFmt, "psd-1"), + } +} + +func TestFromArchitecturePrefetchesOnce(t *testing.T) { + t.Parallel() + + h, srv := newFetchTestHierarchy(t, []string{"pa-pd-3", "pa-pd-1"}, fetchTestAssignments()) + + require.NoError(t, h.FromArchitecture(t.Context(), fetchTestArchitecture, "external", eastUSLocation)) + + assert.Equal(t, []string{"child1", "child2", "grandchild", "root"}, h.ManagementGroupNames()) + assert.Equal(t, map[string]int{"pd-1": 1, "pd-2": 1, "pd-3": 1, "pd-4": 1, "pd-5": 1, "psd-1": 1}, srv.AllCalls()) + assert.Greater(t, srv.PeakInFlight(), 1) +} + +func TestFromArchitectureFetchErrorIsWrapped(t *testing.T) { + t.Parallel() + + h, srv := newFetchTestHierarchy(t, []string{"pa-pd-3", "pa-pd-1"}, fetchTestAssignments()) + srv.FailDefinition("pd-3", http.StatusNotFound) + + err := h.FromArchitecture(t.Context(), fetchTestArchitecture, "external", eastUSLocation) + require.Error(t, err) + assert.Contains( + t, + err.Error(), + "Hierarchy.FromArchitecture: getting policy definitions from Azure for architecture `fetchtest`", + ) + assert.Contains(t, err.Error(), "pd-3") + assert.Empty(t, h.ManagementGroupNames()) +} + +func TestFromArchitectureUnresolvableAssignmentErrorUnchanged(t *testing.T) { + t.Parallel() + + assignments := fetchTestAssignments() + assignments["pa-bad"] = "not-a-resource-id" + + h, srv := newFetchTestHierarchy(t, []string{"pa-pd-3", "pa-bad"}, assignments) + + err := h.FromArchitecture(t.Context(), fetchTestArchitecture, "external", eastUSLocation) + require.Error(t, err) + // The management group reported depends on the order in which recurseAddManagementGroup visits + // children, which is not defined, so only the assignment is asserted. + assert.Contains( + t, + err.Error(), + "Hierarchy.AddManagementGroup: getting referenced policy definition resource ID "+ + "for policy assignment `pa-bad` in management group ", + ) + assert.NotContains(t, err.Error(), "Hierarchy.FromArchitecture: getting policy definitions from Azure") + // The resolvable definitions are still fetched once by the prefetch. + assert.Equal(t, 1, srv.Calls("pd-3")) +} diff --git a/go.mod b/go.mod index b534bf2..98b7f69 100644 --- a/go.mod +++ b/go.mod @@ -18,6 +18,7 @@ require ( github.com/nao1215/markdown v0.6.0 github.com/spf13/cobra v1.8.1 github.com/stretchr/testify v1.11.1 + golang.org/x/sync v0.18.0 gopkg.in/yaml.v3 v3.0.1 ) diff --git a/go.sum b/go.sum index 781093f..34f988f 100644 --- a/go.sum +++ b/go.sum @@ -113,6 +113,8 @@ golang.org/x/crypto v0.45.0 h1:jMBrvKuj23MTlT0bQEOBcAE0mjg8mK9RXFhRH6nyF3Q= golang.org/x/crypto v0.45.0/go.mod h1:XTGrrkGJve7CYK7J8PEww4aY7gM3qMCElcJQ8n8JdX4= golang.org/x/net v0.47.0 h1:Mx+4dIFzqraBXUugkia1OOvlD6LemFo1ALMHjrXDOhY= golang.org/x/net v0.47.0/go.mod h1:/jNxtkgq5yWUGYkaZGqo27cfGZ1c5Nen03aYrrKpVRU= +golang.org/x/sync v0.18.0 h1:kr88TuHDroi+UVf+0hZnirlk8o8T+4MrK6mr60WkH/I= +golang.org/x/sync v0.18.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI= golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.38.0 h1:3yZWxaJjBmCWXqhN1qh02AkOnCQ1poK6oF+a7xWL6Gc= golang.org/x/sys v0.38.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= diff --git a/internal/testutil/fakearm/fakearm.go b/internal/testutil/fakearm/fakearm.go new file mode 100644 index 0000000..b3c06f0 --- /dev/null +++ b/internal/testutil/fakearm/fakearm.go @@ -0,0 +1,354 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Package fakearm provides an in-process fake of the Azure Resource Manager built-in policy APIs. +// It is intended for tests that need deterministic, offline responses and want to observe how many +// requests are made and how many run concurrently. +package fakearm + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "maps" + "net/http" + "strings" + "sync" + "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/resources/armpolicy" +) + +const ( + policyDefinitionIDFmt = "/providers/Microsoft.Authorization/policyDefinitions/%s" + policySetDefinitionIDFmt = "/providers/Microsoft.Authorization/policySetDefinitions/%s" + policyDefinitionsSeg = "policydefinitions" + policySetDefinitionsSeg = "policysetdefinitions" + versionsSeg = "versions" + minPathSegments = 4 + versionsPathSegments = 5 + tokenLifetime = time.Hour +) + +// Reference is a policy definition referenced by a fake policy set definition. +type Reference struct { + // Name is the name of the referenced policy definition. + Name string + // Version is the optional version constraint of the reference, e.g. `1.*.*`. + Version *string +} + +type policyDefinitionEntry struct { + name string + versions []string +} + +type setDefinition struct { + name string + versions []string + refs []Reference +} + +// Server is a fake ARM endpoint for built-in policy (set) definitions. +// It implements [policy.Transporter] and is safe for concurrent use. +type Server struct { + // Latency is added to every request. The request context is honoured while waiting. + Latency time.Duration + + mu sync.Mutex + policyDefinitions map[string]policyDefinitionEntry + policySetDefinitions map[string]setDefinition + failures map[string]int + calls map[string]int + totalCalls int + inFlight int + peakInFlight int +} + +// New returns an empty fake ARM server. +func New() *Server { + return &Server{ + policyDefinitions: make(map[string]policyDefinitionEntry), + policySetDefinitions: make(map[string]setDefinition), + failures: make(map[string]int), + calls: make(map[string]int), + } +} + +// AddPolicyDefinition registers a built-in policy definition with the supplied versions. +// The last version is returned as the latest version. +func (s *Server) AddPolicyDefinition(name string, versions ...string) { + s.mu.Lock() + defer s.mu.Unlock() + + s.policyDefinitions[strings.ToLower(name)] = policyDefinitionEntry{name: name, versions: versions} +} + +// AddPolicySetDefinition registers a built-in policy set definition with the supplied versions and +// referenced policy definitions. The last version is returned as the latest version. +func (s *Server) AddPolicySetDefinition(name string, versions []string, refs ...Reference) { + s.mu.Lock() + defer s.mu.Unlock() + + s.policySetDefinitions[strings.ToLower(name)] = setDefinition{name: name, versions: versions, refs: refs} +} + +// FailDefinition makes every request for the named definition return the supplied HTTP status code. +func (s *Server) FailDefinition(name string, status int) { + s.mu.Lock() + defer s.mu.Unlock() + + s.failures[strings.ToLower(name)] = status +} + +// Calls returns the number of requests made for the named definition. +func (s *Server) Calls(name string) int { + s.mu.Lock() + defer s.mu.Unlock() + + return s.calls[strings.ToLower(name)] +} + +// AllCalls returns a copy of the number of requests made per lower case definition name. +func (s *Server) AllCalls() map[string]int { + s.mu.Lock() + defer s.mu.Unlock() + + return maps.Clone(s.calls) +} + +// TotalCalls returns the total number of requests made. +func (s *Server) TotalCalls() int { + s.mu.Lock() + defer s.mu.Unlock() + + return s.totalCalls +} + +// PeakInFlight returns the highest number of requests that were in flight at the same time. +func (s *Server) PeakInFlight() int { + s.mu.Lock() + defer s.mu.Unlock() + + return s.peakInFlight +} + +// NewClientFactory returns an armpolicy client factory that sends all requests to this server. +// Retries are disabled so that injected failures surface immediately. +func (s *Server) NewClientFactory() (*armpolicy.ClientFactory, error) { + opts := &arm.ClientOptions{ + ClientOptions: policy.ClientOptions{ + Transport: s, + Retry: policy.RetryOptions{MaxRetries: -1}, + }, + } + + cf, err := armpolicy.NewClientFactory("", credential{}, opts) + if err != nil { + return nil, fmt.Errorf("fakearm: creating client factory: %w", err) + } + + return cf, nil +} + +// Do implements [policy.Transporter]. +func (s *Server) Do(req *http.Request) (*http.Response, error) { + segs := strings.Split(strings.Trim(strings.ToLower(req.URL.Path), "/"), "/") + // Expected: providers/microsoft.authorization//[/versions]. + if len(segs) < minPathSegments || segs[0] != "providers" || segs[1] != "microsoft.authorization" { + return jsonResponse(req, http.StatusNotFound, errorBody("PathNotFound", req.URL.Path)) + } + + defType, name := segs[2], segs[3] + listVersions := len(segs) == versionsPathSegments && segs[4] == versionsSeg + + s.begin(name) + defer s.end() + + if s.Latency > 0 { + select { + case <-req.Context().Done(): + return nil, fmt.Errorf("fakearm: %w", req.Context().Err()) + case <-time.After(s.Latency): + } + } + + s.mu.Lock() + status, fail := s.failures[name] + pd, pdOK := s.policyDefinitions[name] + psd, psdOK := s.policySetDefinitions[name] + s.mu.Unlock() + + if fail { + return jsonResponse(req, status, errorBody("InjectedFailure", name)) + } + + switch { + case defType == policyDefinitionsSeg && pdOK: + if listVersions { + return jsonResponse(req, http.StatusOK, policyDefinitionVersionList(pd.name, pd.versions)) + } + + return jsonResponse(req, http.StatusOK, policyDefinition(pd.name, latest(pd.versions))) + case defType == policySetDefinitionsSeg && psdOK: + if listVersions { + return jsonResponse(req, http.StatusOK, policySetDefinitionVersionList(psd)) + } + + return jsonResponse(req, http.StatusOK, policySetDefinition(psd.name, latest(psd.versions), psd.refs)) + } + + return jsonResponse(req, http.StatusNotFound, errorBody("PolicyDefinitionNotFound", name)) +} + +func (s *Server) begin(name string) { + s.mu.Lock() + defer s.mu.Unlock() + + s.calls[name]++ + s.totalCalls++ + s.inFlight++ + s.peakInFlight = max(s.peakInFlight, s.inFlight) +} + +func (s *Server) end() { + s.mu.Lock() + defer s.mu.Unlock() + + s.inFlight-- +} + +func latest(versions []string) string { + if len(versions) == 0 { + return "1.0.0" + } + + return versions[len(versions)-1] +} + +func policyDefinitionProperties(name, version string) map[string]any { + return map[string]any{ + "displayName": name, + "description": name + " description", + "policyType": "BuiltIn", + "mode": "All", + "metadata": map[string]any{}, + "parameters": map[string]any{}, + "policyRule": map[string]any{ + "if": map[string]any{"field": "type", "equals": "Microsoft.Resources/subscriptions"}, + "then": map[string]any{"effect": "audit"}, + }, + "version": version, + } +} + +func policyDefinition(name, version string) map[string]any { + return map[string]any{ + "id": fmt.Sprintf(policyDefinitionIDFmt, name), + "name": name, + "type": "Microsoft.Authorization/policyDefinitions", + "properties": policyDefinitionProperties(name, version), + } +} + +func policyDefinitionVersionList(name string, versions []string) map[string]any { + value := make([]any, 0, len(versions)) + for _, v := range versions { + value = append(value, map[string]any{ + "id": fmt.Sprintf(policyDefinitionIDFmt+"/versions/%s", name, v), + "name": v, + "type": "Microsoft.Authorization/policyDefinitions/versions", + "properties": policyDefinitionProperties(name, v), + }) + } + + return map[string]any{"value": value} +} + +func policySetDefinitionProperties(name, version string, refs []Reference) map[string]any { + pdRefs := make([]any, 0, len(refs)) + + for i, ref := range refs { + r := map[string]any{ + "policyDefinitionId": fmt.Sprintf(policyDefinitionIDFmt, ref.Name), + "policyDefinitionReferenceId": fmt.Sprintf("ref-%d", i), + "parameters": map[string]any{}, + } + if ref.Version != nil { + r["definitionVersion"] = *ref.Version + } + + pdRefs = append(pdRefs, r) + } + + return map[string]any{ + "displayName": name, + "description": name + " description", + "policyType": "BuiltIn", + "metadata": map[string]any{}, + "parameters": map[string]any{}, + "policyDefinitions": pdRefs, + "version": version, + } +} + +func policySetDefinition(name, version string, refs []Reference) map[string]any { + return map[string]any{ + "id": fmt.Sprintf(policySetDefinitionIDFmt, name), + "name": name, + "type": "Microsoft.Authorization/policySetDefinitions", + "properties": policySetDefinitionProperties(name, version, refs), + } +} + +func policySetDefinitionVersionList(psd setDefinition) map[string]any { + value := make([]any, 0, len(psd.versions)) + for _, v := range psd.versions { + value = append(value, map[string]any{ + "id": fmt.Sprintf(policySetDefinitionIDFmt+"/versions/%s", psd.name, v), + "name": v, + "type": "Microsoft.Authorization/policySetDefinitions/versions", + "properties": policySetDefinitionProperties(psd.name, v, psd.refs), + }) + } + + return map[string]any{"value": value} +} + +func errorBody(code, target string) map[string]any { + return map[string]any{ + "error": map[string]any{ + "code": code, + "message": "fakearm: " + target, + }, + } +} + +func jsonResponse(req *http.Request, status int, body any) (*http.Response, error) { + b, err := json.Marshal(body) + if err != nil { + return nil, fmt.Errorf("fakearm: marshalling response: %w", err) + } + + return &http.Response{ + StatusCode: status, + Status: http.StatusText(status), + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(bytes.NewReader(b)), + ContentLength: int64(len(b)), + Request: req, + }, nil +} + +// credential is a fake azcore.TokenCredential that always returns a static token. +type credential struct{} + +// GetToken implements azcore.TokenCredential. +func (credential) GetToken(context.Context, policy.TokenRequestOptions) (azcore.AccessToken, error) { + return azcore.AccessToken{Token: "fake", ExpiresOn: time.Now().Add(tokenLifetime)}, nil +} From 8c42f3f4b64dc44f2b211b517b942f8e2fcfa50f Mon Sep 17 00:00:00 2001 From: Yingjie Date: Wed, 28 Oct 2026 18:08:19 +0800 Subject: [PATCH 2/2] fix: retry built-in definition fetches throttled by Azure Resource Manager --- GNUmakefile | 9 +- alzlib.go | 1 + builtin_fetch.go | 94 ++++++++++++++++++++- builtin_fetch_test.go | 121 +++++++++++++++++++++++++++ internal/testutil/fakearm/fakearm.go | 36 ++++++++ 5 files changed, 257 insertions(+), 4 deletions(-) diff --git a/GNUmakefile b/GNUmakefile index bc7b51a..51ea982 100644 --- a/GNUmakefile +++ b/GNUmakefile @@ -1,4 +1,7 @@ TEST?=$$(go list ./... |grep -v 'vendor') +# Several packages call Azure; run them one at a time so that concurrent CI jobs sharing a tenant +# do not exceed the Azure Resource Manager read limits. +TESTFLAGS?=-p 1 .PHONY: lint lint: @@ -6,17 +9,17 @@ lint: .PHONY: testrace testrace: - go test -test.v -race $(TEST) + go test $(TESTFLAGS) -test.v -race $(TEST) .PHONY: test test: - go test -test.v $(TEST) + go test $(TESTFLAGS) -test.v $(TEST) # Create a test coverage report and launch a browser to view it .PHONY: testcover testcover: if [ -f "coverage.out" ]; then rm coverage.out; fi - go test -v -coverprofile=coverage.out -covermode=count $(TEST) + go test $(TESTFLAGS) -v -coverprofile=coverage.out -covermode=count $(TEST) .PHONY: tool tool: diff --git a/alzlib.go b/alzlib.go index 286147a..5f2e9fe 100644 --- a/alzlib.go +++ b/alzlib.go @@ -70,6 +70,7 @@ type Options struct { // and policy set definitions. // Requests for the same definition name are always made serially. // Set to 1 to fetch all definitions serially. + // Requests throttled by Azure Resource Manager (HTTP 429) are retried after the delay it requests. Parallelism int // UniqueRoleDefinitions indicates whether to update the role definitions to be unique per management group. // If this is not set, you may end up with conflicting role definition names. diff --git a/builtin_fetch.go b/builtin_fetch.go index f0d5fa3..216548b 100644 --- a/builtin_fetch.go +++ b/builtin_fetch.go @@ -5,15 +5,31 @@ package alzlib import ( "context" + "errors" "fmt" + "math/rand/v2" + "net/http" + "strconv" "strings" + "time" "github.com/Azure/alzlib/assets" + "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/resourcemanager/resources/armpolicy" "golang.org/x/sync/errgroup" ) +const ( + // maxThrottleRetries is the number of times a definition is fetched again after Azure Resource + // Manager throttled it (HTTP 429) beyond the retries of the Azure SDK. + maxThrottleRetries = 5 + // throttleBaseDelay is the first back-off delay when a throttled response has no Retry-After. + throttleBaseDelay = 2 * time.Second + // throttleMaxDelay caps every back-off delay, including a requested Retry-After. + throttleMaxDelay = 60 * time.Second +) + // definitionReference is a validated policy definition reference from a built-in policy set definition. type definitionReference struct { setDefinition *assets.PolicySetDefinition @@ -111,6 +127,7 @@ func groupByName[T any](items []T, name func(T) string) [][]T { // forEachGroupConcurrently calls fn for every item, running at most limit groups at once. // Items within a group are processed in order. +// An item throttled by Azure Resource Manager is retried, see withThrottleRetry. // The first error cancels the remaining work and is returned. func forEachGroupConcurrently[T any]( ctx context.Context, @@ -132,7 +149,7 @@ func forEachGroupConcurrently[T any]( return err } - if err := fn(grpCtx, item); err != nil { + if err := withThrottleRetry(grpCtx, func() error { return fn(grpCtx, item) }); err != nil { return err } } @@ -143,3 +160,78 @@ func forEachGroupConcurrently[T any]( return grp.Wait() } + +// withThrottleRetry calls fn and, while it fails because Azure Resource Manager throttled a request +// (HTTP 429), waits and calls it again, up to maxThrottleRetries times. +// The Azure SDK already retries throttled requests, but only a few times. Concurrent fetching, and +// several processes sharing a tenant, can exhaust those retries, so a throttled definition is +// fetched again after the delay requested by Retry-After, or an exponential back-off. +// fn must be safe to call again after a failure: the fetch functions first check whether the +// definition already exists. +func withThrottleRetry(ctx context.Context, fn func() error) error { + for attempt := 0; ; attempt++ { + err := fn() + if err == nil || attempt == maxThrottleRetries || !isThrottled(err) { + return err + } + + timer := time.NewTimer(throttleDelay(err, attempt)) + + select { + case <-ctx.Done(): + timer.Stop() + + return err + case <-timer.C: + } + } +} + +// isThrottled reports whether err is an Azure Resource Manager HTTP 429 response. +func isThrottled(err error) bool { + var respErr *azcore.ResponseError + + return errors.As(err, &respErr) && respErr.StatusCode == http.StatusTooManyRequests +} + +// throttleDelay returns how long to wait before retrying after a throttled response: the delay +// requested by the response, or an exponential back-off with jitter, capped at throttleMaxDelay. +func throttleDelay(err error, attempt int) time.Duration { + var respErr *azcore.ResponseError + if errors.As(err, &respErr) && respErr.RawResponse != nil { + if d, ok := retryAfter(respErr.RawResponse.Header); ok { + return min(d, throttleMaxDelay) + } + } + + backoff := min(throttleBaseDelay<= 0 { + return time.Duration(ms) * time.Millisecond, true + } + } + } + + v := h.Get("Retry-After") + if v == "" { + return 0, false + } + + if s, err := strconv.Atoi(v); err == nil && s >= 0 { + return time.Duration(s) * time.Second, true + } + + if t, err := http.ParseTime(v); err == nil { + return max(time.Until(t), 0), true + } + + return 0, false +} diff --git a/builtin_fetch_test.go b/builtin_fetch_test.go index 0df8e46..fc890e1 100644 --- a/builtin_fetch_test.go +++ b/builtin_fetch_test.go @@ -5,6 +5,7 @@ package alzlib import ( "context" + "errors" "fmt" "net/http" "strings" @@ -223,6 +224,126 @@ func TestGetDefinitionsFromAzureFirstErrorCancelsRemainingWork(t *testing.T) { require.Error(t, err) assert.Contains(t, err.Error(), reqs[0].ResourceID.Name) assert.Less(t, srv.TotalCalls(), count) + assert.Equal(t, 1, srv.Calls(reqs[0].ResourceID.Name), "errors other than HTTP 429 must not be retried") +} + +func TestGetDefinitionsFromAzureRetriesThrottledRequests(t *testing.T) { + t.Parallel() + + // The fake client factory disables the Azure SDK's own retries, so every retry here is alzlib's. + az, srv := newFakeAlzLib(t, 10) + reqs := addFakePolicyDefinitions(t, srv, 20) + + for _, req := range reqs { + srv.ThrottleDefinition(req.ResourceID.Name, 2, "0") + } + + srv.AddPolicyDefinition("ref-versioned", "1.0.0", "1.1.0") + srv.AddPolicyDefinition("ref-latest", "1.0.0") + srv.AddPolicySetDefinition("psd", []string{"1.0.0"}, + fakearm.Reference{Name: "ref-versioned", Version: to.Ptr("1.*.*")}, + fakearm.Reference{Name: "ref-latest"}, + ) + + for _, name := range []string{"psd", "ref-versioned", "ref-latest"} { + srv.ThrottleDefinition(name, maxThrottleRetries, "0") + } + + reqs = append(reqs, policySetDefinitionRequest(t, "psd")) + + require.NoError(t, az.GetDefinitionsFromAzure(t.Context(), reqs)) + + for _, req := range reqs[:20] { + assert.True(t, az.PolicyDefinitionExists(req.ResourceID.Name, nil), req.ResourceID.Name) + assert.Equal(t, 3, srv.Calls(req.ResourceID.Name), req.ResourceID.Name) + } + + assert.True(t, az.PolicySetDefinitionExists("psd", nil)) + assert.True(t, az.PolicyDefinitionExists("ref-versioned", to.Ptr("1.*.*"))) + assert.True(t, az.PolicyDefinitionExists("ref-latest", nil)) + + for _, name := range []string{"psd", "ref-versioned", "ref-latest"} { + assert.Equal(t, maxThrottleRetries+1, srv.Calls(name), name) + } +} + +func TestGetDefinitionsFromAzureGivesUpWhenThrottlingPersists(t *testing.T) { + t.Parallel() + + az, srv := newFakeAlzLib(t, 4) + reqs := addFakePolicyDefinitions(t, srv, 1) + srv.ThrottleDefinition(reqs[0].ResourceID.Name, -1, "0") + + err := az.GetDefinitionsFromAzure(t.Context(), reqs) + require.Error(t, err) + assert.Contains(t, err.Error(), "429") + assert.Equal(t, maxThrottleRetries+1, srv.Calls(reqs[0].ResourceID.Name)) +} + +func TestGetDefinitionsFromAzureThrottleWaitHonoursCancellation(t *testing.T) { + t.Parallel() + + az, srv := newFakeAlzLib(t, 4) + reqs := addFakePolicyDefinitions(t, srv, 1) + srv.ThrottleDefinition(reqs[0].ResourceID.Name, -1, "30") + + ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond) + defer cancel() + + start := time.Now() + err := az.GetDefinitionsFromAzure(ctx, reqs) + + require.Error(t, err) + assert.Less(t, time.Since(start), 5*time.Second) + assert.Equal(t, 1, srv.Calls(reqs[0].ResourceID.Name)) +} + +func TestRetryAfter(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + header http.Header + want time.Duration + ok bool + }{ + {name: "none", header: http.Header{}, ok: false}, + {name: "seconds", header: http.Header{"Retry-After": {"7"}}, want: 7 * time.Second, ok: true}, + {name: "zero", header: http.Header{"Retry-After": {"0"}}, want: 0, ok: true}, + {name: "retry-after-ms", header: http.Header{"Retry-After-Ms": {"250"}}, want: 250 * time.Millisecond, ok: true}, + {name: "x-ms-retry-after-ms", header: http.Header{"X-Ms-Retry-After-Ms": {"1500"}}, want: 1500 * time.Millisecond, ok: true}, + { + name: "milliseconds take precedence", + header: http.Header{"Retry-After": {"9"}, "Retry-After-Ms": {"100"}}, + want: 100 * time.Millisecond, + ok: true, + }, + {name: "past date", header: http.Header{"Retry-After": {"Mon, 02 Jan 2006 15:04:05 GMT"}}, want: 0, ok: true}, + {name: "invalid", header: http.Header{"Retry-After": {"soon"}}, ok: false}, + {name: "negative", header: http.Header{"Retry-After": {"-1"}}, ok: false}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + got, ok := retryAfter(tc.header) + assert.Equal(t, tc.ok, ok) + assert.Equal(t, tc.want, got) + }) + } +} + +func TestThrottleDelayBackoffWithoutRetryAfter(t *testing.T) { + t.Parallel() + + for attempt := range maxThrottleRetries + 1 { + backoff := min(throttleBaseDelay<