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 }