add check service-discovery command (#8970)

Signed-off-by: darshanime <deathbullet@gmail.com>

cmd/promtool/main.go

@@ -63,6 +63,11 @@ func main() {
 
 	checkCmd := app.Command("check", "Check the resources for validity.")
 
+	sdCheckCmd := checkCmd.Command("service-discovery", "Perform service discovery for the given job name and report the results, including relabeling.")
+	sdConfigFile := sdCheckCmd.Arg("config-file", "The prometheus config file.").Required().ExistingFile()
+	sdJobName := sdCheckCmd.Arg("job", "The job to run service discovery for.").Required().String()
+	sdTimeout := sdCheckCmd.Flag("timeout", "The time to wait for discovery results.").Default("30s").Duration()
+
 	checkConfigCmd := checkCmd.Command("config", "Check if the config files are valid or not.")
 	configFiles := checkConfigCmd.Arg(
 		"config-files",
@@ -202,6 +207,9 @@ func main() {
 	}
 
 	switch parsedCmd {
+	case sdCheckCmd.FullCommand():
+		os.Exit(CheckSD(*sdConfigFile, *sdJobName, *sdTimeout))
+
 	case checkConfigCmd.FullCommand():
 		os.Exit(CheckConfig(*agentMode, *configFiles...))
 

cmd/promtool/sd.go

@@ -0,0 +1,148 @@
+// Copyright 2021 The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"os"
+	"reflect"
+	"time"
+
+	"github.com/go-kit/log"
+
+	"github.com/prometheus/prometheus/config"
+	"github.com/prometheus/prometheus/discovery"
+	"github.com/prometheus/prometheus/discovery/targetgroup"
+	"github.com/prometheus/prometheus/pkg/labels"
+	"github.com/prometheus/prometheus/scrape"
+)
+
+type sdCheckResult struct {
+	DiscoveredLabels labels.Labels `json:"discoveredLabels"`
+	Labels           labels.Labels `json:"labels"`
+	Error            error         `json:"error,omitempty"`
+}
+
+// CheckSD performs service discovery for the given job name and reports the results.
+func CheckSD(sdConfigFiles, sdJobName string, sdTimeout time.Duration) int {
+	logger := log.NewLogfmtLogger(log.NewSyncWriter(os.Stderr))
+
+	cfg, err := config.LoadFile(sdConfigFiles, false, false, logger)
+	if err != nil {
+		fmt.Fprintln(os.Stderr, "Cannot load config", err)
+		return 2
+	}
+
+	var scrapeConfig *config.ScrapeConfig
+	jobs := []string{}
+	jobMatched := false
+	for _, v := range cfg.ScrapeConfigs {
+		jobs = append(jobs, v.JobName)
+		if v.JobName == sdJobName {
+			jobMatched = true
+			scrapeConfig = v
+			break
+		}
+	}
+
+	if !jobMatched {
+		fmt.Fprintf(os.Stderr, "Job %s not found. Select one of:\n", sdJobName)
+		for _, job := range jobs {
+			fmt.Fprintf(os.Stderr, "\t%s\n", job)
+		}
+		return 1
+	}
+
+	targetGroupChan := make(chan []*targetgroup.Group)
+	ctx, cancel := context.WithTimeout(context.Background(), sdTimeout)
+	defer cancel()
+
+	for _, cfg := range scrapeConfig.ServiceDiscoveryConfigs {
+		d, err := cfg.NewDiscoverer(discovery.DiscovererOptions{Logger: logger})
+		if err != nil {
+			fmt.Fprintln(os.Stderr, "Could not create new discoverer", err)
+			return 2
+		}
+		go d.Run(ctx, targetGroupChan)
+	}
+
+	var targetGroups []*targetgroup.Group
+	sdCheckResults := make(map[string][]*targetgroup.Group)
+outerLoop:
+	for {
+		select {
+		case targetGroups = <-targetGroupChan:
+			for _, tg := range targetGroups {
+				sdCheckResults[tg.Source] = append(sdCheckResults[tg.Source], tg)
+			}
+		case <-ctx.Done():
+			break outerLoop
+		}
+	}
+	results := []sdCheckResult{}
+	for _, tgs := range sdCheckResults {
+		results = append(results, getSDCheckResult(tgs, scrapeConfig)...)
+	}
+
+	res, err := json.MarshalIndent(results, "", "  ")
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Could not marshal result json: %s", err)
+		return 2
+	}
+
+	fmt.Printf("%s", res)
+	return 0
+}
+
+func getSDCheckResult(targetGroups []*targetgroup.Group, scrapeConfig *config.ScrapeConfig) []sdCheckResult {
+	sdCheckResults := []sdCheckResult{}
+	for _, targetGroup := range targetGroups {
+		for _, target := range targetGroup.Targets {
+			labelSlice := make([]labels.Label, 0, len(target)+len(targetGroup.Labels))
+
+			for name, value := range target {
+				labelSlice = append(labelSlice, labels.Label{Name: string(name), Value: string(value)})
+			}
+
+			for name, value := range targetGroup.Labels {
+				if _, ok := target[name]; !ok {
+					labelSlice = append(labelSlice, labels.Label{Name: string(name), Value: string(value)})
+				}
+			}
+
+			targetLabels := labels.New(labelSlice...)
+			res, orig, err := scrape.PopulateLabels(targetLabels, scrapeConfig)
+			result := sdCheckResult{
+				DiscoveredLabels: orig,
+				Labels:           res,
+				Error:            err,
+			}
+
+			duplicateRes := false
+			for _, sdCheckRes := range sdCheckResults {
+				if reflect.DeepEqual(sdCheckRes, result) {
+					duplicateRes = true
+					break
+				}
+			}
+
+			if !duplicateRes {
+				sdCheckResults = append(sdCheckResults, result)
+			}
+		}
+	}
+	return sdCheckResults
+}

cmd/promtool/sd_test.go

@@ -0,0 +1,66 @@
+// Copyright 2021 The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"testing"
+
+	"github.com/prometheus/common/model"
+	"github.com/prometheus/prometheus/config"
+	"github.com/prometheus/prometheus/discovery/targetgroup"
+	"github.com/prometheus/prometheus/pkg/labels"
+	"github.com/prometheus/prometheus/pkg/relabel"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestSDCheckResult(t *testing.T) {
+	targetGroups := []*targetgroup.Group{{
+		Targets: []model.LabelSet{
+			map[model.LabelName]model.LabelValue{"__address__": "localhost:8080", "foo": "bar"},
+		},
+	}}
+
+	reg, err := relabel.NewRegexp("(.*)")
+	require.Nil(t, err)
+
+	scrapeConfig := &config.ScrapeConfig{
+		RelabelConfigs: []*relabel.Config{{
+			SourceLabels: model.LabelNames{"foo"},
+			Action:       relabel.Replace,
+			TargetLabel:  "newfoo",
+			Regex:        reg,
+			Replacement:  "$1",
+		}},
+	}
+
+	expectedSDCheckResult := []sdCheckResult{
+		sdCheckResult{
+			DiscoveredLabels: labels.Labels{
+				labels.Label{Name: "__address__", Value: "localhost:8080"},
+				labels.Label{Name: "__scrape_interval__", Value: "0s"},
+				labels.Label{Name: "__scrape_timeout__", Value: "0s"},
+				labels.Label{Name: "foo", Value: "bar"}},
+			Labels: labels.Labels{
+				labels.Label{Name: "__address__", Value: "localhost:8080"},
+				labels.Label{Name: "__scrape_interval__", Value: "0s"},
+				labels.Label{Name: "__scrape_timeout__", Value: "0s"},
+				labels.Label{Name: "foo", Value: "bar"},
+				labels.Label{Name: "instance", Value: "localhost:8080"},
+				labels.Label{Name: "newfoo", Value: "bar"}},
+			Error: nil,
+		}}
+
+	require.Equal(t, expectedSDCheckResult, getSDCheckResult(targetGroups, scrapeConfig))
+}

scrape/manager_test.go

@@ -335,7 +335,7 @@ func TestPopulateLabels(t *testing.T) {
 	for _, c := range cases {
 		in := c.in.Copy()
 
-		res, orig, err := populateLabels(c.in, c.cfg)
+		res, orig, err := PopulateLabels(c.in, c.cfg)
 		if c.err != "" {
 			require.EqualError(t, err, c.err)
 		} else {

scrape/target.go

@@ -348,10 +348,10 @@ func (app *timeLimitAppender) Append(ref uint64, lset labels.Labels, t int64, v
 	return ref, nil
 }
 
-// populateLabels builds a label set from the given label set and scrape configuration.
+// PopulateLabels builds a label set from the given label set and scrape configuration.
 // It returns a label set before relabeling was applied as the second return value.
 // Returns the original discovered label set found before relabelling was applied if the target is dropped during relabeling.
-func populateLabels(lset labels.Labels, cfg *config.ScrapeConfig) (res, orig labels.Labels, err error) {
+func PopulateLabels(lset labels.Labels, cfg *config.ScrapeConfig) (res, orig labels.Labels, err error) {
 	// Copy labels into the labelset for the target if they are not set already.
 	scrapeLabels := []labels.Label{
 		{Name: model.JobLabel, Value: cfg.JobName},
@@ -488,7 +488,7 @@ func TargetsFromGroup(tg *targetgroup.Group, cfg *config.ScrapeConfig) ([]*Targe
 
 		lset := labels.New(lbls...)
 
-		lbls, origLabels, err := populateLabels(lset, cfg)
+		lbls, origLabels, err := PopulateLabels(lset, cfg)
 		if err != nil {
 			failures = append(failures, errors.Wrapf(err, "instance %d in group %s", i, tg))
 		}