Skip to content
This repository was archived by the owner on Jun 19, 2026. It is now read-only.

add legacy tenant attribution in router #274

Open
yuchen-db wants to merge 1 commit into
base: db_main
Choose a base branch
from
Open

add legacy tenant attribution in router #274

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
33 changes: 33 additions & 0 deletions cmd/thanos/receive.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -302,6 +302,23 @@ func runReceive(
return errors.Wrap(err, "creating limiter")
}

// Initialize tenant attributor if configured.
var tenantAttributor *receive.TenantAttributor
if conf.tenantRulesPath != "" {
tenantAttributor, err = receive.NewTenantAttributor(
conf.tenantRulesPath,
conf.defaultTenantID,
conf.verifyTenantAttribution,
reg,
log.With(logger, "component", "tenant-attributor"),
)
if err != nil {
return errors.Wrap(err, "creating tenant attributor")
}
} else if conf.verifyTenantAttribution {
level.Warn(logger).Log("msg", "receive.verify-tenant-attribution is set but receive.tenant-rules is not configured, ignoring")
}

webHandler := receive.NewHandler(log.With(logger, "component", "receive-handler"), &receive.Options{
Writer: writer,
ListenAddress: conf.rwAddress,
Expand All @@ -324,6 +341,7 @@ func runReceive(
Limiter: limiter,
AsyncForwardWorkerCount: conf.asyncForwardWorkerCount,
ReplicationProtocol: receive.ReplicationProtocol(conf.replicationProtocol),
TenantAttributor: tenantAttributor,
})

{
Expand Down Expand Up @@ -1004,6 +1022,9 @@ type receiveConfig struct {
writerInterning bool
splitTenantLabelName string

tenantRulesPath string
verifyTenantAttribution bool

hashFunc string

ignoreBlockSize bool
Expand Down Expand Up @@ -1089,6 +1110,18 @@ func (rc *receiveConfig) registerFlag(cmd extkingpin.FlagClause) {

cmd.Flag("receive.split-tenant-label-name", "Label name through which the request will be split into multiple tenants. This takes precedence over the HTTP header.").Default("").StringVar(&rc.splitTenantLabelName)

cmd.Flag("receive.tenant-rules",
"Path to YAML file containing tenant attribution rules. "+
"Rules use M3-style filter syntax (e.g., 'job:prometheus namespace:prod*'). "+
"First matching rule wins. Falls back to --receive.default-tenant-id if no rule matches.").
PlaceHolder("<path>").StringVar(&rc.tenantRulesPath)

cmd.Flag("receive.verify-tenant-attribution",
"When set with --receive.tenant-rules, compares attributed tenant with HTTP header tenant "+
"and exposes metrics (thanos_receive_tenant_attribution_matches_total and "+
"thanos_receive_tenant_attribution_mismatches_total). Does not change routing behavior.").
Default("false").BoolVar(&rc.verifyTenantAttribution)

cmd.Flag("receive.tenant-label-name", "Label name through which the tenant will be announced.").Default(tenancy.DefaultTenantLabel).StringVar(&rc.tenantLabelName)

cmd.Flag("receive.replica-header", "HTTP header specifying the replica number of a write request.").Default(receive.DefaultReplicaHeader).StringVar(&rc.replicaHeader)
Expand Down
39 changes: 33 additions & 6 deletions pkg/receive/handler.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,7 @@ type Options struct {
Limiter *Limiter
AsyncForwardWorkerCount uint
ReplicationProtocol ReplicationProtocol
TenantAttributor *TenantAttributor
}

// Handler serves a Prometheus remote write receiving HTTP endpoint.
Expand Down Expand Up @@ -529,6 +530,9 @@ func (h *Handler) receiveHTTP(w http.ResponseWriter, r *http.Request) {
span.SetTag("receiver.mode", string(h.receiverMode))
defer span.Finish()

// Check if tenant header is present before calling GetTenantFromHTTP
httpHeaderPresent := r.Header.Get(h.options.TenantHeader) != "" || r.Header.Get(tenancy.DefaultTenantHeader) != ""

tenantHTTP, err := tenancy.GetTenantFromHTTP(r, h.options.TenantHeader, h.options.DefaultTenantID, h.options.TenantField)
if err != nil {
level.Error(h.logger).Log("msg", "error getting tenant from HTTP", "err", err)
Expand Down Expand Up @@ -644,7 +648,7 @@ func (h *Handler) receiveHTTP(w http.ResponseWriter, r *http.Request) {
}

responseStatusCode := http.StatusOK
tenantStats, err := h.handleRequest(ctx, rep, tenantHTTP, &wreq)
tenantStats, err := h.handleRequest(ctx, rep, tenantHTTP, httpHeaderPresent, &wreq)
if err != nil {
level.Debug(tLogger).Log("msg", "failed to handle request", "err", err.Error())
switch errors.Cause(err) {
Expand Down Expand Up @@ -703,7 +707,7 @@ type requestStats struct {

type tenantRequestStats map[string]requestStats

func (h *Handler) handleRequest(ctx context.Context, rep uint64, tenantHTTP string, wreq *prompb.WriteRequest) (tenantRequestStats, error) {
func (h *Handler) handleRequest(ctx context.Context, rep uint64, tenantHTTP string, httpHeaderPresent bool, wreq *prompb.WriteRequest) (tenantRequestStats, error) {

@jnyi jnyi Jan 16, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why introduce new variable, can tenantHTTP string tell if a header is present? I think if it isn't, it will be empty?

@jnyi jnyi Jan 16, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is obtained in tenantHTTP, err := tenancy.GetTenantFromHTTP(r, h.options.TenantHeader, h.options.DefaultTenantID, h.options.TenantField)

a few places might be broken for example,

rate limit logic: if !requestLimiter.AllowSizeBytes(tenantHTTP, int64(len(reqBuf))) {
e2e latency calculation: h.writeE2eLatency.WithLabelValues(strconv.Itoa(responseStatusCode), tenantHTTP, strconv.FormatBool(isPreAgged)).Observe(lat)

with new tenant attribution, we might need to fix how those work too, (especially for rate limit which is request scoped, it might be not possible to reject 1 remote write directly anymore)

tLogger := log.With(h.logger, "tenantHTTP", tenantHTTP)

// This replica value is used to detect cycles in cyclic topologies.
Expand Down Expand Up @@ -732,7 +736,7 @@ func (h *Handler) handleRequest(ctx context.Context, rep uint64, tenantHTTP stri
// Forward any time series as necessary. All time series
// destined for the local node will be written to the receiver.
// Time series will be replicated as necessary.
return h.forward(ctx, tenantHTTP, r, wreq)
return h.forward(ctx, tenantHTTP, httpHeaderPresent, r, wreq)
}

// forward accepts a write request, batches its time series by
Expand All @@ -743,7 +747,7 @@ func (h *Handler) handleRequest(ctx context.Context, rep uint64, tenantHTTP stri
// unless the request needs to be replicated.
// The function only returns when all requests have finished
// or the context is canceled.
func (h *Handler) forward(ctx context.Context, tenantHTTP string, r replica, wreq *prompb.WriteRequest) (tenantRequestStats, error) {
func (h *Handler) forward(ctx context.Context, tenantHTTP string, httpHeaderPresent bool, r replica, wreq *prompb.WriteRequest) (tenantRequestStats, error) {
span, ctx := tracing.StartSpan(ctx, "receive_fanout_forward")
defer span.Finish()

Expand All @@ -761,6 +765,7 @@ func (h *Handler) forward(ctx context.Context, tenantHTTP string, r replica, wre
writeRequest: wreq,
replicas: replicas,
alreadyReplicated: r.replicated,
httpHeaderPresent: httpHeaderPresent,
}

return h.fanoutForward(ctx, params)
Expand All @@ -771,6 +776,7 @@ type remoteWriteParams struct {
writeRequest *prompb.WriteRequest
replicas []uint64
alreadyReplicated bool
httpHeaderPresent bool // true if tenant came from HTTP header (not default)
}

func (h *Handler) gatherWriteStats(rf int, writes ...map[endpointReplica]map[string]trackedSeries) tenantRequestStats {
Expand Down Expand Up @@ -831,7 +837,7 @@ func (h *Handler) fanoutForward(ctx context.Context, params remoteWriteParams) (
}
requestLogger := log.With(h.logger, logTags...)

localWrites, remoteWrites, err := h.distributeTimeseriesToReplicas(params.tenant, params.replicas, params.writeRequest.Timeseries)
localWrites, remoteWrites, err := h.distributeTimeseriesToReplicas(params.tenant, params.httpHeaderPresent, params.replicas, params.writeRequest.Timeseries)
if err != nil {
level.Error(requestLogger).Log("msg", "failed to distribute timeseries to replicas", "err", err)
return stats, err
Expand Down Expand Up @@ -914,6 +920,7 @@ func (h *Handler) fanoutForward(ctx context.Context, params remoteWriteParams) (
// series that should be written to remote nodes.
func (h *Handler) distributeTimeseriesToReplicas(
tenantHTTP string,

@jnyi jnyi Jan 16, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

wonder if we can use tenantHTTP to tell the diff, we should set it to empty string if HTTP headers didn't set it and set to default TenantId within this function

httpHeaderPresent bool,
replicas []uint64,
timeseries []prompb.TimeSeries,
) (map[endpointReplica]map[string]trackedSeries, map[endpointReplica]map[string]trackedSeries, error) {
Expand All @@ -935,6 +942,25 @@ func (h *Handler) distributeTimeseriesToReplicas(
}
}

// Tenant attribution logic
if h.options.TenantAttributor != nil {

@jnyi jnyi Jan 16, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

also note this is exclusive to h.splitTenantLabelName != ""

lbls := labelpb.ZLabelsToPromLabels(ts.Labels)
attributedTenant := h.options.TenantAttributor.GetTenantFromLabels(lbls)

if h.options.TenantAttributor.IsVerifyMode() {
// VERIFICATION MODE: compute attribution for metrics only
// Compare attributed tenant with HTTP tenant
h.options.TenantAttributor.RecordVerification(attributedTenant, tenantHTTP)
// DO NOT modify tenant - keep using HTTP tenant for routing
} else {
// ATTRIBUTION MODE: only attribute if no HTTP header was present
if !httpHeaderPresent {
tenant = attributedTenant
}
// If HTTP header present, use that (tenant already set correctly)
}
}

for _, rn := range replicas {
endpoint, err := h.hashring.GetN(tenant, &ts, rn)
if err != nil {
Expand Down Expand Up @@ -1125,7 +1151,8 @@ func (h *Handler) RemoteWrite(ctx context.Context, r *storepb.WriteRequest) (*st
span, ctx := tracing.StartSpan(ctx, "receive_grpc")
defer span.Finish()

_, err := h.handleRequest(ctx, uint64(r.Replica), r.Tenant, &prompb.WriteRequest{Timeseries: r.Timeseries})
// For gRPC requests, tenant is explicitly provided, treat as if HTTP header was present
_, err := h.handleRequest(ctx, uint64(r.Replica), r.Tenant, true, &prompb.WriteRequest{Timeseries: r.Timeseries})
if err != nil {
level.Debug(h.logger).Log("msg", "failed to handle request", "err", err)
}
Expand Down
7 changes: 4 additions & 3 deletions pkg/receive/handler_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1862,6 +1862,7 @@ func TestDistributeSeries(t *testing.T) {

_, remote, err := h.distributeTimeseriesToReplicas(
"foo",
true,
[]uint64{0},
[]prompb.TimeSeries{
{
Expand Down Expand Up @@ -1920,7 +1921,7 @@ func TestHandlerFlippingHashrings(t *testing.T) {
return
}

_, err := h.handleRequest(ctx, 0, "test", &prompb.WriteRequest{
_, err := h.handleRequest(ctx, 0, "test", true, &prompb.WriteRequest{
Timeseries: []prompb.TimeSeries{
{
Labels: labelpb.ZLabelsFromPromLabels(labels.FromStrings("foo", "bar")),
Expand Down Expand Up @@ -2004,7 +2005,7 @@ func TestIngestorRestart(t *testing.T) {
},
}

stats, err := client.handleRequest(ctx, 0, "test", data)
stats, err := client.handleRequest(ctx, 0, "test", true, data)
require.NoError(t, err)
require.Equal(t, tenantRequestStats{
"test": requestStats{timeseries: 1, totalSamples: 1},
Expand All @@ -2019,7 +2020,7 @@ func TestIngestorRestart(t *testing.T) {

iter, errs := 10, 0
for i := 0; i < iter; i++ {
_, err = client.handleRequest(ctx, 0, "test", data)
_, err = client.handleRequest(ctx, 0, "test", true, data)
if err != nil {
require.Error(t, errUnavailable, err)
errs++
Expand Down
Loading
Loading