diff --git a/storage/remote/queue_manager.go b/storage/remote/queue_manager.go index 98ad7d72860..535886ed1e4 100644 --- a/storage/remote/queue_manager.go +++ b/storage/remote/queue_manager.go @@ -36,6 +36,7 @@ import ( "go.uber.org/atomic" "github.com/prometheus/prometheus/config" + "github.com/prometheus/prometheus/model/exemplar" "github.com/prometheus/prometheus/model/histogram" "github.com/prometheus/prometheus/model/labels" "github.com/prometheus/prometheus/model/metadata" @@ -96,9 +97,10 @@ type queueManagerMetrics struct { maxNumShards prometheus.Gauge minNumShards prometheus.Gauge desiredNumShards prometheus.Gauge - sentBytesTotal prometheus.Counter - metadataBytesTotal prometheus.Counter - maxSamplesPerSend prometheus.Gauge + sentBytesTotal prometheus.Counter + metadataBytesTotal prometheus.Counter + maxSamplesPerSend prometheus.Gauge + unmatchedExemplarsDroppedTotal prometheus.Counter } func newQueueManagerMetrics(r prometheus.Registerer, rn, e string) *queueManagerMetrics { @@ -328,6 +330,13 @@ func newQueueManagerMetrics(r prometheus.Registerer, rn, e string) *queueManager Help: "The maximum number of samples to be sent, in a single request, to the remote storage. Note that, when sending of exemplars over remote write is enabled, exemplars count towards this limit.", ConstLabels: constLabels, }) + m.unmatchedExemplarsDroppedTotal = prometheus.NewCounter(prometheus.CounterOpts{ + Namespace: namespace, + Subsystem: subsystem, + Name: "unmatched_exemplars_dropped_total", + Help: "Total number of exemplars dropped due to not matching any sample within the coalescing window.", + ConstLabels: constLabels, + }) return m } @@ -365,6 +374,7 @@ func (m *queueManagerMetrics) register() { m.sentBytesTotal, m.metadataBytesTotal, m.maxSamplesPerSend, + m.unmatchedExemplarsDroppedTotal, ) } } @@ -401,6 +411,7 @@ func (m *queueManagerMetrics) unregister() { m.reg.Unregister(m.sentBytesTotal) m.reg.Unregister(m.metadataBytesTotal) m.reg.Unregister(m.maxSamplesPerSend) + m.reg.Unregister(m.unmatchedExemplarsDroppedTotal) } } @@ -1292,8 +1303,25 @@ func (s *shards) start(n int) { s.qm.metrics.numShards.Set(float64(n)) newQueues := make([]*queue, n) + onDrop := func(ex exemplar.Exemplar) { + s.enqueuedExemplars.Sub(1) + if s.qm.metrics != nil { + if s.qm.metrics.unmatchedExemplarsDroppedTotal != nil { + s.qm.metrics.unmatchedExemplarsDroppedTotal.Inc() + } + if s.qm.metrics.droppedExemplarsTotal != nil { + s.qm.metrics.droppedExemplarsTotal.WithLabelValues("unmatched").Inc() + } + if s.qm.metrics.pendingExemplars != nil { + s.qm.metrics.pendingExemplars.Dec() + } + } + if s.qm.dataDropped != nil { + s.qm.dataDropped.incr(1) + } + } for i := range n { - newQueues[i] = newQueue(s.qm.cfg.MaxSamplesPerSend, s.qm.cfg.Capacity) + newQueues[i] = newQueue(s.qm.cfg.MaxSamplesPerSend, s.qm.cfg.Capacity, s.qm.protoMsg, onDrop) } s.queues = newQueues @@ -1367,6 +1395,7 @@ func (s *shards) enqueue(ref chunks.HeadSeriesRef, data timeSeries) bool { case <-s.softShutdown: return false default: + data.seriesRef = ref appended := s.queues[shard].Append(data) if !appended { return false @@ -1395,6 +1424,9 @@ type queue struct { batch []timeSeries batchQueue chan []timeSeries + coalescer *shardCoalescer + protoMsg remoteapi.WriteMessageType + // Since we know there are a limited number of batches out, using a stack // is easy and safe so a sync.Pool is not necessary. // poolMtx covers adding and removing batches from the batchPool. @@ -1403,6 +1435,7 @@ type queue struct { } type timeSeries struct { + seriesRef chunks.HeadSeriesRef seriesLabels labels.Labels value float64 histogram *histogram.Histogram @@ -1410,6 +1443,7 @@ type timeSeries struct { metadata *metadata.Metadata startTimestamp, timestamp int64 exemplarLabels labels.Labels + exemplars []exemplar.Exemplar // The type of series: sample, exemplar, or histogram. sType seriesType } @@ -1424,7 +1458,7 @@ const ( tMetadata ) -func newQueue(batchSize, capacity int) *queue { +func newQueue(batchSize, capacity int, protoMsg remoteapi.WriteMessageType, onDrop func(exemplar.Exemplar)) *queue { batches := capacity / batchSize // Always create an unbuffered channel even if capacity is configured to be // less than max_samples_per_send. @@ -1437,6 +1471,8 @@ func newQueue(batchSize, capacity int) *queue { // batchPool should have capacity for everything in the channel + 1 for // the batch being processed. batchPool: make([][]timeSeries, 0, batches+1), + coalescer: newShardCoalescer(defaultRingBufferSize, onDrop), + protoMsg: protoMsg, } } @@ -1445,6 +1481,32 @@ func newQueue(batchSize, capacity int) *queue { func (q *queue) Append(datum timeSeries) bool { q.batchMtx.Lock() defer q.batchMtx.Unlock() + + if q.protoMsg == remoteapi.WriteV2MessageType && q.coalescer != nil { + if datum.sType == tExemplar { + ex := exemplar.Exemplar{ + Labels: datum.exemplarLabels, + Value: datum.value, + Ts: datum.timestamp, + HasTs: true, + } + // Try attaching to an existing un-flushed sample/histogram in the current batch. + if q.coalescer.TryAttachToBatch(q.batch, datum.seriesRef, ex) { + return true + } + // Buffer in ring buffer. + q.coalescer.AddPendingExemplar(datum.seriesRef, ex) + return true + } + + if datum.sType == tSample || datum.sType == tHistogram || datum.sType == tFloatHistogram { + // Check if there are any matching pending exemplars in the ring buffer. + if matched := q.coalescer.TryAttachMatchingExemplars(datum.seriesRef, datum.timestamp); len(matched) > 0 { + datum.exemplars = append(datum.exemplars, matched...) + } + } + } + // TODO(cstyan): Check if metadata now means we've reduced the total # of samples // we can batch together here, and if so find a way to not include metadata // in the batch size calculation. @@ -1505,6 +1567,9 @@ loop: q.batchMtx.Lock() defer q.batchMtx.Unlock() + if q.coalescer != nil { + q.coalescer.FlushAndClear() + } q.batch = nil close(q.batchQueue) } @@ -1604,7 +1669,7 @@ func (s *shards) runShard(ctx context.Context, shardID int, queue *queue) { _ = s.sendSamples(ctx, pendingData[:n], nPendingSamples, nPendingExemplars, nPendingHistograms, pBuf, encBuf, compr) case remoteapi.WriteV2MessageType: nPendingSamples, nPendingExemplars, nPendingHistograms, nPendingMetadata, nUnexpectedMetadata := populateV2TimeSeries(&symbolTable, batch, pendingDataV2, s.qm.sendExemplars, s.qm.sendNativeHistograms, s.qm.enableTypeAndUnitLabels) - n := nPendingSamples + nPendingExemplars + nPendingHistograms + n := len(batch) if nUnexpectedMetadata > 0 { s.qm.logger.Warn("unexpected metadata sType in populateV2TimeSeries", "count", nUnexpectedMetadata) } @@ -1678,6 +1743,15 @@ func populateTimeSeries(batch []timeSeries, pendingData []prompb.TimeSeries, sen Value: d.value, Timestamp: d.timestamp, }) + if sendExemplars && len(d.exemplars) > 0 { + for _, ex := range d.exemplars { + pendingData[nPending].Exemplars = append(pendingData[nPending].Exemplars, prompb.Exemplar{ + Labels: prompb.FromLabels(ex.Labels, nil), + Value: ex.Value, + Timestamp: ex.Ts, + }) + } + } nPendingSamples++ case tExemplar: pendingData[nPending].Exemplars = append(pendingData[nPending].Exemplars, prompb.Exemplar{ @@ -1688,9 +1762,27 @@ func populateTimeSeries(batch []timeSeries, pendingData []prompb.TimeSeries, sen nPendingExemplars++ case tHistogram: pendingData[nPending].Histograms = append(pendingData[nPending].Histograms, prompb.FromIntHistogram(d.timestamp, d.histogram)) + if sendExemplars && len(d.exemplars) > 0 { + for _, ex := range d.exemplars { + pendingData[nPending].Exemplars = append(pendingData[nPending].Exemplars, prompb.Exemplar{ + Labels: prompb.FromLabels(ex.Labels, nil), + Value: ex.Value, + Timestamp: ex.Ts, + }) + } + } nPendingHistograms++ case tFloatHistogram: pendingData[nPending].Histograms = append(pendingData[nPending].Histograms, prompb.FromFloatHistogram(d.timestamp, d.floatHistogram)) + if sendExemplars && len(d.exemplars) > 0 { + for _, ex := range d.exemplars { + pendingData[nPending].Exemplars = append(pendingData[nPending].Exemplars, prompb.Exemplar{ + Labels: prompb.FromLabels(ex.Labels, nil), + Value: ex.Value, + Timestamp: ex.Ts, + }) + } + } nPendingHistograms++ } } @@ -2017,19 +2109,51 @@ func populateV2TimeSeries(symbolTable *writev2.SymbolsTable, batch []timeSeries, StartTimestamp: d.startTimestamp, }) nPendingSamples++ + if sendExemplars && len(d.exemplars) > 0 { + for _, ex := range d.exemplars { + pendingData[nPending].Exemplars = append(pendingData[nPending].Exemplars, writev2.Exemplar{ + LabelsRefs: symbolTable.SymbolizeLabels(ex.Labels, nil), + Value: ex.Value, + Timestamp: ex.Ts, + }) + nPendingExemplars++ + } + } case tExemplar: - pendingData[nPending].Exemplars = append(pendingData[nPending].Exemplars, writev2.Exemplar{ - LabelsRefs: symbolTable.SymbolizeLabels(d.exemplarLabels, nil), // TODO: optimize, reuse slice - Value: d.value, - Timestamp: d.timestamp, - }) - nPendingExemplars++ + if sendExemplars { + pendingData[nPending].Exemplars = append(pendingData[nPending].Exemplars, writev2.Exemplar{ + LabelsRefs: symbolTable.SymbolizeLabels(d.exemplarLabels, nil), // TODO: optimize, reuse slice + Value: d.value, + Timestamp: d.timestamp, + }) + nPendingExemplars++ + } case tHistogram: pendingData[nPending].Histograms = append(pendingData[nPending].Histograms, writev2.FromIntHistogram(d.startTimestamp, d.timestamp, d.histogram)) nPendingHistograms++ + if sendExemplars && len(d.exemplars) > 0 { + for _, ex := range d.exemplars { + pendingData[nPending].Exemplars = append(pendingData[nPending].Exemplars, writev2.Exemplar{ + LabelsRefs: symbolTable.SymbolizeLabels(ex.Labels, nil), + Value: ex.Value, + Timestamp: ex.Ts, + }) + nPendingExemplars++ + } + } case tFloatHistogram: pendingData[nPending].Histograms = append(pendingData[nPending].Histograms, writev2.FromFloatHistogram(d.startTimestamp, d.timestamp, d.floatHistogram)) nPendingHistograms++ + if sendExemplars && len(d.exemplars) > 0 { + for _, ex := range d.exemplars { + pendingData[nPending].Exemplars = append(pendingData[nPending].Exemplars, writev2.Exemplar{ + LabelsRefs: symbolTable.SymbolizeLabels(ex.Labels, nil), + Value: ex.Value, + Timestamp: ex.Ts, + }) + nPendingExemplars++ + } + } case tMetadata: nUnexpectedMetadata++ } diff --git a/storage/remote/queue_manager_coalescing_benchmark_test.go b/storage/remote/queue_manager_coalescing_benchmark_test.go new file mode 100644 index 00000000000..663d3ed3ed4 --- /dev/null +++ b/storage/remote/queue_manager_coalescing_benchmark_test.go @@ -0,0 +1,309 @@ +// Copyright The Prometheus Authors +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package remote + +import ( + "context" + "fmt" + "runtime" + "testing" + "time" + + remoteapi "github.com/prometheus/client_golang/exp/api/remote" + "github.com/prometheus/common/model" + "github.com/stretchr/testify/require" + + "github.com/prometheus/prometheus/config" + "github.com/prometheus/prometheus/model/labels" + "github.com/prometheus/prometheus/tsdb/chunks" + "github.com/prometheus/prometheus/tsdb/record" +) + +type noopWriteClient struct{} + +func (c *noopWriteClient) Store(_ context.Context, _ []byte, _ int) (WriteResponseStats, error) { + return WriteResponseStats{}, nil +} +func (c *noopWriteClient) Name() string { return "noop" } +func (c *noopWriteClient) Endpoint() string { return "http://localhost/noop" } + +// BenchmarkQueueManager_PRW2_Baseline measures standard PRW 2.0 sample ingestion without exemplars. +func BenchmarkQueueManager_PRW2_Baseline(b *testing.B) { + dir := b.TempDir() + s := NewStorage(nil, nil, nil, dir, defaultFlushDeadline, nil, false) + defer s.Close() + + queueConfig := config.DefaultQueueConfig + queueConfig.BatchSendDeadline = model.Duration(100 * time.Millisecond) + queueConfig.MaxShards = 4 + queueConfig.MinShards = 4 + queueConfig.Capacity = 10000 + queueConfig.MaxSamplesPerSend = 1000 + + writeConfig := baseRemoteWriteConfig("http://test-storage.com") + writeConfig.QueueConfig = queueConfig + writeConfig.SendExemplars = false + writeConfig.ProtobufMessage = remoteapi.WriteV2MessageType + + conf := &config.Config{ + GlobalConfig: config.DefaultGlobalConfig, + RemoteWriteConfigs: []*config.RemoteWriteConfig{ + writeConfig, + }, + } + require.NoError(b, s.ApplyConfig(conf)) + + hash, err := toHash(writeConfig) + require.NoError(b, err) + qm := s.rws.queues[hash] + qm.SetClient(&noopWriteClient{}) + + numSeries := 1000 + series := make([]record.RefSeries, numSeries) + for i := 0; i < numSeries; i++ { + series[i] = record.RefSeries{ + Ref: chunks.HeadSeriesRef(i + 1), + Labels: labels.FromStrings("__name__", fmt.Sprintf("metric_%d", i), "job", "benchmark"), + } + } + qm.StoreSeries(series, 0) + + samples := make([]record.RefSample, 100) + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + ts := int64(i * 1000) + for j := range samples { + samples[j] = record.RefSample{ + Ref: chunks.HeadSeriesRef((j % numSeries) + 1), + T: ts, + V: float64(i), + } + } + qm.Append(samples) + } +} + +// BenchmarkQueueManager_PRW2_Coalescing measures PRW 2.0 with shard coalescing enabled and samples + exemplars. +func BenchmarkQueueManager_PRW2_Coalescing(b *testing.B) { + dir := b.TempDir() + s := NewStorage(nil, nil, nil, dir, defaultFlushDeadline, nil, false) + defer s.Close() + + queueConfig := config.DefaultQueueConfig + queueConfig.BatchSendDeadline = model.Duration(100 * time.Millisecond) + queueConfig.MaxShards = 4 + queueConfig.MinShards = 4 + queueConfig.Capacity = 10000 + queueConfig.MaxSamplesPerSend = 1000 + + writeConfig := baseRemoteWriteConfig("http://test-storage.com") + writeConfig.QueueConfig = queueConfig + writeConfig.SendExemplars = true + writeConfig.ProtobufMessage = remoteapi.WriteV2MessageType + + conf := &config.Config{ + GlobalConfig: config.DefaultGlobalConfig, + RemoteWriteConfigs: []*config.RemoteWriteConfig{ + writeConfig, + }, + } + require.NoError(b, s.ApplyConfig(conf)) + + hash, err := toHash(writeConfig) + require.NoError(b, err) + qm := s.rws.queues[hash] + qm.SetClient(&noopWriteClient{}) + + numSeries := 1000 + series := make([]record.RefSeries, numSeries) + for i := 0; i < numSeries; i++ { + series[i] = record.RefSeries{ + Ref: chunks.HeadSeriesRef(i + 1), + Labels: labels.FromStrings("__name__", fmt.Sprintf("metric_%d", i), "job", "benchmark"), + } + } + qm.StoreSeries(series, 0) + + samples := make([]record.RefSample, 100) + exemplars := make([]record.RefExemplar, 100) + exLabels := labels.FromStrings("trace_id", "4bf92f3577b34da6a3ce929d0e0e4736") + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + ts := int64(i * 1000) + for j := range samples { + ref := chunks.HeadSeriesRef((j % numSeries) + 1) + samples[j] = record.RefSample{Ref: ref, T: ts, V: float64(i)} + exemplars[j] = record.RefExemplar{Ref: ref, T: ts + 5, V: float64(i), Labels: exLabels} + } + qm.Append(samples) + qm.AppendExemplars(exemplars) + } +} + +// BenchmarkQueueManager_PRW2_SeriesChurn_100k tests memory stability and allocation efficiency under 100k churned series. +func BenchmarkQueueManager_PRW2_SeriesChurn_100k(b *testing.B) { + dir := b.TempDir() + s := NewStorage(nil, nil, nil, dir, defaultFlushDeadline, nil, false) + defer s.Close() + + queueConfig := config.DefaultQueueConfig + queueConfig.BatchSendDeadline = model.Duration(50 * time.Millisecond) + queueConfig.MaxShards = 4 + queueConfig.MinShards = 4 + queueConfig.Capacity = 10000 + queueConfig.MaxSamplesPerSend = 1000 + + writeConfig := baseRemoteWriteConfig("http://test-storage.com") + writeConfig.QueueConfig = queueConfig + writeConfig.SendExemplars = true + writeConfig.ProtobufMessage = remoteapi.WriteV2MessageType + + conf := &config.Config{ + GlobalConfig: config.DefaultGlobalConfig, + RemoteWriteConfigs: []*config.RemoteWriteConfig{ + writeConfig, + }, + } + require.NoError(b, s.ApplyConfig(conf)) + + hash, err := toHash(writeConfig) + require.NoError(b, err) + qm := s.rws.queues[hash] + qm.SetClient(&noopWriteClient{}) + + const totalChurnSeries = 100000 + seriesBatch := make([]record.RefSeries, 1000) + samplesBatch := make([]record.RefSample, 1000) + exemplarsBatch := make([]record.RefExemplar, 1000) + exLabels := labels.FromStrings("trace_id", "trace-churn-benchmark") + + runtime.GC() + var memBefore runtime.MemStats + runtime.ReadMemStats(&memBefore) + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + // Ingest 100,000 unique series in chunks of 1,000 + for offset := 0; offset < totalChurnSeries; offset += 1000 { + ts := time.Now().UnixMilli() + for k := 0; k < 1000; k++ { + ref := chunks.HeadSeriesRef(offset + k + 1) + seriesBatch[k] = record.RefSeries{ + Ref: ref, + Labels: labels.FromStrings("__name__", fmt.Sprintf("churn_metric_%d", offset+k), "cycle", fmt.Sprintf("c%d", i)), + } + samplesBatch[k] = record.RefSample{Ref: ref, T: ts, V: float64(k)} + exemplarsBatch[k] = record.RefExemplar{Ref: ref, T: ts + 2, V: float64(k), Labels: exLabels} + } + + qm.StoreSeries(seriesBatch, 0) + qm.Append(samplesBatch) + qm.AppendExemplars(exemplarsBatch) + } + } + + b.StopTimer() + + runtime.GC() + var memAfter runtime.MemStats + runtime.ReadMemStats(&memAfter) + + heapGrowthMB := float64(memAfter.HeapAlloc-memBefore.HeapAlloc) / (1024 * 1024) + b.Logf("HeapAlloc after 100k churned series: %.2f MB (growth: %.2f MB)", float64(memAfter.HeapAlloc)/(1024*1024), heapGrowthMB) +} + +func TestQueueManager_100kSeriesChurn_HeapStability(t *testing.T) { + if testing.Short() { + t.Skip("skipping 100k churn test in short mode") + } + + dir := t.TempDir() + s := NewStorage(nil, nil, nil, dir, defaultFlushDeadline, nil, false) + defer s.Close() + + queueConfig := config.DefaultQueueConfig + queueConfig.BatchSendDeadline = model.Duration(20 * time.Millisecond) + queueConfig.MaxShards = 4 + queueConfig.MinShards = 4 + queueConfig.Capacity = 5000 + queueConfig.MaxSamplesPerSend = 500 + + writeConfig := baseRemoteWriteConfig("http://test-storage.com") + writeConfig.QueueConfig = queueConfig + writeConfig.SendExemplars = true + writeConfig.ProtobufMessage = remoteapi.WriteV2MessageType + + conf := &config.Config{ + GlobalConfig: config.DefaultGlobalConfig, + RemoteWriteConfigs: []*config.RemoteWriteConfig{ + writeConfig, + }, + } + require.NoError(t, s.ApplyConfig(conf)) + + hash, err := toHash(writeConfig) + require.NoError(t, err) + qm := s.rws.queues[hash] + qm.SetClient(&noopWriteClient{}) + + const totalSeries = 100000 + seriesChunk := make([]record.RefSeries, 1000) + samplesChunk := make([]record.RefSample, 1000) + exemplarsChunk := make([]record.RefExemplar, 1000) + exLabels := labels.FromStrings("trace_id", "trace-churn-stability") + + runtime.GC() + var memStart runtime.MemStats + runtime.ReadMemStats(&memStart) + + for chunk := 0; chunk < totalSeries/1000; chunk++ { + ts := time.Now().UnixMilli() + for i := 0; i < 1000; i++ { + ref := chunks.HeadSeriesRef(chunk*1000 + i + 1) + seriesChunk[i] = record.RefSeries{ + Ref: ref, + Labels: labels.FromStrings("__name__", fmt.Sprintf("churn_%d", ref), "pod", "test"), + } + samplesChunk[i] = record.RefSample{Ref: ref, T: ts, V: float64(i)} + exemplarsChunk[i] = record.RefExemplar{Ref: ref, T: ts + 1, V: float64(i), Labels: exLabels} + } + + qm.StoreSeries(seriesChunk, 0) + qm.Append(samplesChunk) + qm.AppendExemplars(exemplarsChunk) + } + + // Verify all queues process and drain + time.Sleep(500 * time.Millisecond) + + runtime.GC() + var memEnd runtime.MemStats + runtime.ReadMemStats(&memEnd) + + heapMB := float64(memEnd.HeapAlloc) / (1024 * 1024) + t.Logf("HeapAlloc after 100,000 distinct series: %.2f MB (GC cycles: %d)", heapMB, memEnd.NumGC-memStart.NumGC) + + // Ring buffer memory is bounded (2048 slots per shard) + for _, q := range qm.shards.queues { + require.LessOrEqual(t, q.coalescer.PendingCount(), 2048) + } +} diff --git a/storage/remote/queue_manager_coalescing_test.go b/storage/remote/queue_manager_coalescing_test.go new file mode 100644 index 00000000000..cd6fb29a86e --- /dev/null +++ b/storage/remote/queue_manager_coalescing_test.go @@ -0,0 +1,315 @@ +// Copyright The Prometheus Authors +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package remote + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + client_testutil "github.com/prometheus/client_golang/prometheus/testutil" + "github.com/prometheus/common/model" + "github.com/stretchr/testify/require" + + remoteapi "github.com/prometheus/client_golang/exp/api/remote" + "github.com/prometheus/prometheus/config" + "github.com/prometheus/prometheus/model/histogram" + "github.com/prometheus/prometheus/model/labels" + writev2 "github.com/prometheus/prometheus/prompb/io/prometheus/write/v2" + "github.com/prometheus/prometheus/tsdb/chunks" + "github.com/prometheus/prometheus/tsdb/record" + "github.com/prometheus/prometheus/util/compression" +) + +// capturingV2WriteClient captures all deserialized writev2.Request payloads. +type capturingV2WriteClient struct { + mtx sync.Mutex + requests []*writev2.Request + compr compression.Type +} + +func newCapturingV2WriteClient(compr compression.Type) *capturingV2WriteClient { + return &capturingV2WriteClient{ + compr: compr, + } +} + +func (c *capturingV2WriteClient) Store(_ context.Context, req []byte, _ int) (WriteResponseStats, error) { + decomp, err := compression.Decode(c.compr, req, nil) + if err != nil { + return WriteResponseStats{}, err + } + var v2Req writev2.Request + if err := v2Req.Unmarshal(decomp); err != nil { + return WriteResponseStats{}, err + } + + c.mtx.Lock() + c.requests = append(c.requests, &v2Req) + c.mtx.Unlock() + + var numSamples, numHistograms, numExemplars int + for _, ts := range v2Req.Timeseries { + numSamples += len(ts.Samples) + numHistograms += len(ts.Histograms) + numExemplars += len(ts.Exemplars) + } + + return WriteResponseStats{ + Samples: numSamples, + Histograms: numHistograms, + Exemplars: numExemplars, + }, nil +} + +func (c *capturingV2WriteClient) Name() string { return "capturing-v2-client" } +func (c *capturingV2WriteClient) Endpoint() string { return "http://localhost/write" } + +func (c *capturingV2WriteClient) getRequests() []*writev2.Request { + c.mtx.Lock() + defer c.mtx.Unlock() + copied := make([]*writev2.Request, len(c.requests)) + copy(copied, c.requests) + return copied +} + +func TestQueueManager_PRW2_Coalescing(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + s := NewStorage(nil, nil, nil, dir, defaultFlushDeadline, nil, false) + + queueConfig := config.DefaultQueueConfig + queueConfig.BatchSendDeadline = model.Duration(50 * time.Millisecond) + queueConfig.MaxShards = 1 + queueConfig.Capacity = 100 + queueConfig.MaxSamplesPerSend = 50 + + writeConfig := baseRemoteWriteConfig("http://test-storage.com") + writeConfig.QueueConfig = queueConfig + writeConfig.SendExemplars = true + writeConfig.SendNativeHistograms = true + writeConfig.ProtobufMessage = remoteapi.WriteV2MessageType + + conf := &config.Config{ + GlobalConfig: config.DefaultGlobalConfig, + RemoteWriteConfigs: []*config.RemoteWriteConfig{ + writeConfig, + }, + } + require.NoError(t, s.ApplyConfig(conf)) + + hash, err := toHash(writeConfig) + require.NoError(t, err) + qm := s.rws.queues[hash] + + client := newCapturingV2WriteClient(qm.compr) + qm.SetClient(client) + + // Register series 1 to 6 + series := []record.RefSeries{ + {Ref: 1, Labels: labels.FromStrings("__name__", "metric_sample_first", "instance", "localhost")}, + {Ref: 2, Labels: labels.FromStrings("__name__", "metric_exemplar_first", "instance", "localhost")}, + {Ref: 3, Labels: labels.FromStrings("__name__", "metric_histogram", "instance", "localhost")}, + {Ref: 4, Labels: labels.FromStrings("__name__", "metric_float_histogram", "instance", "localhost")}, + {Ref: 5, Labels: labels.FromStrings("__name__", "metric_unmatched_exemplar", "instance", "localhost")}, + {Ref: 6, Labels: labels.FromStrings("__name__", "metric_cross_scrape", "instance", "localhost")}, + } + qm.StoreSeries(series, 0) + + // 1. Sample-First: Enqueue sample at T=1000, then exemplar at T=1010 + qm.Append([]record.RefSample{{Ref: 1, T: 1000, V: 42.0}}) + qm.AppendExemplars([]record.RefExemplar{{Ref: 1, T: 1010, V: 42.0, Labels: labels.FromStrings("trace_id", "trace-sample-first")}}) + + // 2. Exemplar-First: Enqueue exemplar at T=2000, then sample at T=2020 + qm.AppendExemplars([]record.RefExemplar{{Ref: 2, T: 2000, V: 84.0, Labels: labels.FromStrings("trace_id", "trace-exemplar-first")}}) + qm.Append([]record.RefSample{{Ref: 2, T: 2020, V: 84.0}}) + + // 3. Int Histogram: Enqueue exemplar at T=3000, then histogram at T=3015 + h := &histogram.Histogram{Schema: 1, Count: 10, Sum: 25.0} + qm.AppendExemplars([]record.RefExemplar{{Ref: 3, T: 3000, V: 5.0, Labels: labels.FromStrings("trace_id", "trace-histogram")}}) + qm.AppendHistograms([]record.RefHistogramSample{{Ref: 3, T: 3015, H: h}}) + + // 4. Float Histogram: Enqueue exemplar at T=4000, then float histogram at T=4010 + fh := &histogram.FloatHistogram{Schema: 1, Count: 20, Sum: 50.0} + qm.AppendExemplars([]record.RefExemplar{{Ref: 4, T: 4000, V: 7.5, Labels: labels.FromStrings("trace_id", "trace-float-histogram")}}) + qm.AppendFloatHistograms([]record.RefFloatHistogramSample{{Ref: 4, T: 4010, FH: fh}}) + + // 5. Unmatched exemplar (no sample arrives for ref 5) + qm.AppendExemplars([]record.RefExemplar{{Ref: 5, T: 5000, V: 99.0, Labels: labels.FromStrings("trace_id", "trace-unmatched")}}) + + // 6. Cross-scrape rejection (>50ms delta): Sample at T=6000, Exemplar at T=6100 (100ms > 50ms) + qm.Append([]record.RefSample{{Ref: 6, T: 6000, V: 111.0}}) + qm.AppendExemplars([]record.RefExemplar{{Ref: 6, T: 6100, V: 111.0, Labels: labels.FromStrings("trace_id", "trace-cross-scrape")}}) + + // Wait for batches to be flushed by deadline + require.Eventually(t, func() bool { + reqs := client.getRequests() + var totalSeries int + for _, r := range reqs { + totalSeries += len(r.Timeseries) + } + // Expect 5 sent series: refs 1, 2, 3, 4, 6 (ref 5 has no sample so must not be sent) + return totalSeries >= 5 + }, 5*time.Second, 20*time.Millisecond) + + reqs := client.getRequests() + require.NotEmpty(t, reqs) + + var totalSamples, totalHistograms, totalExemplars int + for _, r := range reqs { + for _, ts := range r.Timeseries { + // PRW 2.0 Invariant: Every TimeSeries MUST contain at least one sample or histogram. + // ZERO empty TimeSeries with only exemplars allowed! + hasData := len(ts.Samples) > 0 || len(ts.Histograms) > 0 + require.True(t, hasData, "PRW 2.0 violation: TimeSeries must not be empty of samples/histograms") + + totalSamples += len(ts.Samples) + totalHistograms += len(ts.Histograms) + totalExemplars += len(ts.Exemplars) + + // If exemplars are present, verify they are attached to a valid sample/histogram + if len(ts.Exemplars) > 0 { + require.True(t, len(ts.Samples) > 0 || len(ts.Histograms) > 0) + } + } + } + + // 3 float samples (refs 1, 2, 6) and 2 histograms (refs 3, 4) + require.Equal(t, 3, totalSamples) + require.Equal(t, 2, totalHistograms) + // 4 matched exemplars (refs 1, 2, 3, 4) + require.Equal(t, 4, totalExemplars) + + // Close storage to flush/drain coalescers + s.Close() + + // Metrics validation: + // 4 exemplars sent, 2 dropped (ref 5 unmatched, ref 6 cross-scrape >50ms) + require.Equal(t, float64(4), client_testutil.ToFloat64(qm.metrics.exemplarsTotal)) + require.Equal(t, float64(2), client_testutil.ToFloat64(qm.metrics.unmatchedExemplarsDroppedTotal)) + require.Equal(t, float64(3), client_testutil.ToFloat64(qm.metrics.samplesTotal)) + require.Equal(t, float64(2), client_testutil.ToFloat64(qm.metrics.histogramsTotal)) +} + +func TestQueueManager_Resharding(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + s := NewStorage(nil, nil, nil, dir, defaultFlushDeadline, nil, false) + + queueConfig := config.DefaultQueueConfig + queueConfig.BatchSendDeadline = model.Duration(20 * time.Millisecond) + queueConfig.MinShards = 1 + queueConfig.MaxShards = 16 + queueConfig.Capacity = 200 + queueConfig.MaxSamplesPerSend = 50 + + writeConfig := baseRemoteWriteConfig("http://test-storage.com") + writeConfig.QueueConfig = queueConfig + writeConfig.SendExemplars = true + writeConfig.SendNativeHistograms = true + writeConfig.ProtobufMessage = remoteapi.WriteV2MessageType + + conf := &config.Config{ + GlobalConfig: config.DefaultGlobalConfig, + RemoteWriteConfigs: []*config.RemoteWriteConfig{ + writeConfig, + }, + } + require.NoError(t, s.ApplyConfig(conf)) + + hash, err := toHash(writeConfig) + require.NoError(t, err) + qm := s.rws.queues[hash] + + client := newCapturingV2WriteClient(qm.compr) + qm.SetClient(client) + + numSeries := 100 + series := make([]record.RefSeries, numSeries) + for i := 0; i < numSeries; i++ { + series[i] = record.RefSeries{ + Ref: chunks.HeadSeriesRef(i + 1), + Labels: labels.FromStrings("__name__", fmt.Sprintf("metric_%d", i), "instance", "localhost"), + } + } + qm.StoreSeries(series, 0) + + var ( + stopAppend atomic.Bool + wg sync.WaitGroup + appendedSamples atomic.Int64 + appendedExemplars atomic.Int64 + ) + + // Concurrently append paired samples and exemplars + for worker := 0; worker < 4; worker++ { + wg.Add(1) + go func(w int) { + defer wg.Done() + var seq int64 + for !stopAppend.Load() { + seq++ + ref := chunks.HeadSeriesRef((seq % int64(numSeries)) + 1) + ts := time.Now().UnixMilli() + + qm.Append([]record.RefSample{{Ref: ref, T: ts, V: float64(seq)}}) + appendedSamples.Add(1) + + qm.AppendExemplars([]record.RefExemplar{{ + Ref: ref, + T: ts, + V: float64(seq), + Labels: labels.FromStrings("trace_id", fmt.Sprintf("trace-%d-%d", w, seq)), + }}) + appendedExemplars.Add(1) + + time.Sleep(500 * time.Microsecond) + } + }(worker) + } + + // Concurrently trigger dynamic resharding back and forth + shardTargets := []int{2, 4, 1, 8, 3, 6, 2, 4} + for _, target := range shardTargets { + time.Sleep(50 * time.Millisecond) + qm.reshardChan <- target + } + + time.Sleep(200 * time.Millisecond) + stopAppend.Store(true) + wg.Wait() + + // Wait for flush + require.Eventually(t, func() bool { + return client_testutil.ToFloat64(qm.metrics.pendingSamples) == 0 + }, 10*time.Second, 50*time.Millisecond) + + s.Close() + + reqs := client.getRequests() + require.NotEmpty(t, reqs) + + // Verify all received TimeSeries are strictly valid PRW 2.0 (no empty series) + for _, r := range reqs { + for _, ts := range r.Timeseries { + hasSampleOrHistogram := len(ts.Samples) > 0 || len(ts.Histograms) > 0 + require.True(t, hasSampleOrHistogram, "PRW 2.0 invariant: must contain at least 1 sample or histogram") + } + } +} diff --git a/storage/remote/queue_manager_test.go b/storage/remote/queue_manager_test.go index 90236200ff7..dd87c690b50 100644 --- a/storage/remote/queue_manager_test.go +++ b/storage/remote/queue_manager_test.go @@ -256,6 +256,16 @@ func TestSampleDelivery(t *testing.T) { qm.StoreSeries(series, 0) qm.StoreMetadata(metadata) + if protoMsg == remoteapi.WriteV2MessageType && rc.Name == "exemplars only" { + // In PRW 2.0, standalone exemplars without matching samples/histograms cannot be sent + // and are dropped as unmatched to adhere to PRW 2.0 invariants (0 empty TimeSeries). + qm.AppendExemplars(exemplars) + require.Eventually(t, func() bool { + return client_testutil.ToFloat64(qm.metrics.pendingExemplars) == float64(len(exemplars)) + }, 2*time.Second, 10*time.Millisecond) + return + } + // Send first half of data. c.expectSamples(samples[:len(samples)/2], series) c.expectExemplars(exemplars[:len(exemplars)/2], series) @@ -1743,7 +1753,7 @@ func TestQueueManagerMetrics(t *testing.T) { func TestQueue_FlushAndShutdownDoesNotDeadlock(t *testing.T) { capacity := 100 batchSize := 10 - queue := newQueue(batchSize, capacity) + queue := newQueue(batchSize, capacity, remoteapi.WriteV1MessageType, nil) for i := 0; i < capacity+batchSize; i++ { queue.Append(timeSeries{}) } diff --git a/storage/remote/shard_coalescer.go b/storage/remote/shard_coalescer.go new file mode 100644 index 00000000000..3ab8f380a18 --- /dev/null +++ b/storage/remote/shard_coalescer.go @@ -0,0 +1,276 @@ +// Copyright The Prometheus Authors +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package remote + +import ( + "github.com/prometheus/prometheus/model/exemplar" + "github.com/prometheus/prometheus/tsdb/chunks" +) + +const ( + // defaultRingBufferSize is the fixed capacity of pending exemplars per shard coalescer. + defaultRingBufferSize = 2048 + + // maxCoalescingTimeDeltaMs is the maximum allowed time delta (50ms) between a sample + // and an exemplar for them to be considered part of the same scrape. + maxCoalescingTimeDeltaMs int64 = 50 +) + +type coalescerIndexEntry struct { + slotIdx int + generation uint64 +} + +type coalescerSlot struct { + seriesRef chunks.HeadSeriesRef + exemplar exemplar.Exemplar + generation uint64 + nextSlot int + tombstone bool +} + +// shardCoalescer manages in-memory correlation and coalescing of samples and exemplars +// within a single Remote Write QueueManager shard worker. +// It is single-threaded per shard worker goroutine (lock-free). +type shardCoalescer struct { + capacity int + slots []coalescerSlot + generations []uint64 + seriesIndex map[chunks.HeadSeriesRef]coalescerIndexEntry + head int + onDrop func(exemplar.Exemplar) +} + +func newShardCoalescer(capacity int, onDrop func(exemplar.Exemplar)) *shardCoalescer { + if capacity <= 0 { + capacity = defaultRingBufferSize + } + slots := make([]coalescerSlot, capacity) + for i := range slots { + slots[i].nextSlot = -1 + slots[i].tombstone = true + } + generations := make([]uint64, capacity) + for i := range generations { + generations[i] = 1 + } + return &shardCoalescer{ + capacity: capacity, + slots: slots, + generations: generations, + seriesIndex: make(map[chunks.HeadSeriesRef]coalescerIndexEntry, capacity), + head: 0, + onDrop: onDrop, + } +} + +// TryAttachToBatch checks if there is an un-flushed sample or histogram in the active batch +// matching the exemplar's seriesRef and scrape timestamp (|sample.T - ex.Ts| <= 50ms). +// If found, attaches the exemplar to that batch item and returns true. +func (c *shardCoalescer) TryAttachToBatch(batch []timeSeries, ref chunks.HeadSeriesRef, ex exemplar.Exemplar) bool { + // Search from newest to oldest in the active batch. + for i := len(batch) - 1; i >= 0; i-- { + item := &batch[i] + if item.seriesRef != ref { + continue + } + if item.sType != tSample && item.sType != tHistogram && item.sType != tFloatHistogram { + continue + } + diff := item.timestamp - ex.Ts + if diff < 0 { + diff = -diff + } + if diff <= maxCoalescingTimeDeltaMs { + item.exemplars = append(item.exemplars, ex) + return true + } + } + return false +} + +// AddPendingExemplar inserts an exemplar into the circular ring buffer. +// If the write pointer wraps around and overwrites an active pending exemplar, +// the overwritten exemplar is evicted and onDrop is invoked. +func (c *shardCoalescer) AddPendingExemplar(ref chunks.HeadSeriesRef, ex exemplar.Exemplar) { + slotIdx := c.head + c.head = (c.head + 1) % c.capacity + oldSlot := &c.slots[slotIdx] + + // Check if we are overwriting an active slot due to wrap-around. + if !oldSlot.tombstone && oldSlot.generation == c.generations[slotIdx] && oldSlot.seriesRef != 0 { + if c.onDrop != nil { + c.onDrop(oldSlot.exemplar) + } + oldSlot.tombstone = true + if curEntry, ok := c.seriesIndex[oldSlot.seriesRef]; ok && curEntry.slotIdx == slotIdx && curEntry.generation == oldSlot.generation { + if oldSlot.nextSlot >= 0 && c.isValidSlot(oldSlot.nextSlot, oldSlot.seriesRef) { + c.seriesIndex[oldSlot.seriesRef] = coalescerIndexEntry{ + slotIdx: oldSlot.nextSlot, + generation: c.slots[oldSlot.nextSlot].generation, + } + } else { + delete(c.seriesIndex, oldSlot.seriesRef) + } + } + } + + c.generations[slotIdx]++ + gen := c.generations[slotIdx] + + nextSlot := -1 + if existing, ok := c.seriesIndex[ref]; ok { + if c.isValidSlot(existing.slotIdx, ref) && existing.generation == c.slots[existing.slotIdx].generation { + nextSlot = existing.slotIdx + } + } + + c.slots[slotIdx] = coalescerSlot{ + seriesRef: ref, + exemplar: ex, + generation: gen, + nextSlot: nextSlot, + tombstone: false, + } + c.seriesIndex[ref] = coalescerIndexEntry{ + slotIdx: slotIdx, + generation: gen, + } +} + +func (c *shardCoalescer) isValidSlot(slotIdx int, ref chunks.HeadSeriesRef) bool { + if slotIdx < 0 || slotIdx >= c.capacity { + return false + } + s := &c.slots[slotIdx] + return !s.tombstone && s.seriesRef == ref && s.generation == c.generations[slotIdx] +} + +// TryAttachMatchingExemplars finds and removes all matching pending exemplars in the ring buffer +// for the given seriesRef and sample timestamp (|sampleTs - ex.Ts| <= 50ms). +// Any expired exemplars (sampleTs - ex.Ts > 50ms) are tombstoned, dropped, and notified via onDrop. +func (c *shardCoalescer) TryAttachMatchingExemplars(ref chunks.HeadSeriesRef, sampleTs int64) []exemplar.Exemplar { + entry, ok := c.seriesIndex[ref] + if !ok { + return nil + } + + var matched []exemplar.Exemplar + currIdx := entry.slotIdx + firstActiveSlot := -1 + visited := 0 + + for currIdx >= 0 && currIdx < c.capacity && visited < c.capacity { + visited++ + slot := &c.slots[currIdx] + nextIdx := slot.nextSlot + + // Validate generation, seriesRef, and tombstone. + if slot.generation != c.generations[currIdx] || slot.seriesRef != ref { + break + } + + if !slot.tombstone { + diff := sampleTs - slot.exemplar.Ts + if diff < 0 { + diff = -diff + } + + if diff <= maxCoalescingTimeDeltaMs { + matched = append(matched, slot.exemplar) + slot.tombstone = true + } else if sampleTs > slot.exemplar.Ts+maxCoalescingTimeDeltaMs { + // Exemplar is stale from an earlier scrape interval. + slot.tombstone = true + if c.onDrop != nil { + c.onDrop(slot.exemplar) + } + } else { + // Exemplar timestamp is in the future relative to this sample (|sampleTs - ex.Ts| > 50ms and sampleTs < ex.Ts - 50ms). + // It remains active for a subsequent scrape timestamp. + if firstActiveSlot == -1 { + firstActiveSlot = currIdx + } + } + } + + currIdx = nextIdx + } + + if firstActiveSlot >= 0 { + c.seriesIndex[ref] = coalescerIndexEntry{ + slotIdx: firstActiveSlot, + generation: c.slots[firstActiveSlot].generation, + } + } else { + delete(c.seriesIndex, ref) + } + + return matched +} + +// EvictOlderThan evicts any pending exemplars older than cutoffTimestamp - maxCoalescingTimeDeltaMs. +func (c *shardCoalescer) EvictOlderThan(cutoffTimestamp int64) int { + evicted := 0 + for i := range c.slots { + slot := &c.slots[i] + if !slot.tombstone && slot.generation == c.generations[i] && slot.seriesRef != 0 { + if slot.exemplar.Ts < cutoffTimestamp-maxCoalescingTimeDeltaMs { + slot.tombstone = true + if curEntry, ok := c.seriesIndex[slot.seriesRef]; ok && curEntry.slotIdx == i && curEntry.generation == slot.generation { + delete(c.seriesIndex, slot.seriesRef) + } + if c.onDrop != nil { + c.onDrop(slot.exemplar) + } + evicted++ + } + } + } + return evicted +} + +// FlushAndClear drains all pending exemplars from the coalescer, invoking onDrop for each, +// and resets the ring buffer and index. +func (c *shardCoalescer) FlushAndClear() int { + dropped := 0 + for i := range c.slots { + slot := &c.slots[i] + if !slot.tombstone && slot.generation == c.generations[i] && slot.seriesRef != 0 { + slot.tombstone = true + if c.onDrop != nil { + c.onDrop(slot.exemplar) + } + dropped++ + } + slot.seriesRef = 0 + slot.nextSlot = -1 + slot.tombstone = true + } + clear(c.seriesIndex) + c.head = 0 + return dropped +} + +// PendingCount returns the number of active un-tombstoned exemplars in the ring buffer. +func (c *shardCoalescer) PendingCount() int { + count := 0 + for i := range c.slots { + slot := &c.slots[i] + if !slot.tombstone && slot.generation == c.generations[i] && slot.seriesRef != 0 { + count++ + } + } + return count +} diff --git a/storage/remote/shard_coalescer_test.go b/storage/remote/shard_coalescer_test.go new file mode 100644 index 00000000000..de65e31b234 --- /dev/null +++ b/storage/remote/shard_coalescer_test.go @@ -0,0 +1,375 @@ +// Copyright The Prometheus Authors +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package remote + +import ( + "strconv" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/prometheus/prometheus/model/exemplar" + "github.com/prometheus/prometheus/model/histogram" + "github.com/prometheus/prometheus/model/labels" + "github.com/prometheus/prometheus/tsdb/chunks" +) + +func TestShardCoalescer_SampleFirstArrival(t *testing.T) { + var dropped atomic.Int64 + coalescer := newShardCoalescer(2048, func(e exemplar.Exemplar) { + dropped.Add(1) + }) + + ref := chunks.HeadSeriesRef(100) + ts := int64(1000) + + batch := []timeSeries{ + { + seriesRef: ref, + timestamp: ts, + sType: tSample, + value: 42.0, + }, + } + + ex := exemplar.Exemplar{ + Labels: labels.FromStrings("trace_id", "abc123"), + Value: 42.0, + Ts: ts + 10, // within 50ms window + HasTs: true, + } + + attached := coalescer.TryAttachToBatch(batch, ref, ex) + require.True(t, attached) + require.Len(t, batch[0].exemplars, 1) + require.Equal(t, ex, batch[0].exemplars[0]) + require.Equal(t, 0, coalescer.PendingCount()) + require.Equal(t, int64(0), dropped.Load()) +} + +func TestShardCoalescer_ExemplarFirstArrival(t *testing.T) { + var dropped atomic.Int64 + coalescer := newShardCoalescer(2048, func(e exemplar.Exemplar) { + dropped.Add(1) + }) + + ref := chunks.HeadSeriesRef(200) + ts := int64(1000) + + ex := exemplar.Exemplar{ + Labels: labels.FromStrings("span_id", "xyz789"), + Value: 15.5, + Ts: ts, + HasTs: true, + } + + // No batch item exists yet + var batch []timeSeries + attached := coalescer.TryAttachToBatch(batch, ref, ex) + require.False(t, attached) + + coalescer.AddPendingExemplar(ref, ex) + require.Equal(t, 1, coalescer.PendingCount()) + + // Sample arrives within 50ms + matched := coalescer.TryAttachMatchingExemplars(ref, ts+25) + require.Len(t, matched, 1) + require.Equal(t, ex, matched[0]) + require.Equal(t, 0, coalescer.PendingCount()) + require.Equal(t, int64(0), dropped.Load()) + + // Second sample for same ref should not find anything + matchedAgain := coalescer.TryAttachMatchingExemplars(ref, ts+25) + require.Nil(t, matchedAgain) +} + +func TestShardCoalescer_RingBufferWrapAround_SlotGenerations(t *testing.T) { + var dropped atomic.Int64 + // Small ring buffer to test wrap-around easily + capacity := 4 + coalescer := newShardCoalescer(capacity, func(e exemplar.Exemplar) { + dropped.Add(1) + }) + + // Fill buffer completely (slots 0..3) + for i := 1; i <= 4; i++ { + ex := exemplar.Exemplar{ + Labels: labels.FromStrings("idx", strconv.Itoa(i)), + Value: float64(i), + Ts: 1000, + HasTs: true, + } + coalescer.AddPendingExemplar(chunks.HeadSeriesRef(i), ex) + } + require.Equal(t, 4, coalescer.PendingCount()) + require.Equal(t, int64(0), dropped.Load()) + + // Overwrite slot 0 with ref 5 + ex5 := exemplar.Exemplar{ + Labels: labels.FromStrings("idx", "5"), + Value: 5.0, + Ts: 1000, + HasTs: true, + } + coalescer.AddPendingExemplar(chunks.HeadSeriesRef(5), ex5) + require.Equal(t, 4, coalescer.PendingCount()) + require.Equal(t, int64(1), dropped.Load(), "overwritten slot 0 (ref 1) must be dropped") + + // Sample for ref 1 arrives - should NOT match ref 5 due to generation validation + matched1 := coalescer.TryAttachMatchingExemplars(chunks.HeadSeriesRef(1), 1000) + require.Nil(t, matched1, "ref 1 slot was overwritten and generation changed; must return nil") + + // Sample for ref 5 arrives - should match ref 5 + matched5 := coalescer.TryAttachMatchingExemplars(chunks.HeadSeriesRef(5), 1000) + require.Len(t, matched5, 1) + require.Equal(t, ex5, matched5[0]) + + // Other slots (refs 2, 3, 4) should still be valid + for i := 2; i <= 4; i++ { + matched := coalescer.TryAttachMatchingExemplars(chunks.HeadSeriesRef(i), 1000) + require.Len(t, matched, 1) + require.Equal(t, float64(i), matched[0].Value) + } + require.Equal(t, 0, coalescer.PendingCount()) +} + +func TestShardCoalescer_RejectCrossScrapeExemplarMatching(t *testing.T) { + var dropped atomic.Int64 + coalescer := newShardCoalescer(2048, func(e exemplar.Exemplar) { + dropped.Add(1) + }) + + ref := chunks.HeadSeriesRef(300) + + // Case 1: Exemplar-first, sample arrives > 50ms later + ex1 := exemplar.Exemplar{ + Labels: labels.FromStrings("trace_id", "old"), + Value: 1.0, + Ts: 1000, + HasTs: true, + } + coalescer.AddPendingExemplar(ref, ex1) + + // Sample arrives at T=1051 (51ms later > 50ms) + matched := coalescer.TryAttachMatchingExemplars(ref, 1051) + require.Nil(t, matched, "should reject matching with delta > 50ms") + require.Equal(t, int64(1), dropped.Load(), "stale exemplar must be dropped") + require.Equal(t, 0, coalescer.PendingCount()) + + // Case 2: Exact boundary at 50ms matches + ref2 := chunks.HeadSeriesRef(301) + ex2 := exemplar.Exemplar{ + Labels: labels.FromStrings("trace_id", "boundary"), + Value: 2.0, + Ts: 1000, + HasTs: true, + } + coalescer.AddPendingExemplar(ref2, ex2) + matchedBoundary := coalescer.TryAttachMatchingExemplars(ref2, 1050) + require.Len(t, matchedBoundary, 1, "boundary at exactly 50ms must match") + require.Equal(t, ex2, matchedBoundary[0]) + + // Case 3: Sample-first, exemplar arrives > 50ms later + batch := []timeSeries{ + { + seriesRef: chunks.HeadSeriesRef(302), + timestamp: 1000, + sType: tSample, + value: 3.0, + }, + } + ex3 := exemplar.Exemplar{ + Labels: labels.FromStrings("trace_id", "late"), + Value: 3.0, + Ts: 1060, // 60ms difference + HasTs: true, + } + attached := coalescer.TryAttachToBatch(batch, chunks.HeadSeriesRef(302), ex3) + require.False(t, attached, "should reject attaching exemplar to batch sample with delta > 50ms") + require.Empty(t, batch[0].exemplars) +} + +func TestShardCoalescer_SupportAllMetricTypes(t *testing.T) { + var dropped atomic.Int64 + coalescer := newShardCoalescer(2048, func(e exemplar.Exemplar) { + dropped.Add(1) + }) + + // Float Sample + batchSample := []timeSeries{ + { + seriesRef: 401, + timestamp: 1000, + sType: tSample, + value: 100.0, + }, + } + ex1 := exemplar.Exemplar{Labels: labels.FromStrings("type", "sample"), Value: 100.0, Ts: 1000, HasTs: true} + require.True(t, coalescer.TryAttachToBatch(batchSample, 401, ex1)) + require.Len(t, batchSample[0].exemplars, 1) + + // Int Histogram + h := &histogram.Histogram{Schema: 1, Count: 10} + batchHist := []timeSeries{ + { + seriesRef: 402, + timestamp: 1000, + sType: tHistogram, + histogram: h, + }, + } + ex2 := exemplar.Exemplar{Labels: labels.FromStrings("type", "histogram"), Value: 5.0, Ts: 1000, HasTs: true} + require.True(t, coalescer.TryAttachToBatch(batchHist, 402, ex2)) + require.Len(t, batchHist[0].exemplars, 1) + + // Float Histogram + fh := &histogram.FloatHistogram{Schema: 1, Count: 15} + batchFloatHist := []timeSeries{ + { + seriesRef: 403, + timestamp: 1000, + sType: tFloatHistogram, + floatHistogram: fh, + }, + } + ex3 := exemplar.Exemplar{Labels: labels.FromStrings("type", "floathistogram"), Value: 7.5, Ts: 1000, HasTs: true} + require.True(t, coalescer.TryAttachToBatch(batchFloatHist, 403, ex3)) + require.Len(t, batchFloatHist[0].exemplars, 1) + + // Exemplar-first for all 3 types + coalescer.AddPendingExemplar(404, ex1) + coalescer.AddPendingExemplar(405, ex2) + coalescer.AddPendingExemplar(406, ex3) + + m1 := coalescer.TryAttachMatchingExemplars(404, 1000) + require.Len(t, m1, 1) + m2 := coalescer.TryAttachMatchingExemplars(405, 1000) + require.Len(t, m2, 1) + m3 := coalescer.TryAttachMatchingExemplars(406, 1000) + require.Len(t, m3, 1) + + require.Equal(t, int64(0), dropped.Load()) + require.Equal(t, 0, coalescer.PendingCount()) +} + +func TestShardCoalescer_EvictOlderThanAndFlush(t *testing.T) { + var dropped atomic.Int64 + coalescer := newShardCoalescer(2048, func(e exemplar.Exemplar) { + dropped.Add(1) + }) + + coalescer.AddPendingExemplar(501, exemplar.Exemplar{Ts: 1000, HasTs: true}) + coalescer.AddPendingExemplar(502, exemplar.Exemplar{Ts: 2000, HasTs: true}) + coalescer.AddPendingExemplar(503, exemplar.Exemplar{Ts: 3000, HasTs: true}) + require.Equal(t, 3, coalescer.PendingCount()) + + // Evict older than 2500 (cutoff 2500 - 50 = 2450): Ts 1000 and 2000 will be evicted + evicted := coalescer.EvictOlderThan(2500) + require.Equal(t, 2, evicted) + require.Equal(t, int64(2), dropped.Load()) + require.Equal(t, 1, coalescer.PendingCount()) + + // Remaining 503 can be matched + m := coalescer.TryAttachMatchingExemplars(503, 3000) + require.Len(t, m, 1) + require.Equal(t, 0, coalescer.PendingCount()) + + // Add new and test FlushAndClear + coalescer.AddPendingExemplar(504, exemplar.Exemplar{Ts: 4000, HasTs: true}) + coalescer.AddPendingExemplar(505, exemplar.Exemplar{Ts: 4000, HasTs: true}) + require.Equal(t, 2, coalescer.PendingCount()) + + flushedDropped := coalescer.FlushAndClear() + require.Equal(t, 2, flushedDropped) + require.Equal(t, int64(4), dropped.Load()) + require.Equal(t, 0, coalescer.PendingCount()) +} + +func TestShardCoalescer_MultiExemplarTimestampSeparationAndCycleBounds(t *testing.T) { + var dropped atomic.Int64 + coalescer := newShardCoalescer(16, func(e exemplar.Exemplar) { + dropped.Add(1) + }) + + // Add an exemplar for ref 601 at T=1000, and another for ref 601 at T=5000 (future scrape) + coalescer.AddPendingExemplar(601, exemplar.Exemplar{Ts: 1000, Value: 10, HasTs: true}) + coalescer.AddPendingExemplar(601, exemplar.Exemplar{Ts: 5000, Value: 50, HasTs: true}) + + // Match sample at T=1010 + m1 := coalescer.TryAttachMatchingExemplars(601, 1010) + require.Len(t, m1, 1) + require.Equal(t, float64(10), m1[0].Value) + require.Equal(t, 1, coalescer.PendingCount()) + + // The future exemplar at T=5000 is still preserved in index + m2 := coalescer.TryAttachMatchingExemplars(601, 5010) + require.Len(t, m2, 1) + require.Equal(t, float64(50), m2[0].Value) + require.Equal(t, 0, coalescer.PendingCount()) + require.Equal(t, int64(0), dropped.Load()) +} + +func BenchmarkShardCoalescer_SampleFirst(b *testing.B) { + coalescer := newShardCoalescer(2048, nil) + lbls := labels.FromStrings("trace_id", "1234567890abcdef") + ex := exemplar.Exemplar{Labels: lbls, Value: 123.45, Ts: 1000, HasTs: true} + + batch := make([]timeSeries, 100) + for i := range batch { + batch[i] = timeSeries{ + seriesRef: chunks.HeadSeriesRef(i + 1), + timestamp: 1000, + sType: tSample, + value: float64(i), + } + } + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + ref := chunks.HeadSeriesRef((i % 100) + 1) + coalescer.TryAttachToBatch(batch, ref, ex) + } +} + +func BenchmarkShardCoalescer_ExemplarFirst(b *testing.B) { + coalescer := newShardCoalescer(2048, nil) + lbls := labels.FromStrings("trace_id", "1234567890abcdef") + ex := exemplar.Exemplar{Labels: lbls, Value: 123.45, Ts: 1000, HasTs: true} + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + ref := chunks.HeadSeriesRef((i % 1000) + 1) + coalescer.AddPendingExemplar(ref, ex) + coalescer.TryAttachMatchingExemplars(ref, 1000) + } +} + +func BenchmarkShardCoalescer_RingWrapAround(b *testing.B) { + coalescer := newShardCoalescer(2048, nil) + lbls := labels.FromStrings("trace_id", "1234567890abcdef") + ex := exemplar.Exemplar{Labels: lbls, Value: 123.45, Ts: 1000, HasTs: true} + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + ref := chunks.HeadSeriesRef(i + 1) + coalescer.AddPendingExemplar(ref, ex) + } +}