traces/tempodb/backend/raw.go

package backend

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"path"
	"time"

	"github.com/google/uuid"

	tempo_io "github.com/grafana/tempo/pkg/io"
)

const (
	// JSON
	MetaName          = "meta.json"
	CompactedMetaName = "meta.compacted.json"
	TenantIndexName   = "index.json.gz"

	// Proto
	TenantIndexNamePb = "index.pb.zst"

	// File name for the cluster seed file.
	ClusterSeedFileName = "tempo_cluster_seed.json"

	// File name for the work cache
	WorkFileName = "work.json"

	// File name for the nocompact flag
	NoCompactFileName = "nocompact.flg"
)

// KeyPath is an ordered set of strings that govern where data is read/written
// from the backend
type KeyPath []string

// FindFunc is executed for each object in the backend.  The provided FindMatch
// are used to determine how to handle the object.  Any collection of these
// objects is the callers responsibility.
type FindFunc func(FindMatch)

type FindMatch struct {
	Modified time.Time
	Key      string
}

// RawWriter is a collection of methods to write data to tempodb backends
type RawWriter interface {
	// Write is for in memory data. shouldCache specifies whether or not caching should be attempted.
	Write(ctx context.Context, name string, keypath KeyPath, data io.Reader, size int64, cacheInfo *CacheInfo) error
	// Append starts or continues an Append job. Pass nil to AppendTracker to start a job.
	Append(ctx context.Context, name string, keypath KeyPath, tracker AppendTracker, buffer []byte) (AppendTracker, error)
	// CloseAppend closes any resources associated with the AppendTracker.
	CloseAppend(ctx context.Context, tracker AppendTracker) error
	// Delete deletes a file.
	Delete(ctx context.Context, name string, keypath KeyPath, cacheInfo *CacheInfo) error
}

// RawReader is a collection of methods to read data from tempodb backends
type RawReader interface {
	// List returns all objects one level beneath the provided keypath
	List(ctx context.Context, keypath KeyPath) ([]string, error)
	// ListBlocks returns all blockIDs and compactedBlockIDs for a tenant.
	ListBlocks(ctx context.Context, tenant string) (blockIDs []uuid.UUID, compactedBlockIDs []uuid.UUID, err error)
	// Find executes the FindFunc for each object in the backend starting at the specified keypath.  Collection of these objects is the callers responsibility.
	Find(ctx context.Context, keypath KeyPath, f FindFunc) error
	// Read is for streaming entire objects from the backend.  There will be an attempt to retrieve this from cache if shouldCache is true.
	Read(ctx context.Context, name string, keyPath KeyPath, cacheInfo *CacheInfo) (io.ReadCloser, int64, error)
	// ReadRange is for reading parts of large objects from the backend.
	// There will be an attempt to retrieve this from cache if shouldCache is true. Cache key will be tenantID:blockID:offset:bufferLength
	ReadRange(ctx context.Context, name string, keypath KeyPath, offset uint64, buffer []byte, cacheInfo *CacheInfo) error
	// Shutdown must be called when the Reader is finished and cleans up any associated resources.
	Shutdown()
}

type writer struct {
	w RawWriter
}

// NewWriter returns an object that implements Writer and bridges to a RawWriter
func NewWriter(w RawWriter) Writer {
	return &writer{
		w: w,
	}
}

// TODO: these objects are not raw, so perhaps they could move somewhere else.
// var (
// 	x RawReader = reader{}
// 	y RawWriter = writer{}
// )

// Write implements backend.Writer
func (w *writer) Write(ctx context.Context, name string, blockID uuid.UUID, tenantID string, buffer []byte, cacheInfo *CacheInfo) error {
	return w.w.Write(ctx, name, KeyPathForBlock(blockID, tenantID), bytes.NewReader(buffer), int64(len(buffer)), cacheInfo)
}

// Write implements backend.Writer
func (w *writer) StreamWriter(ctx context.Context, name string, blockID uuid.UUID, tenantID string, data io.Reader, size int64) error {
	return w.w.Write(ctx, name, KeyPathForBlock(blockID, tenantID), data, size, nil)
}

// Write implements backend.Writer
func (w *writer) WriteBlockMeta(ctx context.Context, meta *BlockMeta) error {
	var (
		blockID  = meta.BlockID
		tenantID = meta.TenantID
	)

	bMeta, err := json.Marshal(meta)
	if err != nil {
		return err
	}

	return w.w.Write(ctx, MetaName, KeyPathForBlock((uuid.UUID)(blockID), tenantID), bytes.NewReader(bMeta), int64(len(bMeta)), nil)
}

// Write implements backend.Writer
func (w *writer) Append(ctx context.Context, name string, blockID uuid.UUID, tenantID string, tracker AppendTracker, buffer []byte) (AppendTracker, error) {
	return w.w.Append(ctx, name, KeyPathForBlock(blockID, tenantID), tracker, buffer)
}

// Write implements backend.Writer
func (w *writer) CloseAppend(ctx context.Context, tracker AppendTracker) error {
	return w.w.CloseAppend(ctx, tracker)
}

// Write implements backend.Writer
func (w *writer) WriteTenantIndex(ctx context.Context, tenantID string, meta []*BlockMeta, compactedMeta []*CompactedBlockMeta) error {
	// If meta and compactedMeta are empty, call delete the tenant index.
	if len(meta) == 0 && len(compactedMeta) == 0 {
		// Skip returning an error when the object is already deleted.
		err := w.w.Delete(ctx, TenantIndexName, []string{tenantID}, nil)
		if err != nil && !errors.Is(err, ErrDoesNotExist) {
			return err
		}

		err = w.w.Delete(ctx, TenantIndexNamePb, []string{tenantID}, nil)
		if err != nil && !errors.Is(err, ErrDoesNotExist) {
			return err
		}

		return nil
	}

	b := newTenantIndex(meta, compactedMeta)

	// Marshal and write the proto object.
	indexBytesPb, err := b.marshalPb()
	if err != nil {
		return err
	}

	err = w.w.Write(ctx, TenantIndexNamePb, KeyPath([]string{tenantID}), bytes.NewReader(indexBytesPb), int64(len(indexBytesPb)), nil)
	if err != nil {
		return err
	}

	// Marshal and write the JSON object.
	indexBytesJSON, err := b.marshal()
	if err != nil {
		return err
	}

	return w.w.Write(ctx, TenantIndexName, KeyPath([]string{tenantID}), bytes.NewReader(indexBytesJSON), int64(len(indexBytesJSON)), nil)
}

// Delete implements backend.Writer
func (w *writer) Delete(ctx context.Context, name string, keypath KeyPath) error {
	return w.w.Delete(ctx, name, keypath, nil)
}

// WriteNoCompactFlag implements backend.Writer
func (w *writer) WriteNoCompactFlag(ctx context.Context, blockID uuid.UUID, tenantID string) error {
	return w.w.Write(ctx, NoCompactFileName, KeyPathForBlock(blockID, tenantID), bytes.NewReader([]byte{}), 0, nil)
}

// DeleteNoCompactFlag implements backend.Writer
func (w *writer) DeleteNoCompactFlag(ctx context.Context, blockID uuid.UUID, tenantID string) error {
	return w.w.Delete(ctx, NoCompactFileName, KeyPathForBlock(blockID, tenantID), nil)
}

type reader struct {
	r RawReader
}

// NewReader returns an object that implements Reader and bridges to a RawReader
func NewReader(r RawReader) Reader {
	return &reader{
		r: r,
	}
}

// Read implements backend.Reader
func (r *reader) Read(ctx context.Context, name string, blockID uuid.UUID, tenantID string, cacheInfo *CacheInfo) ([]byte, error) {
	objReader, size, err := r.r.Read(ctx, name, KeyPathForBlock(blockID, tenantID), cacheInfo)
	if err != nil {
		return nil, err
	}
	defer objReader.Close()
	return tempo_io.ReadAllWithEstimate(objReader, size)
}

// StreamReader implements backend.Reader
func (r *reader) StreamReader(ctx context.Context, name string, blockID uuid.UUID, tenantID string) (io.ReadCloser, int64, error) {
	return r.r.Read(ctx, name, KeyPathForBlock(blockID, tenantID), nil)
}

// ReadRange implements backend.Reader
func (r *reader) ReadRange(ctx context.Context, name string, blockID uuid.UUID, tenantID string, offset uint64, buffer []byte, cacheInfo *CacheInfo) error {
	return r.r.ReadRange(ctx, name, KeyPathForBlock(blockID, tenantID), offset, buffer, cacheInfo)
}

// Tenants implements backend.Reader
func (r *reader) Tenants(ctx context.Context) ([]string, error) {
	list, err := r.r.List(ctx, nil)

	// this filter is added to fix a GCS usage stats issue that would result in ""
	var filteredList []string
	for _, tenant := range list {
		if tenant != "" && tenant != ClusterSeedFileName && tenant != WorkFileName {
			filteredList = append(filteredList, tenant)
		}
	}

	return filteredList, err
}

// Blocks implements backend.Reader
func (r *reader) Blocks(ctx context.Context, tenantID string) ([]uuid.UUID, []uuid.UUID, error) {
	return r.r.ListBlocks(ctx, tenantID)
}

// BlockMeta implements backend.Reader
func (r *reader) BlockMeta(ctx context.Context, blockID uuid.UUID, tenantID string) (*BlockMeta, error) {
	reader, size, err := r.r.Read(ctx, MetaName, KeyPathForBlock(blockID, tenantID), nil)
	if err != nil {
		return nil, err
	}
	defer reader.Close()

	bytes, err := tempo_io.ReadAllWithEstimate(reader, size)
	if err != nil {
		return nil, err
	}

	out := &BlockMeta{}
	err = json.Unmarshal(bytes, out)
	if err != nil {
		return nil, err
	}

	return out, nil
}

// TenantIndex implements backend.Reader
func (r *reader) TenantIndex(ctx context.Context, tenantID string) (*TenantIndex, error) {
	ctx, span := tracer.Start(ctx, "reader.TenantIndex")
	defer span.End()

	outPb, err := r.tenantIndexProto(ctx, tenantID)
	if err == nil {
		return outPb, nil
	}

	if !errors.Is(err, ErrDoesNotExist) {
		return nil, err
	}

	span.AddEvent(EventJSONFallback)

	readerJ, size, err := r.r.Read(ctx, TenantIndexName, KeyPath([]string{tenantID}), nil)
	if err != nil {
		return nil, err
	}
	defer readerJ.Close()

	bytesJ, err := tempo_io.ReadAllWithEstimate(readerJ, size)
	if err != nil {
		return nil, err
	}

	i := &TenantIndex{}
	err = i.unmarshal(bytesJ)
	if err != nil {
		return nil, err
	}

	return i, nil
}

func (r *reader) tenantIndexProto(ctx context.Context, tenantID string) (*TenantIndex, error) {
	readerPb, size, err := r.r.Read(ctx, TenantIndexNamePb, KeyPath([]string{tenantID}), nil)
	if err != nil {
		return nil, fmt.Errorf("failed to read tenant index proto: %w", err)
	}
	defer readerPb.Close()

	bytesPb, err := tempo_io.ReadAllWithEstimate(readerPb, size)
	if err != nil {
		return nil, fmt.Errorf("failed to read all with estimate: %w", err)
	}

	out := &TenantIndex{}
	err = out.unmarshalPb(bytesPb)
	if err != nil {
		return nil, fmt.Errorf("failed to unmarshal tenant index proto: %w", err)
	}

	return out, nil
}

// Find implements backend.Reader
func (r *reader) Find(ctx context.Context, keypath KeyPath, f FindFunc) error {
	return r.r.Find(ctx, keypath, f)
}

// Shutdown implements backend.Reader
func (r *reader) Shutdown() {
	r.r.Shutdown()
}

// HasNoCompactFlag implements backend.Reader
func (r *reader) HasNoCompactFlag(ctx context.Context, blockID uuid.UUID, tenantID string) (bool, error) {
	// TODO: can be optimized by not creating a descriptor and just checking file stat
	_, _, err := r.r.Read(ctx, NoCompactFileName, KeyPathForBlock(blockID, tenantID), nil)
	if err != nil {
		if errors.Is(err, ErrDoesNotExist) {
			return false, nil
		}
		return false, err
	}
	return true, nil
}

// KeyPathForBlock returns a correctly ordered keypath given a block id and tenantid
func KeyPathForBlock(blockID uuid.UUID, tenantID string) KeyPath {
	return []string{tenantID, blockID.String()}
}

// ObjectFileName returns a unique identifier for an object in object storage given its name and keypath
func ObjectFileName(keypath KeyPath, name string) string {
	return path.Join(path.Join(keypath...), name)
}

// KeyPathWithPrefix returns a keypath with a prefix
func KeyPathWithPrefix(keypath KeyPath, prefix string) KeyPath {
	if len(prefix) == 0 {
		return keypath
	}

	return append([]string{prefix}, keypath...)
}

// MetaFileName returns the object name for the block meta given a block id and tenantid
func MetaFileName(blockID uuid.UUID, tenantID, prefix string) string {
	return path.Join(prefix, tenantID, blockID.String(), MetaName)
}

// CompactedMetaFileName returns the object name for the compacted block meta given a block id and tenantid
func CompactedMetaFileName(blockID uuid.UUID, tenantID, prefix string) string {
	return path.Join(prefix, tenantID, blockID.String(), CompactedMetaName)
}

// RootPath returns the root path for a block given a block id and tenantid
func RootPath(blockID uuid.UUID, tenantID, prefix string) string {
	return path.Join(prefix, tenantID, blockID.String())
}