Orphan Safety Controller missing VM's due to large resource query (gardener#207)

Author: gagan16k

pkg/azure/access/helpers/fake_resourcegraph.go

@@ -0,0 +1,62 @@
+// SPDX-FileCopyrightText: 2024 SAP SE or an SAP affiliate company and Gardener contributors
+//
+// SPDX-License-Identifier: Apache-2.0
+
+package helpers
+
+import (
+	"context"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resourcegraph/armresourcegraph"
+)
+
+// FakeResourceGraphClient is a fake implementation of ResourceGraphClient for testing.
+type FakeResourceGraphClient struct {
+	// Responses is a list of responses to return in sequence for each call to Resources
+	Responses []armresourcegraph.ClientResourcesResponse
+	// Errors is a list of errors to return in sequence for each call to Resources
+	Errors []error
+	// CallCount tracks how many times Resources has been called
+	CallCount int
+	// RecordedRequests stores all query requests made to the client
+	RecordedRequests []armresourcegraph.QueryRequest
+}
+
+// NewFakeResourceGraphClient creates a new FakeResourceGraphClient for testing.
+func NewFakeResourceGraphClient() *FakeResourceGraphClient {
+	return &FakeResourceGraphClient{
+		Responses:        []armresourcegraph.ClientResourcesResponse{},
+		Errors:           []error{},
+		RecordedRequests: []armresourcegraph.QueryRequest{},
+	}
+}
+
+// Resources implements the ResourceGraphClient interface.
+// Returns the next response/error in the sequence based on CallCount.
+func (f *FakeResourceGraphClient) Resources(_ context.Context, query armresourcegraph.QueryRequest, _ *armresourcegraph.ClientResourcesOptions) (armresourcegraph.ClientResourcesResponse, error) {
+	f.RecordedRequests = append(f.RecordedRequests, query)
+	index := f.CallCount
+	f.CallCount++
+
+	if index < len(f.Errors) && f.Errors[index] != nil {
+		return armresourcegraph.ClientResourcesResponse{}, f.Errors[index]
+	}
+
+	if index < len(f.Responses) {
+		return f.Responses[index], nil
+	}
+
+	return armresourcegraph.ClientResourcesResponse{}, nil
+}
+
+// AddResponse adds a response to be returned by the fake client.
+func (f *FakeResourceGraphClient) AddResponse(response armresourcegraph.ClientResourcesResponse) *FakeResourceGraphClient {
+	f.Responses = append(f.Responses, response)
+	return f
+}
+
+// AddError adds an error to be returned by the fake client.
+func (f *FakeResourceGraphClient) AddError(err error) *FakeResourceGraphClient {
+	f.Errors = append(f.Errors, err)
+	return f
+}

pkg/azure/access/helpers/resourcegraph.go

@@ -12,49 +12,82 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resourcegraph/armresourcegraph"
 	"github.com/gardener/machine-controller-manager-provider-azure/pkg/azure/access/errors"
 	"github.com/gardener/machine-controller-manager-provider-azure/pkg/azure/instrument"
-	"k8s.io/utils/pointer"
+	"k8s.io/klog/v2"
+	"k8s.io/utils/ptr"
 )
 
 const (
 	resourceGraphQueryServiceLabel = "resource_graph_query"
 )
 
+// ResourceGraphClient is an interface for Azure Resource Graph client operations.
+// This allows for easier testing by enabling mock implementations.
+type ResourceGraphClient interface {
+	Resources(ctx context.Context, query armresourcegraph.QueryRequest, options *armresourcegraph.ClientResourcesOptions) (armresourcegraph.ClientResourcesResponse, error)
+}
+
+// Azure client implements interface
+var _ ResourceGraphClient = (*armresourcegraph.Client)(nil)
+
 // MapperFn maps a row of result (represented as map[string]interface{}) to any type T.
 type MapperFn[T any] func(map[string]interface{}) *T
 
 // QueryAndMap fires a resource graph KUSTO query constructing it from queryTemplate and templateArgs.
 // The result of the query are then mapped using a mapperFn and the result or an error is returned.
 // NOTE: All calls to this Azure API are instrumented as prometheus metric.
-func QueryAndMap[T any](ctx context.Context, client *armresourcegraph.Client, subscriptionID string, mapperFn MapperFn[T], queryTemplate string, templateArgs ...any) (results []T, err error) {
+func QueryAndMap[T any](ctx context.Context, client ResourceGraphClient, subscriptionID string, mapperFn MapperFn[T], queryTemplate string, templateArgs ...any) (results []T, err error) {
 	defer instrument.AZAPIMetricRecorderFn(resourceGraphQueryServiceLabel, &err)()
 
 	query := fmt.Sprintf(queryTemplate, templateArgs...)
-	resources, err := client.Resources(ctx,
-		armresourcegraph.QueryRequest{
+	var skipToken *string
+	pageCount := 0
+
+	// Continue fetching results while there is a skipToken
+	for {
+		queryRequest := armresourcegraph.QueryRequest{
 			Query:         to.Ptr(query),
 			Options:       nil,
 			Subscriptions: []*string{to.Ptr(subscriptionID)},
-		}, nil)
+		}
 
-	if err != nil {
-		errors.LogAzAPIError(err, "ResourceGraphQuery failure to execute Query: %s", query)
-		return nil, err
-	}
+		// Set skipToken in options if present for subsequent pages
+		if skipToken != nil {
+			queryRequest.Options = &armresourcegraph.QueryRequestOptions{
+				SkipToken: skipToken,
+			}
+		}
 
-	if resources.TotalRecords == pointer.Int64(0) {
-		return results, nil
-	}
+		resources, err := client.Resources(ctx, queryRequest, nil)
+		if err != nil {
+			errors.LogAzAPIError(err, "ResourceGraphQuery failure to execute Query: %s, with skipToken: %s", query, ptr.Deref(skipToken, "<nil>"))
+			return nil, err
+		}
+		pageCount++
 
-	// resourceResponse.Data is a []interface{}
-	if objSlice, ok := resources.Data.([]interface{}); ok {
-		for _, obj := range objSlice {
-			// Each obj in resourceResponse.Data is a map[string]Interface{}
-			rowElements := obj.(map[string]interface{})
-			result := mapperFn(rowElements)
-			if result != nil {
-				results = append(results, *result)
+		if ptr.Deref(resources.TotalRecords, 0) == 0 {
+			klog.Infof("Query completed: fetched %d pages, no results retrieved", pageCount)
+			return results, nil
+		}
+
+		// resourceResponse.Data is a []interface{}
+		if objSlice, ok := resources.Data.([]interface{}); ok {
+			for _, obj := range objSlice {
+				// Each obj in resourceResponse.Data is a map[string]Interface{}
+				rowElements := obj.(map[string]interface{})
+				result := mapperFn(rowElements)
+				if result != nil {
+					results = append(results, *result)
+				}
 			}
 		}
+		// Check if there are more pages to fetch and set skipToken for next iteration
+		if resources.SkipToken == nil || *resources.SkipToken == "" {
+			break
+		}
+		klog.Infof("Fetching next page (page %d) with skipToken: %s", pageCount+1, *resources.SkipToken)
+		skipToken = resources.SkipToken
 	}
-	return
+
+	klog.Infof("Query completed: fetched %d pages, total results: %d", pageCount, len(results))
+	return results, nil
 }