<commit_before>package profiler

import (
	"fmt"
	"os"
	"runtime"
	"runtime/pprof"
	"time"
)

type Config struct {
	CPU               bool   `json:"cpu"`
	Memory            bool   `json:"memory"`
	Block             bool   `json:"block"`
	Goroutine         bool   `json:"goroutine"`
	Prefix            string `json:"prefix"`
	Interval          string `json:"interval"`
	MemoryProfileRate int    `json:"memory_profile_rate"`
}

type profiler struct {
	conf        Config
	terminateCh chan struct{}
	closers     []func()
}

func NewProfiler(conf Config) *profiler {
	return &profiler{
		conf:        conf,
		terminateCh: make(chan struct{}, 0),
		closers:     make([]func(), 0),
	}
}

func (c Config) isOn() bool {
	return c.CPU || c.Memory || c.Goroutine || c.Block
}

func (p profiler) Run() {

	if p.conf.CPU {
		p.startProfilingCPU()
	}
	if p.conf.Memory {
		runtime.MemProfileRate = p.conf.MemoryProfileRate
	}
	if p.conf.Block {
		runtime.SetBlockProfileRate(1)
	}

	if p.conf.isOn() && "" != p.conf.Interval {
		interval, err := time.ParseDuration(p.conf.Interval)
		if nil != err {
			fmt.Println(err)
			return
		}
		timer := time.NewTimer(interval)
		select {
		case <-timer.C:
			p.TakeSnapshot()
			// start again
			p.Run()
		case <-p.terminateCh:
			p.TakeSnapshot()
		}
	}
}

func (p profiler) TakeSnapshot() {
	if p.conf.CPU {
		p.takeCPUSnapshot()
	}
	if p.conf.Memory {
		p.takeMemorySnapshot()
	}
	if p.conf.Block {
		p.takeBlockSnapshot()
	}
	if p.conf.Goroutine {
		p.takeGoroutineSnapshot()
	}

	for _, c := range p.closers {
		c()
	}
	p.closers = nil
}

func (p profiler) Stop() {
	close(p.terminateCh)
}

func (p profiler) startProfilingCPU() {
	pprofFile := fmt.Sprintf("%scpu.%d.pprof", p.conf.Prefix, time.Now().Unix())
	fmt.Println("Starting new CPU Profiler:", pprofFile)
	f, err := os.Create(pprofFile)
	if err != nil {
		panic(err)
	}
	pprof.StartCPUProfile(f)
	p.closers = append(p.closers, func() {
		f.Close()
	})
}

// collect CPU profiling information
func (p profiler) takeCPUSnapshot() {
	fmt.Println("Stopping CPU Profiler")
	pprof.StopCPUProfile()<commit_msg>make the CPU profiling rate configurable<commit_after>package profiler

import (
	"fmt"
	"os"
	"runtime"
	"runtime/pprof"
	"time"
)

// Config wraps the profiler settings
type Config struct {
	CPU               bool   `json:"cpu"`
	Memory            bool   `json:"memory"`
	Block             bool   `json:"block"`
	Goroutine         bool   `json:"goroutine"`
	Prefix            string `json:"prefix"`
	Interval          string `json:"interval"`
	MemoryProfileRate int    `json:"memory_profile_rate"`
	CPUProfileRate    int    `json:"cpu_profile_rate"`
}

type profiler struct {
	conf        Config
	terminateCh chan struct{}
	closers     []func()
}

// NewProfiler initialises a new instance of a profiler
func NewProfiler(conf Config) *profiler {
	return &profiler{
		conf:        conf,
		terminateCh: make(chan struct{}, 0),
		closers:     make([]func(), 0),
	}
}

func (c Config) isOn() bool {
	return c.CPU || c.Memory || c.Goroutine || c.Block
}

// Run starts the profiler
func (p profiler) Run() {
	if p.conf.CPU {
		if p.conf.CPUProfileRate > 0 {
			runtime.SetCPUProfileRate(p.conf.CPUProfileRate)
		}
		p.startProfilingCPU()
	}
	if p.conf.Memory {
		runtime.MemProfileRate = p.conf.MemoryProfileRate
	}
	if p.conf.Block {
		runtime.SetBlockProfileRate(1)
	}

	if p.conf.isOn() && "" != p.conf.Interval {
		interval, err := time.ParseDuration(p.conf.Interval)
		if nil != err {
			fmt.Println(err)
			return
		}
		timer := time.NewTimer(interval)
		select {
		case <-timer.C:
			p.TakeSnapshot()
			// start again
			p.Run()
		case <-p.terminateCh:
			p.TakeSnapshot()
		}
	}
}

// TakeSnapshot takes a profiling data snapshot for the enabled resources
func (p profiler) TakeSnapshot() {
	if p.conf.CPU {
		p.takeCPUSnapshot()
	}
	if p.conf.Memory {
		p.takeMemorySnapshot()
	}
	if p.conf.Block {
		p.takeBlockSnapshot()
	}
	if p.conf.Goroutine {
		p.takeGoroutineSnapshot()
	}

	for _, c := range p.closers {
		c()
	}
	p.closers = nil
}

// Stop terminates the active profiler(s)
func (p profiler) Stop() {
	close(p.terminateCh)
}

// opens a new output file to collect CPU profiling information
func (p profiler) startProfilingCPU() {
	pprofFile := fmt.Sprintf("%scpu.%d.pprof", p.conf.Prefix, time.Now().Unix())
	fmt.Println("Starting new CPU Profiler:", pprofFile)
	f, err := os.Create(pprofFile)
	if err != nil {
		panic(err)
	}
	pprof.StartCPUProfile(f)
	p.closers = append(p.closers, func() {
		f.Close()
	})
}

// collect CPU profiling information
func (p profiler) takeCPUSnapshot() {
	fmt.Println("Stopping CPU Profiler")
	pprof.StopCPUProfile()
{'fraction_non_alphanumeric': 0.14184549356223175, 'fraction_numerical': 0.0015021459227467812, 'mean_word_length': 3.124778761061947, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '194152', 'n_tokens_mistral': 1811, 'n_tokens_neox': 1731, 'n_words': 391}