Feature: Distributor usage trackers

cmd/tempo/app/modules.go

@@ -243,6 +243,8 @@ func (t *App) initDistributor() (services.Service, error) {
 		t.Server.HTTPRouter().Handle("/distributor/ring", distributor.DistributorRing)
 	}
 
+	t.Server.HTTPRouter().Handle("/usage_metrics", distributor.UsageTrackerHandler())
+
 	return t.distributor, nil
 }
 

modules/distributor/config.go

@@ -8,6 +8,7 @@ import (
 	ring_client "github.com/grafana/dskit/ring/client"
 
 	"github.com/grafana/tempo/modules/distributor/forwarder"
+	"github.com/grafana/tempo/modules/distributor/usage"
 	"github.com/grafana/tempo/pkg/util"
 )
 
@@ -37,8 +38,8 @@ type Config struct {
 	LogReceivedSpans    LogSpansConfig            `yaml:"log_received_spans,omitempty"`
 	LogDiscardedSpans   LogSpansConfig            `yaml:"log_discarded_spans,omitempty"`
 	MetricReceivedSpans MetricReceivedSpansConfig `yaml:"metric_received_spans,omitempty"`
-
-	Forwarders forwarder.ConfigList `yaml:"forwarders"`
+	Forwarders          forwarder.ConfigList      `yaml:"forwarders"`
+	Usage               usage.Config              `yaml:"usage,omitempty"`
 
 	// disables write extension with inactive ingesters. Use this along with ingester.lifecycler.unregister_on_shutdown = true
 	//  note that setting these two config values reduces tolerance to failures on rollout b/c there is always one guaranteed to be failing replica
@@ -80,4 +81,6 @@ func (cfg *Config) RegisterFlagsAndApplyDefaults(prefix string, f *flag.FlagSet)
 	f.BoolVar(&cfg.LogDiscardedSpans.Enabled, util.PrefixConfig(prefix, "log-discarded-spans.enabled"), false, "Enable to log every discarded span to help debug ingestion or calculate span error distributions using the logs.")
 	f.BoolVar(&cfg.LogDiscardedSpans.IncludeAllAttributes, util.PrefixConfig(prefix, "log-discarded-spans.include-attributes"), false, "Enable to include span attributes in the logs.")
 	f.BoolVar(&cfg.LogDiscardedSpans.FilterByStatusError, util.PrefixConfig(prefix, "log-discarded-spans.filter-by-status-error"), false, "Enable to filter out spans without status error.")
+
+	cfg.Usage.RegisterFlagsAndApplyDefaults(prefix, f)
 }

modules/distributor/distributor.go

@@ -5,6 +5,7 @@ import (
 	"encoding/hex"
 	"fmt"
 	"math"
+	"net/http"
 	"sync"
 	"time"
 
@@ -28,6 +29,7 @@ import (
 
 	"github.com/grafana/tempo/modules/distributor/forwarder"
 	"github.com/grafana/tempo/modules/distributor/receiver"
+	"github.com/grafana/tempo/modules/distributor/usage"
 	generator_client "github.com/grafana/tempo/modules/generator/client"
 	ingester_client "github.com/grafana/tempo/modules/ingester/client"
 	"github.com/grafana/tempo/modules/overrides"
@@ -154,6 +156,8 @@ type Distributor struct {
 	subservices        *services.Manager
 	subservicesWatcher *services.FailureWatcher
 
+	usage *usage.Tracker
+
 	logger log.Logger
 }
 
@@ -214,6 +218,14 @@ func New(cfg Config, clientCfg ingester_client.Config, ingestersRing ring.ReadRi
 		logger:               logger,
 	}
 
+	if cfg.Usage.Enabled {
+		usage, err := usage.NewTracker(cfg.Usage, "cost-attribution", o.CostAttributionDimensions)
+		if err != nil {
+			return nil, fmt.Errorf("creating usage tracker: %w", err)
+		}
+		d.usage = usage
+	}
+
 	var generatorsPoolFactory ring_client.PoolAddrFunc = func(addr string) (ring_client.PoolClient, error) {
 		return generator_client.New(addr, generatorClientCfg)
 	}
@@ -328,6 +340,7 @@ func (d *Distributor) PushTraces(ctx context.Context, traces ptrace.Traces) (*te
 		return &tempopb.PushResponse{}, nil
 	}
 	// check limits
+	// todo - usage tracker include discarded bytes?
 	err = d.checkForRateLimits(size, spanCount, userID)
 	if err != nil {
 		return nil, err
@@ -360,6 +373,11 @@ func (d *Distributor) PushTraces(ctx context.Context, traces ptrace.Traces) (*te
 	statBytesReceived.Inc(int64(size))
 	statSpansReceived.Inc(int64(spanCount))
 
+	// Usage tracking
+	if d.usage != nil {
+		d.usage.Observe(userID, batches)
+	}
+
 	keys, rebatchedTraces, err := requestsByTraceID(batches, userID, spanCount)
 	if err != nil {
 		overrides.RecordDiscardedSpans(spanCount, reasonInternalError, userID)
@@ -498,6 +516,14 @@ func (*Distributor) Check(_ context.Context, _ *grpc_health_v1.HealthCheckReques
 	return &grpc_health_v1.HealthCheckResponse{Status: grpc_health_v1.HealthCheckResponse_SERVING}, nil
 }
 
+func (d *Distributor) UsageTrackerHandler() http.Handler {
+	if d.usage != nil {
+		return d.usage.Handler()
+	}
+
+	return nil
+}
+
 // requestsByTraceID takes an incoming tempodb.PushRequest and creates a set of keys for the hash ring
 // and traces to pass onto the ingesters.
 func requestsByTraceID(batches []*v1.ResourceSpans, userID string, spanCount int) ([]uint32, []*rebatchedTrace, error) {

modules/distributor/usage/config.go

@@ -0,0 +1,35 @@
+package usage
+
+import (
+	"flag"
+	"time"
+)
+
+const (
+	defaultMaxCardinality = 1000
+	defaultStaleDuration  = 15 * time.Minute
+	defaultPurgePeriod    = time.Minute
+)
+
+type Config struct {
+	Enabled        bool          `yaml:"enabled,omitempty" json:"enabled,omitempty"`
+	MaxCardinality uint          `yaml:"max_cardinality,omitempty" json:"max_cardinality,omitempty"`
+	StaleDuration  time.Duration `yaml:"stale_duration,omitempty" json:"stale_duration,omitempty"`
+	PurgePeriod    time.Duration `yaml:"purge_period,omitempty" json:"purge_period,omitempty"`
+}
+
+func (c *Config) RegisterFlagsAndApplyDefaults(_ string, _ *flag.FlagSet) {
+	c.Enabled = true
+	c.MaxCardinality = defaultMaxCardinality
+	c.StaleDuration = defaultStaleDuration
+	c.PurgePeriod = defaultPurgePeriod
+}
+
+func DefaultConfig() Config {
+	return Config{
+		Enabled:        true,
+		MaxCardinality: defaultMaxCardinality,
+		StaleDuration:  defaultStaleDuration,
+		PurgePeriod:    defaultPurgePeriod,
+	}
+}