tstor/src/telemetry/setup.go

package telemetry

import (
	"context"
	"fmt"
	"log/slog"
	"os"
	"runtime"

	"git.kmsign.ru/royalcat/tstor/pkg/rlog"
	"github.com/agoda-com/opentelemetry-go/otelslog"
	"github.com/agoda-com/opentelemetry-logs-go/exporters/otlp/otlplogs"
	"github.com/agoda-com/opentelemetry-logs-go/exporters/otlp/otlplogs/otlplogshttp"
	logsdk "github.com/agoda-com/opentelemetry-logs-go/sdk/logs"
	"github.com/google/uuid"
	otelpyroscope "github.com/grafana/otel-profiling-go"
	"github.com/grafana/pyroscope-go"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetrichttp"
	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp"
	"go.opentelemetry.io/otel/exporters/prometheus"
	"go.opentelemetry.io/otel/sdk/metric"
	"go.opentelemetry.io/otel/sdk/resource"
	"go.opentelemetry.io/otel/sdk/trace"
	semconv "go.opentelemetry.io/otel/semconv/v1.26.0"
)

type Client struct {
	log *rlog.Logger

	tracerProvider *trace.TracerProvider
	metricProvider *metric.MeterProvider
	loggerProvider *logsdk.LoggerProvider
}

func (client *Client) Shutdown(ctx context.Context) {
	if client.metricProvider == nil {
		err := client.metricProvider.Shutdown(ctx)
		if err != nil {
			client.log.Error(ctx, "error shutting down metric provider", rlog.Error(err))
		}
	}
	if client.tracerProvider == nil {
		err := client.tracerProvider.Shutdown(ctx)
		if err != nil {
			client.log.Error(ctx, "error shutting down tracer provider", rlog.Error(err))
		}
	}
	if client.loggerProvider == nil {
		err := client.loggerProvider.Shutdown(ctx)
		if err != nil {
			client.log.Error(ctx, "error shutting down logger provider", rlog.Error(err))
		}
	}
}

const appName = "tstor"

func Setup(ctx context.Context, endpoint string) (*Client, error) {
	client := &Client{
		log: rlog.Component("telemetry"),
	}
	otel.SetErrorHandler(otel.ErrorHandlerFunc(func(cause error) {
		client.log.Error(context.Background(), "otel error", rlog.Error(cause))
	}))

	hostName, _ := os.Hostname()

	r, err := resource.Merge(
		resource.Default(),
		resource.NewWithAttributes(
			semconv.SchemaURL,
			semconv.ServiceName(appName),
			semconv.HostName(hostName),
			semconv.ServiceInstanceID(uuid.NewString()),
		),
	)
	if err != nil {
		return nil, err
	}

	meticExporter, err := otlpmetrichttp.New(ctx,
		otlpmetrichttp.WithEndpoint(endpoint),
		otlpmetrichttp.WithRetry(otlpmetrichttp.RetryConfig{
			Enabled: false,
		}),
	)
	if err != nil {
		return nil, err
	}

	promExporter, err := prometheus.New()
	if err != nil {
		return nil, fmt.Errorf("failed to initialize prometheus exporter: %w", err)
	}
	client.metricProvider = metric.NewMeterProvider(
		metric.WithReader(metric.NewPeriodicReader(meticExporter)),
		metric.WithReader(promExporter),
		metric.WithResource(r),
	)
	otel.SetMeterProvider(client.metricProvider)

	var meter = otel.Meter("git.kmsign.ru/royalcat/tstor/pkg/telemetry")
	counter, err := meter.Int64Counter("up")
	if err != nil {
		return nil, err
	}
	counter.Add(ctx, 1)
	client.log.Info(ctx, "metrics provider initialized")

	traceExporter, err := otlptracehttp.New(ctx,
		otlptracehttp.WithEndpoint(endpoint),
		otlptracehttp.WithRetry(otlptracehttp.RetryConfig{
			Enabled: false,
		}),
	)
	if err != nil {
		return nil, err
	}
	client.tracerProvider = trace.NewTracerProvider(
		trace.WithBatcher(traceExporter),
		trace.WithResource(r),
	)
	otel.SetTracerProvider(otelpyroscope.NewTracerProvider(client.tracerProvider))
	client.log.Info(ctx, "tracing provider initialized")

	logExporter, err := otlplogs.NewExporter(ctx,
		otlplogs.WithClient(
			otlplogshttp.NewClient(
				otlplogshttp.WithEndpoint(endpoint),
			),
		),
	)
	if err != nil {
		return nil, err
	}
	client.loggerProvider = logsdk.NewLoggerProvider(
		logsdk.WithBatcher(logExporter),
		logsdk.WithResource(r),
	)
	rlog.AddHandler(otelslog.NewOtelHandler(client.loggerProvider,
		&otelslog.HandlerOptions{
			Level: slog.LevelDebug,
		}),
	)
	client.log.Info(ctx, "logger provider initialized")

	// recreate telemetry logger
	client.log = rlog.Component("telemetry")

	runtime.SetMutexProfileFraction(5)
	runtime.SetBlockProfileRate(5)
	_, err = pyroscope.Start(pyroscope.Config{
		ApplicationName: appName,
		// replace this with the address of pyroscope server
		ServerAddress: "https://pyroscope.kmsign.ru",
		// you can disable logging by setting this to nil
		Logger: newPyroscopeLogger(client.log),
		ProfileTypes: []pyroscope.ProfileType{
			// these profile types are enabled by default:
			pyroscope.ProfileCPU,
			pyroscope.ProfileAllocObjects,
			pyroscope.ProfileAllocSpace,
			pyroscope.ProfileInuseObjects,
			pyroscope.ProfileInuseSpace,
			// these profile types are optional:
			// pyroscope.ProfileGoroutines,
			// pyroscope.ProfileMutexCount,
			// pyroscope.ProfileMutexDuration,
			// pyroscope.ProfileBlockCount,
			// pyroscope.ProfileBlockDuration,
		},
	})
	if err != nil {
		return client, nil
	}

	return client, nil
}

func newPyroscopeLogger(log *rlog.Logger) *pyroscopeLogger {
	return &pyroscopeLogger{
		log: log.WithComponent("pyroscope").Nested(1),
	}
}

type pyroscopeLogger struct {
	log *rlog.Logger
}

var _ pyroscope.Logger = (*pyroscopeLogger)(nil)

// Debugf implements pyroscope.Logger.
func (p *pyroscopeLogger) Debugf(msg string, args ...any) {
	ctx := context.Background()
	p.log.Debug(ctx, fmt.Sprintf(msg, args...))
}

// Errorf implements pyroscope.Logger.
func (p *pyroscopeLogger) Errorf(msg string, args ...any) {
	ctx := context.Background()
	p.log.Error(ctx, fmt.Sprintf(msg, args...))
}

// Infof implements pyroscope.Logger.
func (p *pyroscopeLogger) Infof(msg string, args ...any) {
	ctx := context.Background()
	p.log.Info(ctx, fmt.Sprintf(msg, args...))
}

func functionName() string {
	var pcs [1]uintptr
	runtime.Callers(1, pcs[:])
	return runtime.FuncForPC(pcs[0]).Name()
}