github.com/cloudwego/kitex@v0.9.0/pkg/profiler/profiler.go

/*
 * Copyright 2022 CloudWeGo Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package profiler

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"runtime/pprof"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/google/pprof/profile"

	"github.com/cloudwego/kitex/pkg/klog"
)

type profilerContextKey struct{}

type Profiler interface {
	Run(ctx context.Context) (err error)
	Stop()
	Pause()
	Resume()
	Prepare(ctx context.Context) context.Context
	Tag(ctx context.Context, tags ...string) context.Context
	Untag(ctx context.Context)
	Lookup(ctx context.Context, key string) (string, bool)
}

type Processor func(profiles []*TagsProfile) error

func LogProcessor(profiles []*TagsProfile) error {
	if len(profiles) == 0 {
		return nil
	}
	klog.Infof("KITEX: profiler collect %d records", len(profiles))
	for _, p := range profiles {
		if p.Key != "" {
			klog.Infof("KITEX: profiler - %s %.2f%%", p.Key, p.Percent*100)
		} else {
			klog.Infof("KITEX: profiler - type=default %.2f%%", p.Percent*100)
		}
	}
	klog.Info("---------------------------------")
	return nil
}

func NewProfiler(processor Processor, interval, window time.Duration) *profiler {
	if processor == nil {
		processor = LogProcessor
	}
	return &profiler{
		stateCond: sync.NewCond(&sync.Mutex{}),
		processor: processor,
		interval:  interval,
		window:    window,
	}
}

var _ Profiler = (*profiler)(nil)

const (
	// state changes:
	//   running => pausing => paused => resuming => running
	//           => stopped
	stateRunning  = 0
	statePausing  = 1
	statePaused   = 2
	stateResuming = 3
	stateStopped  = 4
)

type profiler struct {
	data      bytes.Buffer // protobuf
	state     int
	stateCond *sync.Cond
	// settings
	processor Processor
	interval  time.Duration // sleep time between every profiling window
	window    time.Duration // profiling in the window, go pprof collect stack profile every 10ms
}

// Tag current goroutine with tagged ctx
// it's used for reuse goroutine scenario
func Tag(ctx context.Context) {
	if pc, ok := ctx.Value(profilerContextKey{}).(*profilerContext); ok {
		pc.profiler.Tag(ctx)
	}
}

// Untag current goroutine with tagged ctx
// it's used for reuse goroutine scenario
func Untag(ctx context.Context) {
	if pc, ok := ctx.Value(profilerContextKey{}).(*profilerContext); ok {
		pc.profiler.Untag(ctx)
	}
}

type profilerContext struct {
	profiler Profiler
	untagCtx context.Context
	tags     []string
}

func newProfilerContext(profiler Profiler) *profilerContext {
	return &profilerContext{
		profiler: profiler,
		tags:     make([]string, 0, 12),
	}
}

// Prepare the profiler context
func (p *profiler) Prepare(ctx context.Context) context.Context {
	if c := ctx.Value(profilerContextKey{}); c != nil {
		return ctx
	}
	return context.WithValue(ctx, profilerContextKey{}, newProfilerContext(p))
}

// State return current profiler state
func (p *profiler) State() (state int) {
	p.stateCond.L.Lock()
	state = p.state
	p.stateCond.L.Unlock()
	return state
}

// Stop the profiler
func (p *profiler) Stop() {
	if p.State() == stateStopped {
		return
	}
	// stateRunning => stateStopped
	p.stateChange(stateRunning, stateStopped)
}

// Pause the profiler.
// The profiler has been paused when Pause() return
func (p *profiler) Pause() {
	if p.State() == statePaused {
		return
	}
	// stateRunning => statePausing
	p.stateChange(stateRunning, statePausing)
	// => statePaused
	p.stateWait(statePaused)
}

// Resume the profiler.
// The profiler has been resumed when Resume() return
func (p *profiler) Resume() {
	if p.State() == stateRunning {
		return
	}
	// statePaused => stateResuming
	p.stateChange(statePaused, stateResuming)
	// => stateRunning
	p.stateWait(stateRunning)
}

// Run start analyse the pprof data with interval and window settings
func (p *profiler) Run(ctx context.Context) (err error) {
	var profiles []*TagsProfile
	timer := time.NewTimer(0)
	for {
		// check state
		state := p.State()
		switch state {
		case stateRunning: // do nothing
		case statePausing: // pause the loop
			p.stateChange(statePausing, statePaused) // wake up Pause()
			klog.Info("KITEX: profiler paused")
			p.stateChange(stateResuming, stateRunning) // wake up Resume()
			klog.Info("KITEX: profiler resumed")
			continue
		case statePaused, stateResuming: // actually, no such case
			continue
		case stateStopped: // end the loop
			klog.Info("KITEX: profiler stopped")
			return nil
		}

		// wait for an interval time to reduce the cost of profiling
		if p.interval > 0 {
			timer.Reset(p.interval)
			<-timer.C
		}

		// start profiler
		if err = p.startProfile(); err != nil {
			klog.Errorf("KITEX: profiler start profile failed: %v", err)
			// the reason of startProfile failed maybe because of there is a running pprof program.
			// so we need to wait a fixed time for it.
			time.Sleep(time.Second * 30)
			continue
		}

		// wait for a window time to collect pprof data
		if p.window > 0 {
			timer.Reset(p.window)
			<-timer.C
		}

		// stop profiler
		p.stopProfile()

		// analyse and process pprof data
		profiles, err = p.analyse()
		if err != nil {
			klog.Errorf("KITEX: profiler analyse failed: %v", err)
			continue
		}
		err = p.processor(profiles)
		if err != nil {
			klog.Errorf("KITEX: profiler controller process failed: %v", err)
			continue
		}
	}
}

// Tag current goroutine with tags
// If ctx already tagged, append the existed tags
func (p *profiler) Tag(ctx context.Context, tags ...string) context.Context {
	pctx, ok := ctx.Value(profilerContextKey{}).(*profilerContext)
	if !ok {
		pctx = newProfilerContext(p)
		ctx = context.WithValue(ctx, profilerContextKey{}, pctx)
	}
	if pctx.untagCtx == nil {
		pctx.untagCtx = ctx
	}
	pctx.tags = append(pctx.tags, tags...)
	// do not return pprof ctx
	pprof.SetGoroutineLabels(pprof.WithLabels(context.Background(), pprof.Labels(pctx.tags...)))
	return ctx
}

// Untag current goroutine
// Only untag if ctx already tagged, will not clear the goroutine labels if not tagged by profiler
func (p *profiler) Untag(ctx context.Context) {
	if pc, ok := ctx.Value(profilerContextKey{}).(*profilerContext); ok && pc.untagCtx != nil {
		// if ctx have untagCtx, that means the current goroutine created by a tagged goroutine
		// we need to untag the goroutine when finished
		// else, do nothing
		pprof.SetGoroutineLabels(pc.untagCtx)
	}
}

func (p *profiler) Lookup(ctx context.Context, key string) (string, bool) {
	return pprof.Label(ctx, key)
}

func (p *profiler) stateChange(from, to int) {
	p.stateCond.L.Lock()
	for p.state != from { // wait state to from first
		p.stateCond.Wait()
	}
	p.state = to
	p.stateCond.L.Unlock()
	p.stateCond.Broadcast()
}

func (p *profiler) stateWait(to int) {
	p.stateCond.L.Lock()
	for p.state != to {
		p.stateCond.Wait()
	}
	p.stateCond.L.Unlock()
}

func (p *profiler) startProfile() error {
	p.data.Reset()
	return pprof.StartCPUProfile(&p.data)
}

func (p *profiler) stopProfile() {
	pprof.StopCPUProfile()
}

func (p *profiler) analyse() ([]*TagsProfile, error) {
	// parse protobuf data
	pf, err := profile.ParseData(p.data.Bytes())
	if err != nil {
		return nil, err
	}

	// filter cpu value index
	sampleIdx := -1
	for idx, st := range pf.SampleType {
		if st.Type == "cpu" {
			sampleIdx = idx
			break
		}
	}
	if sampleIdx < 0 {
		return nil, errors.New("profiler: sample type not found")
	}

	// calculate every sample expense
	counter := map[string]*TagsProfile{} // map[tagsKey]funcProfile
	var total int64
	for _, sm := range pf.Sample {
		value := sm.Value[sampleIdx]
		tags := labelToTags(sm.Label)
		tagsKey := tagsToKey(tags)
		tp, ok := counter[tagsKey]
		if !ok {
			tp = &TagsProfile{}
			counter[tagsKey] = tp
			tp.Key = tagsKey
			tp.Tags = tags
		}
		tp.Value += value
		total += value
	}

	// compensate value with duration date
	durationRate := float64((p.interval + p.window) / p.window)
	profiles := make([]*TagsProfile, 0, len(counter)) // flat to array
	for _, l := range counter {
		l.Percent = float64(l.Value) / float64(total)
		l.Value = int64(durationRate * float64(l.Value))
		profiles = append(profiles, l)
	}
	return profiles, nil
}

// TagsProfile is the stats result group by tag key
type TagsProfile struct {
	Key     string   // eg: a=1,b=2
	Tags    []string // eg: ["a", "1", "b", "2"]
	Value   int64    // pprof cpu times
	Percent float64  // <= 1.0
}

func labelToTags(label map[string][]string) []string {
	tags := make([]string, 0, len(label)*2)
	for k, v := range label {
		tags = append(tags, k, strings.Join(v, ","))
	}
	return tags
}

func tagsToKey(tags []string) string {
	if len(tags)%2 != 0 {
		return ""
	}
	tagsPair := make([]string, 0, len(tags)/2)
	for i := 0; i < len(tags); i += 2 {
		tagsPair = append(tagsPair, fmt.Sprintf("%s=%s", tags[i], tags[i+1]))
	}
	// sort tags to make it a unique key
	sort.Strings(tagsPair)
	return strings.Join(tagsPair, "|")
}