openzeppelin_monitor/services/blockchain/
pool.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
//! Client pool for managing blockchain clients.
//!
//! This module provides a thread-safe client pooling system that:
//! - Caches blockchain clients by network
//! - Creates clients lazily on first use
//! - Handles both EVM and Stellar clients
//! - Provides type-safe access to clients
//! - Manages client lifecycles automatically
//!
//! The pool uses a fast path for existing clients and a slow path for
//! creating new ones, optimizing performance while maintaining safety.

use crate::{
	models::{BlockChainType, Network},
	services::blockchain::{
		BlockChainClient, BlockFilterFactory, EVMTransportClient, EvmClient, EvmClientTrait,
		StellarClient, StellarClientTrait, StellarTransportClient,
	},
};
use anyhow::Context;
use async_trait::async_trait;
use futures::future::BoxFuture;
use std::{any::Any, collections::HashMap, sync::Arc};
use tokio::sync::RwLock;

/// Trait for the client pool.
#[async_trait]
pub trait ClientPoolTrait: Send + Sync {
	type EvmClient: EvmClientTrait + BlockChainClient + BlockFilterFactory<Self::EvmClient>;
	type StellarClient: StellarClientTrait
		+ BlockChainClient
		+ BlockFilterFactory<Self::StellarClient>;

	async fn get_evm_client(
		&self,
		network: &Network,
	) -> Result<Arc<Self::EvmClient>, anyhow::Error>;
	async fn get_stellar_client(
		&self,
		network: &Network,
	) -> Result<Arc<Self::StellarClient>, anyhow::Error>;
}

/// Generic client storage that can hold any type of blockchain client
///
/// Clients are stored in a thread-safe way using a HashMap and an RwLock.
/// The HashMap is indexed by the network slug and the value is an Arc of the client.
pub struct ClientStorage<T> {
	clients: Arc<RwLock<HashMap<String, Arc<T>>>>,
}

impl<T> ClientStorage<T> {
	pub fn new() -> Self {
		Self {
			clients: Arc::new(RwLock::new(HashMap::new())),
		}
	}
}

/// Main client pool manager that handles multiple blockchain types.
///
/// Provides type-safe access to cached blockchain clients. Clients are created
/// on demand when first requested and then cached for future use. Uses RwLock
/// for thread-safe access and Arc for shared ownership.
pub struct ClientPool {
	/// Map of client storages indexed by client type
	pub storages: HashMap<BlockChainType, Box<dyn Any + Send + Sync>>,
}

impl ClientPool {
	/// Creates a new empty client pool.
	///
	/// Initializes empty hashmaps for both EVM and Stellar clients.
	pub fn new() -> Self {
		let mut pool = Self {
			storages: HashMap::new(),
		};

		// Register client types
		pool.register_client_type::<EvmClient<EVMTransportClient>>(BlockChainType::EVM);
		pool.register_client_type::<StellarClient<StellarTransportClient>>(BlockChainType::Stellar);

		pool
	}

	fn register_client_type<T: 'static + Send + Sync>(&mut self, client_type: BlockChainType) {
		self.storages
			.insert(client_type, Box::new(ClientStorage::<T>::new()));
	}

	/// Internal helper method to get or create a client of any type.
	///
	/// Uses a double-checked locking pattern:
	/// 1. Fast path with read lock to check for existing client
	/// 2. Slow path with write lock to create new client if needed
	///
	/// This ensures thread-safety while maintaining good performance
	/// for the common case of accessing existing clients.
	async fn get_or_create_client<T: BlockChainClient + 'static>(
		&self,
		client_type: BlockChainType,
		network: &Network,
		create_fn: impl Fn(&Network) -> BoxFuture<'static, Result<T, anyhow::Error>>,
	) -> Result<Arc<T>, anyhow::Error> {
		let storage = self
			.storages
			.get(&client_type)
			.and_then(|s| s.downcast_ref::<ClientStorage<T>>())
			.with_context(|| "Invalid client type")?;

		// Fast path: check if client exists
		if let Some(client) = storage.clients.read().await.get(&network.slug) {
			return Ok(client.clone());
		}

		// Slow path: create new client
		let mut clients = storage.clients.write().await;
		let client = Arc::new(create_fn(network).await?);
		clients.insert(network.slug.clone(), client.clone());
		Ok(client)
	}

	/// Get the number of clients for a given client type.
	pub async fn get_client_count<T: 'static>(&self, client_type: BlockChainType) -> usize {
		match self
			.storages
			.get(&client_type)
			.and_then(|s| s.downcast_ref::<ClientStorage<T>>())
		{
			Some(storage) => storage.clients.read().await.len(),
			None => 0,
		}
	}
}

#[async_trait]
impl ClientPoolTrait for ClientPool {
	type EvmClient = EvmClient<EVMTransportClient>;
	type StellarClient = StellarClient<StellarTransportClient>;

	/// Gets or creates an EVM client for the given network.
	///
	/// First checks the cache for an existing client. If none exists,
	/// creates a new client under a write lock.
	async fn get_evm_client(
		&self,
		network: &Network,
	) -> Result<Arc<Self::EvmClient>, anyhow::Error> {
		self.get_or_create_client(BlockChainType::EVM, network, |n| {
			let network = n.clone();
			Box::pin(async move { Self::EvmClient::new(&network).await })
		})
		.await
		.with_context(|| "Failed to get or create EVM client")
	}

	/// Gets or creates a Stellar client for the given network.
	///
	/// First checks the cache for an existing client. If none exists,
	/// creates a new client under a write lock.
	async fn get_stellar_client(
		&self,
		network: &Network,
	) -> Result<Arc<Self::StellarClient>, anyhow::Error> {
		self.get_or_create_client(BlockChainType::Stellar, network, |n| {
			let network = n.clone();
			Box::pin(async move { Self::StellarClient::new(&network).await })
		})
		.await
		.with_context(|| "Failed to get or create Stellar client")
	}
}

impl Default for ClientPool {
	fn default() -> Self {
		Self::new()
	}
}