Fix pagination bug and optimize config caching with RWMutex (#3933)

Signed-off-by: Claude <noreply@anthropic.com>
Co-authored-by: Claude <noreply@anthropic.com>

modules/prometheus-source/pkg/prom/result.go

@@ -107,7 +107,7 @@ func NewQueryResults(query string, queryResult interface{}, resultKeys *source.R
 	}
 
 	// Result vectors from the query
-	var results []*source.QueryResult
+	results := make([]*source.QueryResult, 0, len(resultsData))
 
 	// Parse raw results and into QueryResults
 	for _, val := range resultsData {
@@ -160,6 +160,7 @@ func NewQueryResults(query string, queryResult interface{}, resultKeys *source.R
 				qrs.Error = fmt.Errorf("Values field is improperly formatted")
 				return qrs
 			}
+			vectors = make([]*util.Vector, 0, len(values))
 
 			// Append new data points, log warnings
 			for _, value := range values {

pkg/cloud/alibaba/provider.go

@@ -63,6 +63,9 @@ const (
 var (
 	// sizeRegEx parses a PV capacity string into a numeric part and an optional binary SI suffix (Ki, Mi, Gi, Ti).
 	sizeRegEx = regexp.MustCompile(`^(\d+(?:\.\d+)?)(Ki|Mi|Gi|Ti)?$`)
+
+	// generationRegEx extracts the numeric generation from an instance family name (e.g. 7 from "g7ne").
+	generationRegEx = regexp.MustCompile(`(\d+)`)
 )
 
 // Variable to keep track of instance families that fail in DescribePrice API due improper defaulting of systemDisk if the information is not available
@@ -1163,8 +1166,7 @@ func getInstanceFamilyFromType(instanceType string) string {
 func getInstanceFamilyGenerationFromType(instanceType string) int {
 	// FamilyName format: g7ne or g7 or r7 or r6e,
 	familyName := getInstanceFamilyFromType(instanceType)
-	re := regexp.MustCompile(`(\d+)`)
-	match := re.FindString(familyName)
+	match := generationRegEx.FindString(familyName)
 	if match != "" {
 		generation, err := strconv.Atoi(match)
 		if err != nil {

pkg/cloud/aws/provider.go

@@ -1107,11 +1107,11 @@ func (aws *AWS) DownloadPricingData() error {
 //	}
 //
 // This function uses streaming JSON parsing to handle large pricing files efficiently:
-// 1. Parse "products" section: Extract SKUs and attributes for EC2 instances, EBS volumes, and load balancers
-// 2. Parse "terms" section: Extract on-demand pricing for each SKU. Note: only the
-//    first term-type key is read and processed if it equals "OnDemand", so this
-//    assumes "OnDemand" precedes any "Reserved" terms.
-// 3. Match SKUs to pricing keys and populate the pricing map
+//  1. Parse "products" section: Extract SKUs and attributes for EC2 instances, EBS volumes, and load balancers
+//  2. Parse "terms" section: Extract on-demand pricing for each SKU. Note: only the
+//     first term-type key is read and processed if it equals "OnDemand", so this
+//     assumes "OnDemand" precedes any "Reserved" terms.
+//  3. Match SKUs to pricing keys and populate the pricing map
 func (aws *AWS) populatePricing(resp *http.Response, inputkeys map[string]bool) error {
 	aws.Pricing = make(map[string]*AWSProductTerms)
 	skuToPricingKeyMap := make(map[string]string)

pkg/cloud/gcp/provider.go

@@ -997,7 +997,6 @@ func (gcp *GCP) getBillingAPIClientAndURL(apiKey, currencyCode string) (*http.Cl
 }
 
 func (gcp *GCP) parsePages(inputKeys map[string]models.Key, pvKeys map[string]models.PVKey) (map[string]*GCPPricing, error) {
-	var pages []map[string]*GCPPricing
 	c, err := gcp.GetConfig()
 	if err != nil {
 		return nil, err
@@ -1008,14 +1007,26 @@ func (gcp *GCP) parsePages(inputKeys map[string]models.Key, pvKeys map[string]mo
 		return nil, err
 	}
 
+	return gcp.parsePagesWithClient(httpClient, url, inputKeys, pvKeys)
+}
+
+// parsePagesWithClient pages through the billing API at url using httpClient,
+// parsing each page of SKUs and merging the results.
+func (gcp *GCP) parsePagesWithClient(httpClient *http.Client, url string, inputKeys map[string]models.Key, pvKeys map[string]models.PVKey) (map[string]*GCPPricing, error) {
+	var pages []map[string]*GCPPricing
+
 	var parsePagesHelper func(string) error
 	parsePagesHelper = func(pageToken string) error {
 		if pageToken == "done" {
 			return nil
-		} else if pageToken != "" {
-			url = url + "&pageToken=" + pageToken
 		}
-		resp, err := httpClient.Get(url)
+		// Build the URL per request; appending to the shared url would accumulate
+		// pageToken params across pages.
+		reqURL := url
+		if pageToken != "" {
+			reqURL = url + "&pageToken=" + pageToken
+		}
+		resp, err := httpClient.Get(reqURL)
 		if err != nil {
 			return err
 		}
@@ -1026,7 +1037,7 @@ func (gcp *GCP) parsePages(inputKeys map[string]models.Key, pvKeys map[string]mo
 		pages = append(pages, page)
 		return parsePagesHelper(token)
 	}
-	err = parsePagesHelper("")
+	err := parsePagesHelper("")
 	if err != nil {
 		return nil, err
 	}

pkg/cloud/gcp/provider_test.go

@@ -4,6 +4,8 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
+	"net/http"
+	"net/http/httptest"
 	"net/url"
 	"os"
 	"reflect"
@@ -1096,6 +1098,40 @@ func TestGCP_parsePages(t *testing.T) {
 	assert.Error(t, err) // Expect error due to missing API key
 }
 
+// TestGCP_parsePagesWithClient_Pagination verifies that multi-page traversal
+// sends exactly one pageToken param per request rather than accumulating
+// tokens from earlier pages.
+func TestGCP_parsePagesWithClient_Pagination(t *testing.T) {
+	var pageTokens [][]string
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		tokens := r.URL.Query()["pageToken"]
+		pageTokens = append(pageTokens, tokens)
+		w.Header().Set("Content-Type", "application/json")
+		if len(tokens) > 0 && tokens[0] == "tok2" {
+			fmt.Fprint(w, `{"skus": [], "nextPageToken": ""}`)
+		} else {
+			fmt.Fprint(w, `{"skus": [], "nextPageToken": "tok2"}`)
+		}
+	}))
+	defer srv.Close()
+
+	gcp := &GCP{}
+	_, err := gcp.parsePagesWithClient(srv.Client(), srv.URL+"?currencyCode=USD", map[string]models.Key{}, map[string]models.PVKey{})
+	if err != nil {
+		t.Fatalf("parsePagesWithClient: %v", err)
+	}
+
+	if len(pageTokens) != 2 {
+		t.Fatalf("expected 2 page requests, got %d", len(pageTokens))
+	}
+	if len(pageTokens[0]) != 0 {
+		t.Errorf("first request should have no pageToken, got %v", pageTokens[0])
+	}
+	if len(pageTokens[1]) != 1 || pageTokens[1][0] != "tok2" {
+		t.Errorf("second request should have exactly one pageToken (tok2), got %v", pageTokens[1])
+	}
+}
+
 func TestGCP_DownloadPricingData(t *testing.T) {
 	gcp := &GCP{
 		Config: &mockConfig{},

pkg/cloud/provider/providerconfig.go

@@ -25,7 +25,7 @@ const closedSourceConfigMount = "models/"
 // ProviderConfig is a utility class that provides a thread-safe configuration storage/cache for all Provider
 // implementations
 type ProviderConfig struct {
-	lock            sync.Mutex
+	lock            sync.RWMutex
 	configManager   *config.ConfigFileManager
 	configFile      *config.ConfigFile
 	customPricing   *models.CustomPricing
@@ -137,6 +137,14 @@ func (pc *ProviderConfig) loadConfig(writeIfNotExists bool) (*models.CustomPrici
 
 // ThreadSafe method for retrieving the custom pricing config.
 func (pc *ProviderConfig) GetCustomPricingData() (*models.CustomPricing, error) {
+	// Fast path: once loaded, the config is cached, so readers only need the read lock.
+	pc.lock.RLock()
+	cached := pc.customPricing
+	pc.lock.RUnlock()
+	if cached != nil {
+		return cached, nil
+	}
+
 	pc.lock.Lock()
 	defer pc.lock.Unlock()
 

pkg/costmodel/router_test.go

@@ -20,13 +20,13 @@ func TestAdminAuthMiddleware(t *testing.T) {
 	}
 
 	tests := []struct {
-		name              string
-		setToken          string
-		authHeader        string
-		wantStatus        int
-		wantNextCalled    bool
-		wantBodySubstr    string
-		wantCacheControl  string
+		name             string
+		setToken         string
+		authHeader       string
+		wantStatus       int
+		wantNextCalled   bool
+		wantBodySubstr   string
+		wantCacheControl string
 	}{
 		{
 			name:             "no admin token configured - returns 503",