yyfbaloot/profile/profile.go

package profile

import (
	"os"
	"path/filepath"
	"runtime"
	"runtime/pprof"

	"bet24.com/log"
)

const (
	cpuMode = iota
	memMode
	mutexMode
	blockMode
	traceMode
	goroutineMode
	threadCreateMode
)

type Profile struct {
	// 分析模式
	mode int

	path string

	//内存配置文件采样速率，分析器将会在每分配指定的字节数量后对内存使用情况进行取样
	memProfileRate int

	closer func()

	stopped uint32

	namePrefix string
}

func NamePrefix(name string) func(*Profile) {
	return func(p *Profile) {
		p.namePrefix = name
	}
}

// CPUProfile enables cpu profiling.
// It disables any previous profiling settings.
func CPUProfile(p *Profile) { p.mode = cpuMode }

// DefaultMemProfileRate is the default memory profiling rate.
const DefaultMemProfileRate = 4096

// MemProfile enables memory profiling.
// It disables any previous profiling settings.
func MemProfile(p *Profile) {
	p.memProfileRate = DefaultMemProfileRate
	p.mode = memMode
}

//GoRTProfile enables goroutine profiling
func GoRTProfile(p *Profile) { p.mode = goroutineMode }

//block enables goroutine profiling
func BlockProfile(p *Profile) { p.mode = blockMode }

//MutexProfile enables goroutine profiling
func MutexProfile(p *Profile) { p.mode = mutexMode }

// ProfilePath controls the base path where various profiling
// files are written. If blank, the base path will be generated
// by os.MkdirTemp.
func ProfilePath(path string) func(*Profile) {
	return func(p *Profile) {
		p.path = path
	}
}

// Stop stops the profile and flushes any unwritten data.
func (p *Profile) Stop() {
	p.closer()
}

// Start starts a new profiling session.
// The caller should call the Stop method on the value returned
// to cleanly stop profiling.
func Start(options ...func(*Profile)) interface {
	Stop()
} {
	var prof Profile
	for _, option := range options {
		option(&prof)
	}

	path, err := func() (string, error) {
		if p := prof.path; p != "" {
			return p, os.MkdirAll(p, 0777)
		}
		return os.MkdirTemp("", "profile")
	}()

	if err != nil {
		log.Fatal("profile: could not create initial output directory: %v", err)
	}

	logf := func(format string, args ...interface{}) {
		log.Debug(format, args...)
	}

	switch prof.mode {
	case cpuMode:
		fn := filepath.Join(path, prof.namePrefix+"cpu.pprof")
		f, err := os.Create(fn)
		if err != nil {
			log.Fatal("profile: could not create cpu profile %q: %v", fn, err)
		}
		logf("profile: cpu profiling enabled, %s", fn)
		pprof.StartCPUProfile(f)
		prof.closer = func() {
			pprof.StopCPUProfile()
			f.Close()
			logf("profile: cpu profiling disabled, %s", fn)
		}

	case memMode:
		fn := filepath.Join(path, prof.namePrefix+"mem.pprof")
		f, err := os.Create(fn)
		if err != nil {
			log.Fatal("profile: could not create memory profile %q: %v", fn, err)
		}
		old := runtime.MemProfileRate
		runtime.MemProfileRate = prof.memProfileRate
		logf("profile: memory profiling enabled (rate %d), %s", runtime.MemProfileRate, fn)
		prof.closer = func() {
			pprof.Lookup("heap").WriteTo(f, 0)
			f.Close()
			runtime.MemProfileRate = old
			logf("profile: memory profiling disabled, %s", fn)
		}
	case goroutineMode:
		fn := filepath.Join(path, prof.namePrefix+"goroutine.pprof")
		f, err := os.Create(fn)
		if err != nil {
			log.Fatal("profile: could not create goroutine profile %q: %v", fn, err)
		}

		logf("profile: goroutine profiling enabled, %s", fn)
		prof.closer = func() {
			pprof.Lookup("goroutine").WriteTo(f, 0)
			f.Close()
			logf("profile: goroutine profiling disabled, %s", fn)
		}

	case blockMode:
		fn := filepath.Join(path, prof.namePrefix+"block.pprof")
		f, err := os.Create(fn)
		if err != nil {
			log.Fatal("profile: could not create block profile %q: %v", fn, err)
		}

		logf("profile: block profiling enabled, %s", fn)
		prof.closer = func() {
			pprof.Lookup("block").WriteTo(f, 0)
			f.Close()
			logf("profile: block profiling disabled, %s", fn)
		}

	case mutexMode:
		fn := filepath.Join(path, prof.namePrefix+"mutex.pprof")
		f, err := os.Create(fn)
		if err != nil {
			log.Fatal("profile: could not create mutex profile %q: %v", fn, err)
		}

		logf("profile: mutex profiling enabled, %s", fn)
		prof.closer = func() {
			pprof.Lookup("mutex").WriteTo(f, 0)
			f.Close()
			logf("profile: mutex profiling disabled, %s", fn)
		}
	}

	return &prof
}