From 5fa9c3fb4fe6fac557a3b9f6de33925da2193f45 Mon Sep 17 00:00:00 2001
From: incrypto32 <vpkrishnanand@gmail.com>
Date: Tue, 7 Nov 2023 10:56:19 +0530
Subject: [PATCH] graph: refactor GasCounter for trackiing gas metrics, add a
 new env for gas metrics

---
 core/src/subgraph/instance_manager.rs |  2 +-
 graph/src/components/metrics/gas.rs   | 12 ++++++------
 graph/src/env/mod.rs                  |  7 +++++++
 graph/src/runtime/gas/mod.rs          | 25 +++++++++++++++----------
 4 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/core/src/subgraph/instance_manager.rs b/core/src/subgraph/instance_manager.rs
index 5ef405409a9..a36eb3a3c71 100644
--- a/core/src/subgraph/instance_manager.rs
+++ b/core/src/subgraph/instance_manager.rs
@@ -6,7 +6,7 @@ use std::collections::BTreeSet;
 
 use crate::subgraph::runner::SubgraphRunner;
 use graph::blockchain::block_stream::BlockStreamMetrics;
-use graph::blockchain::{Blockchain, BlockchainKind, DataSource, NodeCapabilities,TriggerFilter};
+use graph::blockchain::{Blockchain, BlockchainKind, DataSource, NodeCapabilities};
 use graph::components::metrics::gas::GasMetrics;
 use graph::components::subgraph::ProofOfIndexingVersion;
 use graph::data::subgraph::{UnresolvedSubgraphManifest, SPEC_VERSION_0_0_6};
diff --git a/graph/src/components/metrics/gas.rs b/graph/src/components/metrics/gas.rs
index 30b537322fb..94bc8550949 100644
--- a/graph/src/components/metrics/gas.rs
+++ b/graph/src/components/metrics/gas.rs
@@ -18,10 +18,10 @@ impl GasMetrics {
                 subgraph_id.as_str(),
                 &["method"],
             )
-            .unwrap_or_else(|_| {
+            .unwrap_or_else(|err| {
                 panic!(
-                    "Failed to register deployment_gas prometheus counter for {}",
-                    subgraph_id
+                    "Failed to register deployment_gas prometheus counter for {}: {}",
+                    subgraph_id, err
                 )
             });
 
@@ -32,10 +32,10 @@ impl GasMetrics {
                 subgraph_id.as_str(),
                 &["method"],
             )
-            .unwrap_or_else(|_| {
+            .unwrap_or_else(|err| {
                 panic!(
-                    "Failed to register deployment_op_count prometheus counter for {}",
-                    subgraph_id
+                    "Failed to register deployment_op_count prometheus counter for {}: {}",
+                    subgraph_id, err
                 )
             });
 
diff --git a/graph/src/env/mod.rs b/graph/src/env/mod.rs
index 1522998d6e9..eb72930e64d 100644
--- a/graph/src/env/mod.rs
+++ b/graph/src/env/mod.rs
@@ -177,6 +177,9 @@ pub struct EnvVars {
     pub subgraph_settings: Option<String>,
     /// Whether to prefer substreams blocks streams over firehose when available.
     pub prefer_substreams_block_streams: bool,
+    /// Set by the flag `GRAPH_ENABLE_GAS_METRICS`. Whether to enable
+    /// gas metrics. Off by default.
+    pub enable_gas_metrics: bool,
 }
 
 impl EnvVars {
@@ -236,6 +239,7 @@ impl EnvVars {
             reorg_threshold: inner.reorg_threshold,
             subgraph_settings: inner.subgraph_settings,
             prefer_substreams_block_streams: inner.prefer_substreams_block_streams,
+            enable_gas_metrics: inner.enable_gas_metrics.0,
         })
     }
 
@@ -361,6 +365,9 @@ struct Inner {
         default = "false"
     )]
     prefer_substreams_block_streams: bool,
+
+    #[envconfig(from = "GRAPH_ENABLE_GAS_METRICS", default = "false")]
+    enable_gas_metrics: EnvVarBoolean,
 }
 
 #[derive(Clone, Debug)]
diff --git a/graph/src/runtime/gas/mod.rs b/graph/src/runtime/gas/mod.rs
index a090192b4c6..e137207d274 100644
--- a/graph/src/runtime/gas/mod.rs
+++ b/graph/src/runtime/gas/mod.rs
@@ -77,16 +77,18 @@ impl Display for Gas {
 }
 
 #[derive(Clone)]
-pub struct GasCounter(Arc<AtomicU64>, GasMetrics);
+pub struct GasCounter {
+    counter: Arc<AtomicU64>,
+    metrics: GasMetrics,
+}
 
 impl CheapClone for GasCounter {}
 
 impl GasCounter {
-    /// Alias of [`Default::default`].
-    pub fn new(gas_metrics: GasMetrics) -> Self {
+    pub fn new(metrics: GasMetrics) -> Self {
         Self {
-            0: Arc::new(AtomicU64::new(0)),
-            1: gas_metrics,
+            counter: Arc::new(AtomicU64::new(0)),
+            metrics,
         }
     }
 
@@ -98,13 +100,16 @@ impl GasCounter {
     ) -> Result<(), DeterministicHostError> {
         amount += costs::HOST_EXPORT_GAS;
 
-        if let Some(method) = method {
-            self.1.track_gas(method, amount.0);
-            self.1.track_operations(method, 1);
+        // If gas metrics are enabled, track the gas used
+        if ENV_VARS.enable_gas_metrics {
+            if let Some(method) = method {
+                self.metrics.track_gas(method, amount.0);
+                self.metrics.track_operations(method, 1);
+            }
         }
 
         let old = self
-            .0
+            .counter
             .fetch_update(SeqCst, SeqCst, |v| Some(v.saturating_add(amount.0)))
             .unwrap();
         let new = old.saturating_add(amount.0);
@@ -131,6 +136,6 @@ impl GasCounter {
     }
 
     pub fn get(&self) -> Gas {
-        Gas(self.0.load(SeqCst))
+        Gas(self.counter.load(SeqCst))
     }
 }