Merge pull request #366 from pyroscope-io/249

server benchmark prep work

benchmark/cmd/command/banner.go

@@ -0,0 +1,26 @@
+package command
+
+import (
+	"github.com/pyroscope-io/pyroscope/pkg/cli"
+)
+
+// made here http://patorjk.com/software/taag/#p=display&f=Doom&t=Pyrobench
+var banner = `
+______                _                     _
+| ___ \              | |                   | |
+| |_/ /   _ _ __ ___ | |__   ___ _ __   ___| |__
+|  __/ | | | '__/ _ \| '_ \ / _ \ '_ \ / __| '_ \
+| |  | |_| | | | (_) | |_) |  __/ | | | (__| | | |
+\_|   \__, |_|  \___/|_.__/ \___|_| |_|\___|_| |_|
+       __/ |
+      |___/
+`
+
+func init() {
+	// removes extra new lines
+	banner = banner[1 : len(banner)-2]
+}
+
+func gradientBanner() string {
+	return cli.GradientBanner(banner)
+}

benchmark/cmd/command/command.go

@@ -0,0 +1,10 @@
+package command
+
+import (
+	"github.com/pyroscope-io/pyroscope/pkg/cli"
+	"github.com/spf13/viper"
+)
+
+func newViper() *viper.Viper {
+	return cli.NewViper("PYROBENCH")
+}

benchmark/cmd/command/loadgen.go

@@ -0,0 +1,22 @@
+package command
+
+import (
+	"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.LoadGen) *cobra.Command {
+	vpr := newViper()
+	loadgenCmd := &cobra.Command{
+		Use:   "loadgen [flags]",
+		Short: "Generates load",
+		RunE: cli.CreateCmdRunFn(cfg, vpr, func(_ *cobra.Command, args []string) error {
+			return loadgen.Cli(cfg)
+		}),
+	}
+
+	cli.PopulateFlagSet(cfg, loadgenCmd.Flags(), vpr)
+	return loadgenCmd
+}

benchmark/cmd/command/root.go

@@ -0,0 +1,66 @@
+package command
+
+import (
+	"fmt"
+	"os"
+	"runtime"
+	"strings"
+
+	"github.com/pyroscope-io/pyroscope/benchmark/config"
+	"github.com/pyroscope-io/pyroscope/pkg/cli"
+	"github.com/sirupsen/logrus"
+	"github.com/spf13/cobra"
+)
+
+func newRootCmd(cfg *config.LoadGen) *cobra.Command {
+	rootCmd := &cobra.Command{
+		Use: "pyrobench [flags] <subcommand>",
+	}
+
+	rootCmd.SetUsageFunc(func(cmd *cobra.Command) error {
+		fmt.Println(gradientBanner())
+		fmt.Println(cli.DefaultUsageFunc(cmd.Flags(), cmd))
+		return nil
+	})
+
+	rootCmd.SetHelpFunc(func(cmd *cobra.Command, a []string) {
+		fmt.Println(gradientBanner())
+		fmt.Println(cli.DefaultUsageFunc(cmd.Flags(), cmd))
+	})
+
+	return rootCmd
+}
+
+// Initialize adds all child commands to the root command and sets flags appropriately
+func Initialize() error {
+	var cfg config.Config
+
+	rootCmd := newRootCmd(&cfg.LoadGen)
+	rootCmd.SilenceErrors = true
+	rootCmd.AddCommand(
+		newLoadGen(&cfg.LoadGen),
+	)
+
+	logrus.SetReportCaller(true)
+	logrus.SetFormatter(&logrus.TextFormatter{
+		TimestampFormat: "2006-01-02T15:04:05.000000",
+		FullTimestamp:   true,
+		CallerPrettyfier: func(f *runtime.Frame) (string, string) {
+			filename := f.File
+			if len(filename) > 38 {
+				filename = filename[38:]
+			}
+			return "", fmt.Sprintf(" %s:%d", filename, f.Line)
+		},
+	})
+
+	args := os.Args[1:]
+	for i, arg := range args {
+		if len(arg) > 2 && strings.HasPrefix(arg, "-") && !strings.HasPrefix(arg, "--") {
+			args[i] = fmt.Sprintf("-%s", arg)
+		}
+	}
+
+	rootCmd.SetArgs(args)
+	return rootCmd.Execute()
+}

benchmark/cmd/logging.go

@@ -0,0 +1,9 @@
+package main
+
+import (
+	"github.com/pyroscope-io/pyroscope/pkg/cli"
+)
+
+func init() {
+	cli.InitLogging()
+}

benchmark/cmd/main.go

@@ -0,0 +1,20 @@
+package main
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/fatih/color"
+	"github.com/pyroscope-io/pyroscope/benchmark/cmd/command"
+)
+
+func main() {
+	if err := command.Initialize(); err != nil {
+		fatalf("%s %v\n\n", color.RedString("Error:"), err)
+	}
+}
+
+func fatalf(format string, args ...interface{}) {
+	_, _ = fmt.Fprintf(os.Stderr, format, args...)
+	os.Exit(1)
+}

benchmark/config/config.go

@@ -0,0 +1,24 @@
+package config
+
+type Config struct {
+	LoadGen LoadGen `skip:"true" mapstructure:",squash"`
+}
+
+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"`
+
+	WaitUntilAvailable bool `def:"true" desc:"wait until endpoint is available"`
+}
+
+// File can be read from file system.
+type File interface{ Path() string }

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)
+}

cmd/pyroscope/command/agent.go

@@ -14,7 +14,7 @@ func newAgentCmd(cfg *config.Agent) *cobra.Command {
 		Short: "Start pyroscope agent",
 
 		DisableFlagParsing: true,
-		RunE: createCmdRunFn(cfg, vpr, func(_ *cobra.Command, _ []string) error {
+		RunE: cli.CreateCmdRunFn(cfg, vpr, func(_ *cobra.Command, _ []string) error {
 			return cli.StartAgent(cfg)
 		}),
 	}