Skip to content

refactor: simplify custom registry usage in metrics tracker #16772

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
parametalol merged 1 commit into from
Sep 12, 2025
Merged

refactor: simplify custom registry usage in metrics tracker #16772

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 2 additions & 3 deletions central/metrics/custom/image_vulnerabilities/tracker.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,23 +5,22 @@ import (
"iter"

deploymentDS "github.com/stackrox/rox/central/deployment/datastore"
"github.com/stackrox/rox/central/metrics"
"github.com/stackrox/rox/central/metrics/custom/tracker"
"github.com/stackrox/rox/generated/storage"
"github.com/stackrox/rox/pkg/sac"
"github.com/stackrox/rox/pkg/sac/resources"
"github.com/stackrox/rox/pkg/search"
)

func New(registryFactory func(string) metrics.CustomRegistry, ds deploymentDS.DataStore) *tracker.TrackerBase[*finding] {
func New(ds deploymentDS.DataStore) *tracker.TrackerBase[*finding] {
return tracker.MakeTrackerBase(
"vulnerabilities",
"CVEs",
lazyLabels,
func(ctx context.Context, md tracker.MetricDescriptors) iter.Seq[*finding] {
return trackVulnerabilityMetrics(ctx, md, ds)
},
registryFactory)
)
}

func trackVulnerabilityMetrics(ctx context.Context, _ tracker.MetricDescriptors, ds deploymentDS.DataStore) iter.Seq[*finding] {
Expand Down
86 changes: 35 additions & 51 deletions central/metrics/custom/image_vulnerabilities/tracker_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,12 @@ package image_vulnerabilities

import (
"context"
"io"
"net/http/httptest"
"testing"

"github.com/prometheus/client_golang/prometheus"
deploymentMockDS "github.com/stackrox/rox/central/deployment/datastore/mocks"
"github.com/stackrox/rox/central/metrics"
"github.com/stackrox/rox/central/metrics/mocks"
v1 "github.com/stackrox/rox/generated/api/v1"
"github.com/stackrox/rox/generated/storage"
"github.com/stackrox/rox/pkg/auth/authproviders"
Expand Down Expand Up @@ -92,11 +92,6 @@ func getTestCVEs(*testing.T) []*storage.EmbeddedVulnerability {
}
}

type labelsTotal struct {
labels prometheus.Labels
total int
}

func TestQueryDeploymentsAndImages(t *testing.T) {
ctrl := gomock.NewController(t)
ds := deploymentMockDS.NewMockDataStore(ctrl)
Expand All @@ -117,23 +112,7 @@ func TestQueryDeploymentsAndImages(t *testing.T) {
Times(1).Return(deploymentImages[deployment.Id], nil)
}

var actual = make(map[string][]*labelsTotal)

mr := mocks.NewMockCustomRegistry(ctrl)
tracker := New(func(string) metrics.CustomRegistry { return mr }, ds)
mr.EXPECT().RegisterMetric(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
Times(3)
mr.EXPECT().SetTotal(gomock.Any(), gomock.Any(), gomock.Any()).
AnyTimes().Do(
func(metric string, labels prometheus.Labels, total int) {
actual[metric] = append(actual[metric], &labelsTotal{labels, total})
},
)
mr.EXPECT().Lock()
mr.EXPECT().Reset("Severity_count")
mr.EXPECT().Reset("Cluster_Namespace_Severity_count")
mr.EXPECT().Reset("Deployment_ImageTag_count")
mr.EXPECT().Unlock()
tracker := New(ds)

cfg, err := tracker.NewConfiguration(
&storage.PrometheusMetrics_Group{
Expand All @@ -154,34 +133,39 @@ func TestQueryDeploymentsAndImages(t *testing.T) {
tracker.Reconfigure(cfg)
tracker.Gather(makeAdminContext(t))

expected := map[string][]*labelsTotal{
"Severity_count": {
{prometheus.Labels{"Severity": "CRITICAL_VULNERABILITY_SEVERITY"}, 3},
{prometheus.Labels{"Severity": "MODERATE_VULNERABILITY_SEVERITY"}, 4},
{prometheus.Labels{"Severity": "LOW_VULNERABILITY_SEVERITY"}, 2},
},
"Cluster_Namespace_Severity_count": {
{prometheus.Labels{"Cluster": "cluster-1", "Namespace": "namespace-1", "Severity": "CRITICAL_VULNERABILITY_SEVERITY"}, 1},
{prometheus.Labels{"Cluster": "cluster-1", "Namespace": "namespace-2", "Severity": "CRITICAL_VULNERABILITY_SEVERITY"}, 2},
{prometheus.Labels{"Cluster": "cluster-1", "Namespace": "namespace-2", "Severity": "MODERATE_VULNERABILITY_SEVERITY"}, 2},
{prometheus.Labels{"Cluster": "cluster-2", "Namespace": "namespace-2", "Severity": "MODERATE_VULNERABILITY_SEVERITY"}, 2},
{prometheus.Labels{"Cluster": "cluster-2", "Namespace": "namespace-2", "Severity": "LOW_VULNERABILITY_SEVERITY"}, 2},
},
"Deployment_ImageTag_count": {
{prometheus.Labels{"Deployment": "D0", "ImageTag": "tag"}, 1},
{prometheus.Labels{"Deployment": "D1", "ImageTag": "tag"}, 3},
{prometheus.Labels{"Deployment": "D2", "ImageTag": "tag"}, 1},
{prometheus.Labels{"Deployment": "D3", "ImageTag": "tag"}, 2},
{prometheus.Labels{"Deployment": "D3", "ImageTag": "latest"}, 2},
},
}
rec := httptest.NewRecorder()
req := httptest.NewRequestWithContext(makeAdminContext(t),
"GET", "/metrics", nil)

for metric := range expected {
assert.Contains(t, actual, metric)
}
for metric, records := range actual {
assert.ElementsMatch(t, expected[metric], records, metric)
}
metrics.GetCustomRegistry("Admin").ServeHTTP(rec, req)

result := rec.Result()
assert.Equal(t, 200, result.StatusCode)
body, err := io.ReadAll(result.Body)
_ = result.Body.Close()
assert.NoError(t, err)
assert.Equal(t,
`# HELP rox_central_Cluster_Namespace_Severity_count The total number of CVEs aggregated by Cluster,Namespace,Severity and gathered every 2h1m0s
# TYPE rox_central_Cluster_Namespace_Severity_count gauge
rox_central_Cluster_Namespace_Severity_count{Cluster="cluster-1",Namespace="namespace-1",Severity="CRITICAL_VULNERABILITY_SEVERITY"} 1
rox_central_Cluster_Namespace_Severity_count{Cluster="cluster-1",Namespace="namespace-2",Severity="CRITICAL_VULNERABILITY_SEVERITY"} 2
rox_central_Cluster_Namespace_Severity_count{Cluster="cluster-1",Namespace="namespace-2",Severity="MODERATE_VULNERABILITY_SEVERITY"} 2
rox_central_Cluster_Namespace_Severity_count{Cluster="cluster-2",Namespace="namespace-2",Severity="LOW_VULNERABILITY_SEVERITY"} 2
rox_central_Cluster_Namespace_Severity_count{Cluster="cluster-2",Namespace="namespace-2",Severity="MODERATE_VULNERABILITY_SEVERITY"} 2
# HELP rox_central_Deployment_ImageTag_count The total number of CVEs aggregated by Deployment,ImageTag and gathered every 2h1m0s
# TYPE rox_central_Deployment_ImageTag_count gauge
rox_central_Deployment_ImageTag_count{Deployment="D0",ImageTag="tag"} 1
rox_central_Deployment_ImageTag_count{Deployment="D1",ImageTag="tag"} 3
rox_central_Deployment_ImageTag_count{Deployment="D2",ImageTag="tag"} 1
rox_central_Deployment_ImageTag_count{Deployment="D3",ImageTag="latest"} 2
rox_central_Deployment_ImageTag_count{Deployment="D3",ImageTag="tag"} 2
# HELP rox_central_Severity_count The total number of CVEs aggregated by Severity and gathered every 2h1m0s
# TYPE rox_central_Severity_count gauge
rox_central_Severity_count{Severity="CRITICAL_VULNERABILITY_SEVERITY"} 3
rox_central_Severity_count{Severity="LOW_VULNERABILITY_SEVERITY"} 2
rox_central_Severity_count{Severity="MODERATE_VULNERABILITY_SEVERITY"} 4
`,
string(body))
}

func Test_forEachImageVuln(t *testing.T) {
Expand Down
5 changes: 2 additions & 3 deletions central/metrics/custom/policy_violations/tracker.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,21 +5,20 @@ import (
"iter"

alertDS "github.com/stackrox/rox/central/alert/datastore"
"github.com/stackrox/rox/central/metrics"
"github.com/stackrox/rox/central/metrics/custom/tracker"
"github.com/stackrox/rox/generated/storage"
"github.com/stackrox/rox/pkg/search"
)

func New(registryFactory func(string) metrics.CustomRegistry, ds alertDS.DataStore) *tracker.TrackerBase[*finding] {
func New(ds alertDS.DataStore) *tracker.TrackerBase[*finding] {
return tracker.MakeTrackerBase(
"alerts",
"policy violations",
lazyLabels,
func(ctx context.Context, _ tracker.MetricDescriptors) iter.Seq[*finding] {
return trackViolations(ctx, ds)
},
registryFactory)
)
}

func trackViolations(ctx context.Context, ds alertDS.DataStore) iter.Seq[*finding] {
Expand Down
6 changes: 3 additions & 3 deletions central/metrics/custom/runner.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,10 +31,10 @@ type RunnerConfiguration struct {
policy_violations *custom.Configuration
}

func makeRunner(registryFactory func(string) metrics.CustomRegistry, dds deploymentDS.DataStore, ads alertDS.DataStore) *aggregatorRunner {
func makeRunner(dds deploymentDS.DataStore, ads alertDS.DataStore) *aggregatorRunner {
return &aggregatorRunner{
image_vulnerabilities: image_vulnerabilities.New(registryFactory, dds),
policy_violations: policy_violations.New(registryFactory, ads),
image_vulnerabilities: image_vulnerabilities.New(dds),
policy_violations: policy_violations.New(ads),
}
}

Expand Down
7 changes: 3 additions & 4 deletions central/metrics/custom/runner_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@ import (
alertDS "github.com/stackrox/rox/central/alert/datastore/mocks"
configDS "github.com/stackrox/rox/central/config/datastore/mocks"
deploymentDS "github.com/stackrox/rox/central/deployment/datastore/mocks"
"github.com/stackrox/rox/central/metrics"
v1 "github.com/stackrox/rox/generated/api/v1"
"github.com/stackrox/rox/generated/storage"
"github.com/stackrox/rox/pkg/auth/authproviders"
Expand All @@ -30,7 +29,7 @@ func TestRunner_makeRunner(t *testing.T) {
Metrics: nil,
},
nil)
runner := makeRunner(metrics.GetCustomRegistry, nil, nil)
runner := makeRunner(nil, nil)
runner.initialize(cds)
assert.NotNil(t, runner)

Expand All @@ -54,7 +53,7 @@ func TestRunner_makeRunner(t *testing.T) {
cds.EXPECT().GetPrivateConfig(gomock.Any()).Times(1).Return(
nil,
errors.New("DB error"))
runner := makeRunner(metrics.GetCustomRegistry, nil, nil)
runner := makeRunner(nil, nil)
assert.NotNil(t, runner)
runner.initialize(cds)

Expand Down Expand Up @@ -142,7 +141,7 @@ func TestRunner_ServeHTTP(t *testing.T) {
}).
Return(nil)

runner := makeRunner(metrics.GetCustomRegistry, dds, ads)
runner := makeRunner(dds, ads)
runner.initialize(cds)
runner.image_vulnerabilities.Gather(makeAdminContext(t))
runner.policy_violations.Gather(makeAdminContext(t))
Expand Down
3 changes: 1 addition & 2 deletions central/metrics/custom/singleton.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@ import (
alertDS "github.com/stackrox/rox/central/alert/datastore"
configDS "github.com/stackrox/rox/central/config/datastore"
deploymentDS "github.com/stackrox/rox/central/deployment/datastore"
"github.com/stackrox/rox/central/metrics"
"github.com/stackrox/rox/generated/storage"
"github.com/stackrox/rox/pkg/logging"
"github.com/stackrox/rox/pkg/sync"
Expand All @@ -29,7 +28,7 @@ type Runner interface {
// initialization. nil runner is safe, but no-op.
func Singleton() Runner {
onceRunner.Do(func() {
runner = makeRunner(metrics.GetCustomRegistry, deploymentDS.Singleton(), alertDS.Singleton())
runner = makeRunner(deploymentDS.Singleton(), alertDS.Singleton())
go runner.initialize(configDS.Singleton())
})
return runner
Expand Down
5 changes: 2 additions & 3 deletions central/metrics/custom/tracker/tracker_base.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ type TrackerBase[Finding WithError] struct {
config *Configuration
metricsConfigMux sync.RWMutex

gatherers sync.Map // map[user ID]tokenGatherer
gatherers sync.Map // map[user ID]*tokenGatherer

registryFactory func(userID string) metrics.CustomRegistry // for mocking in tests.
}
Expand All @@ -99,15 +99,14 @@ func makeGettersMap[Finding WithError](getters []LazyLabel[Finding]) map[Label]f
// configuration. Call Reconfigure to configure the period and the metrics.
func MakeTrackerBase[Finding WithError](category, description string,
getters []LazyLabel[Finding], generator FindingGenerator[Finding],
registryFactory func(string) metrics.CustomRegistry,
) *TrackerBase[Finding] {
return &TrackerBase[Finding]{
category: category,
description: description,
labelOrder: MakeLabelOrderMap(getters),
getters: makeGettersMap(getters),
generator: generator,
registryFactory: registryFactory,
registryFactory: metrics.GetCustomRegistry,
}
}

Expand Down
30 changes: 14 additions & 16 deletions central/metrics/custom/tracker/tracker_base_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,19 +51,15 @@ func makeTestGatherFunc(data []map[Label]string) FindingGenerator[testFinding] {
}

func TestMakeTrackerBase(t *testing.T) {
tracker := MakeTrackerBase("test", "test", testLabelGetters, nilGatherFunc,
func(string) metrics.CustomRegistry { return nil })
tracker := MakeTrackerBase("test", "test", testLabelGetters, nilGatherFunc)
assert.NotNil(t, tracker)
assert.Nil(t, tracker.GetConfiguration())
}

func TestTrackerBase_Reconfigure(t *testing.T) {
ctrl := gomock.NewController(t)
t.Run("nil configuration", func(t *testing.T) {
rf := mocks.NewMockCustomRegistry(ctrl)

tracker := MakeTrackerBase("test", "test", testLabelGetters, nilGatherFunc,
func(string) metrics.CustomRegistry { return rf })
tracker := MakeTrackerBase("test", "test", testLabelGetters, nilGatherFunc)

tracker.Reconfigure(nil)
config := tracker.GetConfiguration()
Expand All @@ -74,9 +70,7 @@ func TestTrackerBase_Reconfigure(t *testing.T) {
})

t.Run("test 0 period", func(t *testing.T) {
rf := mocks.NewMockCustomRegistry(ctrl)
tracker := MakeTrackerBase("test", "test", testLabelGetters, nilGatherFunc,
func(string) metrics.CustomRegistry { return rf })
tracker := MakeTrackerBase("test", "test", testLabelGetters, nilGatherFunc)
cfg0 := &Configuration{}

tracker.Reconfigure(cfg0)
Expand All @@ -90,10 +84,11 @@ func TestTrackerBase_Reconfigure(t *testing.T) {
t.Run("test add -> delete -> stop", func(t *testing.T) {
trackedMetricNames := make([]MetricName, 0)

rf := mocks.NewMockCustomRegistry(ctrl)
tracker := MakeTrackerBase("test", "test", testLabelGetters,
makeTestGatherFunc(testData),
func(string) metrics.CustomRegistry { return rf })
makeTestGatherFunc(testData))

rf := mocks.NewMockCustomRegistry(ctrl)
tracker.registryFactory = func(string) metrics.CustomRegistry { return rf }

var registered, unregistered []MetricName
rf.EXPECT().RegisterMetric(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
Expand Down Expand Up @@ -180,10 +175,11 @@ func TestTrackerBase_Reconfigure(t *testing.T) {
func TestTrackerBase_Track(t *testing.T) {
ctrl := gomock.NewController(t)
rf := mocks.NewMockCustomRegistry(ctrl)

tracker := MakeTrackerBase("test", "test",
testLabelGetters,
makeTestGatherFunc(testData),
func(string) metrics.CustomRegistry { return rf })
makeTestGatherFunc(testData))
tracker.registryFactory = func(string) metrics.CustomRegistry { return rf }

result := make(map[string][]*aggregatedRecord)
rf.EXPECT().SetTotal(gomock.Any(), gomock.Any(), gomock.Any()).
Expand Down Expand Up @@ -253,7 +249,8 @@ func TestTrackerBase_error(t *testing.T) {
}
}
},
func(string) metrics.CustomRegistry { return rf })
)
tracker.registryFactory = func(string) metrics.CustomRegistry { return rf }

tracker.config = &Configuration{
metrics: makeTestMetricDescriptors(t),
Expand All @@ -268,7 +265,8 @@ func TestTrackerBase_Gather(t *testing.T) {
tracker := MakeTrackerBase("test", "test",
testLabelGetters,
makeTestGatherFunc(testData),
func(string) metrics.CustomRegistry { return rf })
)
tracker.registryFactory = func(string) metrics.CustomRegistry { return rf }

result := make(map[string][]*aggregatedRecord)
{ // Capture result with a mock registry.
Expand Down
Loading