github.com/cloudwego/kitex@v0.9.0/pkg/profiler/profiler_test.go (about)

     1  /*
     2   * Copyright 2022 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package profiler
    18  
    19  import (
    20  	"bytes"
    21  	"context"
    22  	"runtime"
    23  	"runtime/pprof"
    24  	"sort"
    25  	"sync/atomic"
    26  	"testing"
    27  	"time"
    28  
    29  	"github.com/cloudwego/kitex/internal/test"
    30  )
    31  
    32  func TestProfiler(t *testing.T) {
    33  	var processed int32
    34  	interval := time.Millisecond * 100
    35  	window := time.Millisecond * 100
    36  	testProcessor := func(pfs []*TagsProfile) error {
    37  		if len(pfs) <= 1 {
    38  			return nil
    39  		}
    40  		defer func() {
    41  			atomic.StoreInt32(&processed, 1)
    42  		}()
    43  		var percent float64
    44  		for _, p := range pfs {
    45  			percent += p.Percent * 100
    46  		}
    47  		test.Assert(t, percent >= 99, percent)
    48  		return nil
    49  	}
    50  
    51  	p := NewProfiler(testProcessor, interval, window)
    52  	ctx := pprof.WithLabels(context.Background(), pprof.Labels("type", "profiler"))
    53  	stopCh := make(chan struct{})
    54  	go func() {
    55  		err := p.Run(ctx)
    56  		test.Assert(t, err == nil, err)
    57  		stopCh <- struct{}{}
    58  	}()
    59  
    60  	time.Sleep(interval) // wait for interval finished
    61  	go func() {
    62  		p.Tag(ctx, []string{"type", "trace"}...)
    63  		defer p.Untag(ctx)
    64  		var sum int
    65  		for atomic.LoadInt32(&processed) == 0 {
    66  			sum++
    67  		}
    68  	}()
    69  
    70  	for atomic.LoadInt32(&processed) == 0 {
    71  		runtime.Gosched()
    72  	}
    73  	p.Stop()
    74  	<-stopCh
    75  }
    76  
    77  func TestProfilerPaused(t *testing.T) {
    78  	interval := time.Millisecond * 50
    79  	window := time.Millisecond * 50
    80  	var count int32
    81  	p := NewProfiler(func(profiles []*TagsProfile) error {
    82  		atomic.AddInt32(&count, 1)
    83  		return nil
    84  	}, interval, window)
    85  	ctx := pprof.WithLabels(context.Background(), pprof.Labels("type", "profiler"))
    86  	stopCh := make(chan struct{})
    87  	go func() {
    88  		err := p.Run(ctx)
    89  		test.Assert(t, err == nil, err)
    90  		close(stopCh)
    91  	}()
    92  	time.Sleep(interval)
    93  
    94  	var data bytes.Buffer
    95  	for i := 0; i < 5; i++ {
    96  		data.Reset()
    97  		p.Pause()
    98  		p.Pause() // pause twice by mistake
    99  		err := pprof.StartCPUProfile(&data)
   100  		test.Assert(t, err == nil, err)
   101  		pprof.StopCPUProfile()
   102  		p.Resume()
   103  		p.Resume() // resume twice by mistake
   104  	}
   105  
   106  	for atomic.LoadInt32(&count) > 5 { // wait for processor finished
   107  		runtime.Gosched()
   108  	}
   109  
   110  	p.Stop()
   111  	p.Stop() // stop twice by mistake
   112  	<-stopCh
   113  }
   114  
   115  func TestLabelToTags(t *testing.T) {
   116  	var labels map[string][]string
   117  	ret := labelToTags(labels)
   118  	test.Assert(t, len(ret) == 0, ret)
   119  
   120  	labels = map[string][]string{}
   121  	ret = labelToTags(labels)
   122  	test.Assert(t, len(ret) == 0, ret)
   123  
   124  	labels["a"] = []string{"b"}
   125  	ret = labelToTags(labels)
   126  	test.DeepEqual(t, ret, []string{"a", "b"})
   127  
   128  	labels["c"] = []string{"d"}
   129  	ret = labelToTags(labels)
   130  	sort.Strings(ret)
   131  	test.DeepEqual(t, ret, []string{"a", "b", "c", "d"})
   132  
   133  	labels = map[string][]string{"a": {"b", "c"}}
   134  	ret = labelToTags(labels)
   135  	test.DeepEqual(t, ret, []string{"a", "b,c"})
   136  }
   137  
   138  func TestTagsToKey(t *testing.T) {
   139  	ret := tagsToKey([]string{})
   140  	test.Assert(t, ret == "")
   141  
   142  	ret = tagsToKey([]string{"a", "b"})
   143  	test.Assert(t, ret == "a=b")
   144  
   145  	ret = tagsToKey([]string{"a", "b,c"})
   146  	test.Assert(t, ret == "a=b,c")
   147  
   148  	ret = tagsToKey([]string{"a", "b", "c"})
   149  	test.Assert(t, ret == "")
   150  
   151  	ret = tagsToKey([]string{"a", "b", "c", "d"})
   152  	test.Assert(t, ret == "a=b|c=d")
   153  
   154  	ret = tagsToKey([]string{"a", "b,c", "d", "e,f"})
   155  	test.Assert(t, ret == "a=b,c|d=e,f")
   156  }