File: envpprof.go

package info (click to toggle)
golang-github-anacrolix-envpprof 1.2.1-2
  • links: PTS, VCS
  • area: main
  • in suites: bookworm, forky, sid, trixie
  • size: 72 kB
  • sloc: makefile: 2
file content (127 lines) | stat: -rw-r--r-- 3,012 bytes parent folder | download
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
package envpprof

import (
	"expvar"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	_ "net/http/pprof"
	"os"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"strings"

	"github.com/anacrolix/log"
)

var (
	pprofDir = filepath.Join(os.Getenv("HOME"), "pprof")
	heap     bool
)

func writeHeapProfile() {
	os.Mkdir(pprofDir, 0750)
	f, err := ioutil.TempFile(pprofDir, "heap")
	if err != nil {
		log.Printf("error creating heap profile file: %s", err)
		return
	}
	defer f.Close()
	pprof.WriteHeapProfile(f)
	log.Printf("wrote heap profile to %q", f.Name())
}

// Stop ends CPU profiling, waiting for writes to complete. If heap profiling is enabled, it also
// writes the heap profile to a file. Stop should be deferred from main if cpu or heap profiling
// are to be used through envpprof.
func Stop() {
	// Should we check if CPU profiling was initiated through this package?
	pprof.StopCPUProfile()
	if heap {
		// Can or should we do this concurrently with stopping CPU profiling?
		writeHeapProfile()
	}
}

func startHTTP(value string) {
	var l net.Listener
	if value == "" {
		for port := uint16(6061); port != 6060; port++ {
			var err error
			l, err = net.Listen("tcp", fmt.Sprintf("localhost:%d", port))
			if err == nil {
				break
			}
		}
		if l == nil {
			log.Print("unable to create envpprof listener for http")
			return
		}
	} else {
		var addr string
		_, _, err := net.SplitHostPort(value)
		if err == nil {
			addr = value
		} else {
			addr = "localhost:" + value
		}
		l, err = net.Listen("tcp", addr)
		if err != nil {
			panic(err)
		}
	}
	log.Printf("(pid=%d) envpprof serving http://%s", os.Getpid(), l.Addr())
	go func() {
		defer l.Close()
		log.Printf("error serving http on envpprof listener: %s", http.Serve(l, nil))
	}()
}

func init() {
	expvar.Publish("numGoroutine", expvar.Func(func() interface{} { return runtime.NumGoroutine() }))
	_var := os.Getenv("GOPPROF")
	if _var == "" {
		return
	}
	for _, item := range strings.Split(os.Getenv("GOPPROF"), ",") {
		equalsPos := strings.IndexByte(item, '=')
		var key, value string
		if equalsPos < 0 {
			key = item
		} else {
			key = item[:equalsPos]
			value = item[equalsPos+1:]
		}
		switch key {
		case "http":
			startHTTP(value)
		case "cpu":
			os.Mkdir(pprofDir, 0750)
			f, err := ioutil.TempFile(pprofDir, "cpu")
			if err != nil {
				log.Printf("error creating cpu pprof file: %s", err)
				break
			}
			err = pprof.StartCPUProfile(f)
			if err != nil {
				log.Printf("error starting cpu profiling: %s", err)
				break
			}
			log.Printf("cpu profiling to file %q", f.Name())
		case "block":
			// Taken from Safe Rate at
			// https://github.com/DataDog/go-profiler-notes/blob/main/guide/README.md#go-profilers.
			runtime.SetBlockProfileRate(10000)
		case "heap":
			heap = true
		case "mutex":
			// Taken from Safe Rate at
			// https://github.com/DataDog/go-profiler-notes/blob/main/guide/README.md#go-profilers.
			runtime.SetMutexProfileFraction(100)
		default:
			log.Printf("unexpected GOPPROF key %q", key)
		}
	}
}