From ccc75ac570d4dc52a5de9c262842f137ffae6c36 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96mer=20Faruk=20IRMAK?= <omerfirmak@gmail.com>
Date: Fri, 9 Feb 2024 15:23:13 +0300
Subject: [PATCH] Switch to using jemalloc

---
 Dockerfile           |  4 ++--
 cmd/juno/juno.go     |  1 +
 jemalloc/jemalloc.go | 28 ++++++++++++++++++++++++++++
 node/metrics.go      | 11 +++++++++++
 node/node.go         |  1 +
 5 files changed, 43 insertions(+), 2 deletions(-)
 create mode 100644 jemalloc/jemalloc.go

diff --git a/Dockerfile b/Dockerfile
index 5f9ecad26d..03940495de 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -5,7 +5,7 @@ ARG VM_DEBUG
 
 # Install Alpine Dependencies
 RUN apt-get update && \
-    apt-get install build-essential cargo git golang upx-ucl -y
+    apt-get install build-essential cargo git golang upx-ucl libjemalloc-dev libjemalloc2 -y
 
 WORKDIR /app
 
@@ -21,7 +21,7 @@ RUN upx-ucl /app/build/juno
 # Stage 2: Build Docker image
 FROM ubuntu:23.10 AS runtime
 
-RUN apt-get update && apt-get install -y ca-certificates curl gawk grep
+RUN apt-get update && apt-get install -y ca-certificates curl gawk grep libjemalloc-dev libjemalloc2
 
 COPY --from=build /app/build/juno /usr/local/bin/
 
diff --git a/cmd/juno/juno.go b/cmd/juno/juno.go
index 0279eb1e50..af1c902b8f 100644
--- a/cmd/juno/juno.go
+++ b/cmd/juno/juno.go
@@ -13,6 +13,7 @@ import (
 	"syscall"
 	"time"
 
+	_ "github.com/NethermindEth/juno/jemalloc"
 	"github.com/NethermindEth/juno/node"
 	"github.com/NethermindEth/juno/utils"
 	"github.com/ethereum/go-ethereum/common"
diff --git a/jemalloc/jemalloc.go b/jemalloc/jemalloc.go
new file mode 100644
index 0000000000..3d3580160e
--- /dev/null
+++ b/jemalloc/jemalloc.go
@@ -0,0 +1,28 @@
+package jemalloc
+
+/*
+// This cgo directive is what actually causes jemalloc to be linked in to the
+// final Go executable
+#cgo pkg-config: jemalloc
+
+#include <jemalloc/jemalloc.h>
+
+void _refresh_jemalloc_stats() {
+	// You just need to pass something not-null into the "epoch" mallctl.
+	size_t random_something = 1;
+	mallctl("epoch", NULL, NULL, &random_something, sizeof(random_something));
+}
+unsigned long long _get_jemalloc_active() {
+	size_t stat, stat_size;
+	stat = 0;
+	stat_size = sizeof(stat);
+	mallctl("stats.active", &stat, &stat_size, NULL, 0);
+	return (unsigned long long)stat;
+}
+*/
+import "C"
+
+func GetActive() C.ulonglong {
+	C._refresh_jemalloc_stats()
+	return C._get_jemalloc_active()
+}
diff --git a/node/metrics.go b/node/metrics.go
index 3849df7774..c857ce4459 100644
--- a/node/metrics.go
+++ b/node/metrics.go
@@ -10,6 +10,7 @@ import (
 	"github.com/NethermindEth/juno/clients/gateway"
 	"github.com/NethermindEth/juno/core"
 	"github.com/NethermindEth/juno/db"
+	"github.com/NethermindEth/juno/jemalloc"
 	"github.com/NethermindEth/juno/jsonrpc"
 	"github.com/NethermindEth/juno/l1"
 	"github.com/NethermindEth/juno/sync"
@@ -301,3 +302,13 @@ func makePebbleMetrics(nodeDB db.DB) {
 	})
 	prometheus.MustRegister(blockCacheSize, blockHitRate, tableCacheSize, tableHitRate)
 }
+
+func makeJeMallocMetrics() {
+	active := prometheus.NewGaugeFunc(prometheus.GaugeOpts{
+		Namespace: "jemalloc",
+		Name:      "active",
+	}, func() float64 {
+		return float64(jemalloc.GetActive())
+	})
+	prometheus.MustRegister(active)
+}
diff --git a/node/node.go b/node/node.go
index cdabba7a9d..ea51455234 100644
--- a/node/node.go
+++ b/node/node.go
@@ -176,6 +176,7 @@ func New(cfg *Config, version string) (*Node, error) { //nolint:gocyclo,funlen
 	}
 	var metricsService service.Service
 	if cfg.Metrics {
+		makeJeMallocMetrics()
 		makePebbleMetrics(database)
 		chain.WithListener(makeBlockchainMetrics())
 		makeJunoMetrics(version)