mirror of
https://github.com/restic/restic.git
synced 2025-03-30 00:00:14 +01:00
126 lines
3.3 KiB
Go
126 lines
3.3 KiB
Go
//go:build debug || profile
|
|
// +build debug profile
|
|
|
|
package main
|
|
|
|
import (
|
|
"fmt"
|
|
"net/http"
|
|
_ "net/http/pprof"
|
|
"os"
|
|
|
|
"github.com/restic/restic/internal/errors"
|
|
"github.com/restic/restic/internal/repository"
|
|
"github.com/spf13/cobra"
|
|
"github.com/spf13/pflag"
|
|
|
|
"github.com/pkg/profile"
|
|
)
|
|
|
|
func registerProfiling(cmd *cobra.Command) {
|
|
var profiler profiler
|
|
|
|
origPreRun := cmd.PersistentPreRunE
|
|
cmd.PersistentPreRunE = func(cmd *cobra.Command, args []string) error {
|
|
if origPreRun != nil {
|
|
if err := origPreRun(cmd, args); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return profiler.Start(profiler.opts)
|
|
}
|
|
|
|
origPostRun := cmd.PersistentPostRunE
|
|
cmd.PersistentPostRunE = func(cmd *cobra.Command, args []string) error {
|
|
profiler.Stop()
|
|
if origPostRun != nil {
|
|
return origPostRun(cmd, args)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
profiler.opts.AddFlags(cmd.PersistentFlags())
|
|
}
|
|
|
|
type profiler struct {
|
|
opts ProfileOptions
|
|
stop interface {
|
|
Stop()
|
|
}
|
|
}
|
|
|
|
type ProfileOptions struct {
|
|
listen string
|
|
memPath string
|
|
cpuPath string
|
|
tracePath string
|
|
blockPath string
|
|
insecure bool
|
|
}
|
|
|
|
func (opts *ProfileOptions) AddFlags(f *pflag.FlagSet) {
|
|
f.StringVar(&opts.listen, "listen-profile", "", "listen on this `address:port` for memory profiling")
|
|
f.StringVar(&opts.memPath, "mem-profile", "", "write memory profile to `dir`")
|
|
f.StringVar(&opts.cpuPath, "cpu-profile", "", "write cpu profile to `dir`")
|
|
f.StringVar(&opts.tracePath, "trace-profile", "", "write trace to `dir`")
|
|
f.StringVar(&opts.blockPath, "block-profile", "", "write block profile to `dir`")
|
|
f.BoolVar(&opts.insecure, "insecure-kdf", false, "use insecure KDF settings")
|
|
}
|
|
|
|
type fakeTestingTB struct{}
|
|
|
|
func (fakeTestingTB) Logf(msg string, args ...interface{}) {
|
|
fmt.Fprintf(os.Stderr, msg, args...)
|
|
}
|
|
|
|
func (p *profiler) Start(profileOpts ProfileOptions) error {
|
|
if profileOpts.listen != "" {
|
|
fmt.Fprintf(os.Stderr, "running profile HTTP server on %v\n", profileOpts.listen)
|
|
go func() {
|
|
err := http.ListenAndServe(profileOpts.listen, nil)
|
|
if err != nil {
|
|
fmt.Fprintf(os.Stderr, "profile HTTP server listen failed: %v\n", err)
|
|
}
|
|
}()
|
|
}
|
|
|
|
profilesEnabled := 0
|
|
if profileOpts.memPath != "" {
|
|
profilesEnabled++
|
|
}
|
|
if profileOpts.cpuPath != "" {
|
|
profilesEnabled++
|
|
}
|
|
if profileOpts.tracePath != "" {
|
|
profilesEnabled++
|
|
}
|
|
if profileOpts.blockPath != "" {
|
|
profilesEnabled++
|
|
}
|
|
|
|
if profilesEnabled > 1 {
|
|
return errors.Fatal("only one profile (memory, CPU, trace, or block) may be activated at the same time")
|
|
}
|
|
|
|
if profileOpts.memPath != "" {
|
|
p.stop = profile.Start(profile.Quiet, profile.NoShutdownHook, profile.MemProfile, profile.ProfilePath(profileOpts.memPath))
|
|
} else if profileOpts.cpuPath != "" {
|
|
p.stop = profile.Start(profile.Quiet, profile.NoShutdownHook, profile.CPUProfile, profile.ProfilePath(profileOpts.cpuPath))
|
|
} else if profileOpts.tracePath != "" {
|
|
p.stop = profile.Start(profile.Quiet, profile.NoShutdownHook, profile.TraceProfile, profile.ProfilePath(profileOpts.tracePath))
|
|
} else if profileOpts.blockPath != "" {
|
|
p.stop = profile.Start(profile.Quiet, profile.NoShutdownHook, profile.BlockProfile, profile.ProfilePath(profileOpts.blockPath))
|
|
}
|
|
|
|
if profileOpts.insecure {
|
|
repository.TestUseLowSecurityKDFParameters(fakeTestingTB{})
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (p *profiler) Stop() {
|
|
if p.stop != nil {
|
|
p.stop.Stop()
|
|
}
|
|
}
|