vendor/github.com/pkg/profile/profile.go

   1 // Package profile provides a simple way to manage runtime/pprof
   2 // profiling of your Go application.
   3 package profile
   4
   5 import (
   6         "io/ioutil"
   7         "log"
   8         "os"
   9         "os/signal"
  10         "path/filepath"
  11         "runtime"
  12         "runtime/pprof"
  13         "sync/atomic"
  14 )
  15
  16 const (
  17         cpuMode = iota
  18         memMode
  19         mutexMode
  20         blockMode
  21         traceMode
  22 )
  23
  24 // Profile represents an active profiling session.
  25 type Profile struct {
  26         // quiet suppresses informational messages during profiling.
  27         quiet bool
  28
  29         // noShutdownHook controls whether the profiling package should
  30         // hook SIGINT to write profiles cleanly.
  31         noShutdownHook bool
  32
  33         // mode holds the type of profiling that will be made
  34         mode int
  35
  36         // path holds the base path where various profiling files are  written.
  37         // If blank, the base path will be generated by ioutil.TempDir.
  38         path string
  39
  40         // memProfileRate holds the rate for the memory profile.
  41         memProfileRate int
  42
  43         // closer holds a cleanup function that run after each profile
  44         closer func()
  45
  46         // stopped records if a call to profile.Stop has been made
  47         stopped uint32
  48 }
  49
  50 // NoShutdownHook controls whether the profiling package should
  51 // hook SIGINT to write profiles cleanly.
  52 // Programs with more sophisticated signal handling should set
  53 // this to true and ensure the Stop() function returned from Start()
  54 // is called during shutdown.
  55 func NoShutdownHook(p *Profile) { p.noShutdownHook = true }
  56
  57 // Quiet suppresses informational messages during profiling.
  58 func Quiet(p *Profile) { p.quiet = true }
  59
  60 // CPUProfile enables cpu profiling.
  61 // It disables any previous profiling settings.
  62 func CPUProfile(p *Profile) { p.mode = cpuMode }
  63
  64 // DefaultMemProfileRate is the default memory profiling rate.
  65 // See also http://golang.org/pkg/runtime/#pkg-variables
  66 const DefaultMemProfileRate = 4096
  67
  68 // MemProfile enables memory profiling.
  69 // It disables any previous profiling settings.
  70 func MemProfile(p *Profile) {
  71         p.memProfileRate = DefaultMemProfileRate
  72         p.mode = memMode
  73 }
  74
  75 // MemProfileRate enables memory profiling at the preferred rate.
  76 // It disables any previous profiling settings.
  77 func MemProfileRate(rate int) func(*Profile) {
  78         return func(p *Profile) {
  79                 p.memProfileRate = rate
  80                 p.mode = memMode
  81         }
  82 }
  83
  84 // MutexProfile enables mutex profiling.
  85 // It disables any previous profiling settings.
  86 //
  87 // Mutex profiling is a no-op before go1.8.
  88 func MutexProfile(p *Profile) { p.mode = mutexMode }
  89
  90 // BlockProfile enables block (contention) profiling.
  91 // It disables any previous profiling settings.
  92 func BlockProfile(p *Profile) { p.mode = blockMode }
  93
  94 // Trace profile controls if execution tracing will be enabled. It disables any previous profiling settings.
  95 func TraceProfile(p *Profile) { p.mode = traceMode }
  96
  97 // ProfilePath controls the base path where various profiling
  98 // files are written. If blank, the base path will be generated
  99 // by ioutil.TempDir.
 100 func ProfilePath(path string) func(*Profile) {
 101         return func(p *Profile) {
 102                 p.path = path
 103         }
 104 }
 105
 106 // Stop stops the profile and flushes any unwritten data.
 107 func (p *Profile) Stop() {
 108         if !atomic.CompareAndSwapUint32(&p.stopped, 0, 1) {
 109                 // someone has already called close
 110                 return
 111         }
 112         p.closer()
 113         atomic.StoreUint32(&started, 0)
 114 }
 115
 116 // started is non zero if a profile is running.
 117 var started uint32
 118
 119 // Start starts a new profiling session.
 120 // The caller should call the Stop method on the value returned
 121 // to cleanly stop profiling.
 122 func Start(options ...func(*Profile)) interface {
 123         Stop()
 124 } {
 125         if !atomic.CompareAndSwapUint32(&started, 0, 1) {
 126                 log.Fatal("profile: Start() already called")
 127         }
 128
 129         var prof Profile
 130         for _, option := range options {
 131                 option(&prof)
 132         }
 133
 134         path, err := func() (string, error) {
 135                 if p := prof.path; p != "" {
 136                         return p, os.MkdirAll(p, 0777)
 137                 }
 138                 return ioutil.TempDir("", "profile")
 139         }()
 140
 141         if err != nil {
 142                 log.Fatalf("profile: could not create initial output directory: %v", err)
 143         }
 144
 145         logf := func(format string, args ...interface{}) {
 146                 if !prof.quiet {
 147                         log.Printf(format, args...)
 148                 }
 149         }
 150
 151         switch prof.mode {
 152         case cpuMode:
 153                 fn := filepath.Join(path, "cpu.pprof")
 154                 f, err := os.Create(fn)
 155                 if err != nil {
 156                         log.Fatalf("profile: could not create cpu profile %q: %v", fn, err)
 157                 }
 158                 logf("profile: cpu profiling enabled, %s", fn)
 159                 pprof.StartCPUProfile(f)
 160                 prof.closer = func() {
 161                         pprof.StopCPUProfile()
 162                         f.Close()
 163                         logf("profile: cpu profiling disabled, %s", fn)
 164                 }
 165
 166         case memMode:
 167                 fn := filepath.Join(path, "mem.pprof")
 168                 f, err := os.Create(fn)
 169                 if err != nil {
 170                         log.Fatalf("profile: could not create memory profile %q: %v", fn, err)
 171                 }
 172                 old := runtime.MemProfileRate
 173                 runtime.MemProfileRate = prof.memProfileRate
 174                 logf("profile: memory profiling enabled (rate %d), %s", runtime.MemProfileRate, fn)
 175                 prof.closer = func() {
 176                         pprof.Lookup("heap").WriteTo(f, 0)
 177                         f.Close()
 178                         runtime.MemProfileRate = old
 179                         logf("profile: memory profiling disabled, %s", fn)
 180                 }
 181
 182         case mutexMode:
 183                 fn := filepath.Join(path, "mutex.pprof")
 184                 f, err := os.Create(fn)
 185                 if err != nil {
 186                         log.Fatalf("profile: could not create mutex profile %q: %v", fn, err)
 187                 }
 188                 enableMutexProfile()
 189                 logf("profile: mutex profiling enabled, %s", fn)
 190                 prof.closer = func() {
 191                         if mp := pprof.Lookup("mutex"); mp != nil {
 192                                 mp.WriteTo(f, 0)
 193                         }
 194                         f.Close()
 195                         disableMutexProfile()
 196                         logf("profile: mutex profiling disabled, %s", fn)
 197                 }
 198
 199         case blockMode:
 200                 fn := filepath.Join(path, "block.pprof")
 201                 f, err := os.Create(fn)
 202                 if err != nil {
 203                         log.Fatalf("profile: could not create block profile %q: %v", fn, err)
 204                 }
 205                 runtime.SetBlockProfileRate(1)
 206                 logf("profile: block profiling enabled, %s", fn)
 207                 prof.closer = func() {
 208                         pprof.Lookup("block").WriteTo(f, 0)
 209                         f.Close()
 210                         runtime.SetBlockProfileRate(0)
 211                         logf("profile: block profiling disabled, %s", fn)
 212                 }
 213
 214         case traceMode:
 215                 fn := filepath.Join(path, "trace.out")
 216                 f, err := os.Create(fn)
 217                 if err != nil {
 218                         log.Fatalf("profile: could not create trace output file %q: %v", fn, err)
 219                 }
 220                 if err := startTrace(f); err != nil {
 221                         log.Fatalf("profile: could not start trace: %v", err)
 222                 }
 223                 logf("profile: trace enabled, %s", fn)
 224                 prof.closer = func() {
 225                         stopTrace()
 226                         logf("profile: trace disabled, %s", fn)
 227                 }
 228         }
 229
 230         if !prof.noShutdownHook {
 231                 go func() {
 232                         c := make(chan os.Signal, 1)
 233                         signal.Notify(c, os.Interrupt)
 234                         <-c
 235
 236                         log.Println("profile: caught interrupt, stopping profiles")
 237                         prof.Stop()
 238
 239                         os.Exit(0)
 240                 }()
 241         }
 242
 243         return &prof
 244 }