move implementation to internal/ directory (#828)

internal/daemon/pprof.go

@@ -0,0 +1,94 @@
+package daemon
+
+import (
+	"net/http"
+	// FIXME: importing this package has the side-effect of poisoning the http.DefaultServeMux
+	// FIXME: with the /debug/pprof endpoints
+	"context"
+	"net"
+	"net/http/pprof"
+
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+	"golang.org/x/net/websocket"
+
+	"github.com/zrepl/zrepl/internal/daemon/job"
+	"github.com/zrepl/zrepl/internal/daemon/logging/trace"
+)
+
+type pprofServer struct {
+	cc       chan PprofServerControlMsg
+	listener net.Listener
+}
+
+type PprofServerControlMsg struct {
+	// Whether the server should listen for requests on the given address
+	Run bool
+	// Must be set if Run is true, undefined otherwise
+	HttpListenAddress string
+}
+
+func NewPProfServer(ctx context.Context) *pprofServer {
+
+	s := &pprofServer{
+		cc: make(chan PprofServerControlMsg),
+	}
+
+	go s.controlLoop(ctx)
+	return s
+}
+
+func (s *pprofServer) controlLoop(ctx context.Context) {
+outer:
+	for {
+
+		var msg PprofServerControlMsg
+		select {
+		case <-ctx.Done():
+			if s.listener != nil {
+				s.listener.Close()
+			}
+			break outer
+		case msg = <-s.cc:
+			// proceed
+		}
+
+		var err error
+		if msg.Run && s.listener == nil {
+
+			s.listener, err = net.Listen("tcp", msg.HttpListenAddress)
+			if err != nil {
+				s.listener = nil
+				continue
+			}
+
+			// FIXME: because net/http/pprof does not provide a mux,
+			mux := http.NewServeMux()
+			mux.Handle("/debug/pprof/", http.HandlerFunc(pprof.Index))
+			mux.Handle("/debug/pprof/cmdline", http.HandlerFunc(pprof.Cmdline))
+			mux.Handle("/debug/pprof/profile", http.HandlerFunc(pprof.Profile))
+			mux.Handle("/debug/pprof/symbol", http.HandlerFunc(pprof.Symbol))
+			mux.Handle("/debug/pprof/trace", http.HandlerFunc(pprof.Trace))
+			mux.Handle("/metrics", promhttp.Handler())
+			mux.Handle("/debug/zrepl/activity-trace", websocket.Handler(trace.ChrometraceClientWebsocketHandler))
+			go func() {
+				err := http.Serve(s.listener, mux)
+				if ctx.Err() != nil {
+					return
+				} else if err != nil {
+					job.GetLogger(ctx).WithError(err).Error("pprof server serve error")
+				}
+			}()
+			continue
+		}
+
+		if !msg.Run && s.listener != nil {
+			s.listener.Close()
+			s.listener = nil
+			continue
+		}
+	}
+}
+
+func (s *pprofServer) Control(msg PprofServerControlMsg) {
+	s.cc <- msg
+}