observer/pkg/tracer/middleware.go

package tracer

import (
	"context"
	"log"
	gelfexporter "maal/tracer/pkg/gelf_exporter"
	"os"

	"github.com/gofiber/fiber/v2"
	fiberOpentelemetry "github.com/psmarcin/fiber-opentelemetry/pkg/fiber-otel"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/exporters/otlp/otlptrace"
	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp"
	"go.opentelemetry.io/otel/exporters/stdout/stdouttrace"
	"go.opentelemetry.io/otel/sdk/resource"
	trc "go.opentelemetry.io/otel/sdk/trace"
	semconv "go.opentelemetry.io/otel/semconv/v1.4.0"
	trace "go.opentelemetry.io/otel/trace"
)

var (
	TracingError error = nil
	TP           trc.TracerProvider
)

type CustomExporter struct {
	jaeger      *otlptrace.Exporter
	stdouttrace *stdouttrace.Exporter
}

type Config struct {
	AppName   string
	JaegerUrl string
	GelfUrl   string
	Version   string
}

func NewCustomExporter(jaegerUrl string) (trc.SpanExporter, error) {
	var jaeg *otlptrace.Exporter
	var outrace *stdouttrace.Exporter
	var err error

	outrace, err = stdouttrace.New(
		stdouttrace.WithWriter(os.Stdout),
		stdouttrace.WithPrettyPrint(),
		stdouttrace.WithoutTimestamps(),
	)
	if err != nil {
		return &CustomExporter{}, err
	}

	if len(jaegerUrl) > 0 {
		jaeg = otlptracehttp.NewUnstarted(otlptracehttp.WithEndpointURL(jaegerUrl))
	}

	return &CustomExporter{
		jaeger:      jaeg,
		stdouttrace: outrace,
	}, nil

}

func (e *CustomExporter) ExportSpans(ctx context.Context, spans []trc.ReadOnlySpan) error {
	if TracingError == nil {
		if e.jaeger != nil {
			err := e.jaeger.ExportSpans(ctx, spans)
			return err
		} else {
			return e.printOnlyOnError(ctx, spans)
		}
	}
	return nil
}

func (e *CustomExporter) printOnlyOnError(ctx context.Context, spans []trc.ReadOnlySpan) error {
	var err error
	for _, s := range spans {
		if s.Status().Code == codes.Error {
			err = e.stdouttrace.ExportSpans(ctx, spans)
			break
		}
	}
	return err
}

func (e *CustomExporter) Shutdown(ctx context.Context) error {
	if e.jaeger != nil {
		e.stdouttrace.Shutdown(ctx)
		return e.jaeger.Shutdown(ctx)
	} else {
		return e.stdouttrace.Shutdown(ctx)
	}
}

func newResource(config Config) *resource.Resource {
	r := resource.NewWithAttributes(
		semconv.SchemaURL,
		semconv.ServiceNameKey.String(config.AppName),
		semconv.ServiceVersionKey.String(config.Version),
		attribute.String("service.provider", "maal"),
	)
	return r
}

func NewTracer(config Config) func(*fiber.Ctx) error {
	l := log.New(os.Stdout, "", 0)

	var tracerProviders []trc.TracerProviderOption

	otlpExporter := otlptracehttp.NewUnstarted(otlptracehttp.WithEndpointURL(config.JaegerUrl))

	gelfExporter, err := gelfexporter.New(gelfexporter.WithGelfUrl(config.GelfUrl))
	if err != nil {
		l.Fatal(err)
	}
	tracerProviders = append(tracerProviders, trc.WithBatcher(otlpExporter))
	tracerProviders = append(tracerProviders, trc.WithBatcher(gelfExporter))

	tracerProviders = append(tracerProviders, trc.WithResource(newResource(config)))

	TP = *trc.NewTracerProvider(tracerProviders...)

	otel.SetTracerProvider(&TP)
	otel.SetErrorHandler(otel.ErrorHandlerFunc(func(err error) {
		if err != TracingError {
			TracingError = err
			log.Println(err)
		}
	}))

	tracer := TP.Tracer("fiber-otel-router")

	return fiberOpentelemetry.New(
		fiberOpentelemetry.Config{
			Tracer:   tracer,
			SpanName: "{{ .Method }} {{ .Path }}",
			TracerStartAttributes: []trace.SpanStartOption{
				trace.WithSpanKind(trace.SpanKindServer),
				trace.WithNewRoot(),
			},
		},
	)
}

func ShutdownTracer() {
	if err := TP.Shutdown(context.Background()); err != nil {
		log.Fatal(err)
	}
}

func Handler(fc *fiber.Ctx) (context.Context, trace.Span) {
	simpleCtx, span := fiberOpentelemetry.Tracer.Start(fc.UserContext(), fc.OriginalURL())
	fc.SetUserContext(simpleCtx)

	span.SetAttributes(attribute.String("service.layer", "handler"))

	return simpleCtx, span
}