From 6a0d342d71e65db73d1a9ef13dfdeb0c90e49a6a Mon Sep 17 00:00:00 2001
From: Calvin Neo <CalvinNeo@users.noreply.github.com>
Date: Wed, 28 Aug 2024 11:31:47 +0800
Subject: [PATCH] Support heap profiling for TiFlash (#1718)

---
 pkg/apiserver/profiling/fetcher.go            | 31 +++++++++++++++++++
 pkg/apiserver/profiling/pprof.go              |  2 +-
 pkg/apiserver/profiling/profile.go            |  4 +--
 pkg/tiflash/client.go                         |  4 +++
 .../apps/ContinuousProfiling/pages/Detail.tsx | 13 ++++++++
 5 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/pkg/apiserver/profiling/fetcher.go b/pkg/apiserver/profiling/fetcher.go
index 62d32adc92..c5341dc458 100644
--- a/pkg/apiserver/profiling/fetcher.go
+++ b/pkg/apiserver/profiling/fetcher.go
@@ -146,6 +146,37 @@ type tiflashFetcher struct {
 }
 
 func (f *tiflashFetcher) fetch(op *fetchOptions) ([]byte, error) {
+	if strings.HasSuffix(op.path, "heap") {
+		scheme := f.client.GetHTTPScheme()
+		cmd := exec.Command("perl", "/dev/stdin", "--raw", scheme+"://"+op.ip+":"+strconv.Itoa(op.port)+op.path) //nolint:gosec
+		cmd.Stdin = strings.NewReader(jeprof)
+		stdout, err := cmd.StdoutPipe()
+		if err != nil {
+			return nil, err
+		}
+		stderr, err := cmd.StderrPipe()
+		if err != nil {
+			return nil, err
+		}
+		// use jeprof to fetch tiflash heap profile
+		err = cmd.Start()
+		if err != nil {
+			return nil, err
+		}
+		data, err := io.ReadAll(stdout)
+		if err != nil {
+			return nil, err
+		}
+		errMsg, err := io.ReadAll(stderr)
+		if err != nil {
+			return nil, err
+		}
+		err = cmd.Wait()
+		if err != nil {
+			return nil, fmt.Errorf("failed to fetch tiflash heap profile: %s", errMsg)
+		}
+		return data, nil
+	}
 	return f.client.WithTimeout(maxProfilingTimeout).AddRequestHeader("Content-Type", "application/protobuf").SendGetRequest(op.ip, op.port, op.path)
 }
 
diff --git a/pkg/apiserver/profiling/pprof.go b/pkg/apiserver/profiling/pprof.go
index 586631248f..f5416008ac 100644
--- a/pkg/apiserver/profiling/pprof.go
+++ b/pkg/apiserver/profiling/pprof.go
@@ -46,7 +46,7 @@ func (f *fetcher) FetchAndWriteToFile(duration uint, fileNameWithoutExt string,
 		fileExtenstion = "*.proto"
 	case ProfilingTypeHeap:
 		url = "/debug/pprof/heap"
-		if f.target.Kind == model.NodeKindTiKV {
+		if f.target.Kind == model.NodeKindTiKV || f.target.Kind == model.NodeKindTiFlash {
 			profilingRawDataType = RawDataTypeJeprof
 			fileExtenstion = "*.prof"
 		} else {
diff --git a/pkg/apiserver/profiling/profile.go b/pkg/apiserver/profiling/profile.go
index 4c87d1141d..bb70ab4d52 100644
--- a/pkg/apiserver/profiling/profile.go
+++ b/pkg/apiserver/profiling/profile.go
@@ -17,8 +17,8 @@ func profileAndWritePprof(_ context.Context, fts *fetchers, target *model.Reques
 		}
 		return fetchPprof(&pprofOptions{duration: profileDurationSecs, fileNameWithoutExt: fileNameWithoutExt, target: target, fetcher: &fts.tikv, profilingType: profilingType})
 	case model.NodeKindTiFlash:
-		// TiFlash only supports CPU Profiling
-		if profilingType != ProfilingTypeCPU {
+		// TiFlash only supports CPU/heap Profiling
+		if profilingType != ProfilingTypeCPU && profilingType != ProfilingTypeHeap {
 			return "", "", ErrUnsupportedProfilingType.NewWithNoMessage()
 		}
 		return fetchPprof(&pprofOptions{duration: profileDurationSecs, fileNameWithoutExt: fileNameWithoutExt, target: target, fetcher: &fts.tiflash, profilingType: profilingType})
diff --git a/pkg/tiflash/client.go b/pkg/tiflash/client.go
index fd9bc7c260..6c0b4b6e7c 100644
--- a/pkg/tiflash/client.go
+++ b/pkg/tiflash/client.go
@@ -49,6 +49,10 @@ func NewTiFlashClient(lc fx.Lifecycle, httpClient *httpc.Client, config *config.
 	return client
 }
 
+func (c Client) GetHTTPScheme() string {
+	return c.httpScheme
+}
+
 func (c Client) WithTimeout(timeout time.Duration) *Client {
 	c.timeout = timeout
 	return &c
diff --git a/ui/packages/tidb-dashboard-lib/src/apps/ContinuousProfiling/pages/Detail.tsx b/ui/packages/tidb-dashboard-lib/src/apps/ContinuousProfiling/pages/Detail.tsx
index 20a70800f6..61f8c8187b 100644
--- a/ui/packages/tidb-dashboard-lib/src/apps/ContinuousProfiling/pages/Detail.tsx
+++ b/ui/packages/tidb-dashboard-lib/src/apps/ContinuousProfiling/pages/Detail.tsx
@@ -106,6 +106,19 @@ export default function Page() {
             break
           default:
         }
+      } else if (component === 'tiflash' && profile_type === 'heap') {
+        switch (action) {
+          case Action.VIEW_FLAMEGRAPH:
+            dataFormat = 'text'
+            break
+          case Action.VIEW_GRAPH:
+            dataFormat = 'svg'
+            break
+          case Action.DOWNLOAD:
+            dataFormat = 'jeprof'
+            break
+          default:
+        }
       } else {
         switch (action) {
           case Action.VIEW_GRAPH: