[#20] Add pprof extension with support for cpu and mem

Signed-off-by: Alejandro Lopez <a.lopez@yadro.com>

frostfs.go

@@ -6,6 +6,7 @@ import (
 	_ "git.frostfs.info/TrueCloudLab/xk6-frostfs/internal/local"
 	_ "git.frostfs.info/TrueCloudLab/xk6-frostfs/internal/logging"
 	_ "git.frostfs.info/TrueCloudLab/xk6-frostfs/internal/native"
+	_ "git.frostfs.info/TrueCloudLab/xk6-frostfs/internal/profile"
 	_ "git.frostfs.info/TrueCloudLab/xk6-frostfs/internal/registry"
 	_ "git.frostfs.info/TrueCloudLab/xk6-frostfs/internal/s3"
 	_ "git.frostfs.info/TrueCloudLab/xk6-frostfs/internal/s3local"

internal/profile/profile.go

@@ -0,0 +1,67 @@
+// Package profile provides an extension to generate profile data from k6 itself.
+//
+// An Output extension is used to leverage the Start and Stop hooks which are
+// otherwise inaccessible in a regular module.
+package profile
+
+import (
+	"fmt"
+	"os"
+	"runtime"
+	"runtime/pprof"
+
+	"go.k6.io/k6/metrics"
+	"go.k6.io/k6/output"
+)
+
+const (
+	cpuProfilePath = "cpu.prof"
+	memProfilePath = "mem.prof"
+)
+
+type profExt struct {
+	cpuFile *os.File
+}
+
+func New(output.Params) (output.Output, error) {
+	return &profExt{}, nil
+}
+
+func (*profExt) Description() string {
+	return "profile"
+}
+
+func (ext *profExt) Start() error {
+	var err error
+	ext.cpuFile, err = os.Create(cpuProfilePath)
+	if err != nil {
+		return fmt.Errorf("creating cpu profile file: %v", err)
+	}
+	if err := pprof.StartCPUProfile(ext.cpuFile); err != nil {
+		return fmt.Errorf("starting cpu profile: %v", err)
+	}
+	return nil
+}
+
+func (ext *profExt) Stop() error {
+	pprof.StopCPUProfile()
+	if err := ext.cpuFile.Close(); err != nil {
+		return fmt.Errorf("closing cpu profile file: %v", err)
+	}
+	f, err := os.Create(memProfilePath)
+	if err != nil {
+		return fmt.Errorf("creating mem profile file: %v", err)
+	}
+	defer f.Close()
+	runtime.GC()
+	if err := pprof.WriteHeapProfile(f); err != nil {
+		return fmt.Errorf("writing mem profile: %v", err)
+	}
+	return nil
+}
+
+func (*profExt) AddMetricSamples([]metrics.SampleContainer) {}
+
+func init() {
+	output.RegisterExtension("profile", New)
+}

scenarios/run_scenarios.md

@@ -19,6 +19,12 @@ Scenarios `grpc.js`, `local.js`, `http.js` and `s3.js` support the following opt
   * `SLEEP_READ` - time interval (in seconds) between reading VU iterations.
   * `SELECTION_SIZE` - size of batch to select for deletion (default: 1000).
 
+Additionally, the profiling extension can be enabled to generate CPU and memory profiles which can be inspected with `go tool pprof file.prof`:
+```shell
+$ ./k6 run --out profile (...)
+```
+The profiles are saved in the current directory as `cpu.prof` and `mem.prof`, respectively.
+
 ## Common options for the local scenarios:
 
   * `DEBUG_LOGGER` - uses a development logger for the local storage engine to aid debugging (default: false).