/
cell.go
131 lines (105 loc) · 3.23 KB
/
cell.go
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
// SPDX-License-Identifier: Apache-2.0
// Copyright Authors of Cilium
package pprof
import (
"errors"
"net"
"net/http"
"net/http/pprof"
_ "net/http/pprof"
"strconv"
"github.com/sirupsen/logrus"
"github.com/spf13/pflag"
"github.com/cilium/cilium/pkg/hive"
"github.com/cilium/cilium/pkg/hive/cell"
)
const (
// Pprof is the flag to enable the registration of pprof HTTP handlers
Pprof = "pprof"
// PprofAddress is the flag to set the address that pprof listens on
PprofAddress = "pprof-address"
// PprofPort is the flag to set the port that pprof listens on
PprofPort = "pprof-port"
)
type Server interface {
// Port returns the port at which the server is listening
Port() int
}
// Cell creates the cell for pprof, that registers its HTTP handlers to serve
// profiling data in the format expected by the pprof visualization tool.
var Cell = cell.Module(
"pprof",
"pprof HTTP server to expose runtime profiling data",
// We don't call cell.Config directly here, because the operator
// uses different flags names for the same pprof config flags.
// Therefore, to register each flag with the correct name, we leave
// the call to cell.Config to the user of the cell.
// Provide coupled with Invoke is used to improve cell testability,
// namely to allow taking a reference to the Server and call Port() on it.
cell.Provide(newServer),
cell.Invoke(func(srv Server) {}),
)
// Config contains the configuration for the pprof cell.
type Config struct {
Pprof bool
PprofAddress string
PprofPort uint16
}
func (def Config) Flags(flags *pflag.FlagSet) {
flags.Bool(Pprof, def.Pprof, "Enable serving pprof debugging API")
flags.String(PprofAddress, def.PprofAddress, "Address that pprof listens on")
flags.Uint16(PprofPort, def.PprofPort, "Port that pprof listens on")
}
func newServer(lc hive.Lifecycle, log logrus.FieldLogger, cfg Config) Server {
if !cfg.Pprof {
return nil
}
srv := &server{
logger: log,
address: cfg.PprofAddress,
port: cfg.PprofPort,
}
lc.Append(srv)
return srv
}
type server struct {
logger logrus.FieldLogger
address string
port uint16
httpSrv *http.Server
listener net.Listener
}
func (s *server) Start(ctx hive.HookContext) error {
listener, err := net.Listen("tcp", net.JoinHostPort(s.address, strconv.FormatUint(uint64(s.port), 10)))
if err != nil {
return err
}
s.listener = listener
s.logger = s.logger.WithFields(logrus.Fields{
"ip": s.listener.Addr().(*net.TCPAddr).IP,
"port": s.listener.Addr().(*net.TCPAddr).Port,
})
mux := http.NewServeMux()
mux.HandleFunc("/debug/pprof/", pprof.Index)
mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
mux.HandleFunc("/debug/pprof/profile", pprof.Profile)
mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
mux.HandleFunc("/debug/pprof/trace", pprof.Trace)
s.httpSrv = &http.Server{
Handler: mux,
}
go func() {
if err := s.httpSrv.Serve(s.listener); !errors.Is(err, http.ErrServerClosed) {
s.logger.WithError(err).Error("server stopped unexpectedly")
}
}()
s.logger.Info("Started pprof server")
return nil
}
func (s *server) Stop(ctx hive.HookContext) error {
s.logger.Info("Stopped pprof server")
return s.httpSrv.Shutdown(ctx)
}
func (s *server) Port() int {
return s.listener.Addr().(*net.TCPAddr).Port
}