opentelemetry-go/oteltest/span.go

// Copyright The OpenTelemetry Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package oteltest

import (
	"context"
	"fmt"
	"reflect"
	"sync"
	"time"

	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/label"
)

const (
	errorTypeKey    = label.Key("error.type")
	errorMessageKey = label.Key("error.message")
	errorEventName  = "error"
)

var _ otel.Span = (*Span)(nil)

// Span is an OpenTelemetry Span used for testing.
type Span struct {
	lock          sync.RWMutex
	tracer        *Tracer
	spanContext   otel.SpanContext
	parentSpanID  otel.SpanID
	ended         bool
	name          string
	startTime     time.Time
	endTime       time.Time
	statusCode    codes.Code
	statusMessage string
	attributes    map[label.Key]label.Value
	events        []Event
	links         map[otel.SpanContext][]label.KeyValue
	spanKind      otel.SpanKind
}

// Tracer returns the Tracer that created s.
func (s *Span) Tracer() otel.Tracer {
	return s.tracer
}

// End ends s. If the Tracer that created s was configured with a
// SpanRecorder, that recorder's OnEnd method is called as the final part of
// this method.
func (s *Span) End(opts ...otel.SpanOption) {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.ended {
		return
	}

	c := otel.NewSpanConfig(opts...)
	s.endTime = time.Now()
	if endTime := c.Timestamp; !endTime.IsZero() {
		s.endTime = endTime
	}

	s.ended = true
	if s.tracer.config.SpanRecorder != nil {
		s.tracer.config.SpanRecorder.OnEnd(s)
	}
}

// RecordError records an error as a Span event.
func (s *Span) RecordError(ctx context.Context, err error, opts ...otel.ErrorOption) {
	if err == nil || s.ended {
		return
	}

	cfg := otel.NewErrorConfig(opts...)

	if cfg.Timestamp.IsZero() {
		cfg.Timestamp = time.Now()
	}

	if cfg.StatusCode != codes.Unset {
		s.SetStatus(cfg.StatusCode, "")
	}

	errType := reflect.TypeOf(err)
	errTypeString := fmt.Sprintf("%s.%s", errType.PkgPath(), errType.Name())
	if errTypeString == "." {
		errTypeString = errType.String()
	}

	s.AddEventWithTimestamp(ctx, cfg.Timestamp, errorEventName,
		errorTypeKey.String(errTypeString),
		errorMessageKey.String(err.Error()),
	)
}

// AddEvent adds an event to s.
func (s *Span) AddEvent(ctx context.Context, name string, attrs ...label.KeyValue) {
	s.AddEventWithTimestamp(ctx, time.Now(), name, attrs...)
}

// AddEventWithTimestamp adds an event that occurred at timestamp to s.
func (s *Span) AddEventWithTimestamp(ctx context.Context, timestamp time.Time, name string, attrs ...label.KeyValue) {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.ended {
		return
	}

	attributes := make(map[label.Key]label.Value)

	for _, attr := range attrs {
		attributes[attr.Key] = attr.Value
	}

	s.events = append(s.events, Event{
		Timestamp:  timestamp,
		Name:       name,
		Attributes: attributes,
	})
}

// IsRecording returns the recording state of s.
func (s *Span) IsRecording() bool {
	return true
}

// SpanContext returns the SpanContext of s.
func (s *Span) SpanContext() otel.SpanContext {
	return s.spanContext
}

// SetStatus sets the status of s in the form of a code and a message.
func (s *Span) SetStatus(code codes.Code, msg string) {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.ended {
		return
	}

	s.statusCode = code
	s.statusMessage = msg
}

// SetName sets the name of s.
func (s *Span) SetName(name string) {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.ended {
		return
	}

	s.name = name
}

// SetAttributes sets attrs as attributes of s.
func (s *Span) SetAttributes(attrs ...label.KeyValue) {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.ended {
		return
	}

	for _, attr := range attrs {
		s.attributes[attr.Key] = attr.Value
	}
}

// Name returns the name most recently set on s, either at or after creation
// time. It cannot be change after End has been called on s.
func (s *Span) Name() string {
	return s.name
}

// ParentSpanID returns the SpanID of the parent Span. If s is a root Span,
// and therefore does not have a parent, the returned SpanID will be invalid
// (i.e., it will contain all zeroes).
func (s *Span) ParentSpanID() otel.SpanID {
	return s.parentSpanID
}

// Attributes returns the attributes set on s, either at or after creation
// time. If the same attribute key was set multiple times, the last call will
// be used. Attributes cannot be changed after End has been called on s.
func (s *Span) Attributes() map[label.Key]label.Value {
	s.lock.RLock()
	defer s.lock.RUnlock()

	attributes := make(map[label.Key]label.Value)

	for k, v := range s.attributes {
		attributes[k] = v
	}

	return attributes
}

// Events returns the events set on s. Events cannot be changed after End has
// been called on s.
func (s *Span) Events() []Event {
	return s.events
}

// Links returns the links set on s at creation time. If multiple links for
// the same SpanContext were set, the last link will be used.
func (s *Span) Links() map[otel.SpanContext][]label.KeyValue {
	links := make(map[otel.SpanContext][]label.KeyValue)

	for sc, attributes := range s.links {
		links[sc] = append([]label.KeyValue{}, attributes...)
	}

	return links
}

// StartTime returns the time at which s was started. This will be the
// wall-clock time unless a specific start time was provided.
func (s *Span) StartTime() time.Time {
	return s.startTime
}

// EndTime returns the time at which s was ended if at has been ended, or
// false otherwise. If the span has been ended, the returned time will be the
// wall-clock time unless a specific end time was provided.
func (s *Span) EndTime() (time.Time, bool) {
	return s.endTime, s.ended
}

// Ended returns whether s has been ended, i.e. whether End has been called at
// least once on s.
func (s *Span) Ended() bool {
	return s.ended
}

// StatusCode returns the code of the status most recently set on s, or
// codes.OK if no status has been explicitly set. It cannot be changed after
// End has been called on s.
func (s *Span) StatusCode() codes.Code {
	return s.statusCode
}

// StatusMessage returns the status message most recently set on s or the
// empty string if no status message was set.
func (s *Span) StatusMessage() string {
	return s.statusMessage
}

// SpanKind returns the span kind of s.
func (s *Span) SpanKind() otel.SpanKind {
	return s.spanKind
}