apm/apm.go

package apm

import (
	"context"
	"encoding/base64"

	"github.com/labstack/echo/v4"
	"go.mongodb.org/mongo-driver/event"
	"go.opentelemetry.io/contrib/instrumentation/github.com/labstack/echo/otelecho"
	"go.opentelemetry.io/contrib/instrumentation/go.mongodb.org/mongo-driver/mongo/otelmongo"
	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp"
	"go.opentelemetry.io/otel/propagation"
	"go.opentelemetry.io/otel/sdk/resource"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	semconv "go.opentelemetry.io/otel/semconv/v1.25.0"
	"go.opentelemetry.io/otel/trace"
	"google.golang.org/grpc"
)

var (
	cfg    OtelConfig
	tp     *sdktrace.TracerProvider
	tracer trace.Tracer
)

type OtelConfig struct {
	Endpoint       string
	Path           string
	AuthUser       string
	AuthPwd        string
	ServiceName    string
	ServiceVersion string
	Env            string
}

func InitOtelAPM(config OtelConfig) error {
	cfg = config
	otel.SetTextMapPropagator(propagation.NewCompositeTextMapPropagator(
		propagation.TraceContext{},
		propagation.Baggage{},
	))
	otlptracehttp.NewClient()

	authHeaders := map[string]string{}
	if config.AuthUser != "" && config.AuthPwd != "" {
		authHeaders["Authorization"] = "Basic " + base64.StdEncoding.EncodeToString([]byte(config.AuthUser+":"+config.AuthPwd))
	}
	otlpHTTPExporter, err := otlptracehttp.New(context.Background(),
		otlptracehttp.WithEndpoint(config.Endpoint),
		otlptracehttp.WithURLPath(config.Path),
		otlptracehttp.WithHeaders(authHeaders),
	)
	if err != nil {
		return err
	}
	res := resource.NewWithAttributes(
		semconv.SchemaURL,
		// the service name used to display traces in backends
		semconv.ServiceNameKey.String(config.ServiceName),
		semconv.ServiceVersionKey.String(config.ServiceVersion),
		attribute.String("environment", config.Env),
	)
	tp = sdktrace.NewTracerProvider(
		sdktrace.WithSampler(sdktrace.AlwaysSample()),
		sdktrace.WithResource(res),
		sdktrace.WithBatcher(otlpHTTPExporter),
	)
	otel.SetTracerProvider(tp)

	return nil
}

func Shutdown(ctx context.Context) error {
	return tp.Shutdown(ctx)
}

func EchoMiddleware() echo.MiddlewareFunc {
	return otelecho.Middleware(cfg.ServiceName)
}

func MongoMonitor() *event.CommandMonitor {
	return otelmongo.NewMonitor(
		otelmongo.WithCommandAttributeDisabled(false),
	)
}

func GRPCServerHandlerOpt() grpc.ServerOption {
	return grpc.StatsHandler(otelgrpc.NewServerHandler())
}

func GRPCClientHandlerOpt() grpc.DialOption {
	return grpc.WithStatsHandler(otelgrpc.NewClientHandler())
}

func TraceStart(ctx context.Context, name string) (context.Context, trace.Span) {
	return getTracer().Start(ctx, name)
}

func getTracer() trace.Tracer {
	if tracer == nil {
		tracer = otel.Tracer(cfg.ServiceName)
	}

	return tracer
}