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

     1  package tree
     2  
     3  import (
     4  	"fmt"
     5  	"math/rand"
     6  
     7  	. "github.com/onsi/ginkgo/v2/dsl/core"
     8  	. "github.com/onsi/gomega"
     9  )
    10  
    11  // StringWithEmpty is used for testing only, hence declared in this file. It's
    12  // different to String() as it includes nodes with zero value.
    13  func (t *Tree) StringWithEmpty() string {
    14  	t.RLock()
    15  	defer t.RUnlock()
    16  
    17  	res := ""
    18  	t.Iterate(func(k []byte, v uint64) {
    19  		if len(k) >= 2 {
    20  			res += fmt.Sprintf("%q %d\n", k[2:], v)
    21  		}
    22  	})
    23  	return res
    24  }
    25  
    26  var _ = Describe("tree package", func() {
    27  	Context("diff", func() {
    28  		Context("similar trees", func() {
    29  			treeA := New()
    30  			treeA.Insert([]byte("a;b"), uint64(1))
    31  			treeA.Insert([]byte("a;c"), uint64(2))
    32  			It("properly sets up tree A", func() {
    33  				Expect(treeA.StringWithEmpty()).To(Equal(treeStr(`"a" 0|"a;b" 1|"a;c" 2|`)))
    34  			})
    35  
    36  			treeB := New()
    37  			treeB.Insert([]byte("a;b"), uint64(4))
    38  			treeB.Insert([]byte("a;c"), uint64(8))
    39  			It("properly sets up tree B", func() {
    40  				Expect(treeB.StringWithEmpty()).To(Equal(treeStr(`"a" 0|"a;b" 4|"a;c" 8|`)))
    41  			})
    42  
    43  			It("properly combine trees", func() {
    44  				CombineTree(treeA, treeB)
    45  
    46  				Expect(treeA.StringWithEmpty()).To(Equal(treeStr(`"a" 0|"a;b" 1|"a;c" 2|`)))
    47  				Expect(treeB.StringWithEmpty()).To(Equal(treeStr(`"a" 0|"a;b" 4|"a;c" 8|`)))
    48  			})
    49  
    50  			It("properly combine trees to flamebearer", func() {
    51  				f := CombineToFlamebearerStruct(treeA, treeB, 1024)
    52  
    53  				Expect(f.Names).To(ConsistOf("total", "a", "b", "c"))
    54  				Expect(f.Levels).To(Equal([][]int{
    55  					// i+0 = x offset, left  tree
    56  					// i+1 = total   , left  tree
    57  					// i+2 = self    , left  tree
    58  					// i+3 = x offset, right tree
    59  					// i+4 = total   , right tree
    60  					// i+5 = self    , right tree
    61  					// i+6 = index in the names array
    62  					{0, 3, 0, 0, 12, 0, 0},
    63  					{0, 3, 0, 0, 12, 0, 1},
    64  					{0, 1, 1, 0, 4, 4, 3, 0, 2, 2, 0, 8, 8, 2},
    65  				}))
    66  				Expect(f.NumTicks).To(Equal(15))
    67  				Expect(f.MaxSelf).To(Equal(8))
    68  			})
    69  		})
    70  
    71  		Context("tree with an extra node", func() {
    72  			treeA := New()
    73  			treeA.Insert([]byte("a;b"), uint64(1))
    74  			treeA.Insert([]byte("a;c"), uint64(2))
    75  			treeA.Insert([]byte("a;e"), uint64(3))
    76  			It("properly sets up tree A", func() {
    77  				Expect(treeA.StringWithEmpty()).To(Equal(treeStr(`"a" 0|"a;b" 1|"a;c" 2|"a;e" 3|`)))
    78  			})
    79  
    80  			treeB := New()
    81  			treeB.Insert([]byte("a;b"), uint64(4))
    82  			treeB.Insert([]byte("a;d"), uint64(8))
    83  			treeB.Insert([]byte("a;e"), uint64(12))
    84  			It("properly sets up tree B", func() {
    85  				Expect(treeB.StringWithEmpty()).To(Equal(treeStr(`"a" 0|"a;b" 4|"a;d" 8|"a;e" 12|`)))
    86  			})
    87  
    88  			It("properly combine trees", func() {
    89  				CombineTree(treeA, treeB)
    90  
    91  				expectedA := `"a" 0|"a;b" 1|"a;c" 2|"a;d" 0|"a;e" 3|`
    92  				expectedB := `"a" 0|"a;b" 4|"a;c" 0|"a;d" 8|"a;e" 12|`
    93  				Expect(treeA.StringWithEmpty()).To(Equal(treeStr(expectedA)))
    94  				Expect(treeB.StringWithEmpty()).To(Equal(treeStr(expectedB)))
    95  			})
    96  
    97  			It("properly combine trees to flamebearer", func() {
    98  				f := CombineToFlamebearerStruct(treeA, treeB, 1024)
    99  
   100  				Expect(f.Names).To(ConsistOf("total", "a", "b", "c", "d", "e"))
   101  				Expect(f.Levels).To(Equal([][]int{
   102  					// i+0 = x offset, left  tree
   103  					// i+1 = total   , left  tree
   104  					// i+2 = self    , left  tree
   105  					// i+3 = x offset, right tree
   106  					// i+4 = total   , right tree
   107  					// i+5 = self    , right tree
   108  					// i+6 = index in the names array
   109  					{0, 6, 0, 0, 24, 0, 0}, // total
   110  					{0, 6, 0, 0, 24, 0, 1}, //    a
   111  					{
   112  						0, 1, 1, 0, 4, 4, 5, //   e
   113  						0, 2, 2, 0, 0, 0, 4, //   d
   114  						0, 0, 0, 0, 8, 8, 3, //   c
   115  						0, 3, 3, 0, 12, 12, 2, // b
   116  					},
   117  				}))
   118  				Expect(f.NumTicks).To(Equal(30))
   119  				Expect(f.MaxSelf).To(Equal(12))
   120  			})
   121  		})
   122  
   123  		Context("tree with many nodes", func() {
   124  			It("groups nodes into a new \"other\" node", func() {
   125  				treeA, treeB := New(), New()
   126  				r := rand.New(rand.NewSource(123))
   127  				for i := 0; i < 2048; i++ {
   128  					treeA.Insert([]byte(fmt.Sprintf("foo;bar%d", i)), uint64(r.Intn(4000)))
   129  					treeB.Insert([]byte(fmt.Sprintf("foo;bar%d", i)), uint64(r.Intn(4000)))
   130  				}
   131  
   132  				f := CombineToFlamebearerStruct(treeA, treeB, 10)
   133  				Expect(f.Names).To(ContainElement("other"))
   134  			})
   135  		})
   136  	})
   137  })