Implement matching objectives with query on ListObjectives

Author: metalmatze

filesystem.go

@@ -7,24 +7,63 @@ import (
 	"log"
 	"net/http"
 	"path/filepath"
+	"sync"
 
 	"github.com/fsnotify/fsnotify"
 	"github.com/oklog/run"
 	monitoringv1 "github.com/prometheus-operator/prometheus-operator/pkg/apis/monitoring/v1"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
+	"github.com/prometheus/prometheus/pkg/labels"
+	"github.com/prometheus/prometheus/promql/parser"
 	"github.com/pyrra-dev/pyrra/kubernetes/api/v1alpha1"
 	"github.com/pyrra-dev/pyrra/openapi"
 	openapiserver "github.com/pyrra-dev/pyrra/openapi/server/go"
 	"github.com/pyrra-dev/pyrra/slo"
 	"sigs.k8s.io/yaml"
 )
 
-var objectives = map[string]slo.Objective{}
+type Objectives struct {
+	mu         sync.RWMutex
+	objectives map[string]slo.Objective
+}
 
-func cmdFilesystem(configFiles, prometheusFolder string) {
-	var gr run.Group
+func (os *Objectives) Set(o slo.Objective) {
+	os.mu.Lock()
+	os.objectives[o.Labels.String()] = o
+	os.mu.Unlock()
+}
+
+func (os *Objectives) Match(ms []*labels.Matcher) []slo.Objective {
+	if ms == nil || len(ms) == 0 {
+		os.mu.RLock()
+		objectives := make([]slo.Objective, 0, len(os.objectives))
+		for _, o := range os.objectives {
+			objectives = append(objectives, o)
+		}
+		os.mu.RUnlock()
+		return objectives
+	}
+
+	os.mu.RLock()
+	defer os.mu.RUnlock()
+
+	var objectives []slo.Objective
 
+	for _, o := range os.objectives {
+		for _, m := range ms {
+			v := o.Labels.Get(m.Name)
+			if !m.Matches(v) {
+				continue
+			}
+			objectives = append(objectives, o)
+		}
+	}
+
+	return objectives
+}
+
+func cmdFilesystem(configFiles, prometheusFolder string) {
 	reg := prometheus.NewRegistry()
 
 	reconcilesTotal := prometheus.NewCounter(prometheus.CounterOpts{
@@ -43,8 +82,10 @@ func cmdFilesystem(configFiles, prometheusFolder string) {
 	)
 
 	ctx, cancel := context.WithCancel(context.Background())
+	objectives := &Objectives{objectives: map[string]slo.Objective{}}
 	files := make(chan string, 16)
 
+	var gr run.Group
 	{
 		gr.Add(func() error {
 			// Initially read all files and send them to be processed and added to the in memory store.
@@ -143,7 +184,7 @@ func cmdFilesystem(configFiles, prometheusFolder string) {
 						return err
 					}
 
-					objectives[objective.Labels.String()] = objective
+					objectives.Set(objective)
 				}
 			}
 		}, func(err error) {
@@ -152,7 +193,9 @@ func cmdFilesystem(configFiles, prometheusFolder string) {
 	}
 	{
 		router := openapiserver.NewRouter(
-			openapiserver.NewObjectivesApiController(&FilesystemObjectiveServer{}),
+			openapiserver.NewObjectivesApiController(&FilesystemObjectiveServer{
+				objectives: objectives,
+			}),
 		)
 		router.Handle("/metrics", promhttp.HandlerFor(reg, promhttp.HandlerOpts{}))
 
@@ -172,12 +215,24 @@ func cmdFilesystem(configFiles, prometheusFolder string) {
 	}
 }
 
-type FilesystemObjectiveServer struct{}
+type FilesystemObjectiveServer struct {
+	objectives *Objectives
+}
 
-func (f FilesystemObjectiveServer) ListObjectives(ctx context.Context) (openapiserver.ImplResponse, error) {
+func (f FilesystemObjectiveServer) ListObjectives(ctx context.Context, query string) (openapiserver.ImplResponse, error) {
+	var matchers []*labels.Matcher
+	if query != "" {
+		var err error
+		matchers, err = parser.ParseMetricSelector(query)
+		if err != nil {
+			return openapiserver.ImplResponse{Code: http.StatusBadRequest}, err
+		}
+	}
+
+	objectives := f.objectives.Match(matchers)
 	list := make([]openapiserver.Objective, 0, len(objectives))
-	for _, objective := range objectives {
-		list = append(list, openapi.ServerFromInternal(objective))
+	for _, o := range objectives {
+		list = append(list, openapi.ServerFromInternal(o))
 	}
 
 	return openapiserver.ImplResponse{
@@ -188,7 +243,10 @@ func (f FilesystemObjectiveServer) ListObjectives(ctx context.Context) (openapis
 
 func (f FilesystemObjectiveServer) GetObjective(ctx context.Context, expr string) (openapiserver.ImplResponse, error) {
 	// TODO: Parse expr to match properly
-	objective, ok := objectives[expr]
+
+	f.objectives.mu.RLock()
+	objective, ok := f.objectives.objectives[expr]
+	f.objectives.mu.RUnlock()
 	if !ok {
 		return openapiserver.ImplResponse{Code: http.StatusNotFound}, nil
 	}

filesystem_test.go

@@ -0,0 +1,49 @@
+package main
+
+import (
+	"testing"
+
+	"github.com/prometheus/prometheus/pkg/labels"
+	"github.com/pyrra-dev/pyrra/slo"
+	"github.com/stretchr/testify/require"
+)
+
+func TestMatchObjectives(t *testing.T) {
+	o1 := slo.Objective{Labels: labels.FromStrings("foo", "bar")}
+	o2 := slo.Objective{Labels: labels.FromStrings("foo", "bar", "ying", "yang")}
+	o3 := slo.Objective{Labels: labels.FromStrings("foo", "baz")}
+
+	objectives := Objectives{objectives: map[string]slo.Objective{}}
+	objectives.Set(o1)
+	objectives.Set(o2)
+	objectives.Set(o3)
+
+	matches := objectives.Match([]*labels.Matcher{
+		labels.MustNewMatcher(labels.MatchEqual, "foo", "foo"),
+	})
+	require.Nil(t, matches)
+
+	matches = objectives.Match([]*labels.Matcher{
+		labels.MustNewMatcher(labels.MatchEqual, "foo", "bar"),
+	})
+	require.Contains(t, matches, o1)
+	require.Contains(t, matches, o2)
+
+	matches = objectives.Match([]*labels.Matcher{
+		labels.MustNewMatcher(labels.MatchEqual, "foo", "baz"),
+	})
+	require.Contains(t, matches, o3)
+
+	matches = objectives.Match([]*labels.Matcher{
+		labels.MustNewMatcher(labels.MatchEqual, "foo", "bar"),
+		labels.MustNewMatcher(labels.MatchEqual, "ying", "yang"),
+	})
+	require.Contains(t, matches, o2)
+
+	matches = objectives.Match([]*labels.Matcher{
+		labels.MustNewMatcher(labels.MatchRegexp, "foo", "ba."),
+	})
+	require.Contains(t, matches, o1)
+	require.Contains(t, matches, o2)
+	require.Contains(t, matches, o3)
+}

main.go

@@ -29,6 +29,7 @@ import (
 	promconfig "github.com/prometheus/common/config"
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/pkg/labels"
+	"github.com/prometheus/prometheus/promql/parser"
 	"github.com/pyrra-dev/pyrra/openapi"
 	openapiclient "github.com/pyrra-dev/pyrra/openapi/client"
 	openapiserver "github.com/pyrra-dev/pyrra/openapi/server/go"
@@ -63,8 +64,6 @@ func main() {
 	case "kubernetes":
 		cmdKubernetes(CLI.Kubernetes.MetricsAddr)
 	}
-
-	return
 }
 
 func cmdAPI(prometheusURL, prometheusExternal, apiURL *url.URL, prometheusBearerTokenPath string) {
@@ -279,8 +278,15 @@ type ObjectivesServer struct {
 	apiclient *openapiclient.APIClient
 }
 
-func (o *ObjectivesServer) ListObjectives(ctx context.Context) (openapiserver.ImplResponse, error) {
-	objectives, _, err := o.apiclient.ObjectivesApi.ListObjectives(ctx).Execute()
+func (o *ObjectivesServer) ListObjectives(ctx context.Context, query string) (openapiserver.ImplResponse, error) {
+	if query != "" {
+		// We'll parse the query matchers already to make sure it's valid before passing on to the backend.
+		if _, err := parser.ParseMetricSelector(query); err != nil {
+			return openapiserver.ImplResponse{Code: http.StatusBadRequest}, err
+		}
+	}
+
+	objectives, _, err := o.apiclient.ObjectivesApi.ListObjectives(ctx, query).Execute()
 	if err != nil {
 		return openapiserver.ImplResponse{Code: http.StatusInternalServerError}, err
 	}
@@ -315,7 +321,7 @@ func (o *ObjectivesServer) GetObjective(ctx context.Context, expr string) (opena
 }
 
 func (o *ObjectivesServer) GetObjectiveStatus(ctx context.Context, expr string) (openapiserver.ImplResponse, error) {
-	clientObjective, _, err := o.apiclient.ObjectivesApi.GetObjective(ctx, expr).Execute()
+	clientObjectives, _, err := o.apiclient.ObjectivesApi.ListObjectives(ctx, expr).Execute()
 	if err != nil {
 		var apiErr openapiclient.GenericOpenAPIError
 		if errors.As(err, &apiErr) {
@@ -326,7 +332,11 @@ func (o *ObjectivesServer) GetObjectiveStatus(ctx context.Context, expr string)
 
 		return openapiserver.ImplResponse{Code: http.StatusInternalServerError}, err
 	}
-	objective := openapi.InternalFromClient(clientObjective)
+	if len(clientObjectives) != 1 {
+		return openapiserver.ImplResponse{Code: http.StatusBadRequest}, fmt.Errorf("expr matches not exactly one SLO, it matches: %d", len(clientObjectives))
+	}
+
+	objective := openapi.InternalFromClient(clientObjectives[0])
 
 	ts := RoundUp(time.Now().UTC(), 5*time.Minute)