github.com/thanos-io/thanos@v0.32.5/pkg/tracing/migration/sampler.go (about)

     1  // Copyright (c) The Thanos Authors.
     2  // Licensed under the Apache License 2.0.
     3  
     4  package migration
     5  
     6  import (
     7  	"fmt"
     8  
     9  	"go.opentelemetry.io/otel/attribute"
    10  	tracesdk "go.opentelemetry.io/otel/sdk/trace"
    11  )
    12  
    13  // ForceTracingAttributeKey is used to signalize a span should be traced.
    14  const ForceTracingAttributeKey = "thanos.force_tracing"
    15  
    16  type samplerWithOverride struct {
    17  	baseSampler tracesdk.Sampler
    18  	overrideKey attribute.Key
    19  }
    20  
    21  // SamplerWithOverride creates a new sampler with the capability to override
    22  // the sampling decision, if the span includes an attribute with the specified key.
    23  // Otherwise the sampler delegates the decision to the wrapped base sampler. This
    24  // is primarily used to enable forced tracing in Thanos components.
    25  // Implements go.opentelemetry.io/otel/sdk/trace.Sampler interface.
    26  func SamplerWithOverride(baseSampler tracesdk.Sampler, overrideKey attribute.Key) tracesdk.Sampler {
    27  	return samplerWithOverride{
    28  		baseSampler,
    29  		overrideKey,
    30  	}
    31  }
    32  
    33  func (s samplerWithOverride) ShouldSample(p tracesdk.SamplingParameters) tracesdk.SamplingResult {
    34  	for _, attr := range p.Attributes {
    35  		if attr.Key == s.overrideKey {
    36  			return tracesdk.SamplingResult{
    37  				Decision: tracesdk.RecordAndSample,
    38  			}
    39  		}
    40  	}
    41  
    42  	return s.baseSampler.ShouldSample(p)
    43  }
    44  
    45  func (s samplerWithOverride) Description() string {
    46  	return fmt.Sprintf("SamplerWithOverride{%s}", string(s.overrideKey))
    47  }