github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/testing/load/tree.go (about)

     1  package load
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/hex"
     6  	"math/rand"
     7  
     8  	"github.com/pyroscope-io/pyroscope/pkg/storage/tree"
     9  )
    10  
    11  type TreeGenerator struct {
    12  	TreeConfig
    13  	seed int
    14  
    15  	trees  []*tree.Tree
    16  	symBuf []byte
    17  	b      *bytes.Buffer
    18  	i      int
    19  }
    20  
    21  type TreeConfig struct {
    22  	MaxSymLen int `yaml:"maxSymLen"`
    23  	MaxDepth  int `yaml:"maxDepth"`
    24  	Width     int `yaml:"width"`
    25  }
    26  
    27  var rootNode = []byte("root")
    28  
    29  const (
    30  	minStackDepth = 2
    31  	minSymLength  = 3
    32  )
    33  
    34  func NewTreeGenerator(seed, trees int, c TreeConfig) *TreeGenerator {
    35  	g := TreeGenerator{
    36  		TreeConfig: c,
    37  		seed:       seed,
    38  		symBuf:     make([]byte, c.MaxSymLen),
    39  		trees:      make([]*tree.Tree, trees),
    40  	}
    41  	g.b = bytes.NewBuffer(make([]byte, 128))
    42  	for i := 0; i < trees; i++ {
    43  		seed++
    44  		g.trees[i] = g.generateTree(newRand(seed))
    45  	}
    46  	return &g
    47  }
    48  
    49  func (g *TreeGenerator) Next() *tree.Tree {
    50  	g.i++
    51  	return g.trees[g.i%len(g.trees)]
    52  }
    53  
    54  func (g *TreeGenerator) generateTree(r *rand.Rand) *tree.Tree {
    55  	t := tree.New()
    56  	for w := 0; w < g.Width; w++ {
    57  		t.Insert(g.generateStack(r), uint64(r.Intn(100)))
    58  	}
    59  	return t
    60  }
    61  
    62  func (g *TreeGenerator) generateStack(r *rand.Rand) []byte {
    63  	g.b.Reset()
    64  	g.b.Write(rootNode)
    65  	e := hex.NewEncoder(g.b)
    66  	d := randInt(r, minStackDepth, g.MaxDepth)
    67  	for i := 0; i < d; i++ {
    68  		l := randInt(r, minSymLength, g.MaxSymLen)
    69  		r.Read(g.symBuf[:l])
    70  		g.b.WriteString(";")
    71  		_, _ = e.Write(g.symBuf[:l])
    72  	}
    73  	s := make([]byte, g.b.Len())
    74  	copy(s, g.b.Bytes())
    75  	return s
    76  }