1package main
2
3import (
4	"github.com/emicklei/go-restful"
5	"io"
6	"log"
7	"os"
8	"runtime/pprof"
9)
10
11// ProfilingService is a WebService that can start/stop a CPU profile and write results to a file
12// 	GET /{rootPath}/start will activate CPU profiling
13//	GET /{rootPath}/stop will stop profiling
14//
15// NewProfileService("/profiler", "ace.prof").AddWebServiceTo(restful.DefaultContainer)
16//
17type ProfilingService struct {
18	rootPath   string   // the base (root) of the service, e.g. /profiler
19	cpuprofile string   // the output filename to write profile results, e.g. myservice.prof
20	cpufile    *os.File // if not nil, then profiling is active
21}
22
23func NewProfileService(rootPath string, outputFilename string) *ProfilingService {
24	ps := new(ProfilingService)
25	ps.rootPath = rootPath
26	ps.cpuprofile = outputFilename
27	return ps
28}
29
30// Add this ProfileService to a restful Container
31func (p ProfilingService) AddWebServiceTo(container *restful.Container) {
32	ws := new(restful.WebService)
33	ws.Path(p.rootPath).Consumes("*/*").Produces(restful.MIME_JSON)
34	ws.Route(ws.GET("/start").To(p.startProfiler))
35	ws.Route(ws.GET("/stop").To(p.stopProfiler))
36	container.Add(ws)
37}
38
39func (p *ProfilingService) startProfiler(req *restful.Request, resp *restful.Response) {
40	if p.cpufile != nil {
41		io.WriteString(resp.ResponseWriter, "[restful] CPU profiling already running")
42		return // error?
43	}
44	cpufile, err := os.Create(p.cpuprofile)
45	if err != nil {
46		log.Fatal(err)
47	}
48	// remember for close
49	p.cpufile = cpufile
50	pprof.StartCPUProfile(cpufile)
51	io.WriteString(resp.ResponseWriter, "[restful] CPU profiling started, writing on:"+p.cpuprofile)
52}
53
54func (p *ProfilingService) stopProfiler(req *restful.Request, resp *restful.Response) {
55	if p.cpufile == nil {
56		io.WriteString(resp.ResponseWriter, "[restful] CPU profiling not active")
57		return // error?
58	}
59	pprof.StopCPUProfile()
60	p.cpufile.Close()
61	p.cpufile = nil
62	io.WriteString(resp.ResponseWriter, "[restful] CPU profiling stopped, closing:"+p.cpuprofile)
63}
64
65func main() {} // exists for example compilation only
66