#4147

1. added logic to add tag `tenant_id` to tracing spans if tenant is defined in context.
2. added additional tracing spans to merge_querier's methods

Signed-off-by: Vladyslav Diachenko <[email protected]>

Author: vlad-diachenko

CHANGELOG.md

@@ -20,6 +20,7 @@
 * [ENHANCEMENT] Distributor: Added distributors ring status section in the admin page. #4151
 * [BUGFIX] Purger: fix `Invalid null value in condition for column range` caused by `nil` value in range for WriteBatch query. #4128
 * [BUGFIX] Ingester: fixed infrequent panic caused by a race condition between TSDB mmap-ed head chunks truncation and queries. #4176
+* [ENHANCEMENT] Added `tenant_id` tag to tracing spans
 
 ## Blocksconvert
 

pkg/querier/tenantfederation/merge_queryable.go

@@ -15,6 +15,7 @@ import (
 
 	"github.com/cortexproject/cortex/pkg/tenant"
 	"github.com/cortexproject/cortex/pkg/util/concurrency"
+	"github.com/cortexproject/cortex/pkg/util/spanlogger"
 )
 
 const (
@@ -96,6 +97,9 @@ type mergeQueryable struct {
 // Querier returns a new mergeQuerier, which aggregates results from multiple
 // underlying queriers into a single result.
 func (m *mergeQueryable) Querier(ctx context.Context, mint int64, maxt int64) (storage.Querier, error) {
+	// TODO: it's necessary to think how to override context inside querier
+	//  to mark spans created inside querier as child of a span created inside
+	//  methods of merged querier.
 	ids, queriers, err := m.callback(ctx, mint, maxt)
 	if err != nil {
 		return nil, err
@@ -133,6 +137,8 @@ type mergeQuerier struct {
 // For the label "original_" + `idLabelName it will return all the values
 // of the underlying queriers for `idLabelName`.
 func (m *mergeQuerier) LabelValues(name string, matchers ...*labels.Matcher) ([]string, storage.Warnings, error) {
+	log, _ := spanlogger.New(m.ctx, "mergeQuerier.labelValues")
+	defer log.Span.Finish()
 	if name == m.idLabelName {
 		return m.ids, nil, nil
 	}
@@ -152,6 +158,8 @@ func (m *mergeQuerier) LabelValues(name string, matchers ...*labels.Matcher) ([]
 // queriers. It also adds the `idLabelName` and if present in the original
 // results the original `idLabelName`.
 func (m *mergeQuerier) LabelNames() ([]string, storage.Warnings, error) {
+	log, _ := spanlogger.New(m.ctx, "mergeQuerier.labelNames")
+	defer log.Span.Finish()
 	labelNames, warnings, err := m.mergeDistinctStringSlice(func(ctx context.Context, q storage.Querier) ([]string, storage.Warnings, error) {
 		return q.LabelNames()
 	})
@@ -272,6 +280,8 @@ type selectJob struct {
 // matching. The forwarded labelSelector is not containing those that operate
 // on `idLabelName`.
 func (m *mergeQuerier) Select(sortSeries bool, hints *storage.SelectHints, matchers ...*labels.Matcher) storage.SeriesSet {
+	log, ctx := spanlogger.New(m.ctx, "mergeQuerier.select")
+	defer log.Span.Finish()
 	matchedValues, filteredMatchers := filterValuesByMatchers(m.idLabelName, m.ids, matchers...)
 	var jobs = make([]interface{}, len(matchedValues))
 	var seriesSets = make([]storage.SeriesSet, len(matchedValues))
@@ -305,7 +315,7 @@ func (m *mergeQuerier) Select(sortSeries bool, hints *storage.SelectHints, match
 		return nil
 	}
 
-	err := concurrency.ForEach(m.ctx, jobs, maxConcurrency, run)
+	err := concurrency.ForEach(ctx, jobs, maxConcurrency, run)
 	if err != nil {
 		return storage.ErrSeriesSet(err)
 	}

pkg/querier/tenantfederation/merge_queryable_test.go

@@ -9,6 +9,9 @@ import (
 	"testing"
 	"time"
 
+	sl "github.com/cortexproject/cortex/pkg/util/spanlogger"
+	"github.com/opentracing/opentracing-go"
+	"github.com/opentracing/opentracing-go/mocktracer"
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/pkg/labels"
 	"github.com/prometheus/prometheus/storage"
@@ -37,7 +40,11 @@ func (m *mockTenantQueryableWithFilter) Querier(ctx context.Context, _, _ int64)
 		return nil, err
 	}
 
-	q := mockTenantQuerier{tenant: tenantIDs[0], extraLabels: m.extraLabels}
+	q := mockTenantQuerier{
+		tenant:      tenantIDs[0],
+		extraLabels: m.extraLabels,
+		ctx:         ctx,
+	}
 
 	// set warning if exists
 	if m.warningsByTenant != nil {
@@ -66,6 +73,7 @@ type mockTenantQuerier struct {
 
 	warnings storage.Warnings
 	queryErr error
+	ctx      context.Context
 }
 
 func (m mockTenantQuerier) matrix() model.Matrix {
@@ -135,6 +143,8 @@ func (m *mockSeriesSet) Warnings() storage.Warnings {
 }
 
 func (m mockTenantQuerier) Select(_ bool, sp *storage.SelectHints, matchers ...*labels.Matcher) storage.SeriesSet {
+	log, _ := sl.New(m.ctx, "mockTenantQuerier.select")
+	defer log.Span.Finish()
 	var matrix model.Matrix
 
 	for _, s := range m.matrix() {
@@ -478,3 +488,64 @@ func TestSetLabelsRetainExisting(t *testing.T) {
 		assert.Equal(t, tc.expected, setLabelsRetainExisting(tc.labels, tc.additionalLabels...))
 	}
 }
+
+func TestTracingMergeQueryable(t *testing.T) {
+	mockTracer := mocktracer.New()
+	opentracing.SetGlobalTracer(mockTracer)
+	ctx := user.InjectOrgID(context.Background(), "team-a|team-b")
+
+	// set a multi tenant resolver
+	tenant.WithDefaultResolver(tenant.NewMultiResolver())
+	filter := mockTenantQueryableWithFilter{}
+	q := NewQueryable(&filter, false)
+	// retrieve querier if set
+	querier, err := q.Querier(ctx, mint, maxt)
+	require.NoError(t, err)
+
+	seriesSet := querier.Select(true, &storage.SelectHints{Start: mint,
+		End: maxt})
+
+	require.NoError(t, seriesSet.Err())
+	spans := mockTracer.FinishedSpans()
+	assertSpanExist(t, spans, "mergeQuerier.select", map[string]string{sl.TenantIdTagName: "team-a|team-b"})
+	assertSpanExist(t, spans, "mockTenantQuerier.select", map[string]string{sl.TenantIdTagName: "team-a"})
+	assertSpanExist(t, spans, "mockTenantQuerier.select", map[string]string{sl.TenantIdTagName: "team-b"})
+}
+
+func assertSpanExist(t *testing.T,
+	actualSpans []*mocktracer.MockSpan,
+	name string,
+	expectedTags map[string]string) {
+	for _, span := range actualSpans {
+		if span.OperationName == name && containsTags(span, expectedTags) {
+			return
+		}
+	}
+	require.FailNow(t, "can not find span matching params",
+		"expected span with name `%v` and with "+
+			"tags %v to be present but it was not. actual spans: %+v",
+		name, expectedTags, extractNameWithTags(actualSpans))
+}
+
+type spanWithTags struct {
+	name string
+	tags map[string]interface{}
+}
+
+func extractNameWithTags(actualSpans []*mocktracer.MockSpan) []spanWithTags {
+	result := make([]spanWithTags, len(actualSpans))
+	for i, span := range actualSpans {
+		result[i] = spanWithTags{span.OperationName, span.Tags()}
+	}
+	return result
+}
+
+func containsTags(span *mocktracer.MockSpan,
+	expectedTags map[string]string) bool {
+	for k, expectedVal := range expectedTags {
+		if span.Tag(k) == expectedVal {
+			return true
+		}
+	}
+	return false
+}

pkg/util/spanlogger/spanlogger.go

@@ -8,12 +8,17 @@ import (
 	opentracing "github.com/opentracing/opentracing-go"
 	"github.com/opentracing/opentracing-go/ext"
 	otlog "github.com/opentracing/opentracing-go/log"
+	"github.com/weaveworks/common/user"
 
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
 )
 
 type loggerCtxMarker struct{}
 
+const (
+	TenantIdTagName = "tenant_id"
+)
+
 var (
 	loggerCtxKey = &loggerCtxMarker{}
 )
@@ -34,6 +39,9 @@ func New(ctx context.Context, method string, kvps ...interface{}) (*SpanLogger,
 // retrieved with FromContext or FromContextWithFallback.
 func NewWithLogger(ctx context.Context, l log.Logger, method string, kvps ...interface{}) (*SpanLogger, context.Context) {
 	span, ctx := opentracing.StartSpanFromContext(ctx, method)
+	if orgId, err := user.ExtractOrgID(ctx); err == nil {
+		span.SetTag(TenantIdTagName, orgId)
+	}
 	logger := &SpanLogger{
 		Logger: log.With(util_log.WithContext(ctx, l), "method", method),
 		Span:   span,

pkg/util/spanlogger/spanlogger_test.go

@@ -5,8 +5,11 @@ import (
 	"testing"
 
 	"github.com/go-kit/kit/log"
+	"github.com/opentracing/opentracing-go"
+	"github.com/opentracing/opentracing-go/mocktracer"
 	"github.com/pkg/errors"
 	"github.com/stretchr/testify/require"
+	"github.com/weaveworks/common/user"
 )
 
 func TestSpanLogger_Log(t *testing.T) {
@@ -44,6 +47,27 @@ func TestSpanLogger_CustomLogger(t *testing.T) {
 	require.Equal(t, expect, logged)
 }
 
+func TestSpanCreatedWithTenantTag(t *testing.T) {
+	mockSpan := createSpan(user.InjectOrgID(context.Background(), "team-a"))
+
+	require.Equal(t, "team-a", mockSpan.Tag(TenantIdTagName))
+}
+
+func TestSpanCreatedWithoutTenantTag(t *testing.T) {
+	mockSpan := createSpan(context.Background())
+
+	_, exist := mockSpan.Tags()[TenantIdTagName]
+	require.False(t, exist)
+}
+
+func createSpan(ctx context.Context) *mocktracer.MockSpan {
+	mockTracer := mocktracer.New()
+	opentracing.SetGlobalTracer(mockTracer)
+
+	logger, _ := New(ctx, "name")
+	return logger.Span.(*mocktracer.MockSpan)
+}
+
 type funcLogger func(keyvals ...interface{}) error
 
 func (f funcLogger) Log(keyvals ...interface{}) error {