feat: Pull cost attribution labels for tenant (#1544)

Following a similar path as limits to pull cost attribution labels for
tenants. Adds a new internal library `cals/cals.go` that provides an
interface that wraps the `TenantProvider` client. Updates all of the
wiring from Updated => Scraper to pass along the new cal client.

Starts the process of using the library in `scraper.go`, which right now
finds the cost attribution labels for a specific tenant. Returns an
empty slice if nothing is found.

This will be dependent upon the API changes being released to get cost
attribution labels defined at the tenant level.

My goal is to have a follow up PR that updates the telemeter bits to track executions by cost attribution labels. 

- relies on grafana/synthetic-monitoring-api#1750

Author: Pokom

cmd/synthetic-monitoring-agent/main.go

@@ -24,6 +24,7 @@ import (
 	"google.golang.org/grpc/grpclog"
 
 	"github.com/grafana/synthetic-monitoring-agent/internal/adhoc"
+	"github.com/grafana/synthetic-monitoring-agent/internal/cals"
 	"github.com/grafana/synthetic-monitoring-agent/internal/checks"
 	"github.com/grafana/synthetic-monitoring-agent/internal/feature"
 	"github.com/grafana/synthetic-monitoring-agent/internal/http"
@@ -315,6 +316,7 @@ func run(args []string, stdout io.Writer) error {
 	publisher := publisherFactory(ctx, tm, zl.With().Str("subsystem", "publisher").Str("version", config.SelectedPublisher).Logger(), promRegisterer)
 	limits := limits.NewTenantLimits(tm)
 	secretProvider := secrets.NewSecretProvider(tm, 60*time.Second, zl.With().Str("subsystem", "secretstore").Logger())
+	cals := cals.NewCostAttributionLabels(tm)
 
 	telemetry := telemetry.NewTelemeter(
 		ctx, uuid.New().String(), time.Duration(config.TelemetryTimeSpan)*time.Minute,
@@ -324,21 +326,23 @@ func run(args []string, stdout io.Writer) error {
 	)
 
 	checksUpdater, err := checks.NewUpdater(checks.UpdaterOptions{
-		Conn:           conn,
-		Logger:         zl.With().Str("subsystem", "updater").Logger(),
-		Backoff:        newConnectionBackoff(),
-		Publisher:      publisher,
-		TenantCh:       tenantCh,
-		IsConnected:    readynessHandler.Set,
-		PromRegisterer: promRegisterer,
-		Features:       features,
-		K6Runner:       k6Runner,
-		ScraperFactory: scraper.New,
-		TenantLimits:   limits,
-		SecretProvider: secretProvider,
-		Telemeter:      telemetry,
-		UsageReporter:  usageReporter,
+		Conn:                  conn,
+		Logger:                zl.With().Str("subsystem", "updater").Logger(),
+		Backoff:               newConnectionBackoff(),
+		Publisher:             publisher,
+		TenantCh:              tenantCh,
+		IsConnected:           readynessHandler.Set,
+		PromRegisterer:        promRegisterer,
+		Features:              features,
+		K6Runner:              k6Runner,
+		ScraperFactory:        scraper.New,
+		TenantLimits:          limits,
+		SecretProvider:        secretProvider,
+		Telemeter:             telemetry,
+		UsageReporter:         usageReporter,
+		CostAttributionLabels: cals,
 	})
+
 	if err != nil {
 		return fmt.Errorf("cannot create checks updater: %w", err)
 	}

internal/cals/cals.go

@@ -0,0 +1,43 @@
+package cals
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	"github.com/grafana/synthetic-monitoring-agent/internal/model"
+	sm "github.com/grafana/synthetic-monitoring-agent/pkg/pb/synthetic_monitoring"
+)
+
+var (
+	ErrTenantProvider = errors.New("fetching tenant data")
+)
+
+type TenantProvider interface {
+	GetTenant(context.Context, *sm.TenantInfo) (*sm.Tenant, error)
+}
+
+// CostAttributionLabels has a TenantProvider that pulls data about a specific tenant
+type CostAttributionLabels struct {
+	provider TenantProvider
+}
+
+// NewCostAttributionLabels is a helper method to create a NewCostAttributionLabels provider
+func NewCostAttributionLabels(provider TenantProvider) *CostAttributionLabels {
+	return &CostAttributionLabels{
+		provider: provider,
+	}
+}
+
+// CostAttributionLabels will call TenantProvider.GetTenant to search for a specific tenant and returns Tenant.CostAttributionLabels
+func (tcal CostAttributionLabels) CostAttributionLabels(ctx context.Context, tenantID model.GlobalID) ([]string, error) {
+	tenant, err := tcal.provider.GetTenant(ctx, &sm.TenantInfo{
+		Id: int64(tenantID),
+	})
+
+	if err != nil {
+		return nil, fmt.Errorf("%w: %v", ErrTenantProvider, err)
+	}
+
+	return tenant.CostAttributionLabels, nil
+}

internal/cals/cals_test.go

@@ -0,0 +1,55 @@
+package cals
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	sm "github.com/grafana/synthetic-monitoring-agent/pkg/pb/synthetic_monitoring"
+	"github.com/stretchr/testify/require"
+)
+
+type mockTenantProvider struct {
+	tenant *sm.Tenant
+	err    error
+}
+
+func (m mockTenantProvider) GetTenant(_ context.Context, _ *sm.TenantInfo) (*sm.Tenant, error) {
+	return m.tenant, m.err
+}
+
+func TestTenantCostAttributionLabels_GetCostAttributionLabels(t *testing.T) {
+	testcases := map[string]struct {
+		tenantProvider                *mockTenantProvider
+		expectedCostAttributionLabels []string
+		expectError                   bool
+	}{
+		"cals should match up": {
+			tenantProvider: &mockTenantProvider{
+				tenant: &sm.Tenant{
+					CostAttributionLabels: []string{"this", "is", "a", "test"},
+				},
+			},
+			expectedCostAttributionLabels: []string{"this", "is", "a", "test"},
+			expectError:                   false,
+		},
+		"Handle returning an error": {
+			tenantProvider: &mockTenantProvider{
+				tenant: &sm.Tenant{},
+				err:    fmt.Errorf("error getting tenant info"),
+			},
+			expectedCostAttributionLabels: []string{},
+			expectError:                   true,
+		},
+	}
+	for name, testcase := range testcases {
+		t.Run(name, func(t *testing.T) {
+			tcal := NewCostAttributionLabels(testcase.tenantProvider)
+			cals, err := tcal.CostAttributionLabels(context.Background(), 1)
+			if testcase.expectError {
+				require.ErrorIs(t, err, ErrTenantProvider)
+			}
+			require.ElementsMatch(t, cals, testcase.expectedCostAttributionLabels)
+		})
+	}
+}

internal/checks/checks.go

@@ -23,6 +23,7 @@ import (
 
 	logproto "github.com/grafana/loki/pkg/push"
 
+	"github.com/grafana/synthetic-monitoring-agent/internal/cals"
 	"github.com/grafana/synthetic-monitoring-agent/internal/error_types"
 	"github.com/grafana/synthetic-monitoring-agent/internal/feature"
 	"github.com/grafana/synthetic-monitoring-agent/internal/k6runner"
@@ -84,6 +85,7 @@ type Updater struct {
 	tenantSecrets  secrets.SecretProvider
 	telemeter      *telemetry.Telemeter
 	usageReporter  usage.Reporter
+	tenantCals     *cals.CostAttributionLabels
 }
 
 type apiInfo struct {
@@ -106,20 +108,21 @@ type (
 )
 
 type UpdaterOptions struct {
-	Conn           *grpc.ClientConn
-	Logger         zerolog.Logger
-	Backoff        Backoffer
-	Publisher      pusher.Publisher
-	TenantCh       chan<- sm.Tenant
-	IsConnected    func(bool)
-	PromRegisterer prometheus.Registerer
-	Features       feature.Collection
-	K6Runner       k6runner.Runner
-	ScraperFactory scraper.Factory
-	TenantLimits   *limits.TenantLimits
-	SecretProvider secrets.SecretProvider
-	Telemeter      *telemetry.Telemeter
-	UsageReporter  usage.Reporter
+	Conn                  *grpc.ClientConn
+	Logger                zerolog.Logger
+	Backoff               Backoffer
+	Publisher             pusher.Publisher
+	TenantCh              chan<- sm.Tenant
+	IsConnected           func(bool)
+	PromRegisterer        prometheus.Registerer
+	Features              feature.Collection
+	K6Runner              k6runner.Runner
+	ScraperFactory        scraper.Factory
+	TenantLimits          *limits.TenantLimits
+	SecretProvider        secrets.SecretProvider
+	Telemeter             *telemetry.Telemeter
+	UsageReporter         usage.Reporter
+	CostAttributionLabels *cals.CostAttributionLabels
 }
 
 func NewUpdater(opts UpdaterOptions) (*Updater, error) {
@@ -253,6 +256,7 @@ func NewUpdater(opts UpdaterOptions) (*Updater, error) {
 			scrapesCounter:      scrapesCounter,
 		},
 		usageReporter: opts.UsageReporter,
+		tenantCals:    opts.CostAttributionLabels,
 	}, nil
 }
 
@@ -935,8 +939,9 @@ func (c *Updater) addAndStartScraperWithLock(ctx context.Context, check model.Ch
 		c.logger,
 		metrics,
 		c.k6Runner,
-		c.tenantLimits, c.telemeter, c.tenantSecrets,
+		c.tenantLimits, c.telemeter, c.tenantSecrets, c.tenantCals,
 	)
+
 	if err != nil {
 		return fmt.Errorf("cannot create new scraper: %w", err)
 	}

internal/checks/checks_test.go

@@ -474,6 +474,7 @@ func testScraperFactory(ctx context.Context, check model.Check, publisher pusher
 	labelsLimiter scraper.LabelsLimiter,
 	telemeter *telemetry.Telemeter,
 	secretStore secrets.SecretProvider,
+	cals scraper.TenantCals,
 ) (*scraper.Scraper, error) {
 	return scraper.NewWithOpts(
 		ctx,

internal/scraper/scraper.go

@@ -58,6 +58,10 @@ type LabelsLimiter interface {
 	LogLabels(ctx context.Context, tenantID model.GlobalID) (int, error)
 }
 
+type TenantCals interface {
+	CostAttributionLabels(ctx context.Context, tenantID model.GlobalID) ([]string, error)
+}
+
 type Telemeter interface {
 	AddExecution(e telemetry.Execution)
 }
@@ -77,6 +81,7 @@ type Scraper struct {
 	summaries     map[uint64]prometheus.Summary
 	histograms    map[uint64]prometheus.Histogram
 	telemeter     Telemeter
+	cals          TenantCals
 }
 
 type Factory func(
@@ -88,6 +93,7 @@ type Factory func(
 	labelsLimiter LabelsLimiter,
 	telemeter *telemetry.Telemeter,
 	secretStore secrets.SecretProvider,
+	cals TenantCals,
 ) (*Scraper, error)
 
 type (
@@ -122,28 +128,31 @@ func New(
 	labelsLimiter LabelsLimiter,
 	telemeter *telemetry.Telemeter,
 	secretStore secrets.SecretProvider,
+	cals TenantCals,
 ) (*Scraper, error) {
 	return NewWithOpts(ctx, check, ScraperOpts{
-		Probe:         probe,
-		Publisher:     publisher,
-		Logger:        logger,
-		Metrics:       metrics,
-		ProbeFactory:  prober.NewProberFactory(k6runner, probe.Id, features, secretStore),
-		LabelsLimiter: labelsLimiter,
-		Telemeter:     telemeter,
+		Probe:                 probe,
+		Publisher:             publisher,
+		Logger:                logger,
+		Metrics:               metrics,
+		ProbeFactory:          prober.NewProberFactory(k6runner, probe.Id, features, secretStore),
+		LabelsLimiter:         labelsLimiter,
+		Telemeter:             telemeter,
+		CostAttributionLabels: cals,
 	})
 }
 
 var _ Factory = New
 
 type ScraperOpts struct {
-	Probe         sm.Probe
-	Publisher     pusher.Publisher
-	Logger        zerolog.Logger
-	Metrics       Metrics
-	ProbeFactory  prober.ProberFactory
-	LabelsLimiter LabelsLimiter
-	Telemeter     Telemeter
+	Probe                 sm.Probe
+	Publisher             pusher.Publisher
+	Logger                zerolog.Logger
+	Metrics               Metrics
+	ProbeFactory          prober.ProberFactory
+	LabelsLimiter         LabelsLimiter
+	Telemeter             Telemeter
+	CostAttributionLabels TenantCals
 }
 
 func NewWithOpts(ctx context.Context, check model.Check, opts ScraperOpts) (*Scraper, error) {
@@ -181,6 +190,7 @@ func NewWithOpts(ctx context.Context, check model.Check, opts ScraperOpts) (*Scr
 		summaries:     make(map[uint64]prometheus.Summary),
 		histograms:    make(map[uint64]prometheus.Histogram),
 		telemeter:     opts.Telemeter,
+		cals:          opts.CostAttributionLabels,
 	}, nil
 }
 
@@ -289,12 +299,26 @@ func (h *scrapeHandler) scrape(ctx context.Context, t time.Time) {
 		})
 	}
 
+	costAttributionLabels, err := h.scraper.cals.CostAttributionLabels(ctx, h.payload.tenantId)
+	if err != nil {
+		// If cals can't be found, do not block
+		h.scraper.logger.Error().
+			Int64("tenantId", int64(h.payload.tenantId)).
+			Msg("Could not load cals")
+	}
+
+	h.scraper.logger.Debug().
+		Int64("tenantId", int64(h.payload.tenantId)).
+		Int("costAttributionLabelsCount", len(costAttributionLabels)).
+		Msgf("Cost Attribution Labels: %v", costAttributionLabels)
+
 	// If we are dropping the data in case of errors, we should not count that execution.
 	h.scraper.telemeter.AddExecution(telemetry.Execution{
 		LocalTenantID: h.scraper.check.TenantId,
 		RegionID:      int32(h.scraper.check.RegionId),
 		CheckClass:    h.scraper.check.Class(),
-		Duration:      duration,
+		// TODO(@pokom): Add cost attribution label bits here
+		Duration: duration,
 	})
 
 	if h.payload != nil {

internal/scraper/scraper_test.go

@@ -1312,6 +1312,14 @@ func (l testLabelsLimiter) LogLabels(ctx context.Context, tenantID model.GlobalI
 	return l.maxLogLabels, nil
 }
 
+type testCalTenants struct {
+	costAttributionLabels []string
+}
+
+func (t testCalTenants) CostAttributionLabels(_ context.Context, tenantID model.GlobalID) ([]string, error) {
+	return t.costAttributionLabels, nil
+}
+
 func TestScraperCollectData(t *testing.T) {
 	const (
 		checkName     = "check name"
@@ -1879,6 +1887,9 @@ func TestScraperRun(t *testing.T) {
 			maxLogLabels:    15,
 		},
 		Telemeter: testTelemeter,
+		CostAttributionLabels: testCalTenants{
+			costAttributionLabels: []string{"testing", "you"},
+		},
 	})
 
 	require.NoError(t, err)