[Tempo CLI] Additional Features

CHANGELOG.md

@@ -6,6 +6,7 @@
 * [BUGFIX] Fix "magic number" errors and other block mishandling when an ingester forcefully shuts down [#937](https://github.com/grafana/tempo/issues/937) (@mdisibio)
 * [BUGFIX] Fix compactor memory leak [#806](https://github.com/grafana/tempo/pull/806) (@mdisibio)
 * [ENHANCEMENT] Added "query blocks" cli option. [#876](https://github.com/grafana/tempo/pull/876) (@joe-elliott)
+* [ENHANCEMENT] Added "search blocks" cli option. [#972](https://github.com/grafana/tempo/pull/972) (@joe-elliott)
 * [ENHANCEMENT] Added traceid to `trace too large message`. [#888](https://github.com/grafana/tempo/pull/888) (@mritunjaysharma394)
 * [ENHANCEMENT] Add support to tempo workloads to `overrides` from single configmap in microservice mode. [#896](https://github.com/grafana/tempo/pull/896) (@kavirajk)
 * [ENHANCEMENT] Make `overrides_config` block name consistent with Loki and Cortex in microservice mode. [#906](https://github.com/grafana/tempo/pull/906) (@kavirajk)

cmd/tempo-cli/cmd-list-block.go

@@ -7,15 +7,32 @@ import (
 	"io"
 	"os"
 	"os/signal"
+	"sort"
+	"strconv"
 	"syscall"
 	"time"
 
 	"github.com/dustin/go-humanize"
 	"github.com/google/uuid"
+	"github.com/grafana/tempo/pkg/model"
+	"github.com/grafana/tempo/pkg/tempopb"
+	v1 "github.com/grafana/tempo/pkg/tempopb/common/v1"
+	"github.com/grafana/tempo/pkg/util"
 	tempodb_backend "github.com/grafana/tempo/tempodb/backend"
 	"github.com/grafana/tempo/tempodb/encoding"
+	"github.com/grafana/tempo/tempodb/encoding/common"
 )
 
+type valueStats struct {
+	count int
+}
+type values struct {
+	all   map[string]valueStats
+	key   string
+	count int
+}
+type kvPairs map[string]values
+
 type listBlockCmd struct {
 	backendOptions
 
@@ -84,14 +101,19 @@ func dumpBlock(r tempodb_backend.Reader, c tempodb_backend.Compactor, tenantID s
 		dupe := 0
 		maxObjSize := 0
 		minObjSize := 0
+		maxObjID := common.ID{}
+		minObjID := common.ID{}
 
+		allKVP := kvPairs{}
 		printStats := func() {
 			fmt.Println()
 			fmt.Println("Scanning results:")
 			fmt.Println("Objects scanned : ", i)
 			fmt.Println("Duplicates      : ", dupe)
-			fmt.Println("Smallest object : ", humanize.Bytes(uint64(minObjSize)))
-			fmt.Println("Largest object  : ", humanize.Bytes(uint64(maxObjSize)))
+			fmt.Println("Smallest object : ", humanize.Bytes(uint64(minObjSize)), " : ", util.TraceIDToHexString(minObjID))
+			fmt.Println("Largest object  : ", humanize.Bytes(uint64(maxObjSize)), " : ", util.TraceIDToHexString(maxObjID))
+			fmt.Println("")
+			printKVPairs(allKVP)
 		}
 
 		// Print stats on ctrl+c
@@ -116,17 +138,32 @@ func dumpBlock(r tempodb_backend.Reader, c tempodb_backend.Compactor, tenantID s
 
 			if len(obj) > maxObjSize {
 				maxObjSize = len(obj)
+				maxObjID = objID
 			}
 
 			if len(obj) < minObjSize || minObjSize == 0 {
 				minObjSize = len(obj)
+				minObjID = objID
 			}
 
 			if bytes.Equal(objID, prevID) {
 				dupe++
 			}
 
 			copy(prevID, objID)
+
+			trace, err := model.Unmarshal(obj, meta.DataEncoding)
+			if err != nil {
+				return err
+			}
+			kvp := extractKVPairs(trace)
+			for k, vs := range kvp {
+				addKey(allKVP, k, 1)
+				for v := range vs.all {
+					addVal(allKVP, k, v, 1)
+				}
+			}
+
 			i++
 			if i%100000 == 0 {
 				fmt.Println("Record: ", i)
@@ -138,3 +175,103 @@ func dumpBlock(r tempodb_backend.Reader, c tempodb_backend.Compactor, tenantID s
 
 	return nil
 }
+
+// helper methods for calculating label stats
+func printKVPairs(kvp kvPairs) {
+	allValues := make([]values, 0, len(kvp))
+	for _, vs := range kvp {
+		allValues = append(allValues, vs)
+	}
+	sort.Slice(allValues, func(i, j int) bool {
+		return relativeValue(allValues[i]) > relativeValue(allValues[j])
+	})
+	for _, vs := range allValues {
+		fmt.Println("key:", vs.key, "count:", vs.count, "len:", len(vs.all), "value:", relativeValue(vs))
+		for a, c := range vs.all {
+			fmt.Printf(" %s:\t%.2f\n", a, float64(c.count)/float64(vs.count))
+		}
+	}
+}
+
+// attempts to calculate the "value" that storing a given label would provide by. currently (number of times appeared)^2 / cardinality
+//  this is not researched and could definitely be improved
+func relativeValue(v values) float64 {
+	return (float64(v.count) * float64(v.count)) / float64(len(v.all))
+}
+func extractKVPairs(t *tempopb.Trace) kvPairs {
+	kvp := kvPairs{}
+	for _, b := range t.Batches {
+		spanCount := 0
+		for _, ils := range b.InstrumentationLibrarySpans {
+			for _, s := range ils.Spans {
+				spanCount++
+				for _, a := range s.Attributes {
+					val, yay := stringVal(a.Value)
+					if !yay {
+						continue
+					}
+					addKey(kvp, a.Key, 1)
+					addVal(kvp, a.Key, val, 1)
+				}
+			}
+		}
+		for _, a := range b.Resource.Attributes {
+			val, yay := stringVal(a.Value)
+			if !yay {
+				continue
+			}
+			addKey(kvp, a.Key, spanCount)
+			addVal(kvp, a.Key, val, spanCount)
+		}
+	}
+	return kvp
+}
+func addKey(kvp kvPairs, key string, count int) {
+	v, ok := kvp[key]
+	if !ok {
+		v = values{
+			all: map[string]valueStats{},
+			key: key,
+		}
+	}
+	v.count += count
+	kvp[key] = v
+}
+func addVal(kvp kvPairs, key string, val string, count int) {
+	v := kvp[key]
+	stats, ok := v.all[val]
+	if !ok {
+		stats = valueStats{
+			count: 0,
+		}
+	}
+	stats.count += count
+	v.all[val] = stats
+	kvp[key] = v
+}
+func stringVal(v *v1.AnyValue) (string, bool) {
+	if sVal, ok := v.Value.(*v1.AnyValue_StringValue); ok {
+		return sVal.StringValue, true
+	}
+	if nVal, ok := v.Value.(*v1.AnyValue_IntValue); ok {
+		return strconv.FormatInt(nVal.IntValue, 10), true
+	}
+	if dVal, ok := v.Value.(*v1.AnyValue_DoubleValue); ok {
+		return fmt.Sprintf("%f", dVal.DoubleValue), true
+		// strconv.FormatFloat()
+	}
+	if bVal, ok := v.Value.(*v1.AnyValue_BoolValue); ok {
+		if bVal.BoolValue {
+			return "true", true
+		}
+		return "false", true
+	}
+	// todo? add support for these?
+	// if kVal, ok := v.Value.(*v1.AnyValue_KvlistValue); ok {
+	//  return fmt.Sprintf("kvval %v", kVal.KvlistValue) // better way?
+	// }
+	// if aVal, ok := v.Value.(*v1.AnyValue_ArrayValue); ok {
+	//  return fmt.Sprintf("arrayval %v", aVal.ArrayValue) // better way?
+	// }
+	return "", false
+}

cmd/tempo-cli/cmd-search.go

@@ -0,0 +1,181 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/grafana/tempo/pkg/boundedwaitgroup"
+	"github.com/grafana/tempo/pkg/model"
+	"github.com/grafana/tempo/pkg/tempopb"
+	"github.com/grafana/tempo/pkg/util"
+	"github.com/grafana/tempo/tempodb/backend"
+	"github.com/grafana/tempo/tempodb/encoding"
+	"github.com/grafana/tempo/tempodb/encoding/common"
+)
+
+const (
+	layoutString   = "2006-01-02T15:04:05"
+	chunkSize      = 10 * 1024 * 1024
+	iteratorBuffer = 10000
+	limit          = 20
+)
+
+type searchBlocksCmd struct {
+	backendOptions
+
+	Name     string `arg:"" help:"attribute name to search for"`
+	Value    string `arg:"" help:"attribute value to search for"`
+	Start    string `arg:"" help:"start of time range to search (YYYY-MM-DDThh:mm:ss)"`
+	End      string `arg:"" help:"end of time range to search (YYYY-MM-DDThh:mm:ss)"`
+	TenantID string `arg:"" help:"tenant ID to search"`
+}
+
+func (cmd *searchBlocksCmd) Run(opts *globalOptions) error {
+	r, _, _, err := loadBackend(&cmd.backendOptions, opts)
+	if err != nil {
+		return err
+	}
+
+	startTime, err := time.Parse(layoutString, cmd.Start)
+	if err != nil {
+		return err
+	}
+	endTime, err := time.Parse(layoutString, cmd.End)
+	if err != nil {
+		return err
+	}
+
+	ctx := context.Background()
+
+	blockIDs, err := r.Blocks(ctx, cmd.TenantID)
+	if err != nil {
+		return err
+	}
+
+	fmt.Println("Total Blocks:", len(blockIDs))
+
+	// Load in parallel
+	wg := boundedwaitgroup.New(20)
+	resultsCh := make(chan *backend.BlockMeta, len(blockIDs))
+	for _, id := range blockIDs {
+		wg.Add(1)
+
+		go func(id2 uuid.UUID) {
+			defer wg.Done()
+
+			// search here
+			meta, err := r.BlockMeta(ctx, id2, cmd.TenantID)
+			if err == backend.ErrDoesNotExist {
+				return
+			}
+			if err != nil {
+				fmt.Println("Error querying block:", err)
+				return
+			}
+			if meta.StartTime.Unix() <= endTime.Unix() &&
+				meta.EndTime.Unix() >= startTime.Unix() {
+				resultsCh <- meta
+			}
+		}(id)
+	}
+
+	wg.Wait()
+	close(resultsCh)
+
+	blockmetas := []*backend.BlockMeta{}
+	for q := range resultsCh {
+		blockmetas = append(blockmetas, q)
+	}
+
+	fmt.Println("Blocks In Range:", len(blockmetas))
+	foundids := []common.ID{}
+	for _, meta := range blockmetas {
+		block, err := encoding.NewBackendBlock(meta, r)
+		if err != nil {
+			return err
+		}
+
+		// todo : graduated chunk sizes will increase throughput. i.e. first request should be small to feed the below parsing faster
+		//  later queries should use larger chunk sizes to be more efficient
+		iter, err := block.Iterator(chunkSize)
+		if err != nil {
+			return err
+		}
+
+		prefetchIter := encoding.NewPrefetchIterator(ctx, iter, iteratorBuffer)
+		ids, err := searchIterator(prefetchIter, meta.DataEncoding, cmd.Name, cmd.Value, limit)
+		prefetchIter.Close()
+		if err != nil {
+			return err
+		}
+
+		foundids = append(foundids, ids...)
+		if len(foundids) >= limit {
+			break
+		}
+	}
+
+	fmt.Println("Matching Traces:", len(foundids))
+	for _, id := range foundids {
+		fmt.Println("  ", util.TraceIDToHexString(id))
+	}
+
+	return nil
+}
+
+func searchIterator(iter encoding.Iterator, dataEncoding string, name string, value string, limit int) ([]common.ID, error) {
+	ctx := context.Background()
+	found := []common.ID{}
+
+	for {
+		id, obj, err := iter.Next(ctx)
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return nil, err
+		}
+
+		// todo : parrallelize unmarshal and search
+		trace, err := model.Unmarshal(obj, dataEncoding)
+		if err != nil {
+			return nil, err
+		}
+
+		if traceContainsKeyValue(trace, name, value) {
+			found = append(found, id)
+		}
+
+		if len(found) >= limit {
+			break
+		}
+	}
+
+	return found, nil
+}
+
+func traceContainsKeyValue(trace *tempopb.Trace, name string, value string) bool {
+	// todo : support other attribute types besides string
+	for _, b := range trace.Batches {
+		for _, a := range b.Resource.Attributes {
+			if a.Key == name && a.Value.GetStringValue() == value {
+				return true
+			}
+		}
+
+		for _, ils := range b.InstrumentationLibrarySpans {
+			for _, s := range ils.Spans {
+				for _, a := range s.Attributes {
+					if a.Key == name && a.Value.GetStringValue() == value {
+						return true
+					}
+				}
+			}
+		}
+	}
+
+	return false
+}

cmd/tempo-cli/main.go

@@ -59,6 +59,10 @@ var cli struct {
 		API    queryCmd       `cmd:"" help:"query tempo http api"`
 		Blocks queryBlocksCmd `cmd:"" help:"query for a traceid directly from backend blocks"`
 	} `cmd:""`
+
+	Search struct {
+		Blocks searchBlocksCmd `cmd:"" help:"search for a traceid directly from backend blocks"`
+	} `cmd:""`
 }
 
 func main() {