xk6-frostfs/internal/profile/profile.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)
}