traces/modules/distributor/receiver/shim.go

package receiver

import (
	"context"
	"fmt"
	"os"
	"time"

	"github.com/go-kit/log/level"
	dslog "github.com/grafana/dskit/log"
	"github.com/grafana/dskit/services"
	zaplogfmt "github.com/jsternberg/zap-logfmt"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/jaegerreceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/kafkareceiver"
	"github.com/open-telemetry/opentelemetry-collector-contrib/receiver/zipkinreceiver"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"go.opentelemetry.io/collector/component"
	"go.opentelemetry.io/collector/confmap"
	"go.opentelemetry.io/collector/consumer"
	"go.opentelemetry.io/collector/exporter"
	"go.opentelemetry.io/collector/exporter/exportertest"
	"go.opentelemetry.io/collector/otelcol"
	"go.opentelemetry.io/collector/pdata/ptrace"
	"go.opentelemetry.io/collector/receiver"
	"go.opentelemetry.io/collector/receiver/otlpreceiver"
	"go.opentelemetry.io/otel"
	otelcodes "go.opentelemetry.io/otel/codes"
	"go.uber.org/multierr"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"google.golang.org/genproto/googleapis/rpc/errdetails"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/types/known/durationpb"

	"github.com/grafana/tempo/pkg/tempopb"
	"github.com/grafana/tempo/pkg/usagestats"
	"github.com/grafana/tempo/pkg/util/log"
)

const (
	logsPerSecond = 10
)

var (
	metricPushDuration = promauto.NewHistogram(prometheus.HistogramOpts{
		Namespace:                       "tempo",
		Name:                            "distributor_push_duration_seconds",
		Help:                            "Records the amount of time to push a batch to the ingester.",
		Buckets:                         prometheus.DefBuckets,
		NativeHistogramBucketFactor:     1.1,
		NativeHistogramMaxBucketNumber:  100,
		NativeHistogramMinResetDuration: 1 * time.Hour,
	})

	statReceiverOtlp   = usagestats.NewInt("receiver_enabled_otlp")
	statReceiverJaeger = usagestats.NewInt("receiver_enabled_jaeger")
	statReceiverZipkin = usagestats.NewInt("receiver_enabled_zipkin")
	statReceiverKafka  = usagestats.NewInt("receiver_enabled_kafka")
)

var tracer = otel.Tracer("modules/distributor/receiver")

type RetryableError struct {
	err error
	st  *status.Status
}

func (r *RetryableError) GRPCStatus() *status.Status {
	return r.st
}

func (r *RetryableError) Error() string {
	return r.err.Error()
}

// wrapErrorIfRetryable wraps the passed in error to meet expectations of the otel collector exporter code:
// https://github.com/open-telemetry/opentelemetry-collector/blob/d7b49df5d9e922df6ce56ad4b64ee1c79f9dbdbe/exporter/otlpexporter/otlp.go#L172
// The otel collector considers some errors retryable and other not. "ResourceExhausted" is special in that it requires a
// RetryInfo detail along with the error code
func wrapErrorIfRetryable(err error, dur *durationpb.Duration, enabled bool) error {
	if dur == nil {
		return err
	}

	s, ok := status.FromError(err)
	if !ok {
		return err
	}

	if s.Code() != codes.ResourceExhausted {
		return err
	}

	// If retryInfo is NOT enabled, return error as is.
	if !enabled {
		return err
	}

	// ignore error. code only errors if Code() == ok
	s, _ = s.WithDetails(&errdetails.RetryInfo{
		RetryDelay: dur,
	})

	return &RetryableError{
		err: err,
		st:  s,
	}
}

type TracesPusher interface {
	PushTraces(ctx context.Context, traces ptrace.Traces) (*tempopb.PushResponse, error)
	RetryInfoEnabled(ctx context.Context) (bool, error)
}

var _ services.Service = (*receiversShim)(nil)

type receiversShim struct {
	services.Service

	retryDelay *durationpb.Duration
	receivers  []receiver.Traces
	pusher     TracesPusher
	logger     *log.RateLimitedLogger
	fatal      chan error
}

var (
	_ consumer.Traces = (*receiversShim)(nil)
	_ component.Host  = (*receiversShim)(nil)
)

func (r *receiversShim) Capabilities() consumer.Capabilities {
	return consumer.Capabilities{MutatesData: false}
}

var _ confmap.Provider = (*mapProvider)(nil)

// mapProvider is a confmap.Provider that returns a single confmap.Retrieved instance with a fixed map.
type mapProvider struct {
	raw map[string]interface{}
}

func (m *mapProvider) Retrieve(context.Context, string, confmap.WatcherFunc) (*confmap.Retrieved, error) {
	return confmap.NewRetrieved(m.raw, []confmap.RetrievedOption{}...)
}

func (m *mapProvider) Scheme() string { return "mock" }

func (m *mapProvider) Shutdown(context.Context) error { return nil }

func New(receiverCfg map[string]interface{}, pusher TracesPusher, middleware Middleware, retryAfterDuration time.Duration, logLevel dslog.Level, reg prometheus.Registerer) (services.Service, error) {
	shim := &receiversShim{
		pusher: pusher,
		logger: log.NewRateLimitedLogger(logsPerSecond, level.Error(log.Logger)),
		fatal:  make(chan error),
	}

	if retryAfterDuration > 0 {
		shim.retryDelay = durationpb.New(retryAfterDuration)
	}

	// shim otel observability
	zapLogger := newLogger(logLevel)

	// load config
	receiverFactories, err := otelcol.MakeFactoryMap(
		jaegerreceiver.NewFactory(),
		zipkinreceiver.NewFactory(),
		otlpreceiver.NewFactory(),
		kafkareceiver.NewFactory(),
	)
	if err != nil {
		return nil, err
	}

	for recv := range receiverCfg {
		switch recv {
		case "otlp":
			statReceiverOtlp.Set(1)
		case "jaeger":
			statReceiverJaeger.Set(1)
		case "zipkin":
			statReceiverZipkin.Set(1)
		case "kafka":
			statReceiverKafka.Set(1)
		}
	}

	receivers := make([]string, 0, len(receiverCfg))
	for k := range receiverCfg {
		receivers = append(receivers, k)
	}

	// Define a factory function to create the mock provider
	mockProviderFactory := confmap.NewProviderFactory(func(confmap.ProviderSettings) confmap.Provider {
		return &mapProvider{
			raw: map[string]interface{}{
				"receivers": receiverCfg,
				"exporters": map[string]interface{}{
					"nop": map[string]interface{}{},
				},
				"service": map[string]interface{}{
					"pipelines": map[string]interface{}{
						"traces": map[string]interface{}{
							"exporters": []string{"nop"}, // nop exporter to avoid errors
							"receivers": receivers,
						},
					},
				},
			},
		}
	})

	// Creates a config provider with the given config map.
	// The provider will be used to retrieve the actual config for the pipeline (although we only need the receivers).
	pro, err := otelcol.NewConfigProvider(otelcol.ConfigProviderSettings{
		ResolverSettings: confmap.ResolverSettings{
			URIs: []string{"mock:/"},
			ProviderFactories: []confmap.ProviderFactory{
				mockProviderFactory,
			},
			DefaultScheme: "mock",
		},
	})
	if err != nil {
		return nil, err
	}

	// Creates the configuration for the pipeline.
	// We only need the receivers, the rest of the configuration is not used.
	conf, err := pro.Get(context.Background(), otelcol.Factories{
		Receivers: receiverFactories,
		Exporters: map[component.Type]exporter.Factory{component.MustNewType("nop"): exportertest.NewNopFactory()}, // nop exporter to avoid errors
	})
	if err != nil {
		return nil, err
	}

	telemetrySettings := component.TelemetrySettings{
		Logger:         zapLogger,
		TracerProvider: otel.GetTracerProvider(),
		MeterProvider:  NewMeterProvider(reg),
	}
	// todo: propagate a real context?  translate our log configuration into zap?
	ctx := context.Background()
	for componentID, cfg := range conf.Receivers {

		factoryBase := receiverFactories[componentID.Type()]
		if factoryBase == nil {
			return nil, fmt.Errorf("receiver factory not found for type: %s", componentID.Type())
		}

		// Make sure that the headers are added to context. Required for Authentication.
		switch componentID.Type().String() {
		case "otlp":
			otlpRecvCfg := cfg.(*otlpreceiver.Config)

			if otlpRecvCfg.HTTP.HasValue() {
				otlpRecvCfg.HTTP.Get().ServerConfig.IncludeMetadata = true
				cfg = otlpRecvCfg
			}

		case "zipkin":
			zipkinRecvCfg := cfg.(*zipkinreceiver.Config)

			zipkinRecvCfg.IncludeMetadata = true
			cfg = zipkinRecvCfg

		case "jaeger":
			jaegerRecvCfg := cfg.(*jaegerreceiver.Config)

			if jaegerRecvCfg.ThriftHTTP.HasValue() {
				jaegerRecvCfg.ThriftHTTP.Get().IncludeMetadata = true
			}

			cfg = jaegerRecvCfg
		}

		params := receiver.Settings{
			ID:                component.NewIDWithName(componentID.Type(), fmt.Sprintf("%s_receiver", componentID.Type().String())),
			TelemetrySettings: telemetrySettings,
		}
		receiver, err := factoryBase.CreateTraces(ctx, params, cfg, middleware.Wrap(shim))
		if err != nil {
			return nil, err
		}

		shim.receivers = append(shim.receivers, receiver)
	}

	shim.Service = services.NewBasicService(shim.starting, shim.running, shim.stopping)

	return shim, nil
}

func (r *receiversShim) running(ctx context.Context) error {
	select {
	case err := <-r.fatal:
		return err
	case <-ctx.Done():
		return ctx.Err()
	}
}

func (r *receiversShim) starting(ctx context.Context) error {
	for _, receiver := range r.receivers {
		err := receiver.Start(ctx, r)
		if err != nil {
			return fmt.Errorf("error starting receiver: %w", err)
		}
	}

	return nil
}

// Called after distributor is asked to stop via StopAsync.
func (r *receiversShim) stopping(_ error) error {
	ctx, cancelFn := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancelFn()

	errs := make([]error, 0)

	for _, receiver := range r.receivers {
		err := receiver.Shutdown(ctx)
		if err != nil {
			errs = append(errs, err)
		}
	}

	if len(errs) > 0 {
		return multierr.Combine(errs...)
	}

	return nil
}

// ConsumeTraces implements consumer.Traces
func (r *receiversShim) ConsumeTraces(ctx context.Context, td ptrace.Traces) error {
	ctx, span := tracer.Start(ctx, "distributor.ConsumeTraces")
	defer span.End()

	var err error

	start := time.Now()
	_, err = r.pusher.PushTraces(ctx, td)
	metricPushDuration.Observe(time.Since(start).Seconds())
	if err != nil {
		r.logger.Log("msg", "pusher failed to consume trace data", "err", err)
		span.SetStatus(otelcodes.Error, err.Error())
		retryInfoEnabled, e := r.pusher.RetryInfoEnabled(ctx)
		if e != nil {
			return e
		}
		err = wrapErrorIfRetryable(err, r.retryDelay, retryInfoEnabled)
	}

	return err
}

// GetExtensions implements component.Host
func (r *receiversShim) GetExtensions() map[component.ID]component.Component { return nil }

// observability shims
func newLogger(level dslog.Level) *zap.Logger {
	zapLevel := zapcore.InfoLevel

	switch level.String() {
	case "error":
		zapLevel = zapcore.ErrorLevel
	case "warn":
		zapLevel = zapcore.WarnLevel
	case "info":
		zapLevel = zapcore.InfoLevel
	case "debug":
		zapLevel = zapcore.DebugLevel
	}

	config := zap.NewProductionEncoderConfig()
	config.EncodeTime = func(ts time.Time, encoder zapcore.PrimitiveArrayEncoder) {
		encoder.AppendString(ts.UTC().Format(time.RFC3339))
	}
	logger := zap.New(zapcore.NewCore(
		zaplogfmt.NewEncoder(config),
		os.Stdout,
		zapLevel,
	))
	logger = logger.With(zap.String("component", "tempo"))
	logger.Info("OTel Shim Logger Initialized")

	return logger
}