profile.go

// Package profile provides a simple way to manage runtime/pprof

// profiling of your Go application.

package profile

import (

    "io/ioutil"

    "log"

    "os"

    "os/signal"

    "path/filepath"

    "runtime"

    "runtime/pprof"

    "sync/atomic"

)

const (

    cpuMode = iota

    memMode

    blockMode

    traceMode

)

// Profile represents an active profiling session.

type Profile struct {

    // quiet suppresses informational messages during profiling.

    quiet bool

    // noShutdownHook controls whether the profiling package should

    // hook SIGINT to write profiles cleanly.

    noShutdownHook bool

    // mode holds the type of profiling that will be made

    mode int

    // path holds the base path where various profiling files are  written.

    // If blank, the base path will be generated by ioutil.TempDir.

    path string

    // memProfileRate holds the rate for the memory profile.

    memProfileRate int

    // closer holds a cleanup function that run after each profile

    closer func()

    // stopped records if a call to profile.Stop has been made

    stopped uint32

}

// NoShutdownHook controls whether the profiling package should

// hook SIGINT to write profiles cleanly.

// Programs with more sophisticated signal handling should set

// this to true and ensure the Stop() function returned from Start()

// is called during shutdown.

func NoShutdownHook(p *Profile) { p.noShutdownHook = true }

// Quiet suppresses informational messages during profiling.

func Quiet(p *Profile) { p.quiet = true }

// CPUProfile enables cpu profiling.

// It disables any previous profiling settings.

func CPUProfile(p *Profile) { p.mode = cpuMode }

// DefaultMemProfileRate is the default memory profiling rate.

// See also http://golang.org/pkg/runtime/#pkg-variables

const DefaultMemProfileRate = 4096

// MemProfile enables memory profiling.

// It disables any previous profiling settings.

func MemProfile(p *Profile) {

    p.memProfileRate = DefaultMemProfileRate

    p.mode = memMode

}

// MemProfileRate enables memory profiling at the preferred rate.

// It disables any previous profiling settings.

func MemProfileRate(rate int) func(*Profile) {

    return func(p *Profile) {

        p.memProfileRate = rate

        p.mode = memMode

    }

}

// BlockProfile enables block (contention) profiling.

// It disables any previous profiling settings.

func BlockProfile(p *Profile) { p.mode = blockMode }

// Trace profile controls if execution tracing will be enabled. It disables any previous profiling settings.

func TraceProfile(p *Profile) { p.mode = traceMode }

// ProfilePath controls the base path where various profiling

// files are written. If blank, the base path will be generated

// by ioutil.TempDir.

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() {

    if !atomic.CompareAndSwapUint32(&p.stopped, 0, 1) {

        // someone has already called close

        return

    }

    p.closer()

    atomic.StoreUint32(&started, 0)

}

// started is non zero if a profile is running.

var started uint32

// 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()

} {

    if !atomic.CompareAndSwapUint32(&started, 0, 1) {

        log.Fatal("profile: Start() already called")

    }

    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 ioutil.TempDir("", "profile")

    }()

    if err != nil {

        log.Fatalf("profile: could not create initial output directory: %v", err)

    }

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

        if !prof.quiet {

            log.Printf(format, args...)

        }

    }

    switch prof.mode {

    case cpuMode:

        fn := filepath.Join(path, "cpu.pprof")

        f, err := os.Create(fn)

        if err != nil {

            log.Fatalf("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, "mem.pprof")

        f, err := os.Create(fn)

        if err != nil {

            log.Fatalf("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 blockMode:

        fn := filepath.Join(path, "block.pprof")

        f, err := os.Create(fn)

        if err != nil {

            log.Fatalf("profile: could not create block profile %q: %v", fn, err)

        }

        runtime.SetBlockProfileRate(1)

        logf("profile: block profiling enabled, %s", fn)

        prof.closer = func() {

            pprof.Lookup("block").WriteTo(f, 0)

            f.Close()

            runtime.SetBlockProfileRate(0)

            logf("profile: block profiling disabled, %s", fn)

        }

    case traceMode:

        fn := filepath.Join(path, "trace.out")

        f, err := os.Create(fn)

        if err != nil {

            log.Fatalf("profile: could not create trace output file %q: %v", fn, err)

        }

        if err := startTrace(f); err != nil {

            log.Fatalf("profile: could not start trace: %v", err)

        }

        logf("profile: trace enabled, %s", fn)

        prof.closer = func() {

            stopTrace()

            logf("profile: trace disabled, %s", fn)

        }

    }

    if !prof.noShutdownHook {

        go func() {

            c := make(chan os.Signal, 1)

            signal.Notify(c, os.Interrupt)

            <-c

            log.Println("profile: caught interrupt, stopping profiles")

            prof.Stop()

            os.Exit(0)

        }()

    }

    return &prof

}