distributor: prevent panics when concurrently calling to forwarder's… (#1422)

* distributor: prevent panics when concurrently calling  to forwarder's queueManager

* Update PR number

* Pin k6 version to v0.37.0

CHANGELOG.md

@@ -1,5 +1,9 @@
 ## main / unreleased
 
+* [BUGFIX] metrics-generator: don't inject X-Scope-OrgID header for single-tenant setups [1417](https://github.com/grafana/tempo/pull/1417) (@kvrhdn)
+* [BUGFIX] compactor: populate `compaction_objects_combined_total` and `tempo_discarded_spans_total{reason="trace_too_large_to_compact"}` metrics again [1420](https://github.com/grafana/tempo/pull/1420) (@mdisibio)
+* [BUGFIX] distributor: prevent panics when concurrently calling `shutdown` to forwarder's queueManager [1422](https://github.com/grafana/tempo/pull/1422) (@mapno)
+
 ## v1.4.0 / 2022-04-28
 
 * [CHANGE] Vulture now exercises search at any point during the block retention to test full backend search.

integration/bench/load_test.go

@@ -13,7 +13,7 @@ import (
 )
 
 const (
-	k6Image = "loadimpact/k6:latest"
+	k6Image = "loadimpact/k6:0.37.0"
 )
 
 func TestAllInOne(t *testing.T) {

modules/distributor/forwarder.go

@@ -278,17 +278,15 @@ func (m *queueManager) forwardRequest(ctx context.Context, req *request) {
 }
 
 func (m *queueManager) shutdown() error {
-	// Already being shutdown
-	if m.readOnly.Load() {
-		return nil
+	// Call to stopWorkers only once
+	if m.readOnly.CAS(false, true) {
+		ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
+		defer cancel()
+
+		return m.stopWorkers(ctx)
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
-	defer cancel()
-
-	m.readOnly.Store(true)
-
-	return m.stopWorkers(ctx)
+	return nil
 }
 
 func (m *queueManager) stopWorkers(ctx context.Context) error {

modules/generator/storage/config_util.go

@@ -7,6 +7,8 @@ import (
 	"github.com/go-kit/log/level"
 	prometheus_config "github.com/prometheus/prometheus/config"
 	"github.com/weaveworks/common/user"
+
+	"github.com/grafana/tempo/pkg/util"
 )
 
 // generateTenantRemoteWriteConfigs creates a copy of the remote write configurations with the
@@ -31,7 +33,9 @@ func generateTenantRemoteWriteConfigs(originalCfgs []prometheus_config.RemoteWri
 		}
 
 		// inject the X-Scope-OrgId header for multi-tenant metrics backends
-		cloneCfg.Headers[user.OrgIDHeaderName] = tenant
+		if tenant != util.FakeTenantID {
+			cloneCfg.Headers[user.OrgIDHeaderName] = tenant
+		}
 
 		cloneCfgs = append(cloneCfgs, cloneCfg)
 	}

modules/generator/storage/config_util_test.go

@@ -9,6 +9,8 @@ import (
 	prometheus_common_config "github.com/prometheus/common/config"
 	prometheus_config "github.com/prometheus/prometheus/config"
 	"github.com/stretchr/testify/assert"
+
+	"github.com/grafana/tempo/pkg/util"
 )
 
 func Test_generateTenantRemoteWriteConfigs(t *testing.T) {
@@ -39,6 +41,23 @@ func Test_generateTenantRemoteWriteConfigs(t *testing.T) {
 	assert.Equal(t, map[string]string{"foo": "bar", "X-Scope-OrgID": "my-tenant"}, result[1].Headers)
 }
 
+func Test_generateTenantRemoteWriteConfigs_singleTenant(t *testing.T) {
+	logger := log.NewLogfmtLogger(log.NewSyncWriter(os.Stdout))
+
+	original := []prometheus_config.RemoteWriteConfig{
+		{
+			URL:     &prometheus_common_config.URL{URL: urlMustParse("http://prometheus-1/api/prom/push")},
+			Headers: map[string]string{},
+		},
+	}
+
+	result := generateTenantRemoteWriteConfigs(original, util.FakeTenantID, logger)
+
+	assert.Equal(t, original[0].URL, result[0].URL)
+	// X-Scope-OrgID has not been injected
+	assert.Empty(t, result[0].Headers)
+}
+
 func Test_copyMap(t *testing.T) {
 	original := map[string]string{
 		"k1": "v1",

tempodb/compactor.go

@@ -127,11 +127,14 @@ func (rw *readerWriter) compact(blockMetas []*backend.BlockMeta, tenantID string
 		return err
 	}
 
-	/*combiner := instrumentedObjectCombiner{
+	compactionLevel := compactionLevelForBlocks(blockMetas)
+	compactionLevelLabel := strconv.Itoa(int(compactionLevel))
+
+	combiner := instrumentedObjectCombiner{
 		tenant:               tenantID,
 		inner:                rw.compactorSharder,
 		compactionLevelLabel: compactionLevelLabel,
-	}*/
+	}
 
 	compactor := enc.NewCompactor()
 	opts := common.DefaultCompactionOptions()
@@ -139,6 +142,7 @@ func (rw *readerWriter) compact(blockMetas []*backend.BlockMeta, tenantID string
 	opts.ChunkSizeBytes = rw.compactorCfg.ChunkSizeBytes
 	opts.FlushSizeBytes = rw.compactorCfg.FlushSizeBytes
 	opts.OutputBlocks = outputBlocks
+	opts.Combiner = combiner
 	newCompactedBlocks, err := compactor.Compact(ctx, rw.logger, rw.r, rw.getWriterForBlock, blockMetas, opts)
 	if err != nil {
 		return err
@@ -147,8 +151,7 @@ func (rw *readerWriter) compact(blockMetas []*backend.BlockMeta, tenantID string
 	// mark old blocks compacted so they don't show up in polling
 	markCompacted(rw, tenantID, blockMetas, newCompactedBlocks)
 
-	compactionLabel := strconv.Itoa(int(newCompactedBlocks[0].CompactionLevel - 1))
-	metrics.MetricCompactionBlocks.WithLabelValues(compactionLabel).Add(float64(len(blockMetas)))
+	metrics.MetricCompactionBlocks.WithLabelValues(compactionLevelLabel).Add(float64(len(blockMetas)))
 
 	return nil
 }
@@ -192,7 +195,19 @@ func measureOutstandingBlocks(tenantID string, blockSelector CompactionBlockSele
 	metrics.MetricCompactionOutstandingBlocks.WithLabelValues(tenantID).Set(float64(totalOutstandingBlocks))
 }
 
-/*type instrumentedObjectCombiner struct {
+func compactionLevelForBlocks(blockMetas []*backend.BlockMeta) uint8 {
+	level := uint8(0)
+
+	for _, m := range blockMetas {
+		if m.CompactionLevel > level {
+			level = m.CompactionLevel
+		}
+	}
+
+	return level
+}
+
+type instrumentedObjectCombiner struct {
 	tenant               string
 	compactionLevelLabel string
 	inner                CompactorSharder
@@ -202,7 +217,7 @@ func measureOutstandingBlocks(tenantID string, blockSelector CompactionBlockSele
 func (i instrumentedObjectCombiner) Combine(dataEncoding string, objs ...[]byte) ([]byte, bool, error) {
 	b, wasCombined, err := i.inner.Combine(dataEncoding, i.tenant, objs...)
 	if wasCombined {
-		metricCompactionObjectsCombined.WithLabelValues(i.compactionLevelLabel).Inc()
+		metrics.MetricCompactionObjectsCombined.WithLabelValues(i.compactionLevelLabel).Inc()
 	}
 	return b, wasCombined, err
-}*/
+}

tempodb/compactor_test.go

@@ -30,6 +30,7 @@ import (
 )
 
 type mockSharder struct {
+	combinerCallCount int
 }
 
 func (m *mockSharder) Owns(hash string) bool {
@@ -37,6 +38,7 @@ func (m *mockSharder) Owns(hash string) bool {
 }
 
 func (m *mockSharder) Combine(dataEncoding string, tenantID string, objs ...[]byte) ([]byte, bool, error) {
+	m.combinerCallCount++
 	return model.StaticCombiner.Combine(dataEncoding, objs...)
 }
 
@@ -231,6 +233,7 @@ func TestSameIDCompaction(t *testing.T) {
 	// make a bunch of sharded requests
 	allReqs := make([][][]byte, 0, recordCount)
 	allIds := make([][]byte, 0, recordCount)
+	sharded := 0
 	for i := 0; i < recordCount; i++ {
 		id := test.ValidTraceID(nil)
 
@@ -246,6 +249,9 @@ func TestSameIDCompaction(t *testing.T) {
 			reqs = append(reqs, buff2)
 		}
 
+		if requestShards > 1 {
+			sharded++
+		}
 		allReqs = append(allReqs, reqs)
 		allIds = append(allIds, id)
 	}
@@ -281,6 +287,9 @@ func TestSameIDCompaction(t *testing.T) {
 	blocks, _ = blockSelector.BlocksToCompact()
 	assert.Len(t, blocks, blockCount)
 
+	combinedStart, err := test.GetCounterVecValue(metrics.MetricCompactionObjectsCombined, "0")
+	require.NoError(t, err)
+
 	err = rw.compact(blocks, testTenantID)
 	require.NoError(t, err)
 
@@ -309,9 +318,12 @@ func TestSameIDCompaction(t *testing.T) {
 
 		expectedBytes, _, err := model.StaticCombiner.Combine(v1.Encoding, allReqs[i]...)
 		require.NoError(t, err)
-
 		require.Equal(t, expectedBytes, b2)
 	}
+
+	combinedEnd, err := test.GetCounterVecValue(metrics.MetricCompactionObjectsCombined, "0")
+	require.NoError(t, err)
+	require.Equal(t, float64(sharded), combinedEnd-combinedStart)
 }
 
 func TestCompactionUpdatesBlocklist(t *testing.T) {
@@ -419,7 +431,7 @@ func TestCompactionMetrics(t *testing.T) {
 
 	// Cut x blocks with y records each
 	blockCount := 5
-	recordCount := 1
+	recordCount := 10
 	cutTestBlocks(t, w, testTenantID, blockCount, recordCount)
 
 	rw := r.(*readerWriter)
@@ -453,6 +465,58 @@ func TestCompactionMetrics(t *testing.T) {
 	assert.Greater(t, bytesEnd, bytesStart) // calculating the exact bytes requires knowledge of the bytes as written in the blocks.  just make sure it goes up
 }
 
+func TestCompactionUsesCombiner(t *testing.T) {
+	tempDir := t.TempDir()
+
+	r, w, c, err := New(&Config{
+		Backend: "local",
+		Pool: &pool.Config{
+			MaxWorkers: 10,
+			QueueDepth: 100,
+		},
+		Local: &local.Config{
+			Path: path.Join(tempDir, "traces"),
+		},
+		Block: &common.BlockConfig{
+			IndexDownsampleBytes: 11,
+			BloomFP:              .01,
+			BloomShardSizeBytes:  100_000,
+			Encoding:             backend.EncNone,
+			IndexPageSizeBytes:   1000,
+		},
+		WAL: &wal.Config{
+			Filepath: path.Join(tempDir, "wal"),
+		},
+		BlocklistPoll: 0,
+	}, log.NewNopLogger())
+	assert.NoError(t, err)
+
+	sharder := &mockSharder{}
+
+	c.EnableCompaction(&CompactorConfig{
+		ChunkSizeBytes:          10,
+		MaxCompactionRange:      24 * time.Hour,
+		BlockRetention:          0,
+		CompactedBlockRetention: 0,
+	}, sharder, &mockOverrides{})
+
+	r.EnablePolling(&mockJobSharder{})
+
+	// Cut x blocks with y records each
+	blockCount := 5
+	recordCount := 7
+	cutTestBlocks(t, w, testTenantID, blockCount, recordCount)
+
+	rw := r.(*readerWriter)
+	rw.pollBlocklist()
+
+	// compact everything
+	err = rw.compact(rw.blocklist.Metas(testTenantID), testTenantID)
+	assert.NoError(t, err)
+
+	require.Equal(t, blockCount*recordCount, sharder.combinerCallCount)
+}
+
 func TestCompactionIteratesThroughTenants(t *testing.T) {
 	tempDir := t.TempDir()
 

tempodb/encoding/common/interfaces.go

@@ -5,6 +5,7 @@ import (
 	"time"
 
 	"github.com/go-kit/log"
+	"github.com/grafana/tempo/pkg/model"
 	"github.com/grafana/tempo/pkg/tempopb"
 	"github.com/grafana/tempo/tempodb/backend"
 )
@@ -41,6 +42,7 @@ type CompactionOptions struct {
 	PrefetchTraceCount int // How many traces to prefetch async.
 	OutputBlocks       uint8
 	BlockConfig        BlockConfig
+	Combiner           model.ObjectCombiner
 }
 
 func DefaultCompactionOptions() CompactionOptions {

tempodb/encoding/v2/compactor.go

@@ -64,12 +64,14 @@ func (*Compactor) Compact(ctx context.Context, l log.Logger, r backend.Reader, w
 		iters = append(iters, iter)
 	}
 
-	//compactionLevelLabel := strconv.Itoa(int(compactionLevel))
 	nextCompactionLevel := compactionLevel + 1
 
 	recordsPerBlock := (totalRecords / int(opts.OutputBlocks))
 
-	combiner := model.StaticCombiner
+	combiner := opts.Combiner
+	if combiner == nil {
+		combiner = model.StaticCombiner
+	}
 
 	var currentBlock *StreamingBlock
 	var tracker backend.AppendTracker