github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/convert/pprof/bench/parser_test.go (about)

     1  package bench
     2  
     3  import (
     4  	"bufio"
     5  	"bytes"
     6  	"compress/gzip"
     7  	"context"
     8  	"encoding/json"
     9  	"fmt"
    10  	"github.com/google/pprof/profile"
    11  	"github.com/pyroscope-io/pyroscope/pkg/convert/pprof"
    12  	"github.com/pyroscope-io/pyroscope/pkg/convert/pprof/streaming"
    13  	"github.com/pyroscope-io/pyroscope/pkg/ingestion"
    14  	"github.com/pyroscope-io/pyroscope/pkg/stackbuilder"
    15  	"github.com/pyroscope-io/pyroscope/pkg/storage"
    16  	"github.com/pyroscope-io/pyroscope/pkg/storage/metadata"
    17  	"github.com/pyroscope-io/pyroscope/pkg/storage/segment"
    18  	"github.com/pyroscope-io/pyroscope/pkg/util/form"
    19  	"golang.org/x/exp/slices"
    20  	"io"
    21  	"io/fs"
    22  	"math/big"
    23  	"mime/multipart"
    24  	"os"
    25  	"sort"
    26  	"strings"
    27  
    28  	"github.com/pyroscope-io/pyroscope/pkg/storage/tree"
    29  	"io/ioutil"
    30  	"net/http"
    31  
    32  	"testing"
    33  	"time"
    34  )
    35  
    36  const benchmarkCorpus = "../../../../../pprof-testdata"
    37  const compareCorpus = "../../../../../pprof-testdata"
    38  
    39  const pprofSmall = benchmarkCorpus +
    40  	"/2022-10-08T00:44:10Z-55903298-d296-4730-a28d-9dcc7c5e25d6.txt"
    41  const pprofBig = benchmarkCorpus +
    42  	"/2022-10-08T00:07:00Z-911c824f-a086-430c-99d7-315a53b58095.txt"
    43  
    44  // GOEXPERIMENT=arenas go test -v -test.count=10 -test.run=none -bench=".*Streaming.*"  ./pkg/convert/pprof/bench
    45  var putter = &MockPutter{}
    46  
    47  const benchWithoutGzip = true
    48  const benchmarkCorpusSize = 5
    49  
    50  var compareCorpusData = readCorpus(compareCorpus, benchWithoutGzip)
    51  
    52  func TestCompare(t *testing.T) {
    53  	if len(compareCorpusData) == 0 {
    54  		t.Skip("empty corpus")
    55  		return
    56  	}
    57  	for _, testType := range streamingTestTypes {
    58  		t.Run(fmt.Sprintf("TestCompare_pool_%v_arenas_%v", testType.pool, testType.arenas), func(t *testing.T) {
    59  			for _, c := range compareCorpusData {
    60  				testCompareOne(t, c, testType)
    61  			}
    62  		})
    63  	}
    64  }
    65  
    66  func TestCompareWriteBatch(t *testing.T) {
    67  	if len(compareCorpusData) == 0 {
    68  		t.Skip("empty corpus")
    69  		return
    70  	}
    71  	for _, c := range compareCorpusData {
    72  		//cur, _ := profile.Parse(bytes.NewReader(c.profile))
    73  		//if c.prev != nil {
    74  		//	prev, _ := profile.Parse(bytes.NewReader(c.prev))
    75  		//	os.WriteFile("p1", []byte(dumpPProfProfile(prev)), 0666)
    76  		//}
    77  		//os.WriteFile("p2", []byte(dumpPProfProfile(cur)), 0666)
    78  		testCompareWriteBatchOne(t, c)
    79  	}
    80  }
    81  
    82  func dumpPProfProfile(p *profile.Profile) string {
    83  	var ls []string
    84  	for _, sample := range p.Sample {
    85  		s := dumpPProfStack(sample, true)
    86  		ls = append(ls, s)
    87  	}
    88  	slices.Sort(ls)
    89  	return strings.Join(ls, "\n")
    90  }
    91  
    92  func dumpPProfStack(sample *profile.Sample, v bool) string {
    93  	sb := strings.Builder{}
    94  	for i := len(sample.Location) - 1; i >= 0; i-- {
    95  		location := sample.Location[i]
    96  		for j := len(location.Line) - 1; j >= 0; j-- {
    97  			line := location.Line[j]
    98  
    99  			sb.WriteString(";")
   100  			//sb.WriteString(fmt.Sprintf("[%x %x] ", location.ID, location.Address))
   101  
   102  			sb.WriteString(line.Function.Name)
   103  		}
   104  	}
   105  	if v {
   106  		sb.WriteString(" ")
   107  		sb.WriteString(fmt.Sprintf("%d", sample.Value[0]))
   108  	}
   109  	s := sb.String()
   110  	return s
   111  }
   112  
   113  func TestIterateWithStackBuilder(t *testing.T) {
   114  	sb := newStackBuilder()
   115  	it := tree.New()
   116  	it.Insert([]byte(""), uint64(43))
   117  	it.Insert([]byte("a"), uint64(42))
   118  	it.Insert([]byte("a;b"), uint64(1))
   119  	it.Insert([]byte("a;c"), uint64(2))
   120  	it.Insert([]byte("a;d;e"), uint64(3))
   121  	it.Insert([]byte("a;d;f"), uint64(4))
   122  
   123  	it.IterateWithStackBuilder(sb, func(stackID uint64, v uint64) {
   124  		sb.stackID2Val[stackID] = v
   125  	})
   126  	sb.expectValue(t, 0, 43)
   127  	sb.expectValue(t, 1, 42)
   128  	sb.expectValue(t, 2, 1)
   129  	sb.expectValue(t, 3, 2)
   130  	sb.expectValue(t, 4, 3)
   131  	sb.expectValue(t, 5, 4)
   132  	sb.expectStack(t, 0, "")
   133  	sb.expectStack(t, 1, "a")
   134  	sb.expectStack(t, 2, "a;b")
   135  	sb.expectStack(t, 3, "a;c")
   136  	sb.expectStack(t, 4, "a;d;e")
   137  	sb.expectStack(t, 5, "a;d;f")
   138  }
   139  
   140  func TestIterateWithStackBuilderEmpty(t *testing.T) {
   141  	it := tree.New()
   142  	sb := newStackBuilder()
   143  	it.IterateWithStackBuilder(sb, func(stackID uint64, v uint64) {
   144  		t.Fatal()
   145  	})
   146  }
   147  
   148  func newStackBuilder() *mockStackBuilder {
   149  	return &mockStackBuilder{
   150  		stackID2Stack:      make(map[uint64]string),
   151  		stackID2Val:        make(map[uint64]uint64),
   152  		stackID2StackBytes: make(map[uint64][][]byte),
   153  	}
   154  }
   155  
   156  func TestTreeIterationCorpus(t *testing.T) {
   157  	corpus := readCorpus(compareCorpus, benchWithoutGzip)
   158  	if len(corpus) == 0 {
   159  		t.Skip("empty corpus")
   160  		return
   161  	}
   162  	for _, c := range corpus {
   163  		key, _ := segment.ParseKey("foo.bar")
   164  		mock1 := &MockPutter{keep: true}
   165  		profile1 := pprof.RawProfile{
   166  			Profile:             c.profile,
   167  			PreviousProfile:     c.prev,
   168  			SampleTypeConfig:    c.config,
   169  			StreamingParser:     true,
   170  			PoolStreamingParser: true,
   171  			ArenasEnabled:       false,
   172  		}
   173  
   174  		err2 := profile1.Parse(context.TODO(), mock1, nil, ingestion.Metadata{Key: key, SpyName: c.spyname})
   175  		if err2 != nil {
   176  			t.Fatal(err2)
   177  		}
   178  		for _, put := range mock1.puts {
   179  			testIterateOne(t, put.ValTree)
   180  		}
   181  	}
   182  }
   183  
   184  func BenchmarkSmallStreaming(b *testing.B) {
   185  	t := readCorpusItemFile(pprofSmall, benchWithoutGzip)
   186  	for _, testType := range streamingTestTypes {
   187  		b.Run(fmt.Sprintf("BenchmarkSmallStreaming_pool_%v_arenas_%v", testType.pool, testType.arenas), func(b *testing.B) {
   188  			benchmarkStreamingOne(b, t, testType)
   189  		})
   190  	}
   191  }
   192  
   193  func BenchmarkBigStreaming(b *testing.B) {
   194  	t := readCorpusItemFile(pprofBig, benchWithoutGzip)
   195  	for _, testType := range streamingTestTypes {
   196  		b.Run(fmt.Sprintf("BenchmarkBigStreaming_pool_%v_arenas_%v", testType.pool, testType.arenas), func(b *testing.B) {
   197  			benchmarkStreamingOne(b, t, testType)
   198  		})
   199  	}
   200  }
   201  
   202  func BenchmarkSmallUnmarshal(b *testing.B) {
   203  	t := readCorpusItemFile(pprofSmall, benchWithoutGzip)
   204  	now := time.Now()
   205  	for i := 0; i < b.N; i++ {
   206  		parser := pprof.NewParser(pprof.ParserConfig{SampleTypes: t.config, Putter: putter})
   207  		err := parser.ParsePprof(context.TODO(), now, now, t.profile, false)
   208  		if err != nil {
   209  			b.Fatal(err)
   210  		}
   211  	}
   212  }
   213  
   214  func BenchmarkBigUnmarshal(b *testing.B) {
   215  	t := readCorpusItemFile(pprofBig, benchWithoutGzip)
   216  	now := time.Now()
   217  	for i := 0; i < b.N; i++ {
   218  		parser := pprof.NewParser(pprof.ParserConfig{SampleTypes: t.config, Putter: putter})
   219  		err := parser.ParsePprof(context.TODO(), now, now, t.profile, false)
   220  		if err != nil {
   221  			b.Fatal(err)
   222  		}
   223  	}
   224  }
   225  
   226  func BenchmarkCorpus(b *testing.B) {
   227  	corpus := readCorpus(benchmarkCorpus, benchWithoutGzip)
   228  	n := benchmarkCorpusSize
   229  	for _, testType := range streamingTestTypes {
   230  		for i := 0; i < n; i++ {
   231  			j := i
   232  			b.Run(fmt.Sprintf("BenchmarkCorpus_%d_pool_%v_arena_%v", j, testType.pool, testType.arenas),
   233  				func(b *testing.B) {
   234  					t := corpus[j]
   235  					benchmarkStreamingOne(b, t, testType)
   236  				})
   237  		}
   238  	}
   239  }
   240  
   241  func TestBugReusingSlices(t *testing.T) {
   242  	profiles := readCorpus(benchmarkCorpus+"/bugs/bug1_slice_reuse", false)
   243  	if len(profiles) == 0 {
   244  		t.Skip()
   245  		return
   246  	}
   247  	for _, p := range profiles {
   248  		parse(t, p, streamingTestType{pool: true, arenas: false})
   249  	}
   250  }
   251  
   252  func parse(t *testing.T, c *testcase, typ streamingTestType) {
   253  	mock := &MockPutter{keep: true}
   254  	key, _ := segment.ParseKey("foo.bar")
   255  	p := pprof.RawProfile{
   256  		Profile:             c.profile,
   257  		PreviousProfile:     c.prev,
   258  		SampleTypeConfig:    c.config,
   259  		StreamingParser:     true,
   260  		PoolStreamingParser: typ.pool,
   261  		ArenasEnabled:       typ.arenas,
   262  	}
   263  	err := p.Parse(context.TODO(), mock, nil, ingestion.Metadata{Key: key, SpyName: c.spyname})
   264  	if err != nil {
   265  		t.Fatal(err)
   266  	}
   267  }
   268  
   269  func benchmarkStreamingOne(b *testing.B, t *testcase, testType streamingTestType) {
   270  	now := time.Now()
   271  	for i := 0; i < b.N; i++ {
   272  		config := t.config
   273  		pConfig := streaming.ParserConfig{SampleTypes: config, Putter: putter, ArenasEnabled: testType.arenas}
   274  		var parser *streaming.VTStreamingParser
   275  		if testType.pool {
   276  			parser = streaming.VTStreamingParserFromPool(pConfig)
   277  		} else {
   278  			parser = streaming.NewStreamingParser(pConfig)
   279  		}
   280  		err := parser.ParsePprof(context.TODO(), now, now, t.profile, false)
   281  		if err != nil {
   282  			b.Fatal(err)
   283  		}
   284  		if testType.pool {
   285  			parser.ResetCache()
   286  			parser.ReturnToPool()
   287  		}
   288  		if testType.arenas {
   289  			parser.FreeArena()
   290  		}
   291  	}
   292  }
   293  
   294  var streamingTestTypes = []streamingTestType{
   295  	{pool: false, arenas: false},
   296  	{pool: true, arenas: false},
   297  	{pool: false, arenas: true},
   298  }
   299  
   300  type streamingTestType struct {
   301  	pool   bool
   302  	arenas bool
   303  }
   304  
   305  type testcase struct {
   306  	profile, prev []byte
   307  	config        map[string]*tree.SampleTypeConfig
   308  	fname         string
   309  	spyname       string
   310  }
   311  
   312  func readCorpus(dir string, doDecompress bool) []*testcase {
   313  	files, err := ioutil.ReadDir(dir)
   314  	if err != nil {
   315  		print(err)
   316  		return nil
   317  	}
   318  	var res []*testcase
   319  	for _, file := range files {
   320  		if strings.HasSuffix(file.Name(), ".txt") {
   321  			res = append(res, readCorpusItem(dir, file, doDecompress))
   322  		}
   323  	}
   324  	return res
   325  }
   326  
   327  func readCorpusItem(dir string, file fs.FileInfo, doDecompress bool) *testcase {
   328  	fname := dir + "/" + file.Name()
   329  	return readCorpusItemFile(fname, doDecompress)
   330  }
   331  
   332  func readCorpusItemFile(fname string, doDecompress bool) *testcase {
   333  	bs, err := ioutil.ReadFile(fname)
   334  	if err != nil {
   335  		panic(err)
   336  	}
   337  	r, err := http.ReadRequest(bufio.NewReader(bytes.NewReader(bs)))
   338  	if err != nil {
   339  		panic(err)
   340  	}
   341  	contentType := r.Header.Get("Content-Type")
   342  	rawData, _ := ioutil.ReadAll(r.Body)
   343  	decompress := func(b []byte) []byte {
   344  		if len(b) < 2 {
   345  			return b
   346  		}
   347  		if b[0] == 0x1f && b[1] == 0x8b {
   348  			gzipr, err := gzip.NewReader(bytes.NewReader(b))
   349  			if err != nil {
   350  				panic(err)
   351  			}
   352  			defer gzipr.Close()
   353  			var buf bytes.Buffer
   354  			if _, err = io.Copy(&buf, gzipr); err != nil {
   355  				panic(err)
   356  			}
   357  			return buf.Bytes()
   358  		}
   359  		return b
   360  	}
   361  
   362  	if contentType == "binary/octet-stream" {
   363  		return &testcase{
   364  			profile: decompress(rawData),
   365  			config:  tree.DefaultSampleTypeMapping,
   366  			fname:   fname,
   367  		}
   368  	}
   369  	boundary, err := form.ParseBoundary(contentType)
   370  	if err != nil {
   371  		panic(err)
   372  	}
   373  
   374  	f, err := multipart.NewReader(bytes.NewReader(rawData), boundary).ReadForm(32 << 20)
   375  	if err != nil {
   376  		panic(err)
   377  	}
   378  	const (
   379  		formFieldProfile          = "profile"
   380  		formFieldPreviousProfile  = "prev_profile"
   381  		formFieldSampleTypeConfig = "sample_type_config"
   382  	)
   383  
   384  	Profile, err := form.ReadField(f, formFieldProfile)
   385  	if err != nil {
   386  		panic(err)
   387  	}
   388  	PreviousProfile, err := form.ReadField(f, formFieldPreviousProfile)
   389  	if err != nil {
   390  		panic(err)
   391  	}
   392  
   393  	stBS, err := form.ReadField(f, formFieldSampleTypeConfig)
   394  	if err != nil {
   395  		panic(err)
   396  	}
   397  	var config map[string]*tree.SampleTypeConfig
   398  	if stBS != nil {
   399  		if err = json.Unmarshal(stBS, &config); err != nil {
   400  			panic(err)
   401  		}
   402  	} else {
   403  		config = tree.DefaultSampleTypeMapping
   404  	}
   405  	_ = Profile
   406  	_ = PreviousProfile
   407  
   408  	if doDecompress {
   409  		Profile = decompress(Profile)
   410  		PreviousProfile = decompress(PreviousProfile)
   411  	}
   412  	elem := &testcase{Profile, PreviousProfile, config, fname, "gospy"}
   413  	return elem
   414  }
   415  
   416  func testCompareOne(t *testing.T, c *testcase, typ streamingTestType) {
   417  	err := pprof.DecodePool(bytes.NewReader(c.profile), func(profile *tree.Profile) error {
   418  		return nil
   419  	})
   420  	fmt.Println(c.fname)
   421  	key, _ := segment.ParseKey("foo.bar")
   422  	mock1 := &MockPutter{keep: true}
   423  	profile1 := pprof.RawProfile{
   424  		Profile:             c.profile,
   425  		PreviousProfile:     c.prev,
   426  		SampleTypeConfig:    c.config,
   427  		StreamingParser:     true,
   428  		PoolStreamingParser: typ.pool,
   429  		ArenasEnabled:       typ.arenas,
   430  	}
   431  
   432  	err2 := profile1.Parse(context.TODO(), mock1, nil, ingestion.Metadata{Key: key, SpyName: c.spyname})
   433  	if err2 != nil {
   434  		t.Fatal(err2)
   435  	}
   436  
   437  	mock2 := &MockPutter{keep: true}
   438  	profile2 := pprof.RawProfile{
   439  		Profile:          c.profile,
   440  		PreviousProfile:  c.prev,
   441  		SampleTypeConfig: c.config,
   442  	}
   443  	err = profile2.Parse(context.TODO(), mock2, nil, ingestion.Metadata{Key: key, SpyName: c.spyname})
   444  	if err != nil {
   445  		t.Fatal(err)
   446  	}
   447  
   448  	if len(mock1.puts) != len(mock2.puts) {
   449  		t.Fatalf("put mismatch %d %d", len(mock1.puts), len(mock2.puts))
   450  	}
   451  	sort.Slice(mock1.puts, func(i, j int) bool {
   452  		return strings.Compare(mock1.puts[i].Key, mock1.puts[j].Key) < 0
   453  	})
   454  	sort.Slice(mock2.puts, func(i, j int) bool {
   455  		return strings.Compare(mock2.puts[i].Key, mock2.puts[j].Key) < 0
   456  	})
   457  	writeGlod := false
   458  	checkGold := true
   459  	trees := map[string]string{}
   460  	gold := c.fname + ".gold.json"
   461  	if checkGold {
   462  		goldBS, err := os.ReadFile(gold)
   463  		if err != nil {
   464  			panic(err)
   465  		}
   466  		err = json.Unmarshal(goldBS, &trees)
   467  		if err != nil {
   468  			panic(err)
   469  		}
   470  	}
   471  	for i := range mock1.puts {
   472  		p1 := mock1.puts[i]
   473  		p2 := mock2.puts[i]
   474  		k1 := p1.Key
   475  		k2 := p2.Key
   476  		if k1 != k2 {
   477  			t.Fatalf("key mismatch %s %s", k1, k2)
   478  		}
   479  		it := p1.Val
   480  		jit := mock2.puts[i].Val
   481  
   482  		if it != jit {
   483  			fmt.Println(key.SegmentKey())
   484  			t.Fatalf("mismatch\n --- actual:\n"+
   485  				"%s\n"+
   486  				" --- exopected\n"+
   487  				"%s\n====", it, jit)
   488  		}
   489  		if checkGold {
   490  			git := trees[k1]
   491  			if it != git {
   492  				t.Fatalf("mismatch ---\n"+
   493  					"%s\n"+
   494  					"---\n"+
   495  					"%s\n====", it, git)
   496  			}
   497  		}
   498  		fmt.Printf("ok %s %d \n", k1, len(it))
   499  		if p1.StartTime != p2.StartTime {
   500  			t.Fatal()
   501  		}
   502  		if p1.EndTime != p2.EndTime {
   503  			t.Fatal()
   504  		}
   505  		if p1.Units != p2.Units {
   506  			t.Fatal()
   507  		}
   508  		if p1.AggregationType != p2.AggregationType {
   509  			t.Fatal()
   510  		}
   511  		if p1.SpyName != p2.SpyName {
   512  			t.Fatal()
   513  		}
   514  		if p1.SampleRate != p2.SampleRate {
   515  			t.Fatal()
   516  		}
   517  		if writeGlod {
   518  			trees[k1] = it
   519  		}
   520  	}
   521  	if writeGlod {
   522  		marshal, err := json.Marshal(trees)
   523  		if err != nil {
   524  			panic(err)
   525  		}
   526  		err = os.WriteFile(gold, marshal, 0666)
   527  		if err != nil {
   528  			panic(err)
   529  		}
   530  	}
   531  }
   532  
   533  func testCompareWriteBatchOne(t *testing.T, c *testcase) {
   534  	fmt.Println(c.fname)
   535  	key, _ := segment.ParseKey("foo.bar")
   536  	profile1 := pprof.RawProfile{
   537  		Profile:          c.profile,
   538  		PreviousProfile:  c.prev,
   539  		SampleTypeConfig: c.config,
   540  	}
   541  	md := ingestion.Metadata{Key: key, SpyName: c.spyname}
   542  	wbf := &mockWriteBatchFactory{}
   543  	err := profile1.ParseWithWriteBatch(context.TODO(), wbf, md)
   544  	if err != nil {
   545  		t.Fatal(err)
   546  	}
   547  
   548  	mock2 := &MockPutter{keep: true}
   549  	profile2 := &pprof.RawProfile{
   550  		Profile:          c.profile,
   551  		PreviousProfile:  c.prev,
   552  		SampleTypeConfig: c.config,
   553  	}
   554  	mergeCumulative(profile2)
   555  
   556  	err = profile2.Parse(context.TODO(), mock2, nil, md)
   557  	if err != nil {
   558  		t.Fatal(err)
   559  	}
   560  
   561  	for _, put := range mock2.puts {
   562  		expectedCollapsed := put.Val
   563  		appenderCollapsed := ""
   564  		var found []*mockSamplesAppender
   565  		for _, batch := range wbf.wbs {
   566  			for _, appender := range batch.appenders {
   567  				labels := make(map[string]string)
   568  				labels["__name__"] = batch.appName
   569  				for _, label := range appender.labels {
   570  					labels[label.Key] = label.Value
   571  				}
   572  				k := segment.NewKey(labels)
   573  				if k.SegmentKey() == put.Key {
   574  					found = append(found, appender)
   575  				}
   576  			}
   577  		}
   578  		if len(found) != 1 {
   579  			if expectedCollapsed == "" {
   580  				continue
   581  			}
   582  			t.Fatalf("not found %s", put.Key)
   583  		}
   584  		appenderCollapsed = found[0].tree.String()
   585  
   586  		if appenderCollapsed != expectedCollapsed {
   587  			os.WriteFile("p3", []byte(expectedCollapsed), 0666)
   588  			os.WriteFile("p4", []byte(appenderCollapsed), 0666)
   589  			t.Fatalf("%s: expected\n%s\ngot\n%s\n failed file:%s\n", put.Key, expectedCollapsed, appenderCollapsed, c.fname)
   590  		}
   591  	}
   592  }
   593  
   594  func mergeCumulative(profile2 *pprof.RawProfile) {
   595  	if profile2.PreviousProfile != nil {
   596  		p1, _ := profile.Parse(bytes.NewReader(profile2.PreviousProfile))
   597  		p2, _ := profile.Parse(bytes.NewReader(profile2.Profile))
   598  		prev := []map[string]int64{
   599  			make(map[string]int64),
   600  			make(map[string]int64),
   601  		}
   602  		for _, sample := range p1.Sample {
   603  			s := dumpPProfStack(sample, false)
   604  			prev[0][s] += sample.Value[0]
   605  			prev[1][s] += sample.Value[1]
   606  		}
   607  		dec := func(s string, i int, v int64) int64 {
   608  			prevV := prev[i][s]
   609  			if v > prevV {
   610  				prev[i][s] = 0
   611  				return v - prevV
   612  			}
   613  			prev[i][s] = prevV - v
   614  			return 0
   615  		}
   616  		for _, sample := range p2.Sample {
   617  			s := dumpPProfStack(sample, false)
   618  			sample.Value[0] = dec(s, 0, sample.Value[0])
   619  			sample.Value[1] = dec(s, 1, sample.Value[1])
   620  		}
   621  
   622  		merged := p2.Compact()
   623  
   624  		bs := bytes.NewBuffer(nil)
   625  		merged.Write(bs)
   626  
   627  		profile2.PreviousProfile = nil
   628  		profile2.Profile = bs.Bytes()
   629  
   630  		sampleTypeConfig := make(map[string]*tree.SampleTypeConfig)
   631  		for k, v := range profile2.SampleTypeConfig {
   632  			vv := *v
   633  			vv.Cumulative = false
   634  			sampleTypeConfig[k] = &vv
   635  		}
   636  		profile2.SampleTypeConfig = sampleTypeConfig
   637  		//os.WriteFile("merged", []byte(dumpPProfProfile(merged)), 0666)
   638  	}
   639  }
   640  
   641  func testIterateOne(t *testing.T, pt *tree.Tree) {
   642  	sb := newStackBuilder()
   643  	var lines []string
   644  	pt.IterateWithStackBuilder(sb, func(stackID uint64, val uint64) {
   645  		lines = append(lines, fmt.Sprintf("%s %d", sb.stackID2Stack[stackID], val))
   646  	})
   647  	s := pt.String()
   648  	s = pt.String()
   649  	var expectedLines []string
   650  	if s != "" {
   651  		expectedLines = strings.Split(strings.Trim(s, "\n"), "\n")
   652  	}
   653  	slices.Sort(lines)
   654  	slices.Sort(expectedLines)
   655  	if !slices.Equal(lines, expectedLines) {
   656  		expected := strings.Join(expectedLines, "\n")
   657  		got := strings.Join(lines, "\n")
   658  		t.Fatalf("expected %v got\n%v", expected, got)
   659  	}
   660  }
   661  
   662  type PutInputCopy struct {
   663  	Val string
   664  	Key string
   665  
   666  	StartTime       time.Time
   667  	EndTime         time.Time
   668  	SpyName         string
   669  	SampleRate      uint32
   670  	Units           metadata.Units
   671  	AggregationType metadata.AggregationType
   672  	ValTree         *tree.Tree
   673  }
   674  
   675  type MockPutter struct {
   676  	keep bool
   677  	puts []PutInputCopy
   678  }
   679  
   680  func (m *MockPutter) Put(_ context.Context, input *storage.PutInput) error {
   681  	if m.keep {
   682  		m.puts = append(m.puts, PutInputCopy{
   683  			Val:             input.Val.String(),
   684  			ValTree:         input.Val.Clone(big.NewRat(1, 1)),
   685  			Key:             input.Key.SegmentKey(),
   686  			StartTime:       input.StartTime,
   687  			EndTime:         input.EndTime,
   688  			SpyName:         input.SpyName,
   689  			SampleRate:      input.SampleRate,
   690  			Units:           input.Units,
   691  			AggregationType: input.AggregationType,
   692  		})
   693  	}
   694  	return nil
   695  }
   696  
   697  type mockStackBuilder struct {
   698  	ss [][]byte
   699  
   700  	stackID2Stack      map[uint64]string
   701  	stackID2StackBytes map[uint64][][]byte
   702  	stackID2Val        map[uint64]uint64
   703  }
   704  
   705  func (s *mockStackBuilder) Push(frame []byte) {
   706  	s.ss = append(s.ss, frame)
   707  }
   708  
   709  func (s *mockStackBuilder) Pop() {
   710  	s.ss = s.ss[0 : len(s.ss)-1]
   711  }
   712  
   713  func (s *mockStackBuilder) Build() (stackID uint64) {
   714  	res := ""
   715  	for _, bs := range s.ss {
   716  		if len(res) != 0 {
   717  			res += ";"
   718  		}
   719  		res += string(bs)
   720  	}
   721  	id := uint64(len(s.stackID2Stack))
   722  	s.stackID2Stack[id] = res
   723  
   724  	bs := make([][]byte, 0, len(s.ss))
   725  	for _, frame := range s.ss {
   726  		bs = append(bs, append([]byte{}, frame...))
   727  	}
   728  	s.stackID2StackBytes[id] = bs
   729  	return id
   730  }
   731  
   732  func (s *mockStackBuilder) Reset() {
   733  	s.ss = s.ss[:0]
   734  }
   735  
   736  func (s *mockStackBuilder) expectValue(t *testing.T, stackID, expected uint64) {
   737  	if s.stackID2Val[stackID] != expected {
   738  		t.Fatalf("expected at %d %d got %d", stackID, expected, s.stackID2Val[stackID])
   739  	}
   740  }
   741  func (s *mockStackBuilder) expectStack(t *testing.T, stackID uint64, expected string) {
   742  	if s.stackID2Stack[stackID] != expected {
   743  		t.Fatalf("expected at %d %s got %s", stackID, expected, s.stackID2Stack[stackID])
   744  	}
   745  }
   746  
   747  type mockWriteBatchFactory struct {
   748  	wbs map[string]*mockWriteBatch
   749  }
   750  
   751  func (m *mockWriteBatchFactory) NewWriteBatch(appName string, _ metadata.Metadata) (stackbuilder.WriteBatch, error) {
   752  	if m.wbs == nil {
   753  		m.wbs = make(map[string]*mockWriteBatch)
   754  	}
   755  	if m.wbs[appName] != nil {
   756  		panic("already exists")
   757  	}
   758  	wb := &mockWriteBatch{
   759  		appName:   appName,
   760  		sb:        newStackBuilder(),
   761  		appenders: make(map[string]*mockSamplesAppender),
   762  	}
   763  	m.wbs[appName] = wb
   764  	return wb, nil
   765  }
   766  
   767  type mockWriteBatch struct {
   768  	appName   string
   769  	sb        *mockStackBuilder
   770  	appenders map[string]*mockSamplesAppender
   771  }
   772  
   773  func (m *mockWriteBatch) StackBuilder() tree.StackBuilder {
   774  	return m.sb
   775  }
   776  
   777  func (m *mockWriteBatch) SamplesAppender(startTime, endTime int64, labels stackbuilder.Labels) stackbuilder.SamplesAppender {
   778  	sLabels, _ := json.Marshal(labels)
   779  	k := fmt.Sprintf("%d-%d-%s", startTime, endTime, sLabels)
   780  	a := m.appenders[k]
   781  	if a != nil {
   782  		return a
   783  	}
   784  	a = &mockSamplesAppender{
   785  		startTime: startTime,
   786  		endTime:   endTime,
   787  		labels:    labels,
   788  		sb:        m.sb,
   789  	}
   790  	m.appenders[k] = a
   791  	return a
   792  }
   793  
   794  func (*mockWriteBatch) Flush() {
   795  
   796  }
   797  
   798  type mockSamplesAppender struct {
   799  	startTime, endTime int64
   800  	labels             stackbuilder.Labels
   801  	stacks             []stackIDToVal
   802  	tree               *tree.Tree
   803  	sb                 *mockStackBuilder
   804  }
   805  
   806  type stackIDToVal struct {
   807  	stackID uint64
   808  	val     uint64
   809  }
   810  
   811  func (m *mockSamplesAppender) Append(stackID, value uint64) {
   812  	m.stacks = append(m.stacks, stackIDToVal{stackID, value})
   813  	stack := m.sb.stackID2StackBytes[stackID]
   814  	if stack == nil {
   815  		panic("not found")
   816  	}
   817  	if m.tree == nil {
   818  		m.tree = tree.New()
   819  	}
   820  	m.tree.InsertStack(stack, value)
   821  }