From 90d809c6391445d6b2038388d604d221ba55583f Mon Sep 17 00:00:00 2001
From: Dmytro Vovk <vovk.dimon@gmail.com>
Date: Thu, 5 Sep 2024 18:26:36 +0100
Subject: [PATCH] Heap (#88)

* added heap profile

* added support to query all pprofs
---
 api/api_handler.go                    | 51 ++++++++++++++-------------
 internal/erigon_node/erigon_client.go |  2 +-
 internal/erigon_node/profile.go       |  6 ++--
 3 files changed, 30 insertions(+), 29 deletions(-)

diff --git a/api/api_handler.go b/api/api_handler.go
index e0a0415..2c763a4 100644
--- a/api/api_handler.go
+++ b/api/api_handler.go
@@ -2,12 +2,14 @@ package api
 
 import (
 	"bytes"
+	"context"
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
 	"net/http"
 	"path"
 	"strconv"
+	"strings"
 
 	"github.com/go-chi/chi/v5"
 
@@ -247,8 +249,6 @@ func (h *APIHandler) findNodeClient(r *http.Request) (erigon_node.Client, error)
 }
 
 func (h *APIHandler) UniversalRequest(w http.ResponseWriter, r *http.Request) {
-	apiStr := chi.URLParam(r, "*")
-
 	client, err := h.findNodeClient(r)
 
 	if err != nil {
@@ -256,42 +256,44 @@ func (h *APIHandler) UniversalRequest(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	response, err := client.GetResponse(r.Context(), apiStr)
+	pprof, data, err := GetResponseData(r.Context(), client, chi.URLParam(r, "*"))
 
 	if err != nil {
 		http.Error(w, fmt.Sprintf("Unable to fetch snapshot files list: %v", err), http.StatusInternalServerError)
 		return
 	}
 
-	jsonData, err := json.Marshal(response)
-
-	if err != nil {
-		api_internal.EncodeError(w, r, err)
+	if pprof {
+		encoded := base64.StdEncoding.EncodeToString(data)
+		data = bytes.NewBufferString(encoded).Bytes()
+	} else {
+		w.Header().Set("Content-Type", "application/json")
 	}
 
-	w.Header().Set("Content-Type", "application/json")
-	w.Write(jsonData)
+	w.Write(data)
 }
 
-func (h *APIHandler) HeapProfile(w http.ResponseWriter, r *http.Request) {
-	client, err := h.findNodeClient(r)
+func GetResponseData(ctx context.Context, client erigon_node.Client, request string) (bool, []byte, error) {
+	if strings.Contains(request, "pprof") {
+		data, err := client.FindProfile(ctx, request)
+		if err != nil {
+			return true, nil, err
+		}
 
-	if err != nil {
-		http.Error(w, err.Error(), http.StatusBadRequest)
-		return
-	}
+		return true, data, nil
+	} else {
+		response, err := client.GetResponse(ctx, request)
+		if err != nil {
+			return false, nil, err
+		}
 
-	heap, err := client.FindHeapProfile(r.Context())
+		jsonData, err := json.Marshal(response)
+		if err != nil {
+			return false, nil, err
+		}
 
-	if err != nil {
-		http.Error(w, fmt.Sprintf("Unable to fetch sync stage progress: %v", err), http.StatusInternalServerError)
-		return
+		return false, jsonData, nil
 	}
-
-	encoded := base64.StdEncoding.EncodeToString(heap)
-
-	//w.Header().Set("Content-Type", "text/plain")
-	w.Write(bytes.NewBufferString(encoded).Bytes())
 }
 
 func NewAPIHandler(
@@ -314,7 +316,6 @@ 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 d0e2794..1532b90 100644
--- a/internal/erigon_node/erigon_client.go
+++ b/internal/erigon_node/erigon_client.go
@@ -93,7 +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)
+	FindProfile(ctx context.Context, profile string) ([]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
index b33460b..60ecbea 100644
--- a/internal/erigon_node/profile.go
+++ b/internal/erigon_node/profile.go
@@ -12,8 +12,8 @@ type ProfileContent struct {
 	Chunk []byte `json:"chunk"`
 }
 
-func (c *NodeClient) FindHeapProfile(ctx context.Context) ([]byte, error) {
-	request, err := c.fetch(ctx, "heap-profile", nil)
+func (c *NodeClient) FindProfile(ctx context.Context, profile string) ([]byte, error) {
+	request, err := c.fetch(ctx, profile, nil)
 
 	if err != nil {
 		return nil, err
@@ -32,7 +32,7 @@ func (c *NodeClient) FindHeapProfile(ctx context.Context) ([]byte, error) {
 	}
 
 	//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")
+	tempFile, err := os.CreateTemp("", "profile-*.pprof")
 	if err != nil {
 		return nil, err
 	}