shared.go

package main

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

	"github.com/gogo/protobuf/jsonpb"
	"github.com/gogo/protobuf/proto"
	"github.com/google/uuid"
	"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
}