Move retry to monitor service

Signed-off-by: constanca <[email protected]>
elastic · Mar 14, 2024 · b4b1dbc · b4b1dbc
1 parent 0d5b521
commit b4b1dbc
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 45 deletions.
diff --git a/x-pack/metricbeat/module/azure/mock_service.go b/x-pack/metricbeat/module/azure/mock_service.go
@@ -35,6 +35,11 @@ func (client *MockService) GetMetricDefinitions(resourceId string, namespace str
 	return args.Get(0).(armmonitor.MetricDefinitionCollection), args.Error(1)
 }
 
+// GetMetricDefinitionsWithRetry is a mock function for the azure service
+func (client *MockService) GetMetricDefinitionsWithRetry(resource *armresources.GenericResourceExpanded, namespace string) (armmonitor.MetricDefinitionCollection, error) {
+	return client.GetMetricDefinitions(*resource.ID, namespace)
+}
+
 // GetMetricNamespaces is a mock function for the azure service
 func (client *MockService) GetMetricNamespaces(resourceId string) (armmonitor.MetricNamespaceCollection, error) {
 	args := client.Called(resourceId)

diff --git a/x-pack/metricbeat/module/azure/monitor/client_helper.go b/x-pack/metricbeat/module/azure/monitor/client_helper.go
@@ -5,13 +5,8 @@
 package monitor
 
 import (
-	"errors"
 	"fmt"
-	"net/http"
 	"strings"
-	"time"
-
-	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/monitor/armmonitor"
 
@@ -22,43 +17,6 @@ import (
 
 const missingNamespace = "no metric definitions were found for resource %s and namespace %s. Verify if the namespace is spelled correctly or if it is supported by the resource in case."
 
-func getMetricsDefinitionsWithRetry(client *azure.Client, resource *armresources.GenericResourceExpanded, namespace string) (armmonitor.MetricDefinitionCollection, error) {
-	for {
-
-		metricDefinitions, err := client.AzureMonitorService.GetMetricDefinitions(*resource.ID, namespace)
-		if err != nil {
-			errorMsg := "no metric definitions were found for resource " + *resource.ID + " and namespace " + namespace
-
-			var respError *azcore.ResponseError
-			ok := errors.As(err, &respError)
-			if !ok {
-				return metricDefinitions, fmt.Errorf("%s, failed to cast error to azcore.ResponseError", errorMsg)
-			}
-			// Check for TooManyRequests error and retry if it is the case
-			if respError.StatusCode != http.StatusTooManyRequests {
-				return metricDefinitions, fmt.Errorf("%s, %w", errorMsg, err)
-			}
-
-			// Check if the error has the header Retry After.
-			// If it is present, then we should try to make this request again.
-			retryAfter := respError.RawResponse.Header.Get("Retry-After")
-			if retryAfter == "" {
-				return metricDefinitions, fmt.Errorf("%s %w, failed to find Retry-After header", errorMsg, err)
-			}
-
-			duration, errD := time.ParseDuration(retryAfter + "s")
-			if errD != nil {
-				return metricDefinitions, fmt.Errorf("%s, failed to parse duration %s from header retry after", errorMsg, retryAfter)
-			}
-
-			client.Log.Infof("%s, metricbeat will try again after %s seconds", errorMsg, retryAfter)
-			time.Sleep(duration)
-		} else {
-			return metricDefinitions, err
-		}
-	}
-}
-
 // mapMetrics should validate and map the metric related configuration to relevant azure monitor api parameters
 func mapMetrics(client *azure.Client, resources []*armresources.GenericResourceExpanded, resourceConfig azure.ResourceConfig) ([]azure.Metric, error) {
 	var metrics []azure.Metric
@@ -73,9 +31,9 @@ func mapMetrics(client *azure.Client, resources []*armresources.GenericResourceE
 
 			var err error
 
-			metricDefinitions, ok := namespaceMetrics[metric.Namespace]
-			if !ok {
-				metricDefinitions, err = getMetricsDefinitionsWithRetry(client, resource, metric.Namespace)
+			metricDefinitions, exists := namespaceMetrics[metric.Namespace]
+			if !exists {
+				metricDefinitions, err = client.AzureMonitorService.GetMetricDefinitionsWithRetry(resource, metric.Namespace)
 				if err != nil {
 					return nil, err
 				}

diff --git a/x-pack/metricbeat/module/azure/monitor_service.go b/x-pack/metricbeat/module/azure/monitor_service.go
@@ -6,8 +6,13 @@ package azure
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"net/http"
 	"strings"
+	"time"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 
 	"github.com/elastic/elastic-agent-libs/logp"
 
@@ -219,6 +224,44 @@ func (service *MonitorService) GetMetricDefinitions(resourceId string, namespace
 	return metricDefinitionCollection, nil
 }
 
+func (service *MonitorService) GetMetricDefinitionsWithRetry(resource *armresources.GenericResourceExpanded, namespace string) (armmonitor.MetricDefinitionCollection, error) {
+	for {
+
+		metricDefinitions, err := service.GetMetricDefinitions(*resource.ID, namespace)
+		if err != nil {
+			errorMsg := "no metric definitions were found for resource " + *resource.ID + " and namespace " + namespace
+
+			var respError *azcore.ResponseError
+			ok := errors.As(err, &respError)
+			if !ok {
+				return metricDefinitions, fmt.Errorf("%s, failed to cast error to azcore.ResponseError", errorMsg)
+			}
+			// Check for TooManyRequests error and retry if it is the case
+			if respError.StatusCode != http.StatusTooManyRequests {
+				return metricDefinitions, fmt.Errorf("%s, %w", errorMsg, err)
+			}
+
+			// Check if the error has the header Retry After.
+			// If it is present, then we should try to make this request again.
+			retryAfter := respError.RawResponse.Header.Get("Retry-After")
+			if retryAfter == "" {
+				return metricDefinitions, fmt.Errorf("%s %w, failed to find Retry-After header", errorMsg, err)
+			}
+
+			duration, errD := time.ParseDuration(retryAfter + "s")
+			if errD != nil {
+				return metricDefinitions, fmt.Errorf("%s, failed to parse duration %s from header retry after", errorMsg, retryAfter)
+			}
+
+			service.log.Infof("%s, metricbeat will try again after %s seconds", errorMsg, retryAfter)
+			time.Sleep(duration)
+			service.log.Infof("%s, metricbeat finished sleeping and will try again now", errorMsg)
+		} else {
+			return metricDefinitions, err
+		}
+	}
+}
+
 // GetMetricValues will return the metric values based on the resource and metric details
 func (service *MonitorService) GetMetricValues(resourceId string, namespace string, timegrain string, timespan string, metricNames []string, aggregations string, filter string) ([]armmonitor.Metric, string, error) {
 	var tg *string

diff --git a/x-pack/metricbeat/module/azure/service_interface.go b/x-pack/metricbeat/module/azure/service_interface.go
@@ -14,6 +14,7 @@ type Service interface {
 	GetResourceDefinitionById(id string) (armresources.GenericResource, error)
 	GetResourceDefinitions(id []string, group []string, rType string, query string) ([]*armresources.GenericResourceExpanded, error)
 	GetMetricDefinitions(resourceId string, namespace string) (armmonitor.MetricDefinitionCollection, error)
+	GetMetricDefinitionsWithRetry(resource *armresources.GenericResourceExpanded, namespace string) (armmonitor.MetricDefinitionCollection, error)
 	GetMetricNamespaces(resourceId string) (armmonitor.MetricNamespaceCollection, error)
 	GetMetricValues(resourceId string, namespace string, timegrain string, timespan string, metricNames []string, aggregations string, filter string) ([]armmonitor.Metric, string, error)
 }