add CLI flag to select specific diag collectors

Author: guseggert

core/commands/profile.go

@@ -21,6 +21,7 @@ type profileResult struct {
 }
 
 const (
+	collectorsOptionName       = "collectors"
 	profileTimeOption          = "profile-time"
 	mutexProfileFractionOption = "mutex-profile-fraction"
 	blockProfileRateOption     = "block-profile-rate"
@@ -72,11 +73,24 @@ However, it could reveal:
 	NoLocal: true,
 	Options: []cmds.Option{
 		cmds.StringOption(outputOptionName, "o", "The path where the output .zip should be stored. Default: ./ipfs-profile-[timestamp].zip"),
+		cmds.DelimitedStringsOption(",", collectorsOptionName, "The list of collectors to use for collecting diagnostic data.").
+			WithDefault([]string{
+				profile.CollectorGoroutinesStack,
+				profile.CollectorGoroutinesPprof,
+				profile.CollectorVersion,
+				profile.CollectorHeap,
+				profile.CollectorBin,
+				profile.CollectorCPU,
+				profile.CollectorMutex,
+				profile.CollectorBlock,
+			}),
 		cmds.StringOption(profileTimeOption, "The amount of time spent profiling. If this is set to 0, then sampling profiles are skipped.").WithDefault("30s"),
 		cmds.IntOption(mutexProfileFractionOption, "The fraction 1/n of mutex contention events that are reported in the mutex profile.").WithDefault(4),
 		cmds.StringOption(blockProfileRateOption, "The duration to wait between sampling goroutine-blocking events for the blocking profile.").WithDefault("1ms"),
 	},
 	Run: func(req *cmds.Request, res cmds.ResponseEmitter, env cmds.Environment) error {
+		collectors := req.Options[collectorsOptionName].([]string)
+
 		profileTimeStr, _ := req.Options[profileTimeOption].(string)
 		profileTime, err := time.ParseDuration(profileTimeStr)
 		if err != nil {
@@ -96,6 +110,7 @@ However, it could reveal:
 		go func() {
 			archive := zip.NewWriter(w)
 			err = profile.WriteProfiles(req.Context, archive, profile.Options{
+				Collectors:           collectors,
 				ProfileDuration:      profileTime,
 				MutexProfileFraction: mutexProfileFraction,
 				BlockProfileRate:     blockProfileRate,

profile/profile.go

@@ -17,19 +17,30 @@ import (
 	"github.com/ipfs/go-log"
 )
 
+const (
+	CollectorGoroutinesStack = "goroutines-stack"
+	CollectorGoroutinesPprof = "goroutines-pprof"
+	CollectorVersion         = "version"
+	CollectorHeap            = "heap"
+	CollectorBin             = "bin"
+	CollectorCPU             = "cpu"
+	CollectorMutex           = "mutex"
+	CollectorBlock           = "block"
+)
+
 var (
 	logger = log.Logger("profile")
 	goos   = runtime.GOOS
 )
 
-type profile struct {
+type collector struct {
 	outputFile   string
 	isExecutable bool
-	profileFunc  func(ctx context.Context, opts Options, writer io.Writer) error
+	collectFunc  func(ctx context.Context, opts Options, writer io.Writer) error
 	enabledFunc  func(opts Options) bool
 }
 
-func (p *profile) fileName() string {
+func (p *collector) outputFileName() string {
 	fName := p.outputFile
 	if p.isExecutable {
 		if goos == "windows" {
@@ -39,51 +50,52 @@ func (p *profile) fileName() string {
 	return fName
 }
 
-var profiles = []profile{
-	{
+var collectors = map[string]collector{
+	CollectorGoroutinesStack: {
 		outputFile:  "goroutines.stacks",
-		profileFunc: goroutineStacksText,
+		collectFunc: goroutineStacksText,
 		enabledFunc: func(opts Options) bool { return true },
 	},
-	{
+	CollectorGoroutinesPprof: {
 		outputFile:  "goroutines.pprof",
-		profileFunc: goroutineStacksProto,
+		collectFunc: goroutineStacksProto,
 		enabledFunc: func(opts Options) bool { return true },
 	},
-	{
+	CollectorVersion: {
 		outputFile:  "version.json",
-		profileFunc: versionInfo,
+		collectFunc: versionInfo,
 		enabledFunc: func(opts Options) bool { return true },
 	},
-	{
+	CollectorHeap: {
 		outputFile:  "heap.pprof",
-		profileFunc: heapProfile,
+		collectFunc: heapProfile,
 		enabledFunc: func(opts Options) bool { return true },
 	},
-	{
+	CollectorBin: {
 		outputFile:   "ipfs",
 		isExecutable: true,
-		profileFunc:  binary,
+		collectFunc:  binary,
 		enabledFunc:  func(opts Options) bool { return true },
 	},
-	{
+	CollectorCPU: {
 		outputFile:  "cpu.pprof",
-		profileFunc: profileCPU,
+		collectFunc: profileCPU,
 		enabledFunc: func(opts Options) bool { return opts.ProfileDuration > 0 },
 	},
-	{
+	CollectorMutex: {
 		outputFile:  "mutex.pprof",
-		profileFunc: mutexProfile,
+		collectFunc: mutexProfile,
 		enabledFunc: func(opts Options) bool { return opts.ProfileDuration > 0 && opts.MutexProfileFraction > 0 },
 	},
-	{
+	CollectorBlock: {
 		outputFile:  "block.pprof",
-		profileFunc: blockProfile,
+		collectFunc: blockProfile,
 		enabledFunc: func(opts Options) bool { return opts.ProfileDuration > 0 && opts.BlockProfileRate > 0 },
 	},
 }
 
 type Options struct {
+	Collectors           []string
 	ProfileDuration      time.Duration
 	MutexProfileFraction int
 	BlockProfileRate     time.Duration
@@ -97,7 +109,7 @@ func WriteProfiles(ctx context.Context, archive *zip.Writer, opts Options) error
 	return p.runProfile(ctx)
 }
 
-// profiler runs the profiles concurrently and writes the results to the zip archive.
+// profiler runs the collectors concurrently and writes the results to the zip archive.
 type profiler struct {
 	archive *zip.Writer
 	opts    Options
@@ -113,22 +125,31 @@ func (p *profiler) runProfile(ctx context.Context) error {
 	ctx, cancelFn := context.WithCancel(ctx)
 	defer cancelFn()
 
-	results := make(chan profileResult, len(profiles))
+	var collectorsToRun []collector
+	for _, name := range p.opts.Collectors {
+		c, ok := collectors[name]
+		if !ok {
+			return fmt.Errorf("unknown collector '%s'", name)
+		}
+		collectorsToRun = append(collectorsToRun, c)
+	}
+
+	results := make(chan profileResult, len(p.opts.Collectors))
 	wg := sync.WaitGroup{}
-	for _, prof := range profiles {
-		if !prof.enabledFunc(p.opts) {
+	for _, c := range collectorsToRun {
+		if !c.enabledFunc(p.opts) {
 			continue
 		}
 
-		fName := prof.fileName()
+		fName := c.outputFileName()
 
 		wg.Add(1)
-		go func(prof profile) {
+		go func(c collector) {
 			defer wg.Done()
 			logger.Infow("collecting profile", "File", fName)
 			defer logger.Infow("profile done", "File", fName)
 			b := bytes.Buffer{}
-			err := prof.profileFunc(ctx, p.opts, &b)
+			err := c.collectFunc(ctx, p.opts, &b)
 			if err != nil {
 				select {
 				case results <- profileResult{err: fmt.Errorf("generating profile data for %q: %w", fName, err)}:
@@ -140,7 +161,7 @@ func (p *profiler) runProfile(ctx context.Context) error {
 			case results <- profileResult{buf: &b, fName: fName}:
 			case <-ctx.Done():
 			}
-		}(prof)
+		}(c)
 	}
 	go func() {
 		wg.Wait()

test/sharness/t0152-profile.sh

@@ -16,7 +16,7 @@ test_expect_success "profiling requires a running daemon" '
 
 test_launch_ipfs_daemon
 
-test_expect_success "test profiling (without samplin)" '
+test_expect_success "test profiling (without sampling)" '
   ipfs diag profile --profile-time=0 > cmd_out
 '
 
@@ -35,14 +35,20 @@ test_expect_success "test profiling with -o" '
 test_expect_success "test that test-profile.zip exists" '
   test -e test-profile.zip
 '
+
+test_expect_success "test profiling with specific collectors" '
+  ipfs diag profile --collectors version,goroutines-stack -o test-profile-small.zip
+'
+
 test_kill_ipfs_daemon
 
 if ! test_have_prereq UNZIP; then
     test_done
 fi
 
 test_expect_success "unpack profiles" '
-  unzip -d profiles test-profile.zip
+  unzip -d profiles test-profile.zip &&
+  unzip -d profiles-small test-profile-small.zip
 '
 
 test_expect_success "cpu profile is valid" '
@@ -69,4 +75,10 @@ test_expect_success "goroutines stacktrace is valid" '
   grep -q "goroutine" "profiles/goroutines.stacks"
 '
 
+test_expect_success "the small profile only contains the requested data" '
+  find profiles-small -type f > actual &&
+  echo -e "profiles-small/goroutines.stacks\nprofiles-small/version.json" > expected &&
+  test_cmp expected actual
+'
+
 test_done