From 4127ac21c7a23503f91a590841f6ea2724bccdde Mon Sep 17 00:00:00 2001
From: Andreas Thomas <dev@chronark.com>
Date: Fri, 28 Jun 2024 10:11:05 +0200
Subject: [PATCH] feat: use a global tracer (#1822)

---
 apps/agent/cmd/agent/agent.go                | 12 ++++-------
 apps/agent/pkg/cache/middleware/tracing.go   | 21 ++++++++++----------
 apps/agent/pkg/connect/middleware_tracing.go |  7 +++----
 apps/agent/pkg/connect/ratelimit.go          |  8 +++-----
 apps/agent/pkg/connect/service.go            |  6 +-----
 apps/agent/pkg/tracing/axiom.go              | 20 ++++++++-----------
 apps/agent/pkg/tracing/trace.go              | 18 +++++++++++++++++
 apps/agent/services/eventrouter/service.go   |  5 +----
 apps/agent/services/ratelimit/middleware.go  | 14 +++++--------
 apps/agent/services/ratelimit/pushpull.go    |  3 ++-
 apps/agent/services/ratelimit/ratelimit.go   |  3 ++-
 apps/agent/services/ratelimit/service.go     |  4 ----
 12 files changed, 57 insertions(+), 64 deletions(-)
 create mode 100644 apps/agent/pkg/tracing/trace.go

diff --git a/apps/agent/cmd/agent/agent.go b/apps/agent/cmd/agent/agent.go
index 816ea7ee3d..134db45137 100644
--- a/apps/agent/cmd/agent/agent.go
+++ b/apps/agent/cmd/agent/agent.go
@@ -73,10 +73,9 @@ func run(c *cli.Context) error {
 
 	logger.Info().Str("file", configFile).Msg("configuration loaded")
 
-	tracer := tracing.NewNoop()
 	{
 		if cfg.Tracing != nil && cfg.Tracing.Axiom != nil {
-			t, closeTracer, err := tracing.New(context.Background(), tracing.Config{
+			closeTracer, err := tracing.Init(context.Background(), tracing.Config{
 				Dataset:     cfg.Tracing.Axiom.Dataset,
 				Application: "agent",
 				Version:     "1.0.0",
@@ -91,7 +90,6 @@ func run(c *cli.Context) error {
 					logger.Error().Err(err).Msg("failed to close tracer")
 				}
 			}()
-			tracer = t
 			logger.Info().Msg("tracing to axiom")
 		}
 	}
@@ -127,7 +125,7 @@ func run(c *cli.Context) error {
 		}
 	}
 
-	srv, err := connect.New(connect.Config{Logger: logger, Tracer: tracer, Image: cfg.Image})
+	srv, err := connect.New(connect.Config{Logger: logger, Image: cfg.Image})
 	if err != nil {
 		return err
 	}
@@ -160,7 +158,6 @@ func run(c *cli.Context) error {
 	if cfg.Services.EventRouter != nil {
 		er, err := eventrouter.New(eventrouter.Config{
 			Logger:        logger,
-			Tracer:        tracer,
 			BatchSize:     cfg.Services.EventRouter.Tinybird.BatchSize,
 			BufferSize:    cfg.Services.EventRouter.Tinybird.BufferSize,
 			FlushInterval: time.Duration(cfg.Services.EventRouter.Tinybird.FlushInterval) * time.Second,
@@ -214,15 +211,14 @@ func run(c *cli.Context) error {
 	if cfg.Services.Ratelimit != nil {
 		rl, err := ratelimit.New(ratelimit.Config{
 			Logger:  logger,
-			Tracer:  tracer,
 			Cluster: clus,
 		})
 		if err != nil {
 			logger.Fatal().Err(err).Msg("failed to create service")
 		}
-		rl = ratelimit.WithTracing(tracer)(rl)
+		rl = ratelimit.WithTracing(rl)
 
-		srv.AddService(connect.NewRatelimitServer(rl, logger, tracer))
+		srv.AddService(connect.NewRatelimitServer(rl, logger))
 		logger.Info().Msg("started ratelimit service")
 	}
 
diff --git a/apps/agent/pkg/cache/middleware/tracing.go b/apps/agent/pkg/cache/middleware/tracing.go
index a4190cfb51..d1280e20c3 100644
--- a/apps/agent/pkg/cache/middleware/tracing.go
+++ b/apps/agent/pkg/cache/middleware/tracing.go
@@ -10,16 +10,15 @@ import (
 )
 
 type tracingMiddleware[T any] struct {
-	next   cache.Cache[T]
-	tracer tracing.Tracer
+	next cache.Cache[T]
 }
 
-func WithTracing[T any](c cache.Cache[T], t tracing.Tracer) cache.Cache[T] {
-	return &tracingMiddleware[T]{next: c, tracer: t}
+func WithTracing[T any](c cache.Cache[T]) cache.Cache[T] {
+	return &tracingMiddleware[T]{next: c}
 }
 
 func (mw *tracingMiddleware[T]) Get(ctx context.Context, key string) (T, cache.CacheHit) {
-	ctx, span := mw.tracer.Start(ctx, "cache.Get", trace.WithAttributes(attribute.String("key", key)))
+	ctx, span := tracing.Start(ctx, "cache.Get", trace.WithAttributes(attribute.String("key", key)))
 	defer span.End()
 
 	value, hit := mw.next.Get(ctx, key)
@@ -29,21 +28,21 @@ func (mw *tracingMiddleware[T]) Get(ctx context.Context, key string) (T, cache.C
 	return value, hit
 }
 func (mw *tracingMiddleware[T]) Set(ctx context.Context, key string, value T) {
-	ctx, span := mw.tracer.Start(ctx, "cache.Set", trace.WithAttributes(attribute.String("key", key)))
+	ctx, span := tracing.Start(ctx, "cache.Set", trace.WithAttributes(attribute.String("key", key)))
 	defer span.End()
 
 	mw.next.Set(ctx, key, value)
 
 }
 func (mw *tracingMiddleware[T]) SetNull(ctx context.Context, key string) {
-	ctx, span := mw.tracer.Start(ctx, "cache.SetNull", trace.WithAttributes(attribute.String("key", key)))
+	ctx, span := tracing.Start(ctx, "cache.SetNull", trace.WithAttributes(attribute.String("key", key)))
 	defer span.End()
 
 	mw.next.SetNull(ctx, key)
 
 }
 func (mw *tracingMiddleware[T]) Remove(ctx context.Context, key string) {
-	ctx, span := mw.tracer.Start(ctx, "cache.Remove", trace.WithAttributes(attribute.String("key", key)))
+	ctx, span := tracing.Start(ctx, "cache.Remove", trace.WithAttributes(attribute.String("key", key)))
 	defer span.End()
 
 	mw.next.Remove(ctx, key)
@@ -51,21 +50,21 @@ func (mw *tracingMiddleware[T]) Remove(ctx context.Context, key string) {
 }
 
 func (mw *tracingMiddleware[T]) Dump(ctx context.Context) ([]byte, error) {
-	ctx, span := mw.tracer.Start(ctx, "cache.Dump")
+	ctx, span := tracing.Start(ctx, "cache.Dump")
 	defer span.End()
 
 	return mw.next.Dump(ctx)
 }
 
 func (mw *tracingMiddleware[T]) Restore(ctx context.Context, data []byte) error {
-	ctx, span := mw.tracer.Start(ctx, "cache.Restore")
+	ctx, span := tracing.Start(ctx, "cache.Restore")
 	defer span.End()
 
 	return mw.next.Restore(ctx, data)
 }
 
 func (mw *tracingMiddleware[T]) Clear(ctx context.Context) {
-	ctx, span := mw.tracer.Start(ctx, "cache.Clear")
+	ctx, span := tracing.Start(ctx, "cache.Clear")
 	defer span.End()
 
 	mw.next.Clear(ctx)
diff --git a/apps/agent/pkg/connect/middleware_tracing.go b/apps/agent/pkg/connect/middleware_tracing.go
index b888771c2b..a20d71084e 100644
--- a/apps/agent/pkg/connect/middleware_tracing.go
+++ b/apps/agent/pkg/connect/middleware_tracing.go
@@ -7,16 +7,15 @@ import (
 )
 
 type tracingMiddleware struct {
-	tracer  tracing.Tracer
 	handler http.Handler
 }
 
-func newTracingMiddleware(handler http.Handler, tracer tracing.Tracer) http.Handler {
-	return &tracingMiddleware{handler: handler, tracer: tracer}
+func newTracingMiddleware(handler http.Handler) http.Handler {
+	return &tracingMiddleware{handler: handler}
 }
 
 func (h *tracingMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	ctx, span := h.tracer.Start(r.Context(), "request")
+	ctx, span := tracing.Start(r.Context(), "request")
 	defer span.End()
 
 	h.handler.ServeHTTP(w, r.WithContext(ctx))
diff --git a/apps/agent/pkg/connect/ratelimit.go b/apps/agent/pkg/connect/ratelimit.go
index 6f47e22891..824e4bed19 100644
--- a/apps/agent/pkg/connect/ratelimit.go
+++ b/apps/agent/pkg/connect/ratelimit.go
@@ -18,16 +18,14 @@ import (
 type ratelimitServer struct {
 	svc    ratelimit.Service
 	logger logging.Logger
-	tracer tracing.Tracer
 	ratelimitv1connect.UnimplementedRatelimitServiceHandler
 }
 
-func NewRatelimitServer(svc ratelimit.Service, logger logging.Logger, tracer tracing.Tracer) *ratelimitServer {
+func NewRatelimitServer(svc ratelimit.Service, logger logging.Logger) *ratelimitServer {
 
 	return &ratelimitServer{
 		svc:    svc,
 		logger: logger,
-		tracer: tracer,
 	}
 
 }
@@ -48,7 +46,7 @@ func (s *ratelimitServer) Ratelimit(
 			Msg("connect.Ratelimit")
 	}()
 
-	ctx, span := s.tracer.Start(ctx, "ratelimit.Ratelimit")
+	ctx, span := tracing.Start(ctx, "ratelimit.Ratelimit")
 	defer span.End()
 	authorization := req.Header().Get("Authorization")
 	err := auth.Authorize(ctx, authorization)
@@ -75,7 +73,7 @@ func (s *ratelimitServer) PushPull(
 			Int64("latency", time.Since(start).Milliseconds()).
 			Msg("connect.PushPull")
 	}()
-	ctx, span := s.tracer.Start(ctx, "ratelimit.PushPull")
+	ctx, span := tracing.Start(ctx, "ratelimit.PushPull")
 	defer span.End()
 	authorization := req.Header().Get("Authorization")
 	err := auth.Authorize(ctx, authorization)
diff --git a/apps/agent/pkg/connect/service.go b/apps/agent/pkg/connect/service.go
index 5ce2356957..f16009114d 100644
--- a/apps/agent/pkg/connect/service.go
+++ b/apps/agent/pkg/connect/service.go
@@ -10,7 +10,6 @@ import (
 	"github.com/bufbuild/connect-go"
 	ratelimitv1 "github.com/unkeyed/unkey/apps/agent/gen/proto/ratelimit/v1"
 	"github.com/unkeyed/unkey/apps/agent/pkg/logging"
-	"github.com/unkeyed/unkey/apps/agent/pkg/tracing"
 	"golang.org/x/net/http2"
 	"golang.org/x/net/http2/h2c"
 )
@@ -22,7 +21,6 @@ type Service interface {
 type Server struct {
 	sync.Mutex
 	logger         logging.Logger
-	tracer         tracing.Tracer
 	mux            *http.ServeMux
 	shutdownC      chan struct{}
 	isShuttingDown bool
@@ -32,7 +30,6 @@ type Server struct {
 
 type Config struct {
 	Logger logging.Logger
-	Tracer tracing.Tracer
 	Image  string
 }
 
@@ -40,7 +37,6 @@ func New(cfg Config) (*Server, error) {
 
 	return &Server{
 		logger:         cfg.Logger,
-		tracer:         cfg.Tracer,
 		isListening:    false,
 		isShuttingDown: false,
 		mux:            http.NewServeMux(),
@@ -54,7 +50,7 @@ func (s *Server) AddService(svc Service) {
 	pattern, handler := svc.CreateHandler()
 	s.logger.Info().Str("pattern", pattern).Msg("adding service")
 
-	h := newTracingMiddleware(newHeaderMiddleware(newLoggingMiddleware(handler, s.logger)), s.tracer)
+	h := newTracingMiddleware(newHeaderMiddleware(newLoggingMiddleware(handler, s.logger)))
 	s.mux.Handle(pattern, h)
 }
 
diff --git a/apps/agent/pkg/tracing/axiom.go b/apps/agent/pkg/tracing/axiom.go
index d5f94f730c..9ef4a8d98a 100644
--- a/apps/agent/pkg/tracing/axiom.go
+++ b/apps/agent/pkg/tracing/axiom.go
@@ -4,14 +4,10 @@ import (
 	"context"
 	"fmt"
 
-	"go.opentelemetry.io/otel"
-	"go.opentelemetry.io/otel/trace"
-
 	axiom "github.com/axiomhq/axiom-go/axiom/otel"
+	"go.opentelemetry.io/otel"
 )
 
-type Tracer trace.Tracer
-
 type Config struct {
 	Dataset     string
 	Application string
@@ -19,7 +15,10 @@ type Config struct {
 	AxiomToken  string
 }
 
-func New(ctx context.Context, config Config) (Tracer, func() error, error) {
+// Coser is a function that closes the global tracer.
+type Closer func() error
+
+func Init(ctx context.Context, config Config) (Closer, error) {
 
 	close, err := axiom.InitTracing(
 		ctx,
@@ -31,12 +30,9 @@ func New(ctx context.Context, config Config) (Tracer, func() error, error) {
 	)
 
 	if err != nil {
-		return nil, nil, fmt.Errorf("unable to init tracing: %w", err)
+		return nil, fmt.Errorf("unable to init tracing: %w", err)
 	}
 
-	return otel.Tracer("main"), close, nil
-}
-
-func NewNoop() Tracer {
-	return trace.NewNoopTracerProvider().Tracer("noop")
+	globalTracer = otel.Tracer("main")
+	return close, nil
 }
diff --git a/apps/agent/pkg/tracing/trace.go b/apps/agent/pkg/tracing/trace.go
new file mode 100644
index 0000000000..0873b94074
--- /dev/null
+++ b/apps/agent/pkg/tracing/trace.go
@@ -0,0 +1,18 @@
+package tracing
+
+import (
+	"context"
+
+	"go.opentelemetry.io/otel/trace"
+	"go.opentelemetry.io/otel/trace/noop"
+)
+
+var globalTracer trace.Tracer
+
+func init() {
+	globalTracer = noop.NewTracerProvider().Tracer("noop")
+}
+
+func Start(ctx context.Context, name string, opts ...trace.SpanStartOption) (context.Context, trace.Span) {
+	return globalTracer.Start(ctx, name, opts...)
+}
diff --git a/apps/agent/services/eventrouter/service.go b/apps/agent/services/eventrouter/service.go
index 3ce9ab775a..75b6e153c5 100644
--- a/apps/agent/services/eventrouter/service.go
+++ b/apps/agent/services/eventrouter/service.go
@@ -26,12 +26,10 @@ type Config struct {
 
 	Tinybird *tinybird.Client
 	Logger   logging.Logger
-	Tracer   tracing.Tracer
 }
 
 type service struct {
 	logger  logging.Logger
-	tracer  tracing.Tracer
 	batcher batch.BatchProcessor[event]
 	tb      *tinybird.Client
 }
@@ -69,7 +67,6 @@ func New(config Config) (*service, error) {
 	return &service{
 		logger:  config.Logger,
 		batcher: *batcher,
-		tracer:  config.Tracer,
 		tb:      config.Tinybird,
 	}, nil
 }
@@ -80,7 +77,7 @@ func (s *service) CreateHandler() (string, http.Handler) {
 		defer func() {
 			s.logger.Info().Str("method", r.Method).Str("path", r.URL.Path).Int64("serviceLatency", time.Since(start).Milliseconds()).Msg("request")
 		}()
-		ctx, span := s.tracer.Start(r.Context(), "events")
+		ctx, span := tracing.Start(r.Context(), "events")
 		defer span.End()
 
 		err := auth.Authorize(ctx, r.Header.Get("Authorization"))
diff --git a/apps/agent/services/ratelimit/middleware.go b/apps/agent/services/ratelimit/middleware.go
index 152b7b0280..83ad1f98f7 100644
--- a/apps/agent/services/ratelimit/middleware.go
+++ b/apps/agent/services/ratelimit/middleware.go
@@ -7,21 +7,17 @@ import (
 	"github.com/unkeyed/unkey/apps/agent/pkg/tracing"
 )
 
-func WithTracing(tracer tracing.Tracer) Middleware {
-
-	return func(svc Service) Service {
-		return &tracingMiddleware{next: svc, tracer: tracer}
-	}
+func WithTracing(svc Service) Service {
+	return &tracingMiddleware{next: svc}
 }
 
 type tracingMiddleware struct {
-	next   Service
-	tracer tracing.Tracer
+	next Service
 }
 
 func (mw *tracingMiddleware) Ratelimit(ctx context.Context, req *ratelimitv1.RatelimitRequest) (res *ratelimitv1.RatelimitResponse, err error) {
 
-	ctx, span := mw.tracer.Start(ctx, tracing.NewSpanName("ratelimit", "Ratelimit"))
+	ctx, span := tracing.Start(ctx, tracing.NewSpanName("ratelimit", "Ratelimit"))
 	defer span.End()
 
 	res, err = mw.next.Ratelimit(ctx, req)
@@ -32,7 +28,7 @@ func (mw *tracingMiddleware) Ratelimit(ctx context.Context, req *ratelimitv1.Rat
 }
 
 func (mw *tracingMiddleware) PushPull(ctx context.Context, req *ratelimitv1.PushPullRequest) (res *ratelimitv1.PushPullResponse, err error) {
-	ctx, span := mw.tracer.Start(ctx, tracing.NewSpanName("ratelimit", "PushPull"))
+	ctx, span := tracing.Start(ctx, tracing.NewSpanName("ratelimit", "PushPull"))
 	defer span.End()
 
 	res, err = mw.next.PushPull(ctx, req)
diff --git a/apps/agent/services/ratelimit/pushpull.go b/apps/agent/services/ratelimit/pushpull.go
index b10692a7ce..6c3256be44 100644
--- a/apps/agent/services/ratelimit/pushpull.go
+++ b/apps/agent/services/ratelimit/pushpull.go
@@ -6,6 +6,7 @@ import (
 
 	ratelimitv1 "github.com/unkeyed/unkey/apps/agent/gen/proto/ratelimit/v1"
 	"github.com/unkeyed/unkey/apps/agent/pkg/ratelimit"
+	"github.com/unkeyed/unkey/apps/agent/pkg/tracing"
 )
 
 func (s *service) PushPull(ctx context.Context, req *ratelimitv1.PushPullRequest) (*ratelimitv1.PushPullResponse, error) {
@@ -15,7 +16,7 @@ func (s *service) PushPull(ctx context.Context, req *ratelimitv1.PushPullRequest
 			Int64("latency", time.Since(start).Milliseconds()).
 			Msg("service.PushPull")
 	}()
-	ctx, span := s.tracer.Start(ctx, "PushPull")
+	ctx, span := tracing.Start(ctx, "PushPull")
 	defer span.End()
 	res := s.ratelimiter.Take(ratelimit.RatelimitRequest{
 		Identifier:     req.Identifier,
diff --git a/apps/agent/services/ratelimit/ratelimit.go b/apps/agent/services/ratelimit/ratelimit.go
index 2b9b4a0e46..51abfd916a 100644
--- a/apps/agent/services/ratelimit/ratelimit.go
+++ b/apps/agent/services/ratelimit/ratelimit.go
@@ -6,6 +6,7 @@ import (
 
 	ratelimitv1 "github.com/unkeyed/unkey/apps/agent/gen/proto/ratelimit/v1"
 	"github.com/unkeyed/unkey/apps/agent/pkg/ratelimit"
+	"github.com/unkeyed/unkey/apps/agent/pkg/tracing"
 )
 
 func (s *service) Ratelimit(ctx context.Context, req *ratelimitv1.RatelimitRequest) (*ratelimitv1.RatelimitResponse, error) {
@@ -15,7 +16,7 @@ func (s *service) Ratelimit(ctx context.Context, req *ratelimitv1.RatelimitReque
 			Int64("latency", time.Since(start).Milliseconds()).
 			Msg("service.Ratelimit")
 	}()
-	_, span := s.tracer.Start(ctx, "Ratelimit")
+	_, span := tracing.Start(ctx, "Ratelimit")
 	defer span.End()
 	res := s.ratelimiter.Take(ratelimit.RatelimitRequest{
 		Identifier:     req.Identifier,
diff --git a/apps/agent/services/ratelimit/service.go b/apps/agent/services/ratelimit/service.go
index dc4413efa7..4b4a7839ba 100644
--- a/apps/agent/services/ratelimit/service.go
+++ b/apps/agent/services/ratelimit/service.go
@@ -4,7 +4,6 @@ import (
 	"github.com/unkeyed/unkey/apps/agent/pkg/cluster"
 	"github.com/unkeyed/unkey/apps/agent/pkg/logging"
 	"github.com/unkeyed/unkey/apps/agent/pkg/ratelimit"
-	"github.com/unkeyed/unkey/apps/agent/pkg/tracing"
 )
 
 type pushPullEvent struct {
@@ -17,7 +16,6 @@ type pushPullEvent struct {
 
 type service struct {
 	logger      logging.Logger
-	tracer      tracing.Tracer
 	ratelimiter ratelimit.Ratelimiter
 	cluster     cluster.Cluster
 
@@ -26,14 +24,12 @@ type service struct {
 
 type Config struct {
 	Logger  logging.Logger
-	Tracer  tracing.Tracer
 	Cluster cluster.Cluster
 }
 
 func New(cfg Config) (Service, error) {
 	s := &service{
 		logger:      cfg.Logger,
-		tracer:      cfg.Tracer,
 		ratelimiter: ratelimit.NewFixedWindow(cfg.Logger.With().Str("ratelimiter", "fixedWindow").Logger()),
 		cluster:     cfg.Cluster,
 	}