Merge pull request #1 from KyberNetwork/sample-rate

Support sample rate for tracer
KyberNetwork · Aug 21, 2023 · a95bd07 · a95bd07
2 parents efae8be + 4256a87
commit a95bd07
Show file tree

Hide file tree

Showing 10 changed files with 72 additions and 27 deletions.
diff --git a/pkg/constant/env.go b/pkg/constant/env.go
@@ -1,10 +1,11 @@
 package constant
 
 const (
-	EnvKeyOTLPEnabled        = "OTLP_ENABLED"
-	EnvKeyOTLPCollectorUrl   = "OTLP_COLLECTOR_URL"
-	EnvKeyOTLPInsecure       = "OTLP_INSECURE"
-	EnvKeyOTLPProtocol       = "OTLP_PROTOCOL"
-	EnvKeyOTLPServiceName    = "OTLP_SERVICE_NAME"
-	EnvKeyOTLPServiceVersion = "OTLP_SERVICE_VERSION"
+	EnvKeyOtelEnabled         = "OTEL_ENABLED"
+	EnvKeyOtelCollectorUrl    = "OTEL_COLLECTOR_URL"
+	EnvKeyOtelInsecure        = "OTEL_INSECURE"
+	EnvKeyOtelProtocol        = "OTEL_PROTOCOL"
+	EnvKeyOtelServiceName     = "OTEL_SERVICE_NAME"
+	EnvKeyOtelServiceVersion  = "OTEL_SERVICE_VERSION"
+	EnvKeyOtelTraceSampleRate = "OTEL_TRACE_SAMPLE_RATE"
 )
diff --git a/pkg/constant/otlp.go b/pkg/constant/otlp.go
@@ -1,8 +1,9 @@
 package constant
 
 const (
-	OTLPDefaultServiceName    = "default-service-name"
-	OTLPDefaultServiceVersion = "0.1.0"
-	OTLPProtocolGRPC          = "grpc"
-	OTLPProtocolHTTP          = "http"
+	OtelDefaultServiceName    = "default-service-name"
+	OtelDefaultServiceVersion = "0.1.0"
+	OtelProtocolGRPC          = "grpc"
+	OtelProtocolHTTP          = "http"
+	OtelDefaultSampleRate     = 0.5
 )
diff --git a/pkg/metric/init.go b/pkg/metric/init.go
@@ -6,7 +6,7 @@ import (
 )
 
 func init() {
-	if env.BoolFromEnv(constant.EnvKeyOTLPEnabled) {
+	if env.BoolFromEnv(constant.EnvKeyOtelEnabled) {
 		InitProvider()
 	}
 }
diff --git a/pkg/metric/meter.go b/pkg/metric/meter.go
@@ -9,5 +9,5 @@ import (
 )
 
 func Meter() metric.Meter {
-	return otel.GetMeterProvider().Meter(env.StringFromEnv(constant.EnvKeyOTLPServiceName, constant.OTLPDefaultServiceName))
+	return otel.GetMeterProvider().Meter(env.StringFromEnv(constant.EnvKeyOtelServiceName, constant.OtelDefaultServiceName))
 }
diff --git a/pkg/metric/provider.go b/pkg/metric/provider.go
@@ -54,12 +54,12 @@ func newOTLPExporter() (metric.Exporter, error) {
 	ctx := context.Background()
 	ctx, cancel := context.WithTimeout(ctx, time.Second)
 	defer cancel()
-	providerServerUrl := env.StringFromEnv(constant.EnvKeyOTLPCollectorUrl, "")
-	isInsecure := env.BoolFromEnv(constant.EnvKeyOTLPInsecure)
-	protocol := env.StringFromEnv(constant.EnvKeyOTLPProtocol, constant.OTLPProtocolGRPC)
+	providerServerUrl := env.StringFromEnv(constant.EnvKeyOtelCollectorUrl, "")
+	isInsecure := env.BoolFromEnv(constant.EnvKeyOtelInsecure)
+	protocol := env.StringFromEnv(constant.EnvKeyOtelProtocol, constant.OtelProtocolGRPC)
 
 	// gRPC
-	if protocol == constant.OTLPProtocolGRPC {
+	if protocol == constant.OtelProtocolGRPC {
 		return newGRPCExporter(ctx, providerServerUrl, isInsecure)
 	}
 
@@ -72,8 +72,8 @@ func newResources() *resource.Resource {
 	// ref: https://opentelemetry.io/docs/instrumentation/go/resources/
 	return resource.NewWithAttributes(
 		semconv.SchemaURL,
-		semconv.ServiceName(env.StringFromEnv(constant.EnvKeyOTLPServiceName, constant.OTLPDefaultServiceName)),
-		semconv.ServiceVersion(env.StringFromEnv(constant.EnvKeyOTLPServiceVersion, constant.OTLPDefaultServiceVersion)),
+		semconv.ServiceName(env.StringFromEnv(constant.EnvKeyOtelServiceName, constant.OtelDefaultServiceName)),
+		semconv.ServiceVersion(env.StringFromEnv(constant.EnvKeyOtelServiceVersion, constant.OtelDefaultServiceVersion)),
 	)
 }
 

diff --git a/pkg/tracer/init.go b/pkg/tracer/init.go
@@ -6,7 +6,7 @@ import (
 )
 
 func init() {
-	if env.BoolFromEnv(constant.EnvKeyOTLPEnabled) {
+	if env.BoolFromEnv(constant.EnvKeyOtelEnabled) {
 		InitProvider()
 	}
 }
diff --git a/pkg/tracer/provider.go b/pkg/tracer/provider.go
@@ -60,12 +60,12 @@ func newOTLPExporter() (*otlptrace.Exporter, error) {
 	ctx := context.Background()
 	ctx, cancel := context.WithTimeout(ctx, time.Second)
 	defer cancel()
-	providerServerUrl := env.StringFromEnv(constant.EnvKeyOTLPCollectorUrl, "")
-	isInsecure := env.BoolFromEnv(constant.EnvKeyOTLPInsecure)
-	protocol := env.StringFromEnv(constant.EnvKeyOTLPProtocol, constant.OTLPProtocolGRPC)
+	providerServerUrl := env.StringFromEnv(constant.EnvKeyOtelCollectorUrl, "")
+	isInsecure := env.BoolFromEnv(constant.EnvKeyOtelInsecure)
+	protocol := env.StringFromEnv(constant.EnvKeyOtelProtocol, constant.OtelProtocolGRPC)
 
 	// gRPC
-	if protocol == constant.OTLPProtocolGRPC {
+	if protocol == constant.OtelProtocolGRPC {
 		return newGRPCExporter(ctx, providerServerUrl, isInsecure)
 	}
 
@@ -78,8 +78,8 @@ func newResources() *resource.Resource {
 	// ref: https://opentelemetry.io/docs/instrumentation/go/resources/
 	return resource.NewWithAttributes(
 		semconv.SchemaURL,
-		semconv.ServiceName(env.StringFromEnv(constant.EnvKeyOTLPServiceName, constant.OTLPDefaultServiceName)),
-		semconv.ServiceVersion(env.StringFromEnv(constant.EnvKeyOTLPServiceVersion, constant.OTLPDefaultServiceVersion)),
+		semconv.ServiceName(env.StringFromEnv(constant.EnvKeyOtelServiceName, constant.OtelDefaultServiceName)),
+		semconv.ServiceVersion(env.StringFromEnv(constant.EnvKeyOtelServiceVersion, constant.OtelDefaultServiceVersion)),
 	)
 }
 
@@ -100,9 +100,11 @@ func InitProvider() {
 	// Register the trace exporter with a TracerProvider, using a batch span processor to aggregate spans before export.
 	bsp := trace.NewBatchSpanProcessor(exporter)
 
+	sampleRate := env.FloatFromEnv(constant.EnvKeyOtelTraceSampleRate, constant.OtelDefaultSampleRate)
+
 	// init tracer provider
 	provider = trace.NewTracerProvider(
-		trace.WithSampler(trace.AlwaysSample()),
+		trace.WithSampler(trace.TraceIDRatioBased(sampleRate)),
 		trace.WithResource(newResources()),
 		trace.WithSpanProcessor(bsp),
 	)

diff --git a/pkg/tracer/tracer.go b/pkg/tracer/tracer.go
@@ -9,5 +9,5 @@ import (
 )
 
 func Tracer() trace.Tracer {
-	return otel.Tracer(env.StringFromEnv(constant.EnvKeyOTLPServiceName, constant.OTLPDefaultServiceName))
+	return otel.Tracer(env.StringFromEnv(constant.EnvKeyOtelServiceName, constant.OtelDefaultServiceName))
 }
diff --git a/pkg/util/env/env.go b/pkg/util/env/env.go
@@ -2,6 +2,7 @@ package env
 
 import (
 	"os"
+	"strconv"
 	"strings"
 )
 
@@ -19,3 +20,16 @@ func BoolFromEnv(key string) bool {
 	v = strings.ToLower(strings.TrimSpace(v))
 	return v != "" && v != "0" && v != "false" && v != "no"
 }
+
+// FloatFromEnv returns the float number for the given key
+// and falls back to the given defaultValue if not set
+func FloatFromEnv(key string, defaultValue float64) float64 {
+	str := ""
+	if str = os.Getenv(key); str != "" {
+		str = strings.TrimSpace(str)
+	}
+	if s, err := strconv.ParseFloat(str, 64); err == nil {
+		return s
+	}
+	return defaultValue
+}
diff --git a/pkg/util/env/env_test.go b/pkg/util/env/env_test.go
@@ -45,3 +45,30 @@ func TestBoolFromEnv(t *testing.T) {
 	assert.Equal(t, true, BoolFromEnv("key_yes"))
 	assert.Equal(t, false, BoolFromEnv("key_no"))
 }
+
+func TestFloatFromEnv(t *testing.T) {
+	// mock
+	_ = os.Setenv("test_float_from_env", "0.75")
+	_ = os.Setenv("test_float_from_env_empty", "")
+	_ = os.Setenv("test_float_from_env_string", "zeropointfive")
+	_ = os.Setenv("test_float_from_env_trim", " 0.75 ")
+	_ = os.Setenv("test_float_from_env_0", "+0.75")
+	_ = os.Setenv("test_float_from_env_1", "-0.75")
+	_ = os.Setenv("test_float_from_env_2", ".75")
+	_ = os.Setenv("test_float_from_env_3", "01.75")
+
+	// Test default value
+	assert.Equal(t, 0.75, FloatFromEnv("not_exists", 0.75))
+	assert.Equal(t, 0.75, FloatFromEnv("test_float_from_env_empty", 0.75))
+	assert.Equal(t, 0.75, FloatFromEnv("test_float_from_env_string", 0.75))
+
+	// Test not default value
+	assert.Equal(t, 0.75, FloatFromEnv("test_float_from_env", 0.5))
+	assert.Equal(t, 0.75, FloatFromEnv("test_float_from_env_0", 0.5))
+	assert.Equal(t, -0.75, FloatFromEnv("test_float_from_env_1", 0.5))
+	assert.Equal(t, 0.75, FloatFromEnv("test_float_from_env_2", 0.5))
+	assert.Equal(t, 1.75, FloatFromEnv("test_float_from_env_3", 0.5))
+
+	// Test trim space
+	assert.Equal(t, 0.75, FloatFromEnv("test_float_from_env_trim", 0.5))
+}