tracegen: add capability to marshal spans to file

Signed-off-by: Samriddha9619 <sumitkumartripathi0@gmail.com>

Author: Samriddha9619

cmd/tracegen/main.go

@@ -8,6 +8,7 @@ import (
 	"errors"
 	"flag"
 	"fmt"
+	"os"
 	"time"
 
 	"github.com/go-logr/zapr"
@@ -71,13 +72,23 @@ func createTracers(cfg *tracegen.Config, logger *zap.Logger) ([]trace.Tracer, fu
 	}
 	var shutdown []func(context.Context) error
 	var tracers []trace.Tracer
+
+	var file *os.File
+	if cfg.MarshalToFile != "" {
+		var err error
+		file, err = os.Create(cfg.MarshalToFile)
+		if err != nil {
+			logger.Sugar().Fatalf("cannot create output file %s: %s", cfg.MarshalToFile, err)
+		}
+	}
+
 	for s := 0; s < cfg.Services; s++ {
 		svc := cfg.Service
 		if cfg.Services > 1 {
 			svc = fmt.Sprintf("%s-%02d", svc, s)
 		}
 
-		exp, err := createOtelExporter(cfg.TraceExporter)
+		exp, err := createOtelExporter(cfg.TraceExporter, file)
 		if err != nil {
 			logger.Sugar().Fatalf("cannot create trace exporter %s: %s", cfg.TraceExporter, err)
 		}
@@ -113,6 +124,12 @@ func createTracers(cfg *tracegen.Config, logger *zap.Logger) ([]trace.Tracer, fu
 		tracers = append(tracers, tp.Tracer(cfg.Service))
 		shutdown = append(shutdown, tp.Shutdown)
 	}
+	if file != nil {
+		shutdown = append(shutdown, func(_ context.Context) error {
+			return file.Close()
+		})
+	}
+
 	return tracers, func(ctx context.Context) error {
 		var errs []error
 		for _, f := range shutdown {
@@ -122,21 +139,24 @@ func createTracers(cfg *tracegen.Config, logger *zap.Logger) ([]trace.Tracer, fu
 	}
 }
 
-func createOtelExporter(exporterType string) (sdktrace.SpanExporter, error) {
+func createOtelExporter(exporterType string, fileWriter *os.File) (sdktrace.SpanExporter, error) {
+	if fileWriter != nil {
+		return stdouttrace.New(
+			stdouttrace.WithWriter(fileWriter),
+			stdouttrace.WithPrettyPrint(),
+		)
+	}
+
 	var exporter sdktrace.SpanExporter
 	var err error
 	switch exporterType {
 	case "jaeger":
 		return nil, errors.New("jaeger exporter is no longer supported, please use otlp")
 	case "otlp", "otlp-http":
-		client := otlptracehttp.NewClient(
-			otlptracehttp.WithInsecure(),
-		)
+		client := otlptracehttp.NewClient(otlptracehttp.WithInsecure())
 		exporter, err = otlptrace.New(context.Background(), client)
 	case "otlp-grpc":
-		client := otlptracegrpc.NewClient(
-			otlptracegrpc.WithInsecure(),
-		)
+		client := otlptracegrpc.NewClient(otlptracegrpc.WithInsecure())
 		exporter, err = otlptrace.New(context.Background(), client)
 	case "stdout":
 		exporter, err = stdouttrace.New()

internal/tracegen/config.go

@@ -30,6 +30,7 @@ type Config struct {
 	Duration      time.Duration
 	Service       string
 	TraceExporter string
+	MarshalToFile string
 }
 
 // Flags registers config flags.
@@ -47,6 +48,7 @@ func (c *Config) Flags(fs *flag.FlagSet) {
 	fs.StringVar(&c.Service, "service", "tracegen", "Service name prefix to use")
 	fs.IntVar(&c.Services, "services", 1, "Number of unique suffixes to add to service name when generating traces, e.g. tracegen-01 (but only one service per trace)")
 	fs.StringVar(&c.TraceExporter, "trace-exporter", "otlp-http", "Trace exporter (otlp/otlp-http|otlp-grpc|stdout). Exporters can be additionally configured via environment variables, see https://github.com/jaegertracing/jaeger/blob/main/cmd/tracegen/README.md")
+	fs.StringVar(&c.MarshalToFile, "marshal-to-file", "", "Filename to write generated spans to (in JSON format) instead of sending to a collector")
 }
 
 // Run executes the test scenario.