From 6628c6dd8f77935b83316abe4b136983b912e4f8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96mer=20Faruk=20IRMAK?= <omerfirmak@gmail.com>
Date: Thu, 15 Feb 2024 12:56:02 +0300
Subject: [PATCH] Add VM throttler metrics

---
 node/metrics.go      | 16 ++++++++++++++++
 node/node.go         |  1 +
 node/throttled_vm.go | 14 ++++++++------
 utils/throttler.go   |  5 +++++
 4 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/node/metrics.go b/node/metrics.go
index c857ce4459..a10ae75a3b 100644
--- a/node/metrics.go
+++ b/node/metrics.go
@@ -312,3 +312,19 @@ func makeJeMallocMetrics() {
 	})
 	prometheus.MustRegister(active)
 }
+
+func makeVMThrottlerMetrics(throttledVM *ThrottledVM) {
+	vmJobs := prometheus.NewGaugeFunc(prometheus.GaugeOpts{
+		Namespace: "vm",
+		Name:      "jobs",
+	}, func() float64 {
+		return float64(throttledVM.JobsRunning())
+	})
+	vmQueue := prometheus.NewGaugeFunc(prometheus.GaugeOpts{
+		Namespace: "vm",
+		Name:      "queue",
+	}, func() float64 {
+		return float64(throttledVM.QueueLen())
+	})
+	prometheus.MustRegister(vmJobs, vmQueue)
+}
diff --git a/node/node.go b/node/node.go
index a3f33ccbf9..0b8f58b264 100644
--- a/node/node.go
+++ b/node/node.go
@@ -207,6 +207,7 @@ func New(cfg *Config, version string) (*Node, error) { //nolint:gocyclo,funlen
 	var metricsService service.Service
 	if cfg.Metrics {
 		makeJeMallocMetrics()
+		makeVMThrottlerMetrics(throttledVM)
 		makePebbleMetrics(database)
 		chain.WithListener(makeBlockchainMetrics())
 		makeJunoMetrics(version)
diff --git a/node/throttled_vm.go b/node/throttled_vm.go
index aaec8ec72b..c0276aa532 100644
--- a/node/throttled_vm.go
+++ b/node/throttled_vm.go
@@ -9,18 +9,21 @@ import (
 
 var _ vm.VM = (*ThrottledVM)(nil)
 
-type ThrottledVM utils.Throttler[vm.VM]
+type ThrottledVM struct {
+	*utils.Throttler[vm.VM]
+}
 
 func NewThrottledVM(res vm.VM, concurrenyBudget uint, maxQueueLen int32) *ThrottledVM {
-	return (*ThrottledVM)(utils.NewThrottler[vm.VM](concurrenyBudget, &res).WithMaxQueueLen(maxQueueLen))
+	return &ThrottledVM{
+		Throttler: utils.NewThrottler[vm.VM](concurrenyBudget, &res).WithMaxQueueLen(maxQueueLen),
+	}
 }
 
 func (tvm *ThrottledVM) Call(contractAddr, classHash, selector *felt.Felt, calldata []felt.Felt, blockNumber,
 	blockTimestamp uint64, state core.StateReader, network *utils.Network, maxSteps uint64,
 ) ([]*felt.Felt, error) {
 	var ret []*felt.Felt
-	throttler := (*utils.Throttler[vm.VM])(tvm)
-	return ret, throttler.Do(func(vm *vm.VM) error {
+	return ret, tvm.Do(func(vm *vm.VM) error {
 		var err error
 		ret, err = (*vm).Call(contractAddr, classHash, selector, calldata, blockNumber, blockTimestamp,
 			state, network, maxSteps)
@@ -34,8 +37,7 @@ func (tvm *ThrottledVM) Execute(txns []core.Transaction, declaredClasses []core.
 ) ([]*felt.Felt, []vm.TransactionTrace, error) {
 	var ret []*felt.Felt
 	var traces []vm.TransactionTrace
-	throttler := (*utils.Throttler[vm.VM])(tvm)
-	return ret, traces, throttler.Do(func(vm *vm.VM) error {
+	return ret, traces, tvm.Do(func(vm *vm.VM) error {
 		var err error
 		ret, traces, err = (*vm).Execute(txns, declaredClasses, blockNumber, blockTimestamp, sequencerAddress,
 			state, network, paidFeesOnL1, skipChargeFee, skipValidate, errOnRevert, gasPriceWEI, gasPriceSTRK, legacyTraceJSON)
diff --git a/utils/throttler.go b/utils/throttler.go
index f4848f26f0..0d54c24e2c 100644
--- a/utils/throttler.go
+++ b/utils/throttler.go
@@ -49,3 +49,8 @@ func (t *Throttler[T]) Do(doer func(resource *T) error) error {
 func (t *Throttler[T]) QueueLen() int {
 	return int(t.queue.Load())
 }
+
+// JobsRunning returns the number of Do calls that are running at the moment
+func (t *Throttler[T]) JobsRunning() int {
+	return len(t.sem)
+}