git-spice/profile_enable.go at main · abhinav/git-spice · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
//go:build profile

package main

import (
	"errors"
	"fmt"
	"io"
	"os"
	"runtime/pprof"
	"runtime/trace"
)

type ProfileFlags struct {
	CPUProfile string `name:"cpuprofile" placeholder:"FILE" help:"Write CPU profile to file"`
	MemProfile string `name:"memprofile" placeholder:"FILE" help:"Write memory profile to file"`
	Trace      string `name:"trace" placeholder:"FILE" help:"Write trace to file"`

	cpuProfileFile io.WriteCloser
	memProfileFile io.WriteCloser
	traceFile      io.WriteCloser
}

func (pf *ProfileFlags) Start() error {
	return errors.Join(
		pf.startCPUProfile(),
		pf.startMemProfile(),
		pf.startTrace(),
	)
}

func (pf *ProfileFlags) Stop() error {
	return errors.Join(
		pf.stopCPUProfile(),
		pf.stopMemProfile(),
		pf.stopTrace(),
	)
}

func (pf *ProfileFlags) startCPUProfile() error {
	path := pf.CPUProfile
	if path == "" {
		return nil
	}

	f, err := os.Create(path)
	if err != nil {
		return fmt.Errorf("create CPU profile file: %w", err)
	}

	if err := pprof.StartCPUProfile(f); err != nil {
		return fmt.Errorf("start CPU profile: %w", err)
	}

	pf.cpuProfileFile = f
	return nil
}

func (pf *ProfileFlags) stopCPUProfile() error {
	f := pf.cpuProfileFile
	if f == nil {
		return nil
	}

	pprof.StopCPUProfile()

	if err := f.Close(); err != nil {
		return fmt.Errorf("close CPU profile file: %w", err)
	}

	return nil
}

func (pf *ProfileFlags) startMemProfile() error {
	path := pf.MemProfile
	if path == "" {
		return nil
	}

	f, err := os.Create(path)
	if err != nil {
		return fmt.Errorf("create memory profile file: %w", err)
	}

	pf.memProfileFile = f
	return nil
}

func (pf *ProfileFlags) stopMemProfile() error {
	f := pf.memProfileFile
	if f == nil {
		return nil
	}

	if err := pprof.Lookup("heap").WriteTo(f, 0); err != nil {
		return fmt.Errorf("write memory profile: %w", err)
	}

	if err := f.Close(); err != nil {
		return fmt.Errorf("close memory profile file: %w", err)
	}

	return nil
}

func (pf *ProfileFlags) startTrace() error {
	path := pf.Trace
	if path == "" {
		return nil
	}

	f, err := os.Create(path)
	if err != nil {
		return fmt.Errorf("create trace file: %w", err)
	}

	if err := trace.Start(f); err != nil {
		return fmt.Errorf("start trace: %w", err)
	}

	pf.traceFile = f
	return nil
}

func (pf *ProfileFlags) stopTrace() error {
	f := pf.traceFile
	if f == nil {
		return nil
	}

	trace.Stop()

	if err := f.Close(); err != nil {
		return fmt.Errorf("close trace file: %w", err)
	}

	return nil
}