add telemetrygen public package

pkg/telemetrygen/config.go

@@ -0,0 +1,95 @@
+// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+// or more contributor license agreements. Licensed under the Elastic License 2.0;
+// you may not use this file except in compliance with the Elastic License 2.0.
+
+package telemetrygen
+
+import (
+	"errors"
+	"fmt"
+	"net/url"
+	"strconv"
+	"strings"
+	"time"
+)
+
+func DefaultConfig() Config {
+	return Config{
+		// default to run 1 replica of 4 agents (one per type).
+		AgentReplicas: 1,
+		// default to expecting a valid TLS certificate.
+		Secure: true,
+		// default to 10 events per second
+		EventRate: RateFlag{Burst: 10, Interval: 1 * time.Second},
+		// default to rewrite ids and timestamp to have new events recorded at the current time.
+		RewriteIDs:        true,
+		RewriteTimestamps: true,
+	}
+}
+
+type Config struct {
+	// number of agents replicas to use, each replica launches 4 agents, one for each type
+	AgentReplicas int
+
+	ServerURL    *url.URL
+	APIKey       string
+	Headers      map[string]string
+	Secure       bool
+	EventRate    RateFlag
+	IgnoreErrors bool
+
+	RewriteIDs                bool
+	RewriteTimestamps         bool
+	RewriteServiceNames       bool
+	RewriteServiceNodeNames   bool
+	RewriteServiceTargetNames bool
+	RewriteSpanNames          bool
+	RewriteTransactionNames   bool
+	RewriteTransactionTypes   bool
+}
+
+func (c Config) Validate() error {
+	errs := []error{}
+
+	if c.ServerURL == nil {
+		errs = append(errs, fmt.Errorf("ServerURL is required"))
+	}
+
+	return errors.Join(errs...)
+}
+
+type RateFlag struct {
+	Burst    int
+	Interval time.Duration
+}
+
+func (f *RateFlag) String() string {
+	return fmt.Sprintf("%d/%s", f.Burst, f.Interval)
+}
+
+func (f *RateFlag) Set(s string) error {
+	before, after, ok := strings.Cut(s, "/")
+	if !ok || before == "" || after == "" {
+		return fmt.Errorf("invalid rate %q, expected format burst/duration", s)
+	}
+
+	burst, err := strconv.Atoi(before)
+	if err != nil {
+		return fmt.Errorf("invalid burst %s in event rate: %w", before, err)
+	}
+
+	if !(after[0] >= '0' && after[0] <= '9') {
+		after = "1" + after
+	}
+	interval, err := time.ParseDuration(after)
+	if err != nil {
+		return fmt.Errorf("invalid interval %q in event rate: %w", after, err)
+	}
+	if interval <= 0 {
+		return fmt.Errorf("invalid interval %q, must be positive", after)
+	}
+
+	f.Burst = burst
+	f.Interval = interval
+	return nil
+}

pkg/telemetrygen/doc.go

@@ -0,0 +1,10 @@
+// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+// or more contributor license agreements. Licensed under the Elastic License 2.0;
+// you may not use this file except in compliance with the Elastic License 2.0.
+
+// Package telemetrygen provides an easy way to generate
+// dummy APM events from a corpus of canned events.
+// Is a public package that is expected to be used outside
+// of this module, as such backward compatibility rules
+// apply.
+package telemetrygen

pkg/telemetrygen/generator.go

@@ -0,0 +1,86 @@
+// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+// or more contributor license agreements. Licensed under the Elastic License 2.0;
+// you may not use this file except in compliance with the Elastic License 2.0.
+
+package telemetrygen
+
+import (
+	"context"
+	cryptorand "crypto/rand"
+	"encoding/binary"
+	"fmt"
+	"math/rand"
+
+	"go.uber.org/zap"
+
+	"github.com/elastic/apm-perf/internal/loadgen"
+
+	"golang.org/x/sync/errgroup"
+	"golang.org/x/time/rate"
+)
+
+type Generator struct {
+	Config Config
+	Logger *zap.Logger
+}
+
+func New(cfg Config) (*Generator, error) {
+	if err := cfg.Validate(); err != nil {
+		return &Generator{}, fmt.Errorf("cannot create generator, configuration is invalid: %w", err)
+	}
+
+	return &Generator{Config: cfg, Logger: zap.NewNop()}, nil
+}
+
+func (g *Generator) RunBlocking(ctx context.Context) error {
+	limiter := loadgen.GetNewLimiter(g.Config.EventRate.Burst, g.Config.EventRate.Interval)
+	errg, gCtx := errgroup.WithContext(ctx)
+
+	var rngseed int64
+	err := binary.Read(cryptorand.Reader, binary.LittleEndian, &rngseed)
+	if err != nil {
+		return fmt.Errorf("failed to generate seed for math/rand: %w", err)
+	}
+
+	for i := 0; i < g.Config.AgentReplicas; i++ {
+		for _, expr := range []string{`apm-go*.ndjson`, `apm-nodejs*.ndjson`, `apm-python*.ndjson`, `apm-ruby*.ndjson`} {
+			expr := expr
+			errg.Go(func() error {
+				rng := rand.New(rand.NewSource(rngseed))
+				return runAgent(gCtx, g.Logger, expr, limiter, rng, g.Config)
+			})
+		}
+	}
+
+	return errg.Wait()
+}
+
+func runAgent(ctx context.Context, l *zap.Logger, expr string, limiter *rate.Limiter, rng *rand.Rand, cfg Config) error {
+	handler, err := loadgen.NewEventHandler(loadgen.EventHandlerParams{
+		Logger:                    l,
+		URL:                       cfg.ServerURL.String(),
+		Path:                      expr,
+		APIKey:                    cfg.APIKey,
+		Limiter:                   limiter,
+		Rand:                      rng,
+		RewriteIDs:                cfg.RewriteIDs,
+		RewriteServiceNames:       cfg.RewriteServiceNames,
+		RewriteServiceNodeNames:   cfg.RewriteServiceNodeNames,
+		RewriteServiceTargetNames: cfg.RewriteServiceTargetNames,
+		RewriteSpanNames:          cfg.RewriteSpanNames,
+		RewriteTransactionNames:   cfg.RewriteTransactionNames,
+		RewriteTransactionTypes:   cfg.RewriteTransactionTypes,
+		RewriteTimestamps:         cfg.RewriteTimestamps,
+		Headers:                   cfg.Headers,
+		Protocol:                  "apm/http",
+	})
+	if err != nil {
+		return err
+	}
+
+	if _, err := handler.SendBatches(ctx); err != nil {
+		return err
+	}
+
+	return nil
+}

pkg/telemetrygen/generator_test.go

@@ -0,0 +1,55 @@
+// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+// or more contributor license agreements. Licensed under the Elastic License 2.0;
+// you may not use this file except in compliance with the Elastic License 2.0.
+
+package telemetrygen_test
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"sync/atomic"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/elastic/apm-perf/pkg/telemetrygen"
+)
+
+func TestGeneration(t *testing.T) {
+	srv, m := newTestServer(t)
+
+	cfg := telemetrygen.DefaultConfig()
+	cfg.Secure = false
+
+	u, err := url.Parse(srv.URL)
+	require.NoError(t, err)
+	cfg.ServerURL = u
+
+	cfg.EventRate.Set("1000/s")
+	g, err := telemetrygen.New(cfg)
+	// g.Logger = zap.Must(zap.NewDevelopment())
+	require.NoError(t, err)
+
+	err = g.RunBlocking(context.Background())
+	require.NoError(t, err)
+	require.Greater(t, m.Load(), int32(0))
+}
+
+func newTestServer(t *testing.T) (*httptest.Server, *atomic.Int32) {
+	t.Helper()
+
+	mux := http.NewServeMux()
+	requestsReceived := &atomic.Int32{}
+
+	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+		requestsReceived.Add(1)
+		w.Write([]byte("ok"))
+	})
+
+	srv := httptest.NewServer(mux)
+	t.Cleanup(srv.Close)
+
+	return srv, requestsReceived
+}