github.com/grafana/pyroscope@v1.18.0/pkg/pyroscope/context/context.go (about)

     1  package context
     2  
     3  import (
     4  	"context"
     5  	"os"
     6  
     7  	"github.com/go-kit/log"
     8  	"github.com/prometheus/client_golang/prometheus"
     9  )
    10  
    11  type contextKey int
    12  
    13  const (
    14  	loggerKey contextKey = iota
    15  	registryKey
    16  	localBucketClient
    17  )
    18  
    19  var (
    20  	defaultLogger = log.NewLogfmtLogger(os.Stderr)
    21  )
    22  
    23  func WithLogger(ctx context.Context, logger log.Logger) context.Context {
    24  	return context.WithValue(ctx, loggerKey, logger)
    25  }
    26  
    27  func Logger(ctx context.Context) log.Logger {
    28  	if logger, ok := ctx.Value(loggerKey).(log.Logger); ok {
    29  		return logger
    30  	}
    31  	return defaultLogger
    32  }
    33  
    34  func WithRegistry(ctx context.Context, registry prometheus.Registerer) context.Context {
    35  	return context.WithValue(ctx, registryKey, registry)
    36  }
    37  
    38  func Registry(ctx context.Context) prometheus.Registerer {
    39  	if registry, ok := ctx.Value(registryKey).(prometheus.Registerer); ok {
    40  		return registry
    41  	}
    42  	return prometheus.NewRegistry()
    43  }
    44  
    45  func WrapTenant(ctx context.Context, tenantID string) context.Context {
    46  	// wrap registry
    47  	reg := Registry(ctx)
    48  	ctx = WithRegistry(ctx, prometheus.WrapRegistererWith(
    49  		prometheus.Labels{"tenant": tenantID},
    50  		reg,
    51  	))
    52  
    53  	// add field to logger
    54  	logger := Logger(ctx)
    55  	ctx = WithLogger(ctx,
    56  		log.With(logger, "tenant", tenantID),
    57  	)
    58  
    59  	return ctx
    60  }