load.go

package main

import (
	"bytes"
	"compress/zlib"
	"context"
	cryptorand "crypto/rand"
	"encoding/binary"
	"flag"
	"fmt"
	"log"
	"math/rand"
	"time"

	"go.elastic.co/apm/model"
	"go.elastic.co/apm/transport"
	"go.elastic.co/apm/transport/transporttest"
	"go.elastic.co/fastjson"
)

var (
	falsy = false
)

func main() {
	dryrun := flag.Bool("n", false, "dryrun")
	transactions := flag.Int("t", 1, "max transactions to generate")
	flag.Parse()

	tr, err := newTransport(!*dryrun)
	if err != nil {
		log.Fatal(err)
	}

	var w fastjson.Writer
	if err := meta(&w); err != nil {
		log.Fatal(err)
	}
	for i := 0; i < *transactions; i++ {
		if err := trans(&w); err != nil {
			log.Fatal(err)
		}
	}
	if err := sendEvents(context.Background(), tr, &w); err != nil {
		log.Fatal(err)
	}
	if t, ok := tr.(*transporttest.RecorderTransport); ok {
		system, process, service, stringMap := t.Metadata()
		fmt.Printf("metadata: %+v %+v %+v %+v\n", system, process, service, stringMap)
		fmt.Printf("payloads: %+v\n", t.Payloads())
	}
}

func meta(w *fastjson.Writer) error {
	service := model.Service{
		Name: "load-apm",
		Agent: &model.Agent{
			Name: "load",
		},
	}
	w.RawString(`{"metadata":{`)
	w.RawString(`"service":`)
	if err := service.MarshalFastJSON(w); err != nil {
		return err
	}
	w.RawString("}}\n")
	return nil
}

func trans(w *fastjson.Writer) error {
	var transactionID model.SpanID
	if _, err := cryptorand.Read(transactionID[:]); err != nil {
		log.Fatal(err)
	}
	var traceID model.TraceID
	binary.LittleEndian.PutUint64(traceID[:8], rand.Uint64())
	binary.LittleEndian.PutUint64(traceID[8:], rand.Uint64())

	var duration = rand.Float64() * 100
	transaction := &model.Transaction{
		// required
		Duration: duration,
		ID:       transactionID,
		Name:     "Do",
		SpanCount: model.SpanCount{
			Dropped: 0,
			Started: 0,
		},
		Timestamp: model.Time(time.Now()),
		TraceID:   traceID,
		Type:      "request",

		// optional
		Sampled: &falsy,
	}
	w.RawString(`{"transaction":`)
	if err := transaction.MarshalFastJSON(w); err != nil {
		return err
	}
	w.RawString("}\n")
	return nil
}

func newTransport(send bool) (transport.Transport, error) {
	if send {
		tr, err := transport.NewHTTPTransport()
		if err != nil {
			return nil, fmt.Errorf("error creating HTTP transport: %v", err)
		}
		return tr, nil
	} else {
		return &transporttest.RecorderTransport{}, nil
	}
}

func sendEvents(ctx context.Context, tr transport.Transport, w *fastjson.Writer) error {
	var buf bytes.Buffer
	zw, err := zlib.NewWriterLevel(&buf, zlib.DefaultCompression)
	if err != nil {
		return err
	}
	if _, err := zw.Write(w.Bytes()); err != nil {
		return err
	}
	if err := zw.Close(); err != nil {
		return err
	}
	if err := tr.SendStream(ctx, &buf); err != nil {
		return err
	}
	return nil
}