From cb80a7aa42fa880ddb5c1d2d769ec0ea53742fe1 Mon Sep 17 00:00:00 2001
From: Ethan Buchman <ethan@coinculture.info>
Date: Fri, 11 Mar 2016 21:45:19 -0500
Subject: [PATCH] rpc: unsafe_write_heap_profile

---
 rpc/core/dev.go             | 19 +++++++++++++++----
 rpc/core/routes.go          |  9 +++++++++
 rpc/core/types/responses.go |  8 +++++---
 3 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/rpc/core/dev.go b/rpc/core/dev.go
index a7a5c54df..2b8dfb8f3 100644
--- a/rpc/core/dev.go
+++ b/rpc/core/dev.go
@@ -36,7 +36,7 @@ func UnsafeSetConfig(typ, key, value string) (*ctypes.ResultUnsafeSetConfig, err
 
 var profFile *os.File
 
-func UnsafeStartCPUProfiler(filename string) (*ctypes.ResultUnsafeCPUProfiler, error) {
+func UnsafeStartCPUProfiler(filename string) (*ctypes.ResultUnsafeProfile, error) {
 	var err error
 	profFile, err = os.Create(filename)
 	if err != nil {
@@ -46,11 +46,22 @@ func UnsafeStartCPUProfiler(filename string) (*ctypes.ResultUnsafeCPUProfiler, e
 	if err != nil {
 		return nil, err
 	}
-	return &ctypes.ResultUnsafeCPUProfiler{}, nil
+	return &ctypes.ResultUnsafeProfile{}, nil
 }
 
-func UnsafeStopCPUProfiler() (*ctypes.ResultUnsafeCPUProfiler, error) {
+func UnsafeStopCPUProfiler() (*ctypes.ResultUnsafeProfile, error) {
 	pprof.StopCPUProfile()
 	profFile.Close()
-	return &ctypes.ResultUnsafeCPUProfiler{}, nil
+	return &ctypes.ResultUnsafeProfile{}, nil
+}
+
+func UnsafeWriteHeapProfile(filename string) (*ctypes.ResultUnsafeProfile, error) {
+	memProfFile, err := os.Create(filename)
+	if err != nil {
+		return nil, err
+	}
+	pprof.WriteHeapProfile(memProfFile)
+	memProfFile.Close()
+
+	return &ctypes.ResultUnsafeProfile{}, nil
 }
diff --git a/rpc/core/routes.go b/rpc/core/routes.go
index b642ad7fe..585a18cd2 100644
--- a/rpc/core/routes.go
+++ b/rpc/core/routes.go
@@ -27,6 +27,7 @@ var Routes = map[string]*rpc.RPCFunc{
 	"unsafe_set_config":         rpc.NewRPCFunc(UnsafeSetConfigResult, "type,key,value"),
 	"unsafe_start_cpu_profiler": rpc.NewRPCFunc(UnsafeStartCPUProfilerResult, "filename"),
 	"unsafe_stop_cpu_profiler":  rpc.NewRPCFunc(UnsafeStopCPUProfilerResult, ""),
+	"unsafe_write_heap_profile": rpc.NewRPCFunc(UnsafeWriteHeapProfileResult, "filename"),
 }
 
 func SubscribeResult(wsCtx rpctypes.WSRPCContext, event string) (ctypes.TMResult, error) {
@@ -164,3 +165,11 @@ func UnsafeStopCPUProfilerResult() (ctypes.TMResult, error) {
 		return r, nil
 	}
 }
+
+func UnsafeWriteHeapProfileResult(filename string) (ctypes.TMResult, error) {
+	if r, err := UnsafeWriteHeapProfile(filename); err != nil {
+		return nil, err
+	} else {
+		return r, nil
+	}
+}
diff --git a/rpc/core/types/responses.go b/rpc/core/types/responses.go
index ee7307a26..d08892e8c 100644
--- a/rpc/core/types/responses.go
+++ b/rpc/core/types/responses.go
@@ -70,7 +70,7 @@ type ResultUnconfirmedTxs struct {
 
 type ResultUnsafeSetConfig struct{}
 
-type ResultUnsafeCPUProfiler struct{}
+type ResultUnsafeProfile struct{}
 
 type ResultSubscribe struct {
 }
@@ -114,6 +114,7 @@ const (
 	ResultTypeUnsafeSetConfig        = byte(0xa0)
 	ResultTypeUnsafeStartCPUProfiler = byte(0xa1)
 	ResultTypeUnsafeStopCPUProfiler  = byte(0xa2)
+	ResultTypeUnsafeWriteHeapProfile = byte(0xa3)
 )
 
 type TMResult interface {
@@ -137,6 +138,7 @@ var _ = wire.RegisterInterface(
 	wire.ConcreteType{&ResultUnsubscribe{}, ResultTypeUnsubscribe},
 	wire.ConcreteType{&ResultEvent{}, ResultTypeEvent},
 	wire.ConcreteType{&ResultUnsafeSetConfig{}, ResultTypeUnsafeSetConfig},
-	wire.ConcreteType{&ResultUnsafeCPUProfiler{}, ResultTypeUnsafeStartCPUProfiler},
-	wire.ConcreteType{&ResultUnsafeCPUProfiler{}, ResultTypeUnsafeStopCPUProfiler},
+	wire.ConcreteType{&ResultUnsafeProfile{}, ResultTypeUnsafeStartCPUProfiler},
+	wire.ConcreteType{&ResultUnsafeProfile{}, ResultTypeUnsafeStopCPUProfiler},
+	wire.ConcreteType{&ResultUnsafeProfile{}, ResultTypeUnsafeWriteHeapProfile},
 )