xk6-frostfs/internal/profile/profile.go
Alejandro Lopez 5a1191a1ab [#20] Add pprof extension with support for cpu and mem
Signed-off-by: Alejandro Lopez <a.lopez@yadro.com>
2023-05-18 15:08:33 +00:00

67 lines
1.4 KiB
Go

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