Merge pull request #47 from blockworks-foundation/block_store

Unified Block store, more metrics and websocket auto fix

Cargo.toml

@@ -13,20 +13,24 @@ members = [
 bench = { path = "./bench" }
 
 [dependencies]
-solana-client = "1.14.13"
+solana-sdk = { git = "https://github.com/blockworks-foundation/solana", branch="lite_rpc" }
-solana-sdk = "1.14.13"
+solana-rpc-client = { git = "https://github.com/blockworks-foundation/solana", branch="lite_rpc" }
-solana-transaction-status = "1.14.13"
+solana-rpc-client-api= { git = "https://github.com/blockworks-foundation/solana", branch="lite_rpc" }
-solana-version = "1.14.13"
+solana-tpu-client= { git = "https://github.com/blockworks-foundation/solana", branch="lite_rpc" }
+solana-quic-client= { git = "https://github.com/blockworks-foundation/solana", branch="lite_rpc" }
+solana-pubsub-client= { git = "https://github.com/blockworks-foundation/solana", branch="lite_rpc" }
+solana-transaction-status = { git = "https://github.com/blockworks-foundation/solana", branch="lite_rpc" }
+solana-version= { git = "https://github.com/blockworks-foundation/solana", branch="lite_rpc" }
 serde = { version = "1.0.152", features = ["derive"] }
-serde_json = "1.0.91"
+serde_json = "1.0.92"
 tokio = { version = "1.25.0", features = ["full"]} 
 bincode = "1.3.3"
 bs58 = "0.4.0"
 base64 = "0.21.0"
 thiserror = "1.0.38"
 futures = "0.3.26"
-bytes = "1.3.0"
+bytes = "1.4.0"
-anyhow = "1.0.68"
+anyhow = "1.0.69"
 log = "0.4.17"
 clap = { version = "4.1.4", features = ["derive"] }
 dashmap = "5.4.0"
@@ -36,4 +40,5 @@ tracing-subscriber = "0.3.16"
 tokio-postgres = "0.7.7"
 native-tls = "0.2.11"
 postgres-native-tls = "0.5.0"
-serde_prometheus = "0.1.6"
+prometheus = "0.13.3"
+lazy_static = "1.4.0"

bench/Cargo.toml

@@ -4,12 +4,12 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
-solana-client = "1.14.13"
+solana-sdk = { git = "https://github.com/blockworks-foundation/solana", branch="lite_rpc" }
-solana-sdk = "1.14.13"
+solana-rpc-client = { git = "https://github.com/blockworks-foundation/solana", branch="lite_rpc" }
 log = "0.4.17"
-anyhow = "1.0.68"
+anyhow = "1.0.69"
 serde = "1.0.152"
-serde_json = "1.0.91"
+serde_json = "1.0.92"
 csv = "1.1.6"
 clap = { version = "4.1.4", features = ["derive"] }
 tokio = { version = "1.25.0", features = ["full", "fs"]} 

bench/src/helpers.rs

@@ -1,7 +1,7 @@
 use std::{ops::Deref, sync::Arc};
 
 use anyhow::Context;
-use solana_client::nonblocking::rpc_client::RpcClient;
+use solana_rpc_client::nonblocking::rpc_client::RpcClient;
 use solana_sdk::{
     commitment_config::CommitmentConfig,
     hash::Hash,

bench/src/main.rs

@@ -11,7 +11,7 @@ use bench::{
 };
 use clap::Parser;
 use log::info;
-use solana_client::{nonblocking::rpc_client::RpcClient, rpc_client::SerializableTransaction};
+use solana_rpc_client::{nonblocking::rpc_client::RpcClient, rpc_client::SerializableTransaction};
 use solana_sdk::{commitment_config::CommitmentConfig, signature::Signature};
 
 #[tokio::main]

src/block_store.rs

@@ -0,0 +1,101 @@
+use std::sync::Arc;
+
+use dashmap::DashMap;
+
+use solana_rpc_client::nonblocking::rpc_client::RpcClient;
+use solana_sdk::commitment_config::CommitmentConfig;
+use tokio::sync::RwLock;
+
+use crate::workers::BlockInformation;
+
+#[derive(Clone)]
+pub struct BlockStore {
+    blocks: Arc<DashMap<String, BlockInformation>>,
+    latest_confirmed_blockhash: Arc<RwLock<String>>,
+    latest_finalized_blockhash: Arc<RwLock<String>>,
+}
+
+impl BlockStore {
+    pub async fn new(rpc_client: &RpcClient) -> anyhow::Result<Self> {
+        let (confirmed_blockhash, confirmed_block) =
+            Self::fetch_latest(rpc_client, CommitmentConfig::confirmed()).await?;
+        let (finalized_blockhash, finalized_block) =
+            Self::fetch_latest(rpc_client, CommitmentConfig::finalized()).await?;
+
+        Ok(Self {
+            latest_confirmed_blockhash: Arc::new(RwLock::new(confirmed_blockhash.clone())),
+            latest_finalized_blockhash: Arc::new(RwLock::new(finalized_blockhash.clone())),
+            blocks: Arc::new({
+                let map = DashMap::new();
+                map.insert(confirmed_blockhash, confirmed_block);
+                map.insert(finalized_blockhash, finalized_block);
+                map
+            }),
+        })
+    }
+
+    pub async fn fetch_latest(
+        rpc_client: &RpcClient,
+        commitment_config: CommitmentConfig,
+    ) -> anyhow::Result<(String, BlockInformation)> {
+        let (latest_block_hash, block_height) = rpc_client
+            .get_latest_blockhash_with_commitment(commitment_config)
+            .await?;
+
+        let latest_block_hash = latest_block_hash.to_string();
+        let slot = rpc_client
+            .get_slot_with_commitment(commitment_config)
+            .await?;
+
+        Ok((latest_block_hash, BlockInformation { slot, block_height }))
+    }
+
+    pub async fn get_block_info(&self, blockhash: &str) -> Option<BlockInformation> {
+        let Some(info) = self.blocks.get(blockhash) else {
+            return None;
+        };
+
+        Some(info.value().to_owned())
+    }
+
+    pub fn get_latest_blockhash(&self, commitment_config: CommitmentConfig) -> Arc<RwLock<String>> {
+        if commitment_config.is_finalized() {
+            self.latest_finalized_blockhash.clone()
+        } else {
+            self.latest_confirmed_blockhash.clone()
+        }
+    }
+
+    pub async fn get_latest_block_info(
+        &self,
+        commitment_config: CommitmentConfig,
+    ) -> (String, BlockInformation) {
+        let blockhash = self
+            .get_latest_blockhash(commitment_config)
+            .read()
+            .await
+            .to_owned();
+
+        let block_info = self
+            .blocks
+            .get(&blockhash)
+            .expect("Race Condition: Latest block not in block store")
+            .value()
+            .to_owned();
+
+        (blockhash, block_info)
+    }
+
+    pub async fn add_block(
+        &self,
+        blockhash: String,
+        block_info: BlockInformation,
+        commitment_config: CommitmentConfig,
+    ) {
+        // Write to block store first in order to prevent
+        // any race condition i.e prevent some one to
+        // ask the map what it doesn't have rn
+        self.blocks.insert(blockhash.clone(), block_info);
+        *self.get_latest_blockhash(commitment_config).write().await = blockhash;
+    }
+}

src/bridge.rs

@@ -1,4 +1,5 @@
 use crate::{
+    block_store::BlockStore,
     configs::{IsBlockHashValidConfig, SendTransactionConfig},
     encoding::BinaryEncoding,
     rpc::LiteRpcServer,
@@ -16,16 +17,15 @@ use anyhow::bail;
 use log::info;
 
 use jsonrpsee::{server::ServerBuilder, types::SubscriptionResult, SubscriptionSink};
-use solana_client::{
+
-    nonblocking::{pubsub_client::PubsubClient, rpc_client::RpcClient},
+use prometheus::{opts, register_counter, Counter};
-    rpc_client::SerializableTransaction,
+use solana_rpc_client::{nonblocking::rpc_client::RpcClient, rpc_client::SerializableTransaction};
-    rpc_config::{RpcContextConfig, RpcRequestAirdropConfig},
+use solana_rpc_client_api::{
-    rpc_response::{Response as RpcResponse, RpcBlockhash, RpcResponseContext, RpcVersionInfo},
+    config::{RpcContextConfig, RpcRequestAirdropConfig, RpcSignatureStatusConfig},
+    response::{Response as RpcResponse, RpcBlockhash, RpcResponseContext, RpcVersionInfo},
 };
 use solana_sdk::{
-    commitment_config::{CommitmentConfig, CommitmentLevel},
+    commitment_config::CommitmentConfig, hash::Hash, pubkey::Pubkey,
-    hash::Hash,
-    pubkey::Pubkey,
     transaction::VersionedTransaction,
 };
 use solana_transaction_status::TransactionStatus;
@@ -35,6 +35,24 @@ use tokio::{
     task::JoinHandle,
 };
 
+lazy_static::lazy_static! {
+    static ref RPC_SEND_TX: Counter =
+        register_counter!(opts!("rpc_send_tx", "RPC call send transaction")).unwrap();
+    static ref RPC_GET_LATEST_BLOCKHASH: Counter =
+        register_counter!(opts!("rpc_get_latest_blockhash", "RPC call to get latest block hash")).unwrap();
+    static ref RPC_IS_BLOCKHASH_VALID: Counter =
+        register_counter!(opts!("rpc_is_blockhash_valid", "RPC call to check if blockhash is vali calld")).unwrap();
+    static ref RPC_GET_SIGNATURE_STATUSES: Counter =
+        register_counter!(opts!("rpc_get_signature_statuses", "RPC call to get signature statuses")).unwrap();
+    static ref RPC_GET_VERSION: Counter =
+        register_counter!(opts!("rpc_get_version", "RPC call to version")).unwrap();
+    static ref RPC_REQUEST_AIRDROP: Counter =
+        register_counter!(opts!("rpc_airdrop", "RPC call to request airdrop")).unwrap();
+    static ref RPC_SIGNATURE_SUBSCRIBE: Counter =
+        register_counter!(opts!("rpc_signature_subscribe", "RPC call to subscribe to signature")).unwrap();
+
+}
+
 /// A bridge between clients and tpu
 pub struct LiteBridge {
     pub rpc_client: Arc<RpcClient>,
@@ -42,56 +60,35 @@ pub struct LiteBridge {
     // None if LiteBridge is not executed
     pub tx_send: Option<UnboundedSender<(String, WireTransaction, u64)>>,
     pub tx_sender: TxSender,
-    pub finalized_block_listener: BlockListener,
+    pub block_listner: BlockListener,
-    pub confirmed_block_listener: BlockListener,
+    pub block_store: BlockStore,
 }
 
 impl LiteBridge {
     pub async fn new(rpc_url: String, ws_addr: String, fanout_slots: u64) -> anyhow::Result<Self> {
         let rpc_client = Arc::new(RpcClient::new(rpc_url.clone()));
 
-        let pub_sub_client = Arc::new(PubsubClient::new(&ws_addr).await?);
-
         let tpu_manager =
             Arc::new(TpuManager::new(rpc_client.clone(), ws_addr, fanout_slots).await?);
 
         let tx_sender = TxSender::new(tpu_manager.clone());
 
-        let finalized_block_listener = BlockListener::new(
+        let block_store = BlockStore::new(&rpc_client).await?;
-            pub_sub_client.clone(),
-            rpc_client.clone(),
-            tx_sender.clone(),
-            CommitmentConfig::finalized(),
-        )
-        .await?;
 
-        let confirmed_block_listener = BlockListener::new(
+        let block_listner = BlockListener::new(tx_sender.clone(), block_store.clone());
-            pub_sub_client,
-            rpc_client.clone(),
-            tx_sender.clone(),
-            CommitmentConfig::confirmed(),
-        )
-        .await?;
 
         Ok(Self {
             rpc_client,
             tpu_manager,
             tx_send: None,
             tx_sender,
-            finalized_block_listener,
+            block_listner,
-            confirmed_block_listener,
+            block_store,
         })
     }
 
-    pub fn get_block_listner(&self, commitment_config: CommitmentConfig) -> BlockListener {
-        if let CommitmentLevel::Finalized = commitment_config.commitment {
-            self.finalized_block_listener.clone()
-        } else {
-            self.confirmed_block_listener.clone()
-        }
-    }
-
     /// List for `JsonRpc` requests
+    #[allow(clippy::too_many_arguments)]
     pub async fn start_services<T: ToSocketAddrs + std::fmt::Debug + 'static + Send + Clone>(
         mut self,
         http_addr: T,
@@ -100,6 +97,7 @@ impl LiteBridge {
         tx_send_interval: Duration,
         clean_interval: Duration,
         enable_postgres: bool,
+        prometheus_addr: T,
     ) -> anyhow::Result<Vec<JoinHandle<anyhow::Result<()>>>> {
         let (postgres, postgres_send) = if enable_postgres {
             let (postgres_send, postgres_recv) = mpsc::unbounded_channel();
@@ -122,24 +120,21 @@ impl LiteBridge {
             postgres_send.clone(),
         );
 
-        let metrics_capture = MetricsCapture::new(self.tx_sender.clone());
+        let metrics_capture = MetricsCapture::new(self.tx_sender.clone()).capture();
-        let prometheus_sync = PrometheusSync::new(metrics_capture.clone()).sync();
+        let prometheus_sync = PrometheusSync.sync(prometheus_addr);
-        let metrics_capture = metrics_capture.capture();
 
         let finalized_block_listener = self
-            .finalized_block_listener
+            .block_listner
             .clone()
-            .listen(postgres_send.clone());
+            .listen(CommitmentConfig::finalized(), postgres_send.clone());
 
-        let confirmed_block_listener = self.confirmed_block_listener.clone().listen(None);
+        let confirmed_block_listener = self
-        let cleaner = Cleaner::new(
+            .block_listner
-            self.tx_sender.clone(),
+            .clone()
-            [
+            .listen(CommitmentConfig::confirmed(), None);
-                self.finalized_block_listener.clone(),
+
-                self.confirmed_block_listener.clone(),
+        let cleaner =
-            ],
+            Cleaner::new(self.tx_sender.clone(), self.block_listner.clone()).start(clean_interval);
-        )
-        .start(clean_interval);
 
         let rpc = self.into_rpc();
 
@@ -198,6 +193,8 @@ impl LiteRpcServer for LiteBridge {
         tx: String,
         send_transaction_config: Option<SendTransactionConfig>,
     ) -> crate::rpc::Result<String> {
+        RPC_SEND_TX.inc();
+
         let SendTransactionConfig {
             encoding,
             max_retries: _,
@@ -219,11 +216,11 @@ impl LiteRpcServer for LiteBridge {
 
         let sig = tx.get_signature();
         let Some(BlockInformation { slot, .. }) = self
-            .confirmed_block_listener
+            .block_store
             .get_block_info(&tx.get_recent_blockhash().to_string())
             .await else {
                 log::warn!("block");
-                return Err(jsonrpsee::core::Error::Custom("Blockhash not found in confirmed block store".to_string()));
+                return Err(jsonrpsee::core::Error::Custom("Blockhash not found in block store".to_string()));
         };
 
         self.tx_send
@@ -237,17 +234,18 @@ impl LiteRpcServer for LiteBridge {
 
     async fn get_latest_blockhash(
         &self,
-        config: Option<solana_client::rpc_config::RpcContextConfig>,
+        config: Option<RpcContextConfig>,
-    ) -> crate::rpc::Result<RpcResponse<solana_client::rpc_response::RpcBlockhash>> {
+    ) -> crate::rpc::Result<RpcResponse<RpcBlockhash>> {
-        let commitment_config = if let Some(RpcContextConfig { commitment, .. }) = config {
+        RPC_GET_LATEST_BLOCKHASH.inc();
-            commitment.unwrap_or_default()
-        } else {
-            CommitmentConfig::default()
-        };
 
-        let block_listner = self.get_block_listner(commitment_config);
+        let commitment_config = config
-        let (blockhash, BlockInformation { slot, block_height }) =
+            .map(|config| config.commitment.unwrap_or_default())
-            block_listner.get_latest_block_info().await;
+            .unwrap_or_default();
+
+        let (blockhash, BlockInformation { slot, block_height }) = self
+            .block_store
+            .get_latest_block_info(commitment_config)
+            .await;
 
         Ok(RpcResponse {
             context: RpcResponseContext {
@@ -266,6 +264,8 @@ impl LiteRpcServer for LiteBridge {
         blockhash: String,
         config: Option<IsBlockHashValidConfig>,
     ) -> crate::rpc::Result<RpcResponse<bool>> {
+        RPC_IS_BLOCKHASH_VALID.inc();
+
         let commitment = config.unwrap_or_default().commitment.unwrap_or_default();
         let commitment = CommitmentConfig { commitment };
 
@@ -276,8 +276,6 @@ impl LiteRpcServer for LiteBridge {
             }
         };
 
-        let block_listner = self.get_block_listner(commitment);
-
         let is_valid = match self
             .rpc_client
             .is_blockhash_valid(&blockhash, commitment)
@@ -289,7 +287,12 @@ impl LiteRpcServer for LiteBridge {
             }
         };
 
-        let slot = block_listner.get_latest_block_info().await.1.slot;
+        let slot = self
+            .block_store
+            .get_latest_block_info(commitment)
+            .await
+            .1
+            .slot;
 
         Ok(RpcResponse {
             context: RpcResponseContext {
@@ -303,8 +306,10 @@ impl LiteRpcServer for LiteBridge {
     async fn get_signature_statuses(
         &self,
         sigs: Vec<String>,
-        _config: Option<solana_client::rpc_config::RpcSignatureStatusConfig>,
+        _config: Option<RpcSignatureStatusConfig>,
     ) -> crate::rpc::Result<RpcResponse<Vec<Option<TransactionStatus>>>> {
+        RPC_GET_SIGNATURE_STATUSES.inc();
+
         let sig_statuses = sigs
             .iter()
             .map(|sig| {
@@ -318,8 +323,8 @@ impl LiteRpcServer for LiteBridge {
         Ok(RpcResponse {
             context: RpcResponseContext {
                 slot: self
-                    .finalized_block_listener
+                    .block_store
-                    .get_latest_block_info()
+                    .get_latest_block_info(CommitmentConfig::finalized())
                     .await
                     .1
                     .slot,
@@ -330,6 +335,8 @@ impl LiteRpcServer for LiteBridge {
     }
 
     fn get_version(&self) -> crate::rpc::Result<RpcVersionInfo> {
+        RPC_GET_VERSION.inc();
+
         let version = solana_version::Version::default();
         Ok(RpcVersionInfo {
             solana_core: version.to_string(),
@@ -343,6 +350,8 @@ impl LiteRpcServer for LiteBridge {
         lamports: u64,
         config: Option<RpcRequestAirdropConfig>,
     ) -> crate::rpc::Result<String> {
+        RPC_REQUEST_AIRDROP.inc();
+
         let pubkey = match Pubkey::from_str(&pubkey_str) {
             Ok(pubkey) => pubkey,
             Err(err) => {
@@ -372,11 +381,11 @@ impl LiteRpcServer for LiteBridge {
         &self,
         mut sink: SubscriptionSink,
         signature: String,
-        commitment_config: CommitmentConfig,
+        _commitment_config: CommitmentConfig,
     ) -> SubscriptionResult {
+        RPC_SIGNATURE_SUBSCRIBE.inc();
         sink.accept()?;
-        self.get_block_listner(commitment_config)
+        self.block_listner.signature_subscribe(signature, sink);
-            .signature_subscribe(signature, sink);
         Ok(())
     }
 }

src/cli.rs

@@ -27,7 +27,10 @@ pub struct Args {
     /// interval between clean
     #[arg(short = 'c', long, default_value_t = DEFAULT_CLEAN_INTERVAL_MS)]
     pub clean_interval_ms: u64,
-    /// addr to postgres
+    /// enable logging to postgres
     #[arg(short = 'p', long)]
-    pub enable_postgres: bool 
+    pub enable_postgres: bool,
+    /// enable metrics to prometheus at addr
+    #[arg(short = 'm', long, default_value_t = String::from("[::]:9091"))]
+    pub prometheus_addr: String,
 }

src/lib.rs

@@ -1,6 +1,7 @@
 use const_env::from_env;
 use solana_transaction_status::TransactionConfirmationStatus;
 
+pub mod block_store;
 pub mod bridge;
 pub mod cli;
 pub mod configs;

src/main.rs

@@ -19,6 +19,7 @@ pub async fn main() -> anyhow::Result<()> {
         clean_interval_ms,
         fanout_size,
         enable_postgres,
+        prometheus_addr,
     } = Args::parse();
 
     let tx_batch_interval_ms = Duration::from_millis(tx_batch_interval_ms);
@@ -34,6 +35,7 @@ pub async fn main() -> anyhow::Result<()> {
             tx_batch_interval_ms,
             clean_interval_ms,
             enable_postgres,
+            prometheus_addr,
         )
         .await?;
 

src/rpc.rs

@@ -1,8 +1,8 @@
 use jsonrpsee::proc_macros::rpc;
-use solana_client::rpc_config::{
+use solana_rpc_client_api::config::{
     RpcContextConfig, RpcRequestAirdropConfig, RpcSignatureStatusConfig,
 };
-use solana_client::rpc_response::{Response as RpcResponse, RpcBlockhash, RpcVersionInfo};
+use solana_rpc_client_api::response::{Response as RpcResponse, RpcBlockhash, RpcVersionInfo};
 use solana_sdk::commitment_config::CommitmentConfig;
 use solana_transaction_status::TransactionStatus;
 

src/tpu_manager.rs

@@ -4,18 +4,19 @@ use std::sync::{
 };
 
 use log::info;
-use solana_client::{
+use solana_quic_client::QuicPool;
-    nonblocking::{rpc_client::RpcClient, tpu_client::TpuClient},
+use solana_rpc_client::nonblocking::rpc_client::RpcClient;
-    tpu_client::TpuClientConfig,
+use solana_tpu_client::{nonblocking::tpu_client::TpuClient, tpu_client::TpuClientConfig};
-};
 use tokio::sync::{RwLock, RwLockReadGuard};
 
+pub type QuicTpuClient = TpuClient<QuicPool>;
+
 #[derive(Clone)]
 pub struct TpuManager {
     error_count: Arc<AtomicU32>,
     rpc_client: Arc<RpcClient>,
-    tpu_client: Arc<RwLock<TpuClient>>,
+    tpu_client: Arc<RwLock<QuicTpuClient>>,
-    ws_addr: String,
+    pub ws_addr: String,
     fanout_slots: u64,
 }
 
@@ -41,7 +42,7 @@ impl TpuManager {
         rpc_client: Arc<RpcClient>,
         ws_addr: &str,
         fanout_slots: u64,
-    ) -> anyhow::Result<TpuClient> {
+    ) -> anyhow::Result<QuicTpuClient> {
         Ok(TpuClient::new(
             rpc_client.clone(),
             ws_addr,
@@ -84,7 +85,7 @@ impl TpuManager {
         }
     }
 
-    pub async fn get_tpu_client(&self) -> RwLockReadGuard<TpuClient> {
+    pub async fn get_tpu_client(&self) -> RwLockReadGuard<QuicTpuClient> {
         self.tpu_client.read().await
     }
 }

src/workers/block_listenser.rs

@@ -4,41 +4,65 @@ use anyhow::{bail, Context};
 use dashmap::DashMap;
 use futures::StreamExt;
 use jsonrpsee::SubscriptionSink;
-use log::info;
+use log::{info, warn};
-use solana_client::{
+use prometheus::{histogram_opts, opts, register_counter, register_histogram, Counter, Histogram};
-    nonblocking::{pubsub_client::PubsubClient, rpc_client::RpcClient},
+use solana_pubsub_client::nonblocking::pubsub_client::PubsubClient;
-    rpc_client::SerializableTransaction,
+use solana_rpc_client::rpc_client::SerializableTransaction;
-    rpc_config::{RpcBlockSubscribeConfig, RpcBlockSubscribeFilter},
+use solana_rpc_client_api::{
-    rpc_response::{Response as RpcResponse, RpcResponseContext},
+    config::{RpcBlockSubscribeConfig, RpcBlockSubscribeFilter},
+    response::{Response as RpcResponse, RpcResponseContext},
 };
+
 use solana_sdk::commitment_config::{CommitmentConfig, CommitmentLevel};
 
 use solana_transaction_status::{
-    option_serializer::OptionSerializer, TransactionConfirmationStatus, TransactionStatus,
+    option_serializer::OptionSerializer, RewardType, TransactionConfirmationStatus,
-    UiConfirmedBlock, UiTransactionStatusMeta,
+    TransactionStatus, UiConfirmedBlock, UiTransactionStatusMeta,
-};
-use tokio::{
-    sync::{mpsc::Sender, RwLock},
-    task::JoinHandle,
 };
+use tokio::{sync::mpsc::Sender, task::JoinHandle};
 
-use crate::workers::{PostgresBlock, PostgresMsg, PostgresUpdateTx};
+use crate::{
+    block_store::BlockStore,
+    workers::{PostgresBlock, PostgresMsg, PostgresUpdateTx},
+};
 
 use super::{PostgresMpscSend, TxProps, TxSender};
 
+lazy_static::lazy_static! {
+    static ref TT_RECV_CON_BLOCK: Histogram = register_histogram!(histogram_opts!(
+        "tt_recv_con_block",
+        "Time to receive confirmed block from block subscribe",
+    ))
+    .unwrap();
+    static ref TT_RECV_FIN_BLOCK: Histogram = register_histogram!(histogram_opts!(
+        "tt_recv_fin_block",
+        "Time to receive finalized block from block subscribe",
+    ))
+    .unwrap();
+    static ref FIN_BLOCKS_RECV: Counter =
+        register_counter!(opts!("fin_blocks_recv", "Number of Finalized Blocks Received")).unwrap();
+    static ref CON_BLOCKS_RECV: Counter =
+        register_counter!(opts!("con_blocks_recv", "Number of Confirmed Blocks Received")).unwrap();
+    static ref INCOMPLETE_FIN_BLOCKS_RECV: Counter =
+        register_counter!(opts!("incomplete_fin_blocks_recv", "Number of Incomplete Finalized Blocks Received")).unwrap();
+    static ref INCOMPLETE_CON_BLOCKS_RECV: Counter =
+        register_counter!(opts!("incomplete_con_blocks_recv", "Number of Incomplete Confirmed Blocks Received")).unwrap();
+    static ref TXS_CONFIRMED: Counter =
+        register_counter!(opts!("txs_confirmed", "Number of Transactions Confirmed")).unwrap();
+    static ref TXS_FINALIZED: Counter =
+        register_counter!(opts!("txs_finalized", "Number of Transactions Finalized")).unwrap();
+}
+
 /// Background worker which listen's to new blocks
 /// and keeps a track of confirmed txs
 #[derive(Clone)]
 pub struct BlockListener {
-    pub_sub_client: Arc<PubsubClient>,
-    commitment_config: CommitmentConfig,
     tx_sender: TxSender,
-    block_store: Arc<DashMap<String, BlockInformation>>,
+    block_store: BlockStore,
-    latest_block_hash: Arc<RwLock<String>>,
     pub signature_subscribers: Arc<DashMap<String, SubscriptionSink>>,
 }
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct BlockInformation {
     pub slot: u64,
     pub block_height: u64,
@@ -50,33 +74,12 @@ pub struct BlockListnerNotificatons {
 }
 
 impl BlockListener {
-    pub async fn new(
+    pub fn new(tx_sender: TxSender, block_store: BlockStore) -> Self {
-        pub_sub_client: Arc<PubsubClient>,
+        Self {
-        rpc_client: Arc<RpcClient>,
-        tx_sender: TxSender,
-        commitment_config: CommitmentConfig,
-    ) -> anyhow::Result<Self> {
-        let (latest_block_hash, block_height) = rpc_client
-            .get_latest_blockhash_with_commitment(commitment_config)
-            .await?;
-
-        let latest_block_hash = latest_block_hash.to_string();
-        let slot = rpc_client
-            .get_slot_with_commitment(commitment_config)
-            .await?;
-
-        Ok(Self {
-            pub_sub_client,
             tx_sender,
-            latest_block_hash: Arc::new(RwLock::new(latest_block_hash.clone())),
+            block_store,
-            block_store: Arc::new({
-                let map = DashMap::new();
-                map.insert(latest_block_hash, BlockInformation { slot, block_height });
-                map
-            }),
-            commitment_config,
             signature_subscribers: Default::default(),
-        })
+        }
     }
 
     pub async fn num_of_sigs_commited(&self, sigs: &[String]) -> usize {
@@ -89,31 +92,6 @@ impl BlockListener {
         num_of_sigs_commited
     }
 
-    pub async fn get_latest_block_info(&self) -> (String, BlockInformation) {
-        let blockhash = &*self.latest_block_hash.read().await;
-
-        (
-            blockhash.to_owned(),
-            self.block_store
-                .get(blockhash)
-                .expect("Internal Error: Block store race condition")
-                .value()
-                .to_owned(),
-        )
-    }
-
-    pub async fn get_block_info(&self, blockhash: &str) -> Option<BlockInformation> {
-        let Some(info) = self.block_store.get(blockhash) else {
-            return None;
-        };
-
-        Some(info.value().to_owned())
-    }
-
-    pub async fn get_latest_blockhash(&self) -> String {
-        self.latest_block_hash.read().await.to_owned()
-    }
-
     pub fn signature_subscribe(&self, signature: String, sink: SubscriptionSink) {
         let _ = self.signature_subscribers.insert(signature, sink);
     }
@@ -122,9 +100,21 @@ impl BlockListener {
         self.signature_subscribers.remove(&signature);
     }
 
-    pub fn listen(self, postgres: Option<PostgresMpscSend>) -> JoinHandle<anyhow::Result<()>> {
+    fn increment_invalid_block_metric(commitment_config: CommitmentConfig) {
-        tokio::spawn(async move {
+        if commitment_config.is_finalized() {
-            let commitment = self.commitment_config.commitment;
+            INCOMPLETE_FIN_BLOCKS_RECV.inc();
+        } else {
+            INCOMPLETE_CON_BLOCKS_RECV.inc();
+        }
+    }
+
+    pub async fn listen_from_pubsub(
+        self,
+        pubsub_client: &PubsubClient,
+        commitment_config: CommitmentConfig,
+        postgres: &Option<PostgresMpscSend>,
+    ) -> anyhow::Result<()> {
+        let commitment = commitment_config.commitment;
 
         let comfirmation_status = match commitment {
             CommitmentLevel::Finalized => TransactionConfirmationStatus::Finalized,
@@ -133,16 +123,13 @@ impl BlockListener {
 
         info!("Subscribing to {commitment:?} blocks");
 
-            let (mut recv, _) = self
+        let (mut recv, _) = pubsub_client
-                .pub_sub_client
             .block_subscribe(
                 RpcBlockSubscribeFilter::All,
                 Some(RpcBlockSubscribeConfig {
-                        commitment: Some(self.commitment_config),
+                    commitment: Some(commitment_config),
                     encoding: None,
-                        transaction_details: Some(
+                    transaction_details: Some(solana_transaction_status::TransactionDetails::Full),
-                            solana_transaction_status::TransactionDetails::Full,
-                        ),
                     show_rewards: None,
                     max_supported_transaction_version: None,
                 }),
@@ -152,33 +139,66 @@ impl BlockListener {
 
         info!("Listening to {commitment:?} blocks");
 
-            while let Some(block) = recv.as_mut().next().await {
+        loop {
-                let slot = block.value.slot;
+            let timer = if commitment_config.is_finalized() {
+                TT_RECV_FIN_BLOCK.start_timer()
+            } else {
+                TT_RECV_CON_BLOCK.start_timer()
+            };
+
+            let Some(block) = recv.as_mut().next().await else {
+                bail!("PubSub broke");
+            };
+
+            timer.observe_duration();
+
+            if commitment_config.is_finalized() {
+                FIN_BLOCKS_RECV.inc();
+            } else {
+                CON_BLOCKS_RECV.inc();
+            };
+
+            let slot = block.context.slot;
 
             let Some(block) = block.value.block else {
+                Self::increment_invalid_block_metric(commitment_config);
                 continue;
             };
 
             let Some(block_height) = block.block_height else {
+                Self::increment_invalid_block_metric(commitment_config);
+                continue;
+            };
+
+            let Some(transactions) = block.transactions else {
+                Self::increment_invalid_block_metric(commitment_config);
                 continue;
             };
 
             let blockhash = block.blockhash;
+            let parent_slot = block.parent_slot;
 
-                let Some(transactions) = block.transactions else {
+            self.block_store
+                .add_block(
+                    blockhash.clone(),
+                    BlockInformation { slot, block_height },
+                    commitment_config,
+                )
+                .await;
+
+            if let Some(postgres) = &postgres {
+                let Some(rewards) = block.rewards else {
                         continue;
                     };
 
-                let parent_slot = block.parent_slot;
+                let Some(leader_reward) = rewards
+                        .iter()
+                        .find(|reward| Some(RewardType::Fee) == reward.reward_type) else {
+                        continue;
+                    };
 
-                // Write to block store first in order to prevent
+                let _leader_id = &leader_reward.pubkey;
-                // any race condition i.e prevent some one to
-                // ask the map what it doesn't have rn
-                self.block_store
-                    .insert(blockhash.clone(), BlockInformation { slot, block_height });
-                *self.latest_block_hash.write().await = blockhash;
 
-                if let Some(postgres) = &postgres {
                 postgres
                     .send(PostgresMsg::PostgresBlock(PostgresBlock {
                         slot: slot as i64,
@@ -202,6 +222,17 @@ impl BlockListener {
                 let sig = tx.get_signature().to_string();
 
                 if let Some(mut tx_status) = self.tx_sender.txs_sent.get_mut(&sig) {
+                    //
+                    // Metrics
+                    //
+                    if status.is_ok() {
+                        if commitment_config.is_finalized() {
+                            TXS_FINALIZED.inc();
+                        } else {
+                            TXS_CONFIRMED.inc();
+                        }
+                    }
+
                     tx_status.value_mut().status = Some(TransactionStatus {
                         slot,
                         confirmations: None,
@@ -245,8 +276,24 @@ impl BlockListener {
                 }
             }
         }
+    }
 
-            bail!("Stopped Listening to {commitment:?} blocks")
+    pub fn listen(
+        self,
+        commitment_config: CommitmentConfig,
+        postgres: Option<PostgresMpscSend>,
+    ) -> JoinHandle<anyhow::Result<()>> {
+        tokio::spawn(async move {
+            loop {
+                let ws_addr = &self.tx_sender.tpu_manager.ws_addr;
+                let pub_sub_client = PubsubClient::new(ws_addr).await?;
+                let err = self
+                    .clone()
+                    .listen_from_pubsub(&pub_sub_client, commitment_config, &postgres)
+                    .await
+                    .unwrap_err();
+                warn!("{commitment_config:?} Block Subscribe error {err}");
+            }
         })
     }
 }

src/workers/cleaner.rs

@@ -7,16 +7,16 @@ use super::{BlockListener, TxSender};
 
 /// Background worker which cleans up memory  
 #[derive(Clone)]
-pub struct Cleaner<const N: usize> {
+pub struct Cleaner {
     tx_sender: TxSender,
-    block_listeners: [BlockListener; N],
+    block_listenser: BlockListener,
 }
 
-impl<const N: usize> Cleaner<N> {
+impl Cleaner {
-    pub fn new(tx_sender: TxSender, block_listeners: [BlockListener; N]) -> Self {
+    pub fn new(tx_sender: TxSender, block_listenser: BlockListener) -> Self {
         Self {
             tx_sender,
-            block_listeners,
+            block_listenser,
         }
     }
 
@@ -38,22 +38,20 @@ impl<const N: usize> Cleaner<N> {
 
     /// Clean Signature Subscribers from Block Listeners
     pub fn clean_block_listeners(&self) {
-        for block_listenser in &self.block_listeners {
         let mut to_remove = vec![];
 
-            for subscriber in block_listenser.signature_subscribers.iter() {
+        for subscriber in self.block_listenser.signature_subscribers.iter() {
             if subscriber.value().is_closed() {
                 to_remove.push(subscriber.key().to_owned());
             }
         }
 
         for to_remove in &to_remove {
-                block_listenser.signature_subscribers.remove(to_remove);
+            self.block_listenser.signature_subscribers.remove(to_remove);
         }
 
         info!("Cleaned {} Signature Subscribers", to_remove.len());
     }
-    }
 
     pub fn start(self, ttl_duration: Duration) -> JoinHandle<anyhow::Result<()>> {
         let mut ttl = tokio::time::interval(ttl_duration);

src/workers/mod.rs

@@ -2,12 +2,12 @@ mod block_listenser;
 mod cleaner;
 mod metrics_capture;
 mod postgres;
-mod prometheus;
+mod prometheus_sync;
 mod tx_sender;
 
 pub use block_listenser::*;
 pub use cleaner::*;
 pub use metrics_capture::*;
 pub use postgres::*;
-pub use prometheus::*;
+pub use prometheus_sync::*;
 pub use tx_sender::*;

src/workers/prometheus_sync.rs

@@ -1,23 +1,13 @@
-use std::collections::HashMap;
+use prometheus::{Encoder, TextEncoder};
-
 use tokio::{
     io::AsyncWriteExt,
-    net::{TcpListener, TcpStream},
+    net::{TcpListener, TcpStream, ToSocketAddrs},
     task::JoinHandle,
 };
 
-use super::MetricsCapture;
+pub struct PrometheusSync;
-
-#[derive(Clone)]
-pub struct PrometheusSync {
-    metrics_capture: MetricsCapture,
-}
 
 impl PrometheusSync {
-    pub fn new(metrics_capture: MetricsCapture) -> Self {
-        Self { metrics_capture }
-    }
-
     fn create_response(payload: &str) -> String {
         format!(
             "HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n{}",
@@ -27,10 +17,16 @@ impl PrometheusSync {
     }
 
     async fn handle_stream(&self, stream: &mut TcpStream) -> anyhow::Result<()> {
-        let metrics = self.metrics_capture.get_metrics().await;
+        let mut metrics_buffer = Vec::new();
-        let metrics = serde_prometheus::to_string(&metrics, Some("literpc"), HashMap::new())?;
+        let encoder = TextEncoder::new();
 
-        let response = Self::create_response(&metrics);
+        let metric_families = prometheus::gather();
+        encoder
+            .encode(&metric_families, &mut metrics_buffer)
+            .unwrap();
+
+        let metrics_buffer = String::from_utf8(metrics_buffer).unwrap();
+        let response = Self::create_response(&metrics_buffer);
 
         stream.writable().await?;
         stream.write_all(response.as_bytes()).await?;
@@ -40,10 +36,10 @@ impl PrometheusSync {
         Ok(())
     }
 
-    pub fn sync(self) -> JoinHandle<anyhow::Result<()>> {
+    pub fn sync(self, addr: impl ToSocketAddrs + Send + 'static) -> JoinHandle<anyhow::Result<()>> {
         #[allow(unreachable_code)]
         tokio::spawn(async move {
-            let listener = TcpListener::bind("[::]:9091").await?;
+            let listener = TcpListener::bind(addr).await?;
 
             loop {
                 let Ok((mut stream, _addr)) =  listener.accept().await else {

src/workers/tx_sender.rs

@@ -7,6 +7,7 @@ use anyhow::bail;
 use dashmap::DashMap;
 use log::{info, warn};
 
+use prometheus::{register_counter, Counter};
 use solana_transaction_status::TransactionStatus;
 use tokio::{
     sync::mpsc::{error::TryRecvError, UnboundedReceiver},
@@ -20,6 +21,11 @@ use crate::{
 
 use super::PostgresMpscSend;
 
+lazy_static::lazy_static! {
+    static ref TXS_SENT: Counter =
+        register_counter!("txs_sent", "Number of transactions forwarded to tpu").unwrap();
+}
+
 pub type WireTransaction = Vec<u8>;
 
 /// Retry transactions to a maximum of `u16` times, keep a track of confirmed transactions
@@ -28,7 +34,7 @@ pub struct TxSender {
     /// Tx(s) forwarded to tpu
     pub txs_sent: Arc<DashMap<String, TxProps>>,
     /// TpuClient to call the tpu port
-    tpu_manager: Arc<TpuManager>,
+    pub tpu_manager: Arc<TpuManager>,
 }
 
 /// Transaction Properties
@@ -77,6 +83,9 @@ impl TxSender {
                     for (sig, _) in &sigs_and_slots {
                         txs_sent.insert(sig.to_owned(), TxProps::default());
                     }
+                    // metrics
+                    TXS_SENT.inc_by(sigs_and_slots.len() as f64);
+
                     1
                 }
                 Err(err) => {
@@ -86,15 +95,17 @@ impl TxSender {
             };
 
             if let Some(postgres) = postgres {
+                let forwarded_slot = tpu_client.get_tpu_client().await.estimated_current_slot();
+
                 for (sig, recent_slot) in sigs_and_slots {
                     postgres
                         .send(PostgresMsg::PostgresTx(PostgresTx {
                             signature: sig.clone(),
                             recent_slot: recent_slot as i64,
-                            forwarded_slot: 0,    // FIX: figure this out
+                            forwarded_slot: forwarded_slot as i64,
-                            processed_slot: None, // FIX: figure this out
+                            processed_slot: None,
-                            cu_consumed: None,    // FIX: figure this out
+                            cu_consumed: None,
-                            cu_requested: None,   // FIX: figure this out
+                            cu_requested: None,
                             quic_response,
                         }))
                         .expect("Error writing to postgres service");

tests/client.rs

@@ -3,7 +3,7 @@ use std::sync::Arc;
 use bench::helpers::BenchHelper;
 use lite_rpc::DEFAULT_LITE_RPC_ADDR;
 use log::info;
-use solana_client::{nonblocking::rpc_client::RpcClient, rpc_client::SerializableTransaction};
+use solana_rpc_client::{nonblocking::rpc_client::RpcClient, rpc_client::SerializableTransaction};
 use solana_sdk::commitment_config::CommitmentConfig;
 
 const AMOUNT: usize = 5;

tests/diff.rs

@@ -0,0 +1,34 @@
+use lite_rpc::{DEFAULT_LITE_RPC_ADDR, DEFAULT_RPC_ADDR};
+use solana_rpc_client::nonblocking::rpc_client::RpcClient;
+use solana_sdk::commitment_config::CommitmentConfig;
+
+#[tokio::test]
+async fn diff_rpc() -> anyhow::Result<()> {
+    let rpc_client = RpcClient::new(DEFAULT_RPC_ADDR.to_string());
+    let lite_rpc_client = RpcClient::new(DEFAULT_LITE_RPC_ADDR.to_string());
+
+    check_block_hash(&rpc_client, &lite_rpc_client, CommitmentConfig::confirmed()).await?;
+    check_block_hash(&rpc_client, &lite_rpc_client, CommitmentConfig::finalized()).await?;
+
+    Ok(())
+}
+
+async fn check_block_hash(
+    rpc_client: &RpcClient,
+    lite_rpc_client: &RpcClient,
+    commitment_config: CommitmentConfig,
+) -> anyhow::Result<()> {
+    let rpc_blockhash = rpc_client
+        .get_latest_blockhash_with_commitment(commitment_config)
+        .await?;
+    let lite_blockhash = lite_rpc_client
+        .get_latest_blockhash_with_commitment(commitment_config)
+        .await?;
+
+    println!("{commitment_config:?} {rpc_blockhash:?} {lite_blockhash:?}");
+
+    assert_eq!(rpc_blockhash.0, lite_blockhash.0);
+    assert_eq!(rpc_blockhash.1, lite_blockhash.1);
+
+    Ok(())
+}

tests/workers.rs

@@ -3,13 +3,14 @@ use std::{sync::Arc, time::Duration};
 use bench::helpers::BenchHelper;
 use futures::future::try_join_all;
 use lite_rpc::{
+    block_store::BlockStore,
     encoding::BinaryEncoding,
     tpu_manager::TpuManager,
     workers::{BlockListener, TxSender},
     DEFAULT_LITE_RPC_ADDR, DEFAULT_RPC_ADDR, DEFAULT_TX_BATCH_INTERVAL_MS, DEFAULT_TX_BATCH_SIZE,
     DEFAULT_WS_ADDR,
 };
-use solana_client::nonblocking::{pubsub_client::PubsubClient, rpc_client::RpcClient};
+use solana_rpc_client::nonblocking::rpc_client::RpcClient;
 
 use solana_sdk::commitment_config::CommitmentConfig;
 use solana_transaction_status::TransactionConfirmationStatus;
@@ -31,23 +32,17 @@ async fn send_and_confirm_txs() {
         .unwrap(),
     );
 
-    let pub_sub_client = Arc::new(PubsubClient::new(DEFAULT_WS_ADDR).await.unwrap());
-
     let tx_sender = TxSender::new(tpu_client);
+    let block_store = BlockStore::new(&rpc_client).await.unwrap();
 
-    let block_listener = BlockListener::new(
+    let block_listener = BlockListener::new(tx_sender.clone(), block_store);
-        pub_sub_client.clone(),
-        rpc_client.clone(),
-        tx_sender.clone(),
-        CommitmentConfig::confirmed(),
-    )
-    .await
-    .unwrap();
 
     let (tx_send, tx_recv) = mpsc::unbounded_channel();
 
     let services = try_join_all(vec![
-        block_listener.clone().listen(None),
+        block_listener
+            .clone()
+            .listen(CommitmentConfig::confirmed(), None),
         tx_sender.clone().execute(
             tx_recv,
             DEFAULT_TX_BATCH_SIZE,