traces/cmd/tempo-cli/shared.go

package main

import (
	"context"
	"crypto/x509"
	"errors"
	"fmt"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/gogo/protobuf/jsonpb"
	"github.com/gogo/protobuf/proto"
	"github.com/google/uuid"
	"github.com/prometheus/common/model"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/insecure"

	"github.com/grafana/tempo/pkg/boundedwaitgroup"
	"github.com/grafana/tempo/tempodb/backend"
)

type unifiedBlockMeta struct {
	backend.BlockMeta
	backend.CompactedBlockMeta

	window    int64
	compacted bool
}

func getMeta(meta *backend.BlockMeta, compactedMeta *backend.CompactedBlockMeta, windowRange time.Duration) unifiedBlockMeta {
	if meta != nil {
		return unifiedBlockMeta{
			BlockMeta: *meta,
			window:    meta.EndTime.Unix() / int64(windowRange/time.Second),
			compacted: false,
		}
	}
	if compactedMeta != nil {
		return unifiedBlockMeta{
			BlockMeta:          compactedMeta.BlockMeta,
			CompactedBlockMeta: *compactedMeta,
			window:             compactedMeta.EndTime.Unix() / int64(windowRange/time.Second),
			compacted:          true,
		}
	}

	return unifiedBlockMeta{
		BlockMeta: backend.BlockMeta{
			BlockID:         backend.UUID{},
			CompactionLevel: 0,
			TotalObjects:    -1,
		},
		window:    -1,
		compacted: false,
	}
}

type blockStats struct {
	unifiedBlockMeta
}

func loadBucket(r backend.Reader, c backend.Compactor, tenantID string, windowRange time.Duration, includeCompacted bool) ([]blockStats, error) {
	blockIDs, compactedBlockIDs, err := r.Blocks(context.Background(), tenantID)
	if err != nil {
		return nil, err
	}

	blockIDs = append(blockIDs, compactedBlockIDs...)

	fmt.Println("total blocks: ", len(blockIDs))

	// Load in parallel
	wg := boundedwaitgroup.New(20)
	resultsCh := make(chan blockStats, len(blockIDs))

	for blockNum, id := range blockIDs {
		wg.Add(1)

		go func(id2 backend.UUID, blockNum2 int) {
			defer wg.Done()

			b, err := loadBlock(r, c, tenantID, id2, blockNum2, windowRange, includeCompacted)
			if err != nil {
				fmt.Println("Error loading block:", id2, err)
				return
			}

			if b != nil {
				resultsCh <- *b
			}
		}(backend.UUID(id), blockNum)
	}

	wg.Wait()
	close(resultsCh)

	results := make([]blockStats, 0)
	for b := range resultsCh {
		results = append(results, b)
	}

	sort.Slice(results, func(i, j int) bool {
		return results[i].EndTime.Before(results[j].EndTime)
	})

	return results, nil
}

func loadBlock(r backend.Reader, c backend.Compactor, tenantID string, id backend.UUID, blockNum int, windowRange time.Duration, includeCompacted bool) (*blockStats, error) {
	fmt.Print(".")
	if blockNum%100 == 0 {
		fmt.Print(strconv.Itoa(blockNum))
	}

	meta, err := r.BlockMeta(context.Background(), (uuid.UUID)(id), tenantID)
	if errors.Is(err, backend.ErrDoesNotExist) && !includeCompacted {
		return nil, nil
	} else if err != nil && !errors.Is(err, backend.ErrDoesNotExist) {
		return nil, err
	}

	compactedMeta, err := c.CompactedBlockMeta((uuid.UUID)(id), tenantID)
	if err != nil && !errors.Is(err, backend.ErrDoesNotExist) {
		return nil, err
	}

	return &blockStats{
		unifiedBlockMeta: getMeta(meta, compactedMeta, windowRange),
	}, nil
}

func printAsJSON(pb proto.Message) error {
	m := jsonpb.Marshaler{}
	traceJSON, err := m.MarshalToString(pb)
	if err != nil {
		return err
	}

	fmt.Println(string(traceJSON))
	return nil
}

func httpScheme(secure bool) string {
	if secure {
		return "https"
	}
	return "http"
}

func grpcTransportCredentials(secure bool) (opt grpc.DialOption, err error) {
	var creds credentials.TransportCredentials
	if secure {
		certPool, err := x509.SystemCertPool()
		if err != nil {
			return nil, err
		}
		creds = credentials.NewClientTLSFromCert(certPool, "")
	} else {
		creds = insecure.NewCredentials()
	}

	return grpc.WithTransportCredentials(creds), nil
}

// parseTime parses a time string that can be:
// - relative: "now", "now-1h", "now-30m", "now-3h30m"
// - RFC3339: "2024-01-01T00:00:00Z"
func parseTime(s string) (time.Time, error) {
	s = strings.TrimSpace(s)

	if strings.HasPrefix(s, "now") {
		return parseRelativeTime(s)
	}

	if t, err := time.Parse(time.RFC3339, s); err == nil {
		return t, nil
	}

	return time.Time{}, fmt.Errorf("failed to parse time: %q use relative (now, now-1h) or absolute RFC3339 (2006-01-02T15:04:05Z07:00) format", s)
}

func parseRelativeTime(s string) (time.Time, error) {
	now := time.Now()
	s = strings.TrimSpace(s)

	if s == "now" {
		return now, nil
	}

	// expect "now-<duration>"
	if !strings.HasPrefix(s, "now-") {
		return time.Time{}, fmt.Errorf("invalid relative time %q: expected format like now-1h", s)
	}

	durStr := s[len("now-"):]
	d, err := model.ParseDuration(durStr)
	if err != nil {
		return time.Time{}, fmt.Errorf("failed to parse duration %q: %w", durStr, err)
	}

	return now.Add(-time.Duration(d)), nil
}