feat(tempo-cli): add benchmark profile command

.chloggen/benchmark-profile.yaml

@@ -0,0 +1,5 @@
+change_type: feature
+component: tempo-cli
+note: Add `tempo-cli benchmark profile` to profile a block for read-path benchmarking.
+issues: []
+user: zhxiaogg

cmd/tempo-cli/cmd-benchmark-profile.go

@@ -0,0 +1,62 @@
+package main
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"os"
+	"strconv"
+	"strings"
+
+	"github.com/grafana/tempo/pkg/benchmark"
+)
+
+type benchmarkProfileCmd struct {
+	Block string `arg:"" help:"path to the block directory, laid out as <bucket>/<tenant-id>/<block-id>"`
+
+	TraceIDs string `name:"trace-ids" help:"number of present trace IDs to sample, or 'all' to enumerate at run time; one absent ID is derived per present ID. 0 skips the scan they need" default:"10000"`
+	Out      string `short:"o" help:"file to write the profile to, instead of stdout" default:""`
+}
+
+func (cmd *benchmarkProfileCmd) Run(_ *globalOptions) error {
+	numTraceIDs, err := parseTraceIDCount(cmd.TraceIDs)
+	if err != nil {
+		return err
+	}
+
+	ctx := context.Background()
+	meta, r, err := benchmark.LoadLocalBlock(ctx, cmd.Block)
+	if err != nil {
+		return err
+	}
+
+	profile, err := benchmark.ProfileBlock(ctx, meta, r, benchmark.ProfileOptions{NumTraceIDs: numTraceIDs})
+	if err != nil {
+		return err
+	}
+
+	out := os.Stdout
+	if cmd.Out != "" {
+		f, err := os.Create(cmd.Out)
+		if err != nil {
+			return err
+		}
+		defer f.Close()
+		out = f
+	}
+	return profile.Write(out)
+}
+
+func parseTraceIDCount(s string) (int, error) {
+	if strings.EqualFold(strings.TrimSpace(s), "all") {
+		return benchmark.TraceIDsAll, nil
+	}
+	n, err := strconv.Atoi(s)
+	if err != nil {
+		return 0, fmt.Errorf("--trace-ids must be a number or 'all': %w", err)
+	}
+	if n < 0 {
+		return 0, errors.New("--trace-ids must not be negative")
+	}
+	return n, nil
+}

cmd/tempo-cli/cmd-benchmark-profile_test.go

@@ -0,0 +1,36 @@
+package main
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/tempo/pkg/benchmark"
+)
+
+func TestParseTraceIDCount(t *testing.T) {
+	for _, tc := range []struct {
+		in      string
+		want    int
+		wantErr bool
+	}{
+		{in: "0", want: 0},
+		{in: "10000", want: 10000},
+		{in: "all", want: benchmark.TraceIDsAll},
+		{in: "ALL", want: benchmark.TraceIDsAll},
+		{in: " all ", want: benchmark.TraceIDsAll},
+		{in: "-1", wantErr: true},
+		{in: "lots", wantErr: true},
+		{in: "", wantErr: true},
+	} {
+		t.Run(tc.in, func(t *testing.T) {
+			got, err := parseTraceIDCount(tc.in)
+			if tc.wantErr {
+				require.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+			require.Equal(t, tc.want, got)
+		})
+	}
+}

cmd/tempo-cli/main.go

@@ -56,6 +56,10 @@ var cli struct {
 		Schema viewSchemaCmd `cmd:"" help:"View parquet schema"`
 	} `cmd:""`
 
+	Benchmark struct {
+		Profile benchmarkProfileCmd `cmd:"" help:"Profile a block for read-path benchmarking"`
+	} `cmd:""`
+
 	Gen struct {
 		AttrIndex attrIndexCmd `cmd:"" help:"Generate an attribute index for a parquet block (EXPERIMENTAL)"`
 	} `cmd:""`

docs/sources/tempo/operations/tempo_cli.md

@@ -487,6 +487,40 @@ Example:
 tempo-cli view schema -c ./tempo.yaml single-tenant ca314fba-efec-4852-ba3f-8d2b0bbf69f1
 ```
 
+## Benchmark profile
+
+Profile a local block for read-path benchmarking. Writes a JSON file recording
+what had to be measured from the block — its metadata, its row-group count, and
+present and absent trace IDs to look up — so that a benchmark run does not have
+to inspect the block, and every variant of an experiment works from the same
+measurements.
+
+```bash
+tempo-cli benchmark profile <block-path>
+```
+
+Arguments:
+
+- `block-path` Path to the block directory on local disk, laid out as
+  `<bucket>/<tenant-id>/<block-id>`.
+
+Options:
+
+- `--trace-ids` Number of present trace IDs to sample, or `all` to enumerate
+  every ID at run time rather than embedding them. Defaults to `10000`. One
+  absent ID is derived per present ID. Pass `0` to skip trace IDs, and with them
+  the full scan they require.
+- `-o`, `--out` File to write the profile to. Defaults to stdout.
+
+Profiles built from a customer block embed real trace IDs. Treat them as local
+artifacts.
+
+Example:
+
+```bash
+tempo-cli benchmark profile /data/traces/single-tenant/ca314fba-efec-4852-ba3f-8d2b0bbf69f1 --trace-ids=10000 -o profile.json
+```
+
 ## Query search command
 
 Search blocks in a given time range for a specific key/value pair.

pkg/benchmark/helper_test.go

@@ -0,0 +1,114 @@
+package benchmark
+
+import (
+	"bytes"
+	"context"
+	"io"
+	"math/rand"
+	"slices"
+	"sort"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/tempo/pkg/tempopb"
+	"github.com/grafana/tempo/pkg/util/test"
+	"github.com/grafana/tempo/tempodb/backend"
+	"github.com/grafana/tempo/tempodb/backend/local"
+	"github.com/grafana/tempo/tempodb/encoding"
+	"github.com/grafana/tempo/tempodb/encoding/common"
+)
+
+// testBlock writes a block of numTraces traces, with enough row groups to
+// exercise per-row-group work.
+func testBlock(t *testing.T, numTraces int) (*backend.BlockMeta, backend.Reader, string) {
+	t.Helper()
+
+	// Real trace IDs are random across all 16 bytes, and code under test hashes
+	// them, so a fixed seed keeps the fixture both realistic and reproducible.
+	rng := rand.New(rand.NewSource(0xbeef))
+
+	ids := make([][]byte, 0, numTraces)
+	for range numTraces {
+		id := make([]byte, 16)
+		_, err := rng.Read(id)
+		require.NoError(t, err)
+		ids = append(ids, test.ValidTraceID(id))
+	}
+
+	meta, r, bucket := testBlockWithTraceIDs(t, ids)
+
+	rowGroups, err := rowGroupCount(context.Background(), meta, r)
+	require.NoError(t, err)
+	require.Greater(t, rowGroups, 1, "need more than one row group to exercise per-row-group work")
+
+	return meta, r, bucket
+}
+
+// testBlockWithTraceIDs writes a block holding exactly the given trace IDs,
+// using a small row group size. It returns the bucket root alongside the block,
+// for tests that address it by path.
+func testBlockWithTraceIDs(t *testing.T, ids [][]byte) (*backend.BlockMeta, backend.Reader, string) {
+	t.Helper()
+
+	bucket := t.TempDir()
+	rawR, rawW, _, err := local.New(&local.Config{Path: bucket})
+	require.NoError(t, err)
+	r, w := backend.NewReader(rawR), backend.NewWriter(rawW)
+
+	ids = slices.Clone(ids)
+	// vParquet blocks are sorted by trace ID.
+	sort.Slice(ids, func(i, j int) bool { return bytes.Compare(ids[i], ids[j]) < 0 })
+
+	iter := &sliceIterator{}
+	for _, id := range ids {
+		iter.add(id, test.MakeTraceWithSpanCount(1, 4, id))
+	}
+
+	// test.MakeSpan timestamps spans at time.Now(), so the block's window has
+	// to bracket that or every query filters them all out.
+	now := time.Now()
+
+	enc := encoding.LatestEncoding()
+	meta := backend.NewBlockMeta("test-tenant", uuid.New(), enc.Version())
+	meta.TotalObjects = int64(len(ids))
+	meta.StartTime = now.Add(-time.Hour)
+	meta.EndTime = now.Add(time.Hour)
+
+	cfg := &common.BlockConfig{
+		BloomFP:             0.01,
+		BloomShardSizeBytes: 100 * 1024,
+		RowGroupSizeBytes:   8 * 1024,
+		Version:             enc.Version(),
+	}
+
+	out, err := enc.CreateBlock(context.Background(), cfg, meta, iter, r, w)
+	require.NoError(t, err)
+
+	return out, r, bucket
+}
+
+type sliceIterator struct {
+	ids    []common.ID
+	traces []*tempopb.Trace
+}
+
+var _ common.Iterator = (*sliceIterator)(nil)
+
+func (i *sliceIterator) add(id common.ID, tr *tempopb.Trace) {
+	i.ids = append(i.ids, id)
+	i.traces = append(i.traces, tr)
+}
+
+func (i *sliceIterator) Next(context.Context) (common.ID, *tempopb.Trace, error) {
+	if len(i.ids) == 0 {
+		return nil, nil, io.EOF
+	}
+	id, tr := i.ids[0], i.traces[0]
+	i.ids, i.traces = i.ids[1:], i.traces[1:]
+	return id, tr, nil
+}
+
+func (i *sliceIterator) Close() {}

pkg/benchmark/profile.go

@@ -0,0 +1,109 @@
+// Package benchmark profiles a Tempo block for read-path benchmarking.
+//
+// A profile records what had to be measured from the block: its metadata, the
+// authoritative row-group count, and a sample of trace IDs. Everything a run
+// derives from those — shards, query windows, steps — is left to the runner, so
+// one profile serves every variant of an experiment.
+package benchmark
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"slices"
+	"time"
+
+	"github.com/grafana/tempo/pkg/util"
+	"github.com/grafana/tempo/tempodb/backend"
+)
+
+const ProfileSchemaVersion = 1
+
+// traceIDHexLen is the length of a padded 16-byte trace ID in hex.
+const traceIDHexLen = 32
+
+const (
+	TraceIDModeSample = "sample"
+	TraceIDModeAll    = "all"
+)
+
+// TraceIDProfile describes which trace IDs a trace-by-ID benchmark can use.
+type TraceIDProfile struct {
+	Mode    string   `json:"mode"`
+	Present []string `json:"present,omitempty"`
+	// Absent IDs are verified missing, so the miss path is measured against a
+	// real negative rather than a guess.
+	Absent []string `json:"absent,omitempty"`
+}
+
+type BuildInfo struct {
+	TempoVersion string `json:"tempoVersion,omitempty"`
+	GitSHA       string `json:"gitSHA,omitempty"`
+}
+
+// BlockProfile is what a benchmark needs to know about a block. Block is
+// recorded in full because two runs are only comparable if they agree on every
+// property of the block except the one under test.
+type BlockProfile struct {
+	SchemaVersion int                `json:"schemaVersion"`
+	GeneratedAt   time.Time          `json:"generatedAt"`
+	GeneratedBy   BuildInfo          `json:"generatedBy"`
+	Block         *backend.BlockMeta `json:"block"`
+	// RowGroups comes from the parquet footer, not Block.TotalRecords, which
+	// Tempo's own sharding calls an estimate.
+	RowGroups int            `json:"rowGroups"`
+	TraceIDs  TraceIDProfile `json:"traceIDs"`
+}
+
+func (p *BlockProfile) Write(w io.Writer) error {
+	enc := json.NewEncoder(w)
+	enc.SetIndent("", "  ")
+	return enc.Encode(p)
+}
+
+func LoadProfile(r io.Reader) (*BlockProfile, error) {
+	var p BlockProfile
+	if err := json.NewDecoder(r).Decode(&p); err != nil {
+		return nil, err
+	}
+	if p.SchemaVersion != ProfileSchemaVersion {
+		return nil, fmt.Errorf("profile schema version %d is not supported, expected %d", p.SchemaVersion, ProfileSchemaVersion)
+	}
+	if err := p.Validate(); err != nil {
+		return nil, err
+	}
+	return &p, nil
+}
+
+// Validate runs on load and after profiling, so an unusable profile is
+// rejected where it is produced rather than midway through a run.
+func (p *BlockProfile) Validate() error {
+	if p.Block == nil {
+		return errors.New("profile has no block metadata")
+	}
+	if p.RowGroups <= 0 {
+		return fmt.Errorf("profile reports %d row groups", p.RowGroups)
+	}
+
+	switch p.TraceIDs.Mode {
+	case TraceIDModeSample, TraceIDModeAll:
+	default:
+		return fmt.Errorf("unknown trace ID mode %q", p.TraceIDs.Mode)
+	}
+	if p.TraceIDs.Mode == TraceIDModeAll && len(p.TraceIDs.Present) > 0 {
+		return errors.New(`trace ID mode is "all" but present IDs are embedded`)
+	}
+
+	for _, id := range slices.Concat(p.TraceIDs.Present, p.TraceIDs.Absent) {
+		// Require the padded form: HexStringToTraceID accepts a short string,
+		// so a truncated ID would otherwise pass and decode to the wrong trace.
+		if len(id) != traceIDHexLen {
+			return fmt.Errorf("trace ID %q is %d characters, expected %d", id, len(id), traceIDHexLen)
+		}
+		if _, err := util.HexStringToTraceID(id); err != nil {
+			return fmt.Errorf("invalid trace ID %q: %w", id, err)
+		}
+	}
+	return nil
+}