expose downloaded data size in query hints (thanos-io#6409)

Signed-off-by: Ben Ye <[email protected]>
databricks · Jun 27, 2023 · cfedce8 · cfedce8
1 parent 540f95c
commit cfedce8
Show file tree

Hide file tree

Showing 4 changed files with 86 additions and 41 deletions.
diff --git a/pkg/store/bucket.go b/pkg/store/bucket.go
@@ -2393,6 +2393,7 @@ func (r *bucketIndexReader) fetchPostings(ctx context.Context, keys []labels.Lab
 		if err := bytesLimiter.Reserve(uint64(len(dataFromCache))); err != nil {
 			return nil, closeFns, httpgrpc.Errorf(int(codes.ResourceExhausted), "exceeded bytes limit while loading postings from index cache: %s", err)
 		}
+		r.stats.DataDownloadedSizeSum += units.Base2Bytes(len(dataFromCache))
 	}
 
 	// Iterate over all groups and fetch posting from cache.
@@ -2466,6 +2467,7 @@ func (r *bucketIndexReader) fetchPostings(ctx context.Context, keys []labels.Lab
 		if err := bytesLimiter.Reserve(uint64(length)); err != nil {
 			return nil, closeFns, httpgrpc.Errorf(int(codes.ResourceExhausted), "exceeded bytes limit while fetching postings: %s", err)
 		}
+		r.stats.DataDownloadedSizeSum += units.Base2Bytes(length)
 	}
 
 	g, ctx := errgroup.WithContext(ctx)
@@ -2626,6 +2628,7 @@ func (r *bucketIndexReader) PreloadSeries(ctx context.Context, ids []storage.Ser
 		if err := bytesLimiter.Reserve(uint64(len(b))); err != nil {
 			return httpgrpc.Errorf(int(codes.ResourceExhausted), "exceeded bytes limit while loading series from index cache: %s", err)
 		}
+		r.stats.DataDownloadedSizeSum += units.Base2Bytes(len(b))
 	}
 
 	parts := r.block.partitioner.Partition(len(ids), func(i int) (start, end uint64) {
@@ -2651,6 +2654,7 @@ func (r *bucketIndexReader) loadSeries(ctx context.Context, ids []storage.Series
 		if err := bytesLimiter.Reserve(uint64(end - start)); err != nil {
 			return httpgrpc.Errorf(int(codes.ResourceExhausted), "exceeded bytes limit while fetching series: %s", err)
 		}
+		r.stats.DataDownloadedSizeSum += units.Base2Bytes(end - start)
 	}
 
 	b, err := r.block.readIndexRange(ctx, int64(start), int64(end-start))
@@ -2923,6 +2927,7 @@ func (r *bucketChunkReader) load(ctx context.Context, res []seriesEntry, aggrs [
 			if err := bytesLimiter.Reserve(uint64(p.End - p.Start)); err != nil {
 				return httpgrpc.Errorf(int(codes.ResourceExhausted), "exceeded bytes limit while fetching chunks: %s", err)
 			}
+			r.stats.DataDownloadedSizeSum += units.Base2Bytes(p.End - p.Start)
 		}
 
 		for _, p := range parts {
@@ -3040,6 +3045,7 @@ func (r *bucketChunkReader) loadChunks(ctx context.Context, res []seriesEntry, a
 		if err := bytesLimiter.Reserve(uint64(chunkLen)); err != nil {
 			return httpgrpc.Errorf(int(codes.ResourceExhausted), "exceeded bytes limit while fetching chunks: %s", err)
 		}
+		r.stats.DataDownloadedSizeSum += units.Base2Bytes(chunkLen)
 		nb, err := r.block.readChunkRange(ctx, seq, int64(pIdx.offset), int64(chunkLen), []byteRange{{offset: 0, length: chunkLen}})
 		if err != nil {
 			return errors.Wrapf(err, "preloaded chunk too small, expecting %d, and failed to fetch full chunk", chunkLen)
@@ -3147,6 +3153,8 @@ type queryStats struct {
 	mergedSeriesCount int
 	mergedChunksCount int
 	MergeDuration     time.Duration
+
+	DataDownloadedSizeSum units.Base2Bytes
 }
 
 func (s queryStats) merge(o *queryStats) *queryStats {
@@ -3187,6 +3195,8 @@ func (s queryStats) merge(o *queryStats) *queryStats {
 	s.mergedChunksCount += o.mergedChunksCount
 	s.MergeDuration += o.MergeDuration
 
+	s.DataDownloadedSizeSum += o.DataDownloadedSizeSum
+
 	return &s
 }
 
@@ -3211,6 +3221,7 @@ func (s queryStats) toHints() *hintspb.QueryStats {
 		ChunksFetchCount:       int64(s.chunksFetchCount),
 		MergedSeriesCount:      int64(s.mergedSeriesCount),
 		MergedChunksCount:      int64(s.mergedChunksCount),
+		DataDownloadedSizeSum:  int64(s.DataDownloadedSizeSum),
 	}
 }
 

diff --git a/pkg/store/hintspb/custom.go b/pkg/store/hintspb/custom.go
@@ -27,6 +27,7 @@ func (m *QueryStats) Merge(other *QueryStats) {
 	m.BlocksQueried += other.BlocksQueried
 	m.MergedSeriesCount += other.MergedSeriesCount
 	m.MergedChunksCount += other.MergedChunksCount
+	m.DataDownloadedSizeSum += other.DataDownloadedSizeSum
 
 	m.PostingsFetched += other.PostingsFetched
 	m.PostingsToFetch += other.PostingsToFetch

diff --git a/pkg/store/hintspb/hints.pb.go b/pkg/store/hintspb/hints.pb.go
diff --git a/pkg/store/hintspb/hints.proto b/pkg/store/hintspb/hints.proto
@@ -88,4 +88,6 @@ message QueryStats {
     int64 chunks_fetched = 17;
     int64 chunks_fetched_size_sum = 18;
     int64 chunks_fetch_count = 19;
+
+    int64 data_downloaded_size_sum = 20;
 }