configurable span formatter

Signed-off-by: Daniel Rammer <[email protected]>
flyteorg · Jan 11, 2024 · 41b47de · 41b47de
1 parent b0db4d6
commit 41b47de
Show file tree

Hide file tree

Showing 12 changed files with 239 additions and 178 deletions.
diff --git a/flytestdlib/go.mod b/flytestdlib/go.mod
@@ -16,6 +16,7 @@ require (
 	github.com/go-test/deep v1.0.7
 	github.com/golang/protobuf v1.5.3
 	github.com/hashicorp/golang-lru v0.5.4
+	github.com/hashicorp/golang-lru/v2 v2.0.7
 	github.com/jackc/pgconn v1.14.1
 	github.com/jackc/pgx/v5 v5.4.3
 	github.com/magiconair/properties v1.8.6

diff --git a/flytestdlib/go.sum b/flytestdlib/go.sum
@@ -245,6 +245,8 @@ github.com/hashicorp/golang-lru v0.5.0/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ
 github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
 github.com/hashicorp/golang-lru v0.5.4 h1:YDjusn29QI/Das2iO9M0BHnIbxPeyuCHsjMW+lJfyTc=
 github.com/hashicorp/golang-lru v0.5.4/go.mod h1:iADmTwqILo4mZ8BN3D2Q6+9jd8WM5uGBxy+E8yxSoD4=
+github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k=
+github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
 github.com/hashicorp/hcl v1.0.0 h1:0Anlzjpi4vEasTeNFn2mLJgTSwt0+6sfsiTG8qcWGx4=
 github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ=
 github.com/ianlancetaylor/demangle v0.0.0-20181102032728-5e5cf60278f6/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=

diff --git a/flytestdlib/otelutils/config.go b/flytestdlib/otelutils/config.go
@@ -12,19 +12,24 @@ import (
 const configSectionKey = "otel"
 
 type ExporterType = string
+type SpanFormat = string
 
 const (
 	NoopExporter   ExporterType = "noop"
 	FileExporter   ExporterType = "file"
 	JaegerExporter ExporterType = "jaeger"
+
+	OtelSpanFormat          SpanFormat = "otel"
+	OtelCollectorSpanFormat SpanFormat = "otel_collector"
 )
 
 var (
 	ConfigSection = config.MustRegisterSection(configSectionKey, defaultConfig)
 	defaultConfig = &Config{
 		ExporterType: NoopExporter,
 		FileConfig: FileConfig{
-			Filename: "/tmp/otel.json",
+			Filename:   "/tmp/otel.json",
+			SpanFormat: OtelSpanFormat,
 		},
 		JaegerConfig: JaegerConfig{
 			Endpoint: "http://localhost:14268/api/traces",
@@ -39,7 +44,8 @@ type Config struct {
 }
 
 type FileConfig struct {
-	Filename string `json:"filename" pflag:",Filename to store exported telemetry traces"`
+	Filename   string     `json:"filename" pflag:",Filename to store exported telemetry traces"`
+	SpanFormat SpanFormat `json:"spanFormat" pflag:",Format of the telemetry traces to export [otel/otel_collector]"`
 }
 
 type JaegerConfig struct {

diff --git a/flytestdlib/otelutils/config_flags.go b/flytestdlib/otelutils/config_flags.go
diff --git a/flytestdlib/otelutils/config_flags_test.go b/flytestdlib/otelutils/config_flags_test.go
diff --git a/flytestdlib/otelutils/exporter.go b/flytestdlib/otelutils/exporter.go
diff --git a/flytestdlib/otelutils/factory.go b/flytestdlib/otelutils/factory.go
@@ -7,7 +7,6 @@ import (
 
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/exporters/jaeger"
-	//"go.opentelemetry.io/otel/exporters/stdout/stdouttrace"
 	"go.opentelemetry.io/otel/sdk/resource"
 	"go.opentelemetry.io/otel/sdk/trace"
 	semconv "go.opentelemetry.io/otel/semconv/v1.21.0"
@@ -29,11 +28,11 @@ const (
 	K8sClientTracer         = "k8s-client"
 )
 
-var tracerProviders = make(map[string]*trace.TracerProvider)
-var noopTracerProvider = rawtrace.NewNoopTracerProvider()
-
-//var fileExporter *stdouttrace.Exporter
-var fileExporter trace.SpanExporter
+var (
+	fileExporter       trace.SpanExporter
+	noopTracerProvider = rawtrace.NewNoopTracerProvider()
+	tracerProviders    = make(map[string]*trace.TracerProvider)
+)
 
 func RegisterTracerProvider(serviceName string, config *Config) error {
 	if config == nil {
@@ -46,21 +45,26 @@ func RegisterTracerProvider(serviceName string, config *Config) error {
 		return nil
 	case FileExporter:
 		if fileExporter == nil {
-			// configure file exporter
+			// initialie file exporter and reuse between all services
 			f, err := os.Create(config.FileConfig.Filename)
 			if err != nil {
 				return err
 			}
 
-			/*fileExporter, err = stdouttrace.New(
-				stdouttrace.WithWriter(f),
-				//stdouttrace.WithPrettyPrint(), // TODO @hamersaw - remove?
-			)
-			if err != nil {
-				return err
-			}*/
+			var spanFormatter spanFormatter
+			switch config.FileConfig.SpanFormat {
+			case OtelSpanFormat:
+				spanFormatter = NewOTELSpanFormatter()
+			case OtelCollectorSpanFormat:
+				spanFormatter, err = NewOTELCollectorSpanFormatter()
+				if err != nil {
+					return err
+				}
+			default:
+				return fmt.Errorf("unknown span format [%v]", config.FileConfig.SpanFormat)
+			}
 
-			fileExporter, err = NewExporter(f)
+			fileExporter, err = NewExporter(spanFormatter, f)
 			if err != nil {
 				return err
 			}

diff --git a/flytestdlib/otelutils/factory_test.go b/flytestdlib/otelutils/factory_test.go
@@ -23,8 +23,7 @@ func TestRegisterTracerProvider(t *testing.T) {
 		FileConfig: FileConfig{
 			Filename: "/dev/null",
 		},
-		JaegerConfig: JaegerConfig{
-		},
+		JaegerConfig: JaegerConfig{},
 	}
 	err = RegisterTracerProvider(serviceName, &fullConfig)
 	assert.Nil(t, err)