github.com/annwntech/go-micro/v2@v2.9.5/debug/profile/pprof/pprof.go (about)

     1  // Package pprof provides a pprof profiler
     2  package pprof
     3  
     4  import (
     5  	"os"
     6  	"path/filepath"
     7  	"runtime"
     8  	"runtime/pprof"
     9  	"sync"
    10  	"time"
    11  
    12  	"github.com/annwntech/go-micro/v2/debug/profile"
    13  )
    14  
    15  type profiler struct {
    16  	opts profile.Options
    17  
    18  	sync.Mutex
    19  	running bool
    20  	exit    chan bool
    21  
    22  	// where the cpu profile is written
    23  	cpuFile *os.File
    24  	// where the mem profile is written
    25  	memFile *os.File
    26  }
    27  
    28  func (p *profiler) writeHeap(f *os.File) {
    29  	defer f.Close()
    30  
    31  	t := time.NewTicker(time.Second * 30)
    32  	defer t.Stop()
    33  
    34  	for {
    35  		select {
    36  		case <-t.C:
    37  			runtime.GC()
    38  			pprof.WriteHeapProfile(f)
    39  		case <-p.exit:
    40  			return
    41  		}
    42  	}
    43  }
    44  
    45  func (p *profiler) Start() error {
    46  	p.Lock()
    47  	defer p.Unlock()
    48  
    49  	if p.running {
    50  		return nil
    51  	}
    52  
    53  	// create exit channel
    54  	p.exit = make(chan bool)
    55  
    56  	cpuFile := filepath.Join("/tmp", "cpu.pprof")
    57  	memFile := filepath.Join("/tmp", "mem.pprof")
    58  
    59  	if len(p.opts.Name) > 0 {
    60  		cpuFile = filepath.Join("/tmp", p.opts.Name+".cpu.pprof")
    61  		memFile = filepath.Join("/tmp", p.opts.Name+".mem.pprof")
    62  	}
    63  
    64  	f1, err := os.Create(cpuFile)
    65  	if err != nil {
    66  		return err
    67  	}
    68  
    69  	f2, err := os.Create(memFile)
    70  	if err != nil {
    71  		return err
    72  	}
    73  
    74  	// start cpu profiling
    75  	if err := pprof.StartCPUProfile(f1); err != nil {
    76  		return err
    77  	}
    78  
    79  	// write the heap periodically
    80  	go p.writeHeap(f2)
    81  
    82  	// set cpu file
    83  	p.cpuFile = f1
    84  	// set mem file
    85  	p.memFile = f2
    86  
    87  	p.running = true
    88  
    89  	return nil
    90  }
    91  
    92  func (p *profiler) Stop() error {
    93  	p.Lock()
    94  	defer p.Unlock()
    95  
    96  	select {
    97  	case <-p.exit:
    98  		return nil
    99  	default:
   100  		close(p.exit)
   101  		pprof.StopCPUProfile()
   102  		p.cpuFile.Close()
   103  		p.running = false
   104  		p.cpuFile = nil
   105  		p.memFile = nil
   106  		return nil
   107  	}
   108  }
   109  
   110  func (p *profiler) String() string {
   111  	return "pprof"
   112  }
   113  
   114  func NewProfile(opts ...profile.Option) profile.Profile {
   115  	var options profile.Options
   116  	for _, o := range opts {
   117  		o(&options)
   118  	}
   119  	p := new(profiler)
   120  	p.opts = options
   121  	return p
   122  }