crdbspan.go

// Copyright 2021 The Cockroach Authors.
//
// Use of this software is governed by the Business Source License
// included in the file licenses/BSL.txt.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0, included in the file
// licenses/APL.txt.

package tracing

import (
	"fmt"
	"sort"
	"sync"
	"sync/atomic"
	"time"

	"github.com/cockroachdb/cockroach/pkg/util/ring"
	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
	"github.com/cockroachdb/cockroach/pkg/util/tracing/tracingpb"
	"github.com/cockroachdb/logtags"
	"github.com/gogo/protobuf/types"
	"github.com/opentracing/opentracing-go"
)

// crdbSpan is a span for internal crdb usage. This is used to power SQL session
// tracing.
type crdbSpan struct {
	traceID      uint64 // probabilistically unique
	spanID       uint64 // probabilistically unique
	parentSpanID uint64
	goroutineID  uint64

	operation string
	startTime time.Time

	// logTags are set to the log tags that were available when this Span was
	// created, so that there's no need to eagerly copy all of those log tags
	// into this Span's tags. If the Span's tags are actually requested, these
	// logTags will be copied out at that point.
	//
	// Note that these tags have not gone through the log tag -> Span tag
	// remapping procedure; tagName() needs to be called before exposing each
	// tag's key to a user.
	logTags *logtags.Buffer

	mu crdbSpanMu
}

type crdbSpanMu struct {
	syncutil.Mutex
	// duration is initialized to -1 and set on Finish().
	duration time.Duration

	recording struct {
		// recordingType is the recording type of the ongoing recording, if any.
		// Its 'load' method may be called without holding the surrounding mutex,
		// but its 'swap' method requires the mutex.
		recordingType atomicRecordingType

		logBytes int64
		logs     ring.Buffer // of tracingpb.LogRecords

		structuredBytes int64
		structured      ring.Buffer // of Structured events

		// dropped is true if the span has capped out it's memory limits for
		// logs and structured events, and has had to drop some.
		dropped bool

		// children contains the list of child spans started after this Span
		// started recording.
		children []*crdbSpan
		// remoteSpan contains the list of remote child span recordings that
		// were manually imported.
		remoteSpans []tracingpb.RecordedSpan
	}

	// tags are only set when recording. These are tags that have been added to
	// this Span, and will be appended to the tags in logTags when someone
	// needs to actually observe the total set of tags that is a part of this
	// Span.
	// TODO(radu): perhaps we want a recording to capture all the tags (even
	// those that were set before recording started)?
	tags opentracing.Tags

	// The Span's associated baggage.
	baggage map[string]string
}

func (s *crdbSpan) recordingType() RecordingType {
	if s == nil {
		return RecordingOff
	}
	return s.mu.recording.recordingType.load()
}

// enableRecording start recording on the Span. From now on, log events and
// child spans will be stored.
//
// If parent != nil, the Span will be registered as a child of the respective
// parent. If nil, the parent's recording will not include this child.
func (s *crdbSpan) enableRecording(parent *crdbSpan, recType RecordingType) {
	if parent != nil {
		parent.addChild(s)
	}
	if recType == RecordingOff {
		return
	}

	s.mu.Lock()
	defer s.mu.Unlock()
	s.mu.recording.recordingType.swap(recType)
	if recType == RecordingVerbose {
		s.setBaggageItemLocked(verboseTracingBaggageKey, "1")
	}
}

// resetRecording clears any previously recorded info.
//
// NB: This is needed by SQL SessionTracing, who likes to start and stop
// recording repeatedly on the same Span, and collect the (separate) recordings
// every time.
func (s *crdbSpan) resetRecording() {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.mu.recording.logs.Reset()
	s.mu.recording.logBytes = 0
	s.mu.recording.structured.Reset()
	s.mu.recording.structuredBytes = 0
	s.mu.recording.dropped = false

	s.mu.recording.children = nil
	s.mu.recording.remoteSpans = nil
}

func (s *crdbSpan) disableRecording() {
	s.mu.Lock()
	defer s.mu.Unlock()
	oldRecType := s.mu.recording.recordingType.swap(RecordingOff)
	// We test the duration as a way to check if the Span has been finished. If it
	// has, we don't want to do the call below as it might crash (at least if
	// there's a netTr).
	if (s.mu.duration == -1) && (oldRecType == RecordingVerbose) {
		// Clear the verboseTracingBaggageKey baggage item, assuming that it was set by
		// enableRecording().
		s.setBaggageItemLocked(verboseTracingBaggageKey, "")
	}
}

func (s *crdbSpan) getRecording(everyoneIsV211 bool, wantTags bool) Recording {
	if s == nil {
		return nil // noop span
	}

	s.mu.Lock()

	if !everyoneIsV211 {
		// The cluster may contain nodes that are running v20.2. Unfortunately that
		// version can easily crash when a peer returns a recording that that node
		// did not expect would get created. To circumvent this, retain the v20.2
		// behavior of eliding recordings when verbosity is off until we're sure
		// that v20.2 is not around any longer.
		//
		// TODO(tbg): remove this in the v21.2 cycle.
		if s.recordingType() == RecordingOff {
			s.mu.Unlock()
			return nil
		}
	}

	// The capacity here is approximate since we don't know how many grandchildren
	// there are.
	result := make(Recording, 0, 1+len(s.mu.recording.children)+len(s.mu.recording.remoteSpans))
	// Shallow-copy the children so we can process them without the lock.
	children := s.mu.recording.children
	result = append(result, s.getRecordingLocked(wantTags))
	result = append(result, s.mu.recording.remoteSpans...)
	s.mu.Unlock()

	for _, child := range children {
		result = append(result, child.getRecording(everyoneIsV211, wantTags)...)
	}

	// Sort the spans by StartTime, except the first Span (the root of this
	// recording) which stays in place.
	toSort := sortPool.Get().(*Recording) // avoids allocations in sort.Sort
	*toSort = result[1:]
	sort.Sort(toSort)
	*toSort = nil
	sortPool.Put(toSort)
	return result
}

func (s *crdbSpan) importRemoteSpans(remoteSpans []tracingpb.RecordedSpan) error {
	// Change the root of the remote recording to be a child of this Span. This is
	// usually already the case, except with DistSQL traces where remote
	// processors run in spans that FollowFrom an RPC Span that we don't collect.
	remoteSpans[0].ParentSpanID = s.spanID

	s.mu.Lock()
	s.mu.recording.remoteSpans = append(s.mu.recording.remoteSpans, remoteSpans...)
	s.mu.Unlock()
	return nil
}

func (s *crdbSpan) setTagLocked(key string, value interface{}) {
	if s.mu.tags == nil {
		s.mu.tags = make(opentracing.Tags)
	}
	s.mu.tags[key] = value
}

func (s *crdbSpan) record(msg string) {
	if s.recordingType() != RecordingVerbose {
		return
	}

	logRecord := tracingpb.LogRecord{
		Time: time.Now(),
		Fields: []tracingpb.LogRecord_Field{
			{Key: tracingpb.LogMessageField, Value: msg},
		},
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	s.mu.recording.logBytes += int64(logRecord.Size())
	if s.mu.recording.logBytes > maxLogBytesPerSpan {
		s.mu.recording.dropped = true
	}
	for s.mu.recording.logBytes > maxLogBytesPerSpan {
		first := s.mu.recording.logs.GetFirst().(tracingpb.LogRecord)
		s.mu.recording.logs.RemoveFirst()
		s.mu.recording.logBytes -= int64(first.Size())
	}
	s.mu.recording.logs.AddLast(logRecord)
}

func (s *crdbSpan) recordStructured(item Structured) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.mu.recording.structuredBytes += int64(item.Size())
	if s.mu.recording.structuredBytes > maxStructuredBytesPerSpan {
		s.mu.recording.dropped = true
	}
	for s.mu.recording.structuredBytes > maxStructuredBytesPerSpan {
		last := s.mu.recording.structured.GetLast().(Structured)
		s.mu.recording.structured.RemoveLast()
		s.mu.recording.structuredBytes -= int64(last.Size())
	}

	s.mu.recording.structured.AddFirst(item)
}

func (s *crdbSpan) setBaggageItemAndTag(restrictedKey, value string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.setBaggageItemLocked(restrictedKey, value)
	// Don't set the tag if this is the special cased baggage item indicating
	// span verbosity, as it is named nondescriptly and the recording knows
	// how to display its verbosity independently.
	if restrictedKey != verboseTracingBaggageKey {
		s.setTagLocked(restrictedKey, value)
	}
}

func (s *crdbSpan) setBaggageItemLocked(restrictedKey, value string) {
	if oldVal, ok := s.mu.baggage[restrictedKey]; ok && oldVal == value {
		// No-op.
		return
	}
	if s.mu.baggage == nil {
		s.mu.baggage = make(map[string]string)
	}
	s.mu.baggage[restrictedKey] = value
}

// getRecordingLocked returns the Span's recording. This does not include
// children.
//
// When wantTags is false, no tags will be added. This is a performance
// optimization as stringifying the tag values can be expensive.
func (s *crdbSpan) getRecordingLocked(wantTags bool) tracingpb.RecordedSpan {
	rs := tracingpb.RecordedSpan{
		TraceID:      s.traceID,
		SpanID:       s.spanID,
		ParentSpanID: s.parentSpanID,
		GoroutineID:  s.goroutineID,
		Operation:    s.operation,
		StartTime:    s.startTime,
		Duration:     s.mu.duration,
	}

	if rs.Duration == -1 {
		// -1 indicates an unfinished Span. For a recording it's better to put some
		// duration in it, otherwise tools get confused. For example, we export
		// recordings to Jaeger, and spans with a zero duration don't look nice.
		rs.Duration = timeutil.Now().Sub(rs.StartTime)
		rs.Finished = false
	} else {
		rs.Finished = true
	}

	addTag := func(k, v string) {
		if rs.Tags == nil {
			rs.Tags = make(map[string]string)
		}
		rs.Tags[k] = v
	}

	if wantTags {
		if s.mu.duration == -1 {
			addTag("_unfinished", "1")
		}
		if s.mu.recording.recordingType.load() == RecordingVerbose {
			addTag("_verbose", "1")
		}
		if s.mu.recording.dropped {
			addTag("_dropped", "1")
		}
	}

	if numEvents := s.mu.recording.structured.Len(); numEvents != 0 {
		rs.InternalStructured = make([]*types.Any, 0, numEvents)
		for i := 0; i < numEvents; i++ {
			event := s.mu.recording.structured.Get(i).(Structured)
			item, err := types.MarshalAny(event)
			if err != nil {
				// An error here is an error from Marshal; these
				// are unlikely to happen.
				continue
			}
			rs.InternalStructured = append(rs.InternalStructured, item)
		}
	}

	if len(s.mu.baggage) > 0 {
		rs.Baggage = make(map[string]string)
		for k, v := range s.mu.baggage {
			rs.Baggage[k] = v
		}
	}
	if wantTags {
		if s.logTags != nil {
			setLogTags(s.logTags.Get(), func(remappedKey string, tag *logtags.Tag) {
				addTag(remappedKey, tag.ValueStr())
			})
		}
		if len(s.mu.tags) > 0 {
			for k, v := range s.mu.tags {
				// We encode the tag values as strings.
				addTag(k, fmt.Sprint(v))
			}
		}
	}

	if numLogs := s.mu.recording.logs.Len(); numLogs != 0 {
		rs.Logs = make([]tracingpb.LogRecord, numLogs)
		for i := 0; i < numLogs; i++ {
			rs.Logs[i] = s.mu.recording.logs.Get(i).(tracingpb.LogRecord)
		}
	}

	return rs
}

func (s *crdbSpan) addChild(child *crdbSpan) {
	s.mu.Lock()
	// Only record the child if the parent still has room.
	if len(s.mu.recording.children) < maxChildrenPerSpan {
		s.mu.recording.children = append(s.mu.recording.children, child)
	}
	s.mu.Unlock()
}

var sortPool = sync.Pool{
	New: func() interface{} {
		return &Recording{}
	},
}

// Less implements sort.Interface.
func (r Recording) Less(i, j int) bool {
	return r[i].StartTime.Before(r[j].StartTime)
}

// Swap implements sort.Interface.
func (r Recording) Swap(i, j int) {
	r[i], r[j] = r[j], r[i]
}

// Len implements sort.Interface.
func (r Recording) Len() int {
	return len(r)
}

type atomicRecordingType RecordingType

// load returns the recording type.
func (art *atomicRecordingType) load() RecordingType {
	return RecordingType(atomic.LoadInt32((*int32)(art)))
}

// swap stores the new recording type and returns the old one.
func (art *atomicRecordingType) swap(recType RecordingType) RecordingType {
	return RecordingType(atomic.SwapInt32((*int32)(art), int32(recType)))
}