tempodb_segmentio.go

package tempodb

import (
	"fmt"
	"io"
	"math"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/grafana/tempo/pkg/tempopb"
	pq "github.com/segmentio/parquet-go"
)

type instrumentedFile struct {
	f          *os.File
	n          int64
	reads      int
	seeks      int
	writes     int
	writeBytes int64
}

var _ io.Seeker = (*instrumentedFile)(nil)
var _ io.ReaderAt = (*instrumentedFile)(nil)
var _ io.Writer = (*instrumentedFile)(nil)

func (i *instrumentedFile) Seek(offset int64, whence int) (int64, error) {
	i.seeks++
	return i.f.Seek(offset, whence)
}

func (i *instrumentedFile) ReadAt(b []byte, offset int64) (int, error) {
	//fmt.Printf("Read: %d len %d\n", offset, len(b))

	n, err := i.f.ReadAt(b, offset)
	i.n += int64(n)
	i.reads++
	return n, err
}

func (i *instrumentedFile) Write(p []byte) (n int, err error) {
	n, err = i.f.Write(p)
	i.writes++
	i.writeBytes += int64(n)
	return n, err
}

// bufferedReader implements io.ReaderAt but extends and buffers reads up the given buffer size.
// Subsequent reads are returned from the previously buffered data when possiblve. Additionally
// it supports concurrent readers by maintaining multiple buffers at different offsets, and
// matching up reads with existing buffers where possible. When needed the least-recently-used
// buffer is overwritten with new reads.
type bufferedReader struct {
	ra    io.ReaderAt
	rasz  int64
	rdsz  int64
	mtx   sync.Mutex
	count int64

	buffers []readerBuffer
}

type readerBuffer struct {
	buf   []byte
	off   int64
	count int64
}

var _ io.ReaderAt = (*bufferedReader)(nil)

func NewBufferedReaderAt(ra io.ReaderAt, readerSize int64, bufSize int64, bufCount int) *bufferedReader {
	r := &bufferedReader{
		ra:      ra,
		rasz:    readerSize,
		rdsz:    bufSize,
		buffers: make([]readerBuffer, bufCount),
	}

	return r
}

func (r *bufferedReader) canRead(buf *readerBuffer, offset, length int64) bool {
	return offset >= buf.off && (offset+length <= buf.off+int64(len(buf.buf)))
}

func (r *bufferedReader) read(buf *readerBuffer, b []byte, offset int64) {
	start := offset - buf.off
	copy(b, buf.buf[start:start+int64(len(b))])
}

func (r *bufferedReader) populate(buf *readerBuffer, offset, length int64) (int, error) {

	// Increase to minimim read size
	sz := r.rdsz
	if sz < length {
		sz = length
	}

	// Don't read past end of reader
	if offset+sz >= r.rasz {
		sz = r.rasz - offset
	}

	// Realloc?
	if int64(cap(buf.buf)) < sz {
		buf.buf = make([]byte, sz)
	}
	buf.buf = buf.buf[:sz]

	// Read
	buf.off = offset
	n, err := r.ra.ReadAt(buf.buf, offset)
	return n, err
}

func (r *bufferedReader) ReadAt(b []byte, offset int64) (int, error) {
	r.mtx.Lock()
	defer r.mtx.Unlock()

	// Least-recently-used tracking
	r.count++
	lruCnt := int64(math.MaxInt64)
	var lruIdx int

	for i := range r.buffers {
		buf := &r.buffers[i]
		if r.canRead(buf, offset, int64(len(b))) {
			r.read(buf, b, offset)
			r.buffers[i].count = r.count
			return len(b), nil
		}

		if r.buffers[i].count < lruCnt {
			lruCnt = r.buffers[i].count
			lruIdx = i
		}
	}

	// Need to read, overwrite least-recently-used
	buf := &r.buffers[lruIdx]
	if n, err := r.populate(buf, offset, int64(len(b))); err != nil {
		return n, err
	}

	r.buffers[lruIdx].count = r.count
	r.read(buf, b, offset)
	return len(b), nil
}

type bufferedWriter struct {
	w   io.Writer
	buf []byte
}

var _ io.WriteCloser = (*bufferedWriter)(nil)

func (b *bufferedWriter) Write(p []byte) (n int, err error) {
	b.buf = append(b.buf, p...)

	if len(b.buf) > 30_000_000 {
		n, err = b.w.Write(b.buf)
		b.buf = b.buf[:len(b.buf)-n]
		if err != nil {
			// What length makes sense here..?
			return 0, err
		}
	}

	return len(p), nil
}

func (b *bufferedWriter) Close() error {
	if len(b.buf) > 0 {
		_, err := b.w.Write(b.buf)
		b.buf = nil
		return err
	}
	return nil
}

func SearchParquetPlease(req *tempopb.SearchRequest) (*tempopb.SearchResponse, error) {
	path := "/Users/marty/src/tempo/tempodb/88d9f6c2-5f52-4fa6-b146-f3079553cb11-segment.parquet"

	ff, err := os.OpenFile(path, os.O_RDONLY, 0)
	if err != nil {
		return nil, err
	}

	fo, err := ff.Stat()
	if err != nil {
		return nil, err
	}

	f := &instrumentedFile{f: ff}

	pf, err := pq.OpenFile(f, fo.Size())
	if err != nil {
		return nil, err
	}

	results := searchParquetFile(pf, req, false)
	return results, nil
}

func getColumnIndexByPath(pf *pq.File, s string) (index, depth int) {
	colSelector := strings.Split(s, ".")
	n := pf.Root()
	for len(colSelector) > 0 {
		n = n.Column(colSelector[0])
		if n == nil {
			return -1, -1
		}

		colSelector = colSelector[1:]
		depth++
	}

	return int(n.Index()), depth
}

func makePipeline(req *tempopb.SearchRequest, pf *pq.File) (iterator, parquetSearchMetrics) {

	// Get list of row groups to inspect. Ideally we use predicate pushdown
	// here to keep only row groups that satisfy all known conditions, but
	// don't have it figured out yet.
	var rgs []pq.RowGroup
	for i := 0; i < pf.NumRowGroups(); i++ {
		rgs = append(rgs, pf.RowGroup(i))
	}

	makeIter := func(name string, predicate parquetPredicate, selectAs string) iterator {
		index, _ := getColumnIndexByPath(pf, name)
		if index == -1 {
			panic("column not found in parquet file:" + name)
		}
		return newColumnIterator(rgs, index, 1000, predicate, selectAs)
	}

	// Wire up iterators
	var resourceIters []iterator
	var traceIters []iterator

	otherAttrConditions := map[string]string{}

	for k, v := range req.Tags {
		switch k {
		case "cluster":
			resourceIters = append(resourceIters, makeIter("rs.Resource.Cluster", NewSubstringPredicate(v), ""))
		case "service.name":
			resourceIters = append(resourceIters, makeIter("rs.Resource.ServiceName", NewSubstringPredicate(v), ""))
		case "namespace":
			resourceIters = append(resourceIters, makeIter("rs.Resource.Namespace", NewSubstringPredicate(v), ""))
		case "pod":
			resourceIters = append(resourceIters, makeIter("rs.Resource.Pod", NewSubstringPredicate(v), ""))
		case "container":
			resourceIters = append(resourceIters, makeIter("rs.Resource.Container", NewSubstringPredicate(v), ""))
		case "name":
			resourceIters = append(resourceIters, makeIter("rs.ils.Spans.Name", NewSubstringPredicate(v), ""))
		default:
			otherAttrConditions[k] = v
		}
	}

	// Generic attribute conditions?
	if len(otherAttrConditions) > 0 {
		// We are looking for one or more foo=bar attributes that aren't
		// projected to their own columns, they are in the generic Key/Value
		// columns at the resource or span levels.  We want to search
		// both locations. But we also only want to read the columns once.

		var keys []string
		var vals []string
		for k, v := range otherAttrConditions {
			keys = append(keys, k)
			vals = append(vals, v)
		}

		keyPred := NewStringInPredicate(keys)
		valPred := NewStringInPredicate(vals)

		// This iterator combines the results from the resource
		// and span searches, and checks if all conditions were satisfied
		// on each ResourceSpans.  This is a single-pass over the attribute columns.
		j := NewUnionIterator(DefinitionLevelResourceSpans, []iterator{
			// This iterator finds all keys/values at the resource level
			NewJoinIterator(DefinitionLevelResourceAttrs, []iterator{
				makeIter("rs.Resource.Attrs.Key", keyPred, "keys"),
				makeIter("rs.Resource.Attrs.Value", valPred, "values"),
			}, nil),
			// This iterator finds all keys/values at the span level
			NewJoinIterator(DefinitionLevelResourceSpansILSSpan, []iterator{
				makeIter("rs.ils.Spans.Attrs.Key", keyPred, "keys"),
				makeIter("rs.ils.Spans.Attrs.Value", valPred, "values"),
			}, nil),
		}, NewKeyValueGroupPredicate(keys, vals))

		resourceIters = append(resourceIters, j)
	}

	// Multiple resource-level filters get joined and wrapped
	// up to trace-level. A single filter can be used as-is
	if len(resourceIters) == 1 {
		traceIters = append(traceIters, resourceIters[0])
	}
	if len(resourceIters) > 1 {
		traceIters = append(traceIters, NewJoinIterator(DefinitionLevelTrace, resourceIters, nil))
	}

	// We always pull back duration for the search results, but it also
	// has a predicate when bounded by the request
	var durFilter parquetPredicate
	if req.MinDurationMs > 0 {
		durFilter = NewIntGreaterThanOrEqualToPredicate((time.Millisecond * time.Duration(req.MinDurationMs)).Nanoseconds())
	}
	traceIters = append(traceIters, makeIter("DurationNanos", durFilter, "Duration"))

	// Join in values for search results. These have
	// no filters so they will always be in the results.
	traceIDMetrics := &instrumentedPredicate{}
	traceIters = append(traceIters, makeIter("TraceID", traceIDMetrics, "TraceID"))
	traceIters = append(traceIters, makeIter("RootServiceName", nil, "RootServiceName"))
	traceIters = append(traceIters, makeIter("RootSpanName", nil, "RootSpanName"))
	traceIters = append(traceIters, makeIter("StartTimeUnixNano", nil, "StartTime"))

	return NewJoinIterator(DefinitionLevelTrace, traceIters, nil), parquetSearchMetrics{
		pTraceID: traceIDMetrics,
	}
}

func makePipelineDynamic(req *tempopb.SearchRequest, pf *pq.File) (iterator, parquetSearchMetrics) {

	// Get list of row groups to inspect. Ideally we use predicate pushdown
	// here to keep only row groups that satisfy all known conditions, but
	// don't have it figured out yet.
	var rgs []pq.RowGroup
	for i := 0; i < pf.NumRowGroups(); i++ {
		rgs = append(rgs, pf.RowGroup(i))
	}

	makeIter := func(name string, predicate parquetPredicate, selectAs string) iterator {
		index, _ := getColumnIndexByPath(pf, name)
		if index == -1 {
			panic("column not found in parquet file:" + name)
		}
		return newColumnIterator(rgs, index, 1000, predicate, selectAs)
	}

	// Wire up iterators
	var resourceIters []iterator
	var traceIters []iterator

	for k, v := range req.Tags {
		colName := normalizeAttributeName(k)
		resourceCol, _ := getColumnIndexByPath(pf, fmt.Sprintf("rs.resource.%s", colName))
		var resourceIter iterator
		if resourceCol >= 0 {
			resourceIter = newColumnIterator(rgs, resourceCol, 1000, NewSubstringPredicate(v), "")
		}

		var spanIter iterator
		spanCol, _ := getColumnIndexByPath(pf, fmt.Sprintf("rs.ils.spans.%s", colName))
		if spanCol >= 0 {
			spanIter = newColumnIterator(rgs, spanCol, 1000, NewSubstringPredicate(v), "")
		}

		if resourceIter != nil && spanIter != nil {
			// Search both resource and span attributes for the value
			// Union them together means it can match in either location
			resourceIters = append(resourceIters, NewUnionIterator(DefinitionLevelResourceSpans, []iterator{
				resourceIter,
				spanIter,
			}, nil))
		} else if resourceIter != nil {
			resourceIters = append(resourceIters, resourceIter)
		} else if spanIter != nil {
			resourceIters = append(resourceIters, spanIter)
		} else {
			// No dedicated column for this key.
			// Ignore for now, data isn't searchable?
			return nil, parquetSearchMetrics{}
		}
	}
	// Multiple resource-level filters get joined and wrapped
	// up to trace-level. A single filter can be used as-is
	if len(resourceIters) == 1 {
		traceIters = append(traceIters, resourceIters[0])
	}
	if len(resourceIters) > 1 {
		traceIters = append(traceIters, NewJoinIterator(DefinitionLevelTrace, resourceIters, nil))
	}

	// We always pull back duration for the search results, but it also
	// has a predicate when bounded by the request
	var durFilter parquetPredicate
	if req.MinDurationMs > 0 {
		durFilter = NewIntGreaterThanOrEqualToPredicate((time.Millisecond * time.Duration(req.MinDurationMs)).Nanoseconds())
	}
	traceIters = append(traceIters, makeIter("DurationNanos", durFilter, "Duration"))

	// Join in values for search results. These have
	// no filters so they will always be in the results.
	traceIDMetrics := &instrumentedPredicate{}
	traceIters = append(traceIters, makeIter("TraceID", traceIDMetrics, "TraceID"))
	traceIters = append(traceIters, makeIter("RootServiceName", nil, "RootServiceName"))
	traceIters = append(traceIters, makeIter("RootSpanName", nil, "RootSpanName"))
	traceIters = append(traceIters, makeIter("StartTimeUnixNano", nil, "StartTime"))

	return NewJoinIterator(DefinitionLevelTrace, traceIters, nil), parquetSearchMetrics{
		pTraceID: traceIDMetrics,
	}
}

func searchParquetFile(pf *pq.File, req *tempopb.SearchRequest, dynamic bool) *tempopb.SearchResponse {
	results := []*tempopb.TraceSearchMetadata{}

	var iter iterator
	var metrics parquetSearchMetrics
	if dynamic {
		iter, metrics = makePipelineDynamic(req, pf)
	} else {
		iter, metrics = makePipeline(req, pf)
	}
	if iter == nil {
		panic("make pipeline failed")
	}
	defer iter.Close()

	for {
		match := iter.Next()
		if match == nil {
			break
		}

		result := &tempopb.TraceSearchMetadata{
			TraceID:           match.value["TraceID"][0].String(),
			RootServiceName:   match.value["RootServiceName"][0].String(),
			RootTraceName:     match.value["RootSpanName"][0].String(),
			StartTimeUnixNano: uint64(match.value["StartTime"][0].Int64()),
			DurationMs:        uint32(match.value["Duration"][0].Int64() / int64(time.Millisecond)),
		}
		results = append(results, result)

		if req.Limit > 0 && len(results) >= int(req.Limit) {
			break
		}
	}

	return &tempopb.SearchResponse{
		Traces:  results,
		Metrics: metrics.ToProto(),
	}
}

type parquetSearchMetrics struct {
	pTraceID *instrumentedPredicate
}

func (p *parquetSearchMetrics) ToProto() *tempopb.SearchMetrics {
	return &tempopb.SearchMetrics{
		InspectedTraces: uint32(p.pTraceID.inspectedValues),
	}
}

func normalizeAttributeName(val string) string {
	val = strings.ReplaceAll(val, ".", "_")
	val = strings.ReplaceAll(val, "-", "_")
	val = strings.ReplaceAll(val, "/", "_")
	val = strings.ToUpper(val)
	return val
}