Skip to content

Commit

Permalink
fix: more on throttling #249
Browse files Browse the repository at this point in the history
  • Loading branch information
cmendible committed Aug 22, 2024
1 parent 602dd2a commit 8fc33c5
Show file tree
Hide file tree
Showing 3 changed files with 21 additions and 3 deletions.
7 changes: 7 additions & 0 deletions internal/aprl_scanner.go
Original file line number Diff line number Diff line change
Expand Up @@ -188,6 +188,7 @@ func (sc AprlScanner) graphScan(ctx context.Context, graphClient *graph.GraphQue
subs = append(subs, &s)
}

sentQueries := 0
for _, rule := range rules {
if rule.GraphQuery != "" {
result := graphClient.Query(ctx, rule.GraphQuery, subs)
Expand Down Expand Up @@ -258,6 +259,12 @@ func (sc AprlScanner) graphScan(ctx context.Context, graphClient *graph.GraphQue
})
}
}
sentQueries++
if sentQueries == 2 {
// Staggering queries to avoid throttling. Max 10 queries each 5 seconds.
// https://learn.microsoft.com/en-us/azure/governance/resource-graph/concepts/guidance-for-throttled-requests#staggering-queries
time.Sleep(1 * time.Second)
}
}
}

Expand Down
1 change: 1 addition & 0 deletions internal/scanner.go
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,7 @@ func (sc Scanner) Scan(params *ScanParams) {
if err != nil {
log.Fatal().Err(err).Msg("Failed to initialize diagnostic settings scanner")
}

diagResults = diagnosticsScanner.Scan(reportData.ResourceIDs())
}

Expand Down
16 changes: 13 additions & 3 deletions internal/scanners/diagnostics_settings.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import (
"net/http"
"strings"
"sync"
"time"

"github.com/Azure/azqr/internal/azqr"
"github.com/Azure/azure-sdk-for-go/sdk/azcore"
Expand Down Expand Up @@ -59,20 +60,31 @@ func (d *DiagnosticSettingsScanner) ListResourcesWithDiagnosticSettings(resource

// Start workers
// Based on: https://medium.com/insiderengineering/concurrent-http-requests-in-golang-best-practices-and-techniques-f667e5a19dea
numWorkers := 200 // Define the number of workers in the pool
numWorkers := 100 // Define the number of workers in the pool
for w := 0; w < numWorkers; w++ {
go d.worker(jobs, ch, &wg)
}
wg.Add(batches)

// Split resources into batches of 20 items.
batchSize := 20
batchCount := 0
for i := 0; i < len(resources); i += batchSize {
j := i + batchSize
if j > len(resources) {
j = len(resources)
}
jobs <- resources[i:j]

batchCount++
if batchCount == numWorkers {
log.Debug().Msgf("all %d workers are running. Sleeping for 4 seconds to avoid throttling", numWorkers)
batchCount = 0
// there are more batches to process
// Staggering queries to avoid throttling. Max 15 queries each 5 seconds.
// https://learn.microsoft.com/en-us/azure/governance/resource-graph/concepts/guidance-for-throttled-requests#staggering-queries
time.Sleep(4 * time.Second)
}
}

// Wait for all workers to finish
Expand Down Expand Up @@ -127,7 +139,6 @@ func (d *DiagnosticSettingsScanner) restCall(ctx context.Context, resourceIds []

for _, resourceId := range resourceIds {
batch.Requests = append(batch.Requests, ArmBatchRequestItem{
Name: *resourceId,
HttpMethod: http.MethodGet,
RelativeUrl: *resourceId + "/providers/microsoft.insights/diagnosticSettings?api-version=2021-05-01-preview",
})
Expand Down Expand Up @@ -167,7 +178,6 @@ type (
}

ArmBatchRequestItem struct {
Name string `json:"name"`
HttpMethod string `json:"httpMethod"`
RelativeUrl string `json:"relativeUrl"`
}
Expand Down

0 comments on commit 8fc33c5

Please sign in to comment.