From d9be11429b439365393ac3efc38acab8dd343d57 Mon Sep 17 00:00:00 2001
From: Dmytro Vovk <vovk.dimon@gmail.com>
Date: Mon, 19 Aug 2024 09:27:57 +0100
Subject: [PATCH] added heap profile (#87)

---
 api/api_handler.go                    | 24 +++++++++++
 internal/erigon_node/erigon_client.go |  2 +
 internal/erigon_node/profile.go       | 62 +++++++++++++++++++++++++++
 3 files changed, 88 insertions(+)
 create mode 100644 internal/erigon_node/profile.go

diff --git a/api/api_handler.go b/api/api_handler.go
index f4b3e76..e0a0415 100644
--- a/api/api_handler.go
+++ b/api/api_handler.go
@@ -1,6 +1,8 @@
 package api
 
 import (
+	"bytes"
+	"encoding/base64"
 	"encoding/json"
 	"fmt"
 	"net/http"
@@ -271,6 +273,27 @@ func (h *APIHandler) UniversalRequest(w http.ResponseWriter, r *http.Request) {
 	w.Write(jsonData)
 }
 
+func (h *APIHandler) HeapProfile(w http.ResponseWriter, r *http.Request) {
+	client, err := h.findNodeClient(r)
+
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusBadRequest)
+		return
+	}
+
+	heap, err := client.FindHeapProfile(r.Context())
+
+	if err != nil {
+		http.Error(w, fmt.Sprintf("Unable to fetch sync stage progress: %v", err), http.StatusInternalServerError)
+		return
+	}
+
+	encoded := base64.StdEncoding.EncodeToString(heap)
+
+	//w.Header().Set("Content-Type", "text/plain")
+	w.Write(bytes.NewBufferString(encoded).Bytes())
+}
+
 func NewAPIHandler(
 	sessions sessions.CacheService,
 	erigonNode erigon_node.Client,
@@ -291,6 +314,7 @@ func NewAPIHandler(
 	r.Get("/sessions/{sessionId}/nodes/{nodeId}/headers/download-summary", r.HeadersDownload)
 	r.Get("/sessions/{sessionId}/nodes/{nodeId}/sync-stages", r.SyncStages)
 	r.Get("/v2/sessions/{sessionId}/nodes/{nodeId}/*", r.UniversalRequest)
+	r.Get("/sessions/{sessionId}/nodes/{nodeId}/profile/heap", r.HeapProfile)
 
 	return r
 }
diff --git a/internal/erigon_node/erigon_client.go b/internal/erigon_node/erigon_client.go
index a06eafe..d0e2794 100644
--- a/internal/erigon_node/erigon_client.go
+++ b/internal/erigon_node/erigon_client.go
@@ -93,5 +93,7 @@ type Client interface {
 	BodiesDownload(ctx context.Context, w http.ResponseWriter)
 	HeadersDownload(ctx context.Context, w http.ResponseWriter)
 
+	FindHeapProfile(ctx context.Context) ([]byte, error)
+
 	fetch(ctx context.Context, method string, params url.Values) (*NodeRequest, error)
 }
diff --git a/internal/erigon_node/profile.go b/internal/erigon_node/profile.go
new file mode 100644
index 0000000..b33460b
--- /dev/null
+++ b/internal/erigon_node/profile.go
@@ -0,0 +1,62 @@
+package erigon_node
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"os"
+	"os/exec"
+)
+
+type ProfileContent struct {
+	Chunk []byte `json:"chunk"`
+}
+
+func (c *NodeClient) FindHeapProfile(ctx context.Context) ([]byte, error) {
+	request, err := c.fetch(ctx, "heap-profile", nil)
+
+	if err != nil {
+		return nil, err
+	}
+
+	_, result, err := request.nextResult(ctx)
+
+	if err != nil {
+		return nil, err
+	}
+
+	var content ProfileContent
+
+	if err := json.Unmarshal(result, &content); err != nil {
+		return nil, err
+	}
+
+	//result is a file content so I need to save it to file and return the file path
+	tempFile, err := os.CreateTemp("", "heap-profile-*.pprof")
+	if err != nil {
+		return nil, err
+	}
+
+	defer func() {
+		if err := tempFile.Close(); err != nil {
+			fmt.Printf("Error closing temporary file: %v\n", err)
+		}
+
+		// Remove the file after closing it
+		if err := os.Remove(tempFile.Name()); err != nil {
+			fmt.Printf("Error removing temporary file: %v\n", err)
+		}
+	}()
+
+	if _, err := tempFile.Write(content.Chunk); err != nil {
+		return nil, err
+	}
+
+	cmd := exec.Command("go", "tool", "pprof", "-png", tempFile.Name())
+	svgOutput, err := cmd.Output()
+	if err != nil {
+		return nil, err
+	}
+
+	return svgOutput, nil
+}