Skip to content

Commit

Permalink
feat: add prometheus timeout attribute
Browse files Browse the repository at this point in the history
Signed-off-by: AhmedGrati <[email protected]>
  • Loading branch information
TessaIO committed Jul 23, 2023
1 parent 45e549d commit af6439a
Show file tree
Hide file tree
Showing 4 changed files with 65 additions and 19 deletions.
2 changes: 1 addition & 1 deletion metricproviders/metricproviders.go
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ func (f *ProviderFactory) NewProvider(logCtx log.Entry, metric v1alpha1.Metric)
if err != nil {
return nil, err
}
return prometheus.NewPrometheusProvider(api, logCtx), nil
return prometheus.NewPrometheusProvider(api, logCtx, metric)
case job.ProviderType:
return job.NewJobProvider(logCtx, f.KubeClient, f.JobLister), nil
case kayenta.ProviderType:
Expand Down
25 changes: 20 additions & 5 deletions metricproviders/prometheus/prometheus.go
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,9 @@ const (

// Provider contains all the required components to run a prometheus query
type Provider struct {
api v1.API
logCtx log.Entry
api v1.API
logCtx log.Entry
Timeout time.Duration
}

// Type indicates provider is a prometheus provider
Expand All @@ -59,7 +60,7 @@ func (p *Provider) Run(run *v1alpha1.AnalysisRun, metric v1alpha1.Metric) v1alph
}

//TODO(dthomson) make timeout configurable
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
ctx, cancel := context.WithTimeout(context.Background(), p.Timeout)
defer cancel()

response, warnings, err := p.api.Query(ctx, metric.Provider.Prometheus.Query, time.Now())
Expand Down Expand Up @@ -138,11 +139,25 @@ func (p *Provider) processResponse(metric v1alpha1.Metric, response model.Value)
}

// NewPrometheusProvider Creates a new Prometheus client
func NewPrometheusProvider(api v1.API, logCtx log.Entry) *Provider {
return &Provider{
func NewPrometheusProvider(api v1.API, logCtx log.Entry, metric v1alpha1.Metric) (*Provider, error) {
provider := &Provider{
logCtx: logCtx,
api: api,
}

if metric.Provider.Prometheus == nil || metric.Provider.Prometheus.Timeout == nil {
provider.Timeout = 30 * time.Second
return provider, nil
}

metricTimeout := metric.Provider.Prometheus.Timeout

if *metricTimeout < 0 {
return nil, errors.New("prometheus timeout should not be negative")
}

provider.Timeout = time.Duration(*metricTimeout * int(time.Second))
return provider, nil
}

// NewPrometheusAPI generates a prometheus API from the metric configuration
Expand Down
54 changes: 41 additions & 13 deletions metricproviders/prometheus/prometheus_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"math"
"os"
"testing"
"time"

"github.com/argoproj/argo-rollouts/pkg/apis/rollouts/v1alpha1"
v1 "github.com/prometheus/client_golang/api/prometheus/v1"
Expand All @@ -30,16 +31,29 @@ func TestType(t *testing.T) {
mock := mockAPI{
value: newScalar(10),
}
p := NewPrometheusProvider(mock, e)
timeout := 5
metric := v1alpha1.Metric{
Name: "foo",
SuccessCondition: "result == 10",
FailureCondition: "result != 10",
Provider: v1alpha1.MetricProvider{
Prometheus: &v1alpha1.PrometheusMetric{
Query: "test",
Timeout: &timeout,
},
},
}
p, err := NewPrometheusProvider(mock, e, metric)
assert.NoError(t, err)
assert.Equal(t, ProviderType, p.Type())
assert.Equal(t, p.Timeout, time.Duration(timeout*int(time.Second)))
}

func TestRunSuccessfully(t *testing.T) {
e := log.Entry{}
mock := mockAPI{
value: newScalar(10),
}
p := NewPrometheusProvider(mock, e)
metric := v1alpha1.Metric{
Name: "foo",
SuccessCondition: "result == 10",
Expand All @@ -50,8 +64,11 @@ func TestRunSuccessfully(t *testing.T) {
},
},
}
p, err := NewPrometheusProvider(mock, e, metric)

measurement := p.Run(newAnalysisRun(), metric)
assert.NotNil(t, measurement.StartedAt)
assert.NoError(t, err)
assert.Equal(t, "10", measurement.Value)
assert.NotNil(t, measurement.FinishedAt)
assert.Equal(t, v1alpha1.AnalysisPhaseSuccessful, measurement.Phase)
Expand All @@ -64,7 +81,6 @@ func TestRunSuccessfullyWithEnv(t *testing.T) {
}
address := "http://127.0.0.1:9090"
os.Setenv(EnvVarArgoRolloutsPrometheusAddress, address)
p := NewPrometheusProvider(mock, e)
metric := v1alpha1.Metric{
Name: "foo",
SuccessCondition: "result == 10",
Expand All @@ -75,8 +91,10 @@ func TestRunSuccessfullyWithEnv(t *testing.T) {
},
},
}
p, err := NewPrometheusProvider(mock, e, metric)
measurement := p.Run(newAnalysisRun(), metric)
assert.NotNil(t, measurement.StartedAt)
assert.NoError(t, err)
assert.Equal(t, "10", measurement.Value)
assert.NotNil(t, measurement.FinishedAt)
assert.Equal(t, v1alpha1.AnalysisPhaseSuccessful, measurement.Phase)
Expand All @@ -88,7 +106,6 @@ func TestRunSuccessfullyWithWarning(t *testing.T) {
value: newScalar(10),
warnings: v1.Warnings([]string{"warning", "warning2"}),
}
p := NewPrometheusProvider(mock, *e)
metric := v1alpha1.Metric{
Name: "foo",
SuccessCondition: "result == 10",
Expand All @@ -99,8 +116,10 @@ func TestRunSuccessfullyWithWarning(t *testing.T) {
},
},
}
p, err := NewPrometheusProvider(mock, *e, metric)
measurement := p.Run(newAnalysisRun(), metric)
assert.NotNil(t, measurement.StartedAt)
assert.NoError(t, err)
assert.Equal(t, "10", measurement.Value)
assert.NotNil(t, measurement.FinishedAt)
assert.Equal(t, `"warning", "warning2"`, measurement.Metadata["warnings"])
Expand All @@ -113,7 +132,6 @@ func TestRunSuccessfullyWithWarningWithEnv(t *testing.T) {
value: newScalar(10),
warnings: v1.Warnings([]string{"warning", "warning2"}),
}
p := NewPrometheusProvider(mock, *e)
metric := v1alpha1.Metric{
Name: "foo",
SuccessCondition: "result == 10",
Expand All @@ -124,8 +142,10 @@ func TestRunSuccessfullyWithWarningWithEnv(t *testing.T) {
},
},
}
p, err := NewPrometheusProvider(mock, *e, metric)
measurement := p.Run(newAnalysisRun(), metric)
assert.NotNil(t, measurement.StartedAt)
assert.NoError(t, err)
assert.Equal(t, "10", measurement.Value)
assert.NotNil(t, measurement.FinishedAt)
assert.Equal(t, `"warning", "warning2"`, measurement.Metadata["warnings"])
Expand All @@ -138,7 +158,6 @@ func TestRunWithQueryError(t *testing.T) {
mock := mockAPI{
err: expectedErr,
}
p := NewPrometheusProvider(mock, *e)
metric := v1alpha1.Metric{
Name: "foo",
SuccessCondition: "result == 10",
Expand All @@ -149,8 +168,10 @@ func TestRunWithQueryError(t *testing.T) {
},
},
}
p, err := NewPrometheusProvider(mock, *e, metric)
measurement := p.Run(newAnalysisRun(), metric)
assert.Equal(t, expectedErr.Error(), measurement.Message)
assert.NoError(t, err)
assert.NotNil(t, measurement.StartedAt)
assert.Equal(t, "", measurement.Value)
assert.NotNil(t, measurement.FinishedAt)
Expand All @@ -163,7 +184,6 @@ func TestRunWithResolveArgsError(t *testing.T) {
mock := mockAPI{
err: expectedErr,
}
p := NewPrometheusProvider(mock, e)
metric := v1alpha1.Metric{
Name: "foo",
Provider: v1alpha1.MetricProvider{
Expand All @@ -172,8 +192,10 @@ func TestRunWithResolveArgsError(t *testing.T) {
},
},
}
p, err := NewPrometheusProvider(mock, e, metric)
measurement := p.Run(newAnalysisRun(), metric)
assert.Equal(t, expectedErr.Error(), measurement.Message)
assert.NoError(t, err)
assert.NotNil(t, measurement.StartedAt)
assert.Equal(t, "", measurement.Value)
assert.NotNil(t, measurement.FinishedAt)
Expand All @@ -183,7 +205,6 @@ func TestRunWithResolveArgsError(t *testing.T) {
func TestGetStatusReturnsResolvedQuery(t *testing.T) {
e := log.Entry{}
mock := mockAPI{}
p := NewPrometheusProvider(mock, e)
metric := v1alpha1.Metric{
Name: "foo",
Provider: v1alpha1.MetricProvider{
Expand All @@ -192,15 +213,16 @@ func TestGetStatusReturnsResolvedQuery(t *testing.T) {
},
},
}
p, err := NewPrometheusProvider(mock, e, metric)
metricsMetadata := p.GetMetadata(metric)
assert.NotNil(t, metricsMetadata)
assert.NoError(t, err)
assert.Equal(t, "resolved-query", metricsMetadata["ResolvedPrometheusQuery"])
}

func TestRunWithEvaluationError(t *testing.T) {
e := log.WithField("", "")
mock := mockAPI{}
p := NewPrometheusProvider(mock, *e)
metric := v1alpha1.Metric{
Name: "foo",
SuccessCondition: "result == 10",
Expand All @@ -211,8 +233,10 @@ func TestRunWithEvaluationError(t *testing.T) {
},
},
}
p, err := NewPrometheusProvider(mock, *e, metric)
measurement := p.Run(newAnalysisRun(), metric)
assert.Equal(t, "Prometheus metric type not supported", measurement.Message)
assert.NoError(t, err)
assert.NotNil(t, measurement.StartedAt)
assert.Equal(t, "", measurement.Value)
assert.NotNil(t, measurement.FinishedAt)
Expand All @@ -222,7 +246,6 @@ func TestRunWithEvaluationError(t *testing.T) {
func TestResume(t *testing.T) {
e := log.WithField("", "")
mock := mockAPI{}
p := NewPrometheusProvider(mock, *e)
metric := v1alpha1.Metric{
Name: "foo",
SuccessCondition: "result == 10",
Expand All @@ -233,6 +256,8 @@ func TestResume(t *testing.T) {
},
},
}
p, err := NewPrometheusProvider(mock, *e, metric)
assert.NoError(t, err)
now := metav1.Now()
previousMeasurement := v1alpha1.Measurement{
StartedAt: &now,
Expand All @@ -245,8 +270,9 @@ func TestResume(t *testing.T) {
func TestTerminate(t *testing.T) {
e := log.NewEntry(log.New())
mock := mockAPI{}
p := NewPrometheusProvider(mock, *e)
metric := v1alpha1.Metric{}
p, err := NewPrometheusProvider(mock, *e, metric)
assert.NoError(t, err)
now := metav1.Now()
previousMeasurement := v1alpha1.Measurement{
StartedAt: &now,
Expand All @@ -259,8 +285,10 @@ func TestTerminate(t *testing.T) {
func TestGarbageCollect(t *testing.T) {
e := log.NewEntry(log.New())
mock := mockAPI{}
p := NewPrometheusProvider(mock, *e)
err := p.GarbageCollect(nil, v1alpha1.Metric{}, 0)
metric := v1alpha1.Metric{}
p, err := NewPrometheusProvider(mock, *e, metric)
assert.NoError(t, err)
err = p.GarbageCollect(nil, v1alpha1.Metric{}, 0)
assert.NoError(t, err)
}

Expand Down
3 changes: 3 additions & 0 deletions pkg/apis/rollouts/v1alpha1/analysis_types.go
Original file line number Diff line number Diff line change
Expand Up @@ -213,6 +213,9 @@ type PrometheusMetric struct {
// Sigv4 Config is the aws SigV4 configuration to use for SigV4 signing if using Amazon Managed Prometheus
// +optional
Authentication PrometheusAuth `json:"authentication,omitempty" protobuf:"bytes,3,opt,name=authentication"`
// Timeout represents the duration within which a prometheus query should complete. It is expressed in seconds.
// +optional
Timeout *int `json:"timeout,omitempty" protobuf:"bytes,4,opt,name=timeout"`
}

// PrometheusMetric defines the prometheus query to perform canary analysis
Expand Down

0 comments on commit af6439a

Please sign in to comment.