Add congestion info Rust script

Cargo.toml

@@ -78,6 +78,7 @@ members = [
     "test-utils/testlib",
     "tools/database",
     "tools/chainsync-loadtest",
+    "tools/congestion-metrics",
     "tools/congestion-model",
     "tools/fork-network",
     "tools/indexer/example",

chain/jsonrpc/client/src/lib.rs

@@ -5,6 +5,9 @@ use near_jsonrpc_primitives::message::{from_slice, Message};
 use near_jsonrpc_primitives::types::changes::{
     RpcStateChangesInBlockByTypeRequest, RpcStateChangesInBlockByTypeResponse,
 };
+use near_jsonrpc_primitives::types::congestion::{
+    RpcCongestionLevelRequest, RpcCongestionLevelResponse,
+};
 use near_jsonrpc_primitives::types::transactions::{
     RpcTransactionResponse, RpcTransactionStatusRequest,
 };
@@ -234,6 +237,14 @@ impl JsonRpcClient {
         call_method(&self.client, &self.server_addr, "EXPERIMENTAL_changes", request)
     }
 
+    #[allow(non_snake_case)]
+    pub fn EXPERIMENTAL_congestion_level(
+        &self,
+        request: RpcCongestionLevelRequest,
+    ) -> RpcRequest<RpcCongestionLevelResponse> {
+        call_method(&self.client, &self.server_addr, "EXPERIMENTAL_congestion_level", request)
+    }
+
     #[allow(non_snake_case)]
     pub fn EXPERIMENTAL_validators_ordered(
         &self,

tools/congestion-metrics/Cargo.toml

@@ -0,0 +1,44 @@
+[package]
+name = "congestion-metrics"
+version.workspace = true
+authors.workspace = true
+edition.workspace = true
+rust-version.workspace = true
+repository.workspace = true
+license.workspace = true
+
+[dependencies]
+actix-rt.workspace = true
+actix.workspace = true
+anyhow.workspace = true
+clap.workspace = true
+futures.workspace = true
+pin-project.workspace = true
+rand.workspace = true
+rayon.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+tempfile.workspace = true
+tokio.workspace = true
+tracing.workspace = true
+
+near-actix-test-utils.workspace = true
+near-time.workspace = true
+near-chain.workspace = true
+near-chain-configs.workspace = true
+near-client.workspace = true
+near-crypto.workspace = true
+near-chunks.workspace = true
+near-epoch-manager.workspace = true
+near-jsonrpc.workspace = true
+near-jsonrpc-client.workspace = true
+near-network.workspace = true
+near-store.workspace = true
+near-o11y.workspace = true
+near-telemetry.workspace = true
+near-performance-metrics.workspace = true
+near-primitives.workspace = true
+nearcore.workspace = true
+
+[lints]
+workspace = true

tools/congestion-metrics/src/main.rs

@@ -0,0 +1,46 @@
+use near_jsonrpc::primitives::types::chunks::ChunkReference;
+use near_jsonrpc::primitives::types::congestion::RpcCongestionLevelRequest;
+use near_jsonrpc_client::new_client;
+use tokio::fs::File;
+use tokio::io::AsyncWriteExt;
+
+use near_primitives::types::{BlockId, BlockReference};
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let mut file = File::create("data.csv").await?;
+    file.write_all(b"block_height,shard_id,congestion_level\n").await?;
+
+    let client = new_client("https://archival-rpc.mainnet.near.org");
+    let start_height: u64 = 133453990;
+    let end_height: u64 = 133454000;
+    let mut current_height = end_height;
+
+    while current_height >= start_height {
+        let block =
+            client.block(BlockReference::BlockId(BlockId::Height(current_height))).await.unwrap();
+        for chunk in block.chunks {
+            let congestion_level = client
+                .EXPERIMENTAL_congestion_level(RpcCongestionLevelRequest {
+                    chunk_reference: ChunkReference::BlockShardId {
+                        block_id: BlockId::Height(current_height),
+                        shard_id: chunk.shard_id,
+                    },
+                })
+                .await
+                .unwrap();
+            file.write_all(
+                format!(
+                    "{},{},{}\n",
+                    current_height, chunk.shard_id, congestion_level.congestion_level
+                )
+                .as_bytes(),
+            )
+            .await?;
+        }
+        current_height = block.header.prev_height.unwrap();
+    }
+
+    file.flush().await?;
+    Ok(())
+}