From 1be209c8aad5656759f88e05792c45220e47dcab Mon Sep 17 00:00:00 2001
From: Parham Alvani <parham.alvani@gmail.com>
Date: Mon, 16 Sep 2024 21:38:21 +0000
Subject: [PATCH] feat: add trace sampler with configuration

---
 internal/infra/config/default.go    | 3 ++-
 internal/infra/telemetry/config.go  | 5 +++--
 internal/infra/telemetry/provide.go | 6 +++++-
 3 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/internal/infra/config/default.go b/internal/infra/config/default.go
index 49f7a38..9a08ddd 100644
--- a/internal/infra/config/default.go
+++ b/internal/infra/config/default.go
@@ -9,7 +9,7 @@ import (
 
 // Default return default configuration.
 func Default() Config {
-	// nolint: exhaustruct
+	// nolint: exhaustruct, mnd
 	return Config{
 		Logger: logger.Config{
 			Level: "debug",
@@ -24,6 +24,7 @@ func Default() Config {
 			Trace: telemetry.Trace{
 				Enabled:  true,
 				Endpoint: "127.0.0.1:4317",
+				Ratio:    0.1,
 			},
 		},
 		NATS: cmq.Config{
diff --git a/internal/infra/telemetry/config.go b/internal/infra/telemetry/config.go
index 09a2000..35f9e2d 100644
--- a/internal/infra/telemetry/config.go
+++ b/internal/infra/telemetry/config.go
@@ -13,6 +13,7 @@ type Meter struct {
 }
 
 type Trace struct {
-	Enabled  bool   `json:"enabled,omitempty"  koanf:"enabled"`
-	Endpoint string `json:"endpoint,omitempty" koanf:"endpoint"`
+	Enabled  bool    `json:"enabled,omitempty"  koanf:"enabled"`
+	Endpoint string  `json:"endpoint,omitempty" koanf:"endpoint"`
+	Ratio    float64 `json:"ratio,omitempty"    koanf:"ratio"`
 }
diff --git a/internal/infra/telemetry/provide.go b/internal/infra/telemetry/provide.go
index 01b7dea..1fb9a96 100644
--- a/internal/infra/telemetry/provide.go
+++ b/internal/infra/telemetry/provide.go
@@ -98,7 +98,11 @@ func Provide(lc fx.Lifecycle, cfg Config) Telemetery {
 	}
 
 	bsp := trace.NewBatchSpanProcessor(exporter)
-	tp := trace.NewTracerProvider(trace.WithSpanProcessor(bsp), trace.WithResource(res))
+	tp := trace.NewTracerProvider(
+		trace.WithSpanProcessor(bsp),
+		trace.WithSampler(trace.ParentBased(trace.TraceIDRatioBased(cfg.Trace.Ratio))),
+		trace.WithResource(res),
+	)
 
 	mp := metric.NewMeterProvider(metric.WithResource(res), metric.WithReader(reader))