From 5a1191a1ab3561d82a2d3503018ba913837528b4 Mon Sep 17 00:00:00 2001
From: Alejandro Lopez <a.lopez@yadro.com>
Date: Tue, 2 May 2023 13:58:57 +0300
Subject: [PATCH] [#20] Add pprof extension with support for cpu and mem

Signed-off-by: Alejandro Lopez <a.lopez@yadro.com>
---
 frostfs.go                  |  1 +
 internal/profile/profile.go | 67 +++++++++++++++++++++++++++++++++++++
 scenarios/run_scenarios.md  |  6 ++++
 3 files changed, 74 insertions(+)
 create mode 100644 internal/profile/profile.go

diff --git a/frostfs.go b/frostfs.go
index a702e22..b6fb527 100644
--- a/frostfs.go
+++ b/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"
diff --git a/internal/profile/profile.go b/internal/profile/profile.go
new file mode 100644
index 0000000..9f8b351
--- /dev/null
+++ b/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)
+}
diff --git a/scenarios/run_scenarios.md b/scenarios/run_scenarios.md
index c39739b..57a4c37 100644
--- a/scenarios/run_scenarios.md
+++ b/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).