pyrobench: basic load generation

benchmark/cmd/command/loadgen.go

@@ -1,22 +1,19 @@
 package command
 
 import (
-	"fmt"
-
 	"github.com/pyroscope-io/pyroscope/benchmark/config"
+	"github.com/pyroscope-io/pyroscope/benchmark/loadgen"
 	"github.com/pyroscope-io/pyroscope/pkg/cli"
 	"github.com/spf13/cobra"
 )
 
-func newLoadGen(cfg *config.Config) *cobra.Command {
+func newLoadGen(cfg *config.LoadGen) *cobra.Command {
 	vpr := newViper()
 	loadgenCmd := &cobra.Command{
 		Use:   "loadgen [flags]",
 		Short: "Generates load",
 		RunE: createCmdRunFn(cfg, vpr, false, func(cmd *cobra.Command, args []string, logger config.LoggerFunc) error {
-			fmt.Println("address", cfg.ServerAddress)
-
-			return nil
+			return loadgen.Cli(cfg)
 		}),
 	}
 

benchmark/cmd/command/root.go

@@ -11,7 +11,7 @@ import (
 	"github.com/spf13/cobra"
 )
 
-func newRootCmd(cfg *config.Config) *cobra.Command {
+func newRootCmd(cfg *config.LoadGen) *cobra.Command {
 	rootCmd := &cobra.Command{
 		Use: "pyrobench [flags] <subcommand>",
 	}
@@ -32,7 +32,7 @@ func newRootCmd(cfg *config.Config) *cobra.Command {
 
 // Initialize adds all child commands to the root command and sets flags appropriately
 func Initialize() error {
-	var cfg config.Config
+	var cfg config.LoadGen
 
 	rootCmd := newRootCmd(&cfg)
 	rootCmd.SilenceErrors = true

benchmark/cmd/logging.go

@@ -0,0 +1,23 @@
+// Copied as is from github.com/pyroscope-io/pyroscope/cmd/logging.go
+package main
+
+import (
+	"log"
+	"os"
+	"runtime"
+
+	"github.com/fatih/color"
+	"github.com/sirupsen/logrus"
+)
+
+func init() {
+	log.SetFlags(log.Lshortfile | log.Ldate | log.Ltime)
+
+	logrus.SetFormatter(&logrus.TextFormatter{})
+	logrus.SetOutput(os.Stdout)
+	logrus.SetLevel(logrus.DebugLevel)
+
+	if runtime.GOOS == "windows" {
+		color.NoColor = true
+	}
+}

benchmark/config/config.go

@@ -1,9 +1,17 @@
 package config
 
-type Config struct {
-	Version bool `mapstructure:"version"`
-}
+type LoadGen struct {
+	LogLevel string `def:"info" desc:"log level: debug|info|warn|error" mapstructure:"log-level"`
+
+	ServerAddress       string `def:"http://localhost:4040" desc:"address of the pyroscope instance being attacked" mapstructure:"server-address"`
+	RandSeed            int    `def:"23061912" desc:""`
+	ProfileWidth        int    `def:"20"`
+	ProfileDepth        int    `def:"20"`
+	ProfileSymbolLength int    `def:"30"`
+	Fixtures            int    `def:"30" desc:"how many different profiles to generate per app"`
+	Apps                int    `def:"20" desc:"how many pyroscope apps to emulate"`
+	Clients             int    `def:"20" desc:"how many pyroscope clients to emulate"`
+	Requests            int    `def:"10000" desc:"how many requests each clients should make"`
 
-type LoggerFunc func(s string)
-type LoggerConfiger interface{ InitializeLogging() LoggerFunc }
-type FileConfiger interface{ ConfigFilePath() string }
+	WaitUntilAvailable bool `def:"true" desc:"wait until endpoint is available"`
+}

benchmark/config/interface.go

@@ -0,0 +1,18 @@
+package config
+
+import (
+	"github.com/sirupsen/logrus"
+)
+
+type FileConfiger interface{ ConfigFilePath() string }
+
+type LoggerFunc func(s string)
+type LoggerConfiger interface{ InitializeLogging() LoggerFunc }
+
+func (cfg LoadGen) InitializeLogging() LoggerFunc {
+	if l, err := logrus.ParseLevel(cfg.LogLevel); err == nil {
+		logrus.SetLevel(l)
+	}
+
+	return nil
+}

benchmark/loadgen/loadgen.go

@@ -0,0 +1,170 @@
+package loadgen
+
+import (
+	"encoding/hex"
+	"fmt"
+	"math/rand"
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/pyroscope-io/pyroscope/benchmark/config"
+	"github.com/pyroscope-io/pyroscope/pkg/agent/upstream"
+	"github.com/pyroscope-io/pyroscope/pkg/agent/upstream/remote"
+	"github.com/pyroscope-io/pyroscope/pkg/structs/transporttrie"
+	"github.com/sirupsen/logrus"
+)
+
+// how many retries to check the pyroscope server is up
+const MaxReadinessRetries = 10
+
+type Fixtures [][]*transporttrie.Trie
+
+type LoadGen struct {
+	Config    *config.LoadGen
+	Rand      *rand.Rand
+	SymbolBuf []byte
+}
+
+func Cli(cfg *config.LoadGen) error {
+	r := rand.New(rand.NewSource(int64(cfg.RandSeed)))
+	l := &LoadGen{
+		Config:    cfg,
+		Rand:      r,
+		SymbolBuf: make([]byte, cfg.ProfileSymbolLength),
+	}
+
+	return l.Run(cfg)
+}
+
+func (l *LoadGen) Run(cfg *config.LoadGen) error {
+	logrus.Info("checking server is available...")
+	err := waitUntilEndpointReady(cfg.ServerAddress)
+	if err != nil {
+		return err
+	}
+
+	logrus.Info("generating fixtures")
+	fixtures := l.generateFixtures()
+	logrus.Debug("done generating fixtures.")
+
+	logrus.Info("starting sending requests")
+	wg := sync.WaitGroup{}
+	wg.Add(l.Config.Apps * l.Config.Clients)
+	appNameBuf := make([]byte, 25)
+
+	for i := 0; i < l.Config.Apps; i++ {
+		// generate a random app name
+		l.Rand.Read(appNameBuf)
+		appName := hex.EncodeToString(appNameBuf)
+		for j := 0; j < l.Config.Clients; j++ {
+			go l.startClientThread(appName, &wg, fixtures[i])
+		}
+	}
+	wg.Wait()
+
+	logrus.Debug("done sending requests")
+	return nil
+}
+
+func (l *LoadGen) generateFixtures() Fixtures {
+	var f Fixtures
+
+	for i := 0; i < l.Config.Apps; i++ {
+		f = append(f, []*transporttrie.Trie{})
+
+		randomGen := rand.New(rand.NewSource(int64(l.Config.RandSeed + i)))
+		p := l.generateProfile(randomGen)
+		for j := 0; j < l.Config.Fixtures; j++ {
+			f[i] = append(f[i], p)
+		}
+	}
+
+	return f
+}
+
+func (l *LoadGen) startClientThread(appName string, wg *sync.WaitGroup, appFixtures []*transporttrie.Trie) {
+	rc := remote.RemoteConfig{
+		UpstreamThreads:        1,
+		UpstreamAddress:        l.Config.ServerAddress,
+		UpstreamRequestTimeout: 10 * time.Second,
+	}
+	r, err := remote.New(rc, logrus.StandardLogger())
+	if err != nil {
+		panic(err)
+	}
+
+	requestsCount := l.Config.Requests
+
+	threadStartTime := time.Now().Truncate(10 * time.Second)
+	threadStartTime = threadStartTime.Add(time.Duration(-1*requestsCount) * (10 * time.Second))
+
+	st := threadStartTime
+
+	for i := 0; i < requestsCount; i++ {
+		t := appFixtures[i%len(appFixtures)]
+
+		st = st.Add(10 * time.Second)
+		et := st.Add(10 * time.Second)
+		err := r.UploadSync(&upstream.UploadJob{
+			Name:            appName + "{}",
+			StartTime:       st,
+			EndTime:         et,
+			SpyName:         "gospy",
+			SampleRate:      100,
+			Units:           "samples",
+			AggregationType: "sum",
+			Trie:            t,
+		})
+		if err != nil {
+			// TODO(eh-am): calculate errors
+			time.Sleep(time.Second)
+		} else {
+			// TODO(eh-am): calculate success
+		}
+	}
+
+	wg.Done()
+}
+
+func (l *LoadGen) generateProfile(randomGen *rand.Rand) *transporttrie.Trie {
+	t := transporttrie.New()
+
+	for w := 0; w < l.Config.ProfileWidth; w++ {
+		symbol := []byte("root")
+		for d := 0; d < 2+l.Rand.Intn(l.Config.ProfileDepth); d++ {
+			randomGen.Read(l.SymbolBuf)
+			symbol = append(symbol, byte(';'))
+			symbol = append(symbol, []byte(hex.EncodeToString(l.SymbolBuf))...)
+			if l.Rand.Intn(100) <= 20 {
+				t.Insert(symbol, uint64(l.Rand.Intn(100)), true)
+			}
+		}
+
+		t.Insert(symbol, uint64(l.Rand.Intn(100)), true)
+	}
+	return t
+}
+
+// TODO(eh-am) exponential backoff and whatnot
+func waitUntilEndpointReady(url string) error {
+	client := http.Client{Timeout: 10 * time.Second}
+	retries := 0
+
+	for {
+		_, err := client.Get(url)
+
+		// all good?
+		if err == nil {
+			return nil
+		}
+		if retries >= MaxReadinessRetries {
+			break
+		}
+
+		time.Sleep(time.Second)
+		retries++
+	}
+
+	return fmt.Errorf("maximum retries exceeded ('%d') waiting for server ('%s') to respond", retries, url)
+}