From ce080d76c643c592f4f5ac7e06df62a979a8f2dc Mon Sep 17 00:00:00 2001
From: Asim Aslam <asim@aslam.me>
Date: Wed, 6 Nov 2019 19:36:04 +0000
Subject: [PATCH] add debug/profile package (#920)

* add debug/profile package

* set service+version for profile
---
 config/cmd/cmd.go            |  17 +++--
 debug/profile/pprof/pprof.go | 118 +++++++++++++++++++++++++++++++++++
 debug/profile/profile.go     |  23 +++++++
 service.go                   |  16 +++++
 4 files changed, 168 insertions(+), 6 deletions(-)
 create mode 100644 debug/profile/pprof/pprof.go
 create mode 100644 debug/profile/profile.go

diff --git a/config/cmd/cmd.go b/config/cmd/cmd.go
index e66836c1..18f12528 100644
--- a/config/cmd/cmd.go
+++ b/config/cmd/cmd.go
@@ -71,12 +71,6 @@ var (
 	DefaultCmd = newCmd()
 
 	DefaultFlags = []cli.Flag{
-		cli.StringFlag{
-			Name:   "runtime",
-			Usage:  "Micro runtime",
-			EnvVar: "MICRO_RUNTIME",
-			Value:  "local",
-		},
 		cli.StringFlag{
 			Name:   "client",
 			EnvVar: "MICRO_CLIENT",
@@ -161,6 +155,11 @@ var (
 			EnvVar: "MICRO_BROKER_ADDRESS",
 			Usage:  "Comma-separated list of broker addresses",
 		},
+		cli.StringFlag{
+			Name:   "profile",
+			Usage:  "Debug profiler for cpu and memory stats",
+			EnvVar: "MICRO_DEBUG_PROFILE",
+		},
 		cli.StringFlag{
 			Name:   "registry",
 			EnvVar: "MICRO_REGISTRY",
@@ -171,6 +170,12 @@ var (
 			EnvVar: "MICRO_REGISTRY_ADDRESS",
 			Usage:  "Comma-separated list of registry addresses",
 		},
+		cli.StringFlag{
+			Name:   "runtime",
+			Usage:  "Runtime for building and running services e.g local, kubernetes",
+			EnvVar: "MICRO_RUNTIME",
+			Value:  "local",
+		},
 		cli.StringFlag{
 			Name:   "selector",
 			EnvVar: "MICRO_SELECTOR",
diff --git a/debug/profile/pprof/pprof.go b/debug/profile/pprof/pprof.go
new file mode 100644
index 00000000..b1f126d9
--- /dev/null
+++ b/debug/profile/pprof/pprof.go
@@ -0,0 +1,118 @@
+// Package pprof provides a pprof profiler
+package pprof
+
+import (
+	"os"
+	"path/filepath"
+	"runtime"
+	"runtime/pprof"
+	"sync"
+	"time"
+
+	"github.com/micro/go-micro/debug/profile"
+)
+
+type profiler struct {
+	opts profile.Options
+
+	sync.Mutex
+	running bool
+	exit    chan bool
+
+	// where the cpu profile is written
+	cpuFile *os.File
+	// where the mem profile is written
+	memFile *os.File
+}
+
+func (p *profiler) writeHeap(f *os.File) {
+	defer f.Close()
+
+	t := time.NewTicker(time.Second * 30)
+	defer t.Stop()
+
+	for {
+		select {
+		case <-t.C:
+			runtime.GC()
+			pprof.WriteHeapProfile(f)
+		case <-p.exit:
+			return
+		}
+	}
+}
+
+func (p *profiler) Start() error {
+	p.Lock()
+	defer p.Unlock()
+
+	if p.running {
+		return nil
+	}
+
+	// create exit channel
+	p.exit = make(chan bool)
+
+	cpuFile := filepath.Join("/tmp", "cpu.pprof")
+	memFile := filepath.Join("/tmp", "mem.pprof")
+
+	if len(p.opts.Name) > 0 {
+		cpuFile = filepath.Join("/tmp", p.opts.Name+".cpu.pprof")
+		memFile = filepath.Join("/tmp", p.opts.Name+".mem.pprof")
+	}
+
+	f1, err := os.Create(cpuFile)
+	if err != nil {
+		return err
+	}
+
+	f2, err := os.Create(memFile)
+	if err != nil {
+		return err
+	}
+
+	// start cpu profiling
+	if err := pprof.StartCPUProfile(f1); err != nil {
+		return err
+	}
+
+	// write the heap periodically
+	go p.writeHeap(f2)
+
+	// set cpu file
+	p.cpuFile = f1
+	// set mem file
+	p.memFile = f2
+
+	p.running = true
+
+	return nil
+}
+
+func (p *profiler) Stop() error {
+	p.Lock()
+	defer p.Unlock()
+
+	select {
+	case <-p.exit:
+		return nil
+	default:
+		close(p.exit)
+		pprof.StopCPUProfile()
+		p.cpuFile.Close()
+		p.running = false
+		p.cpuFile = nil
+		p.memFile = nil
+		return nil
+	}
+}
+
+func NewProfile(opts ...profile.Option) profile.Profile {
+	var options profile.Options
+	for _, o := range opts {
+		o(&options)
+	}
+	p := new(profiler)
+	p.opts = options
+	return p
+}
diff --git a/debug/profile/profile.go b/debug/profile/profile.go
new file mode 100644
index 00000000..2dbe1e13
--- /dev/null
+++ b/debug/profile/profile.go
@@ -0,0 +1,23 @@
+// Package profile is for profilers
+package profile
+
+type Profile interface {
+	// Start the profiler
+	Start() error
+	// Stop the profiler
+	Stop() error
+}
+
+type Options struct {
+	// Name to use for the profile
+	Name string
+}
+
+type Option func(o *Options)
+
+// Name of the profile
+func Name(n string) Option {
+	return func(o *Options) {
+		o.Name = n
+	}
+}
diff --git a/service.go b/service.go
index a0f9817c..4073ff8d 100644
--- a/service.go
+++ b/service.go
@@ -10,6 +10,8 @@ import (
 	"github.com/micro/go-micro/client"
 	"github.com/micro/go-micro/config/cmd"
 	"github.com/micro/go-micro/debug/handler"
+	"github.com/micro/go-micro/debug/profile"
+	"github.com/micro/go-micro/debug/profile/pprof"
 	"github.com/micro/go-micro/metadata"
 	"github.com/micro/go-micro/plugin"
 	"github.com/micro/go-micro/server"
@@ -147,6 +149,20 @@ func (s *service) Run() error {
 		),
 	)
 
+	// start the profiler
+	// TODO: set as an option to the service, don't just use pprof
+	if prof := os.Getenv("MICRO_DEBUG_PROFILE"); len(prof) > 0 {
+		service := s.opts.Server.Options().Name
+		version := s.opts.Server.Options().Version
+		profiler := pprof.NewProfile(
+			profile.Name(service + "." + version),
+		)
+		if err := profiler.Start(); err != nil {
+			return err
+		}
+		defer profiler.Stop()
+	}
+
 	if err := s.Start(); err != nil {
 		return err
 	}