Skip to main content

surfpool_sdk/
surfnet.rs

1use std::{
2    net::TcpListener,
3    thread::sleep,
4    time::{Duration, Instant},
5};
6
7use crossbeam_channel::{Receiver, Sender};
8use solana_commitment_config::CommitmentConfig;
9use solana_keypair::Keypair;
10use solana_pubkey::Pubkey;
11use solana_rpc_client::rpc_client::RpcClient;
12use solana_signer::Signer;
13use surfpool_core::surfnet::{
14    locker::SurfnetSvmLocker,
15    svm::{SurfnetSvm, SurfnetSvmConfig},
16};
17use surfpool_types::{
18    BlockProductionMode, RpcConfig, SimnetCommand, SimnetConfig, SimnetEvent, SurfpoolConfig,
19};
20
21use crate::{
22    Cheatcodes,
23    error::{SurfnetError, SurfnetResult},
24};
25
26/// Builder for configuring a [`Surfnet`] instance before starting it.
27///
28/// ```rust
29/// use surfpool_sdk::{Surfnet, BlockProductionMode};
30///
31/// # async fn example() {
32/// let surfnet = Surfnet::builder()
33///     .offline(true)
34///     .block_production_mode(BlockProductionMode::Transaction)
35///     .airdrop_sol(10_000_000_000)
36///     .start()
37///     .await
38///     .unwrap();
39/// # }
40/// ```
41pub struct SurfnetBuilder {
42    offline_mode: bool,
43    remote_rpc_url: Option<String>,
44    block_production_mode: BlockProductionMode,
45    slot_time_ms: u64,
46    airdrop_addresses: Vec<Pubkey>,
47    airdrop_lamports: u64,
48    payer: Option<Keypair>,
49}
50
51impl Default for SurfnetBuilder {
52    fn default() -> Self {
53        Self {
54            offline_mode: true,
55            remote_rpc_url: None,
56            block_production_mode: BlockProductionMode::Transaction,
57            slot_time_ms: 1,
58            airdrop_addresses: vec![],
59            airdrop_lamports: 10_000_000_000, // 10 SOL
60            payer: None,
61        }
62    }
63}
64
65impl SurfnetBuilder {
66    /// Run in offline mode (no mainnet RPC fallback). Default: `true`.
67    pub fn offline(mut self, offline: bool) -> Self {
68        self.offline_mode = offline;
69        self
70    }
71
72    /// Set a remote RPC URL for account fallback (implies `offline(false)`).
73    pub fn remote_rpc_url(mut self, url: impl Into<String>) -> Self {
74        self.remote_rpc_url = Some(url.into());
75        self.offline_mode = false;
76        self
77    }
78
79    /// How blocks are produced. Default: `Transaction` (advance on each tx).
80    pub fn block_production_mode(mut self, mode: BlockProductionMode) -> Self {
81        self.block_production_mode = mode;
82        self
83    }
84
85    /// Slot time in milliseconds. Default: `1` (fast for tests).
86    pub fn slot_time_ms(mut self, ms: u64) -> Self {
87        self.slot_time_ms = ms;
88        self
89    }
90
91    /// Additional addresses to airdrop SOL to at startup.
92    pub fn airdrop_addresses(mut self, addresses: Vec<Pubkey>) -> Self {
93        self.airdrop_addresses = addresses;
94        self
95    }
96
97    /// Amount of lamports to airdrop to the payer (and additional addresses) at startup.
98    /// Default: 10 SOL.
99    pub fn airdrop_sol(mut self, lamports: u64) -> Self {
100        self.airdrop_lamports = lamports;
101        self
102    }
103
104    /// Use a specific keypair as the payer. If not set, a random one is generated.
105    pub fn payer(mut self, keypair: Keypair) -> Self {
106        self.payer = Some(keypair);
107        self
108    }
109
110    /// Start the surfnet with the configured options.
111    pub async fn start(self) -> SurfnetResult<Surfnet> {
112        let payer = self.payer.unwrap_or_else(Keypair::new);
113        let airdrop_lamports = self.airdrop_lamports;
114
115        let bind_port = get_free_port()?;
116        let ws_port = get_free_port()?;
117        let bind_host = "127.0.0.1".to_string();
118
119        let mut startup_airdrop_addresses = vec![payer.pubkey()];
120        startup_airdrop_addresses.extend(self.airdrop_addresses);
121
122        let surfpool_config = SurfpoolConfig {
123            simnets: vec![SimnetConfig {
124                offline_mode: self.offline_mode,
125                remote_rpc_url: self.remote_rpc_url,
126                slot_time: self.slot_time_ms,
127                block_production_mode: self.block_production_mode,
128                airdrop_addresses: startup_airdrop_addresses.clone(),
129                airdrop_token_amount: airdrop_lamports,
130                ..Default::default()
131            }],
132            rpc: RpcConfig {
133                bind_host: bind_host.clone(),
134                bind_port,
135                ws_port,
136                ..Default::default()
137            },
138            ..Default::default()
139        };
140
141        let rpc_url = format!("http://{bind_host}:{bind_port}");
142        let ws_url = format!("ws://{bind_host}:{ws_port}");
143
144        let svm_config = SurfnetSvmConfig {
145            surfnet_id: surfpool_config.simnets[0].surfnet_id.clone(),
146            slot_time: surfpool_config.simnets[0].slot_time,
147            instruction_profiling_enabled: surfpool_config.simnets[0].instruction_profiling_enabled,
148            max_profiles: surfpool_config.simnets[0].max_profiles,
149            log_bytes_limit: surfpool_config.simnets[0].log_bytes_limit,
150            ..SurfnetSvmConfig::default()
151        };
152        let (surfnet_svm, simnet_events_rx, geyser_events_rx) = SurfnetSvm::new(svm_config)
153            .map_err(|e| SurfnetError::Runtime(format!("failed to initialize Surfnet SVM: {e}")))?;
154        let (simnet_commands_tx, simnet_commands_rx) = crossbeam_channel::unbounded();
155
156        let svm_locker = SurfnetSvmLocker::new(surfnet_svm);
157        let svm_locker_clone = svm_locker.clone();
158        let simnet_commands_tx_clone = simnet_commands_tx.clone();
159
160        let _handle = std::thread::Builder::new()
161            .name("surfnet-sdk".into())
162            .spawn(move || {
163                let future = surfpool_core::runloops::start_local_surfnet_runloop(
164                    svm_locker_clone,
165                    surfpool_config,
166                    simnet_commands_tx_clone,
167                    simnet_commands_rx,
168                    geyser_events_rx,
169                );
170                if let Err(e) = hiro_system_kit::nestable_block_on(future) {
171                    log::error!("Surfnet exited with error: {e}");
172                }
173            })
174            .map_err(|e| SurfnetError::Runtime(e.to_string()))?;
175
176        // Wait for the runtime to signal ready
177        wait_for_ready(&simnet_events_rx)?;
178        wait_for_startup_airdrops(&rpc_url, &startup_airdrop_addresses, airdrop_lamports)?;
179
180        Ok(Surfnet {
181            rpc_url,
182            ws_url,
183            payer,
184            simnet_commands_tx,
185            simnet_events_rx,
186            svm_locker,
187            instance_id: uuid::Uuid::new_v4().to_string(),
188        })
189    }
190}
191
192/// A running Surfpool instance with RPC/WS endpoints on dynamic ports.
193///
194/// Provides:
195/// - Pre-funded payer keypair
196/// - [`RpcClient`] connected to the local instance
197/// - [`Cheatcodes`] for direct state manipulation (fund accounts, set token balances, etc.)
198///
199/// The instance is shut down when dropped.
200pub struct Surfnet {
201    rpc_url: String,
202    ws_url: String,
203    payer: Keypair,
204    simnet_commands_tx: Sender<SimnetCommand>,
205    simnet_events_rx: Receiver<SimnetEvent>,
206    #[allow(dead_code)] // retained for future direct profiling access
207    svm_locker: SurfnetSvmLocker,
208    instance_id: String,
209}
210
211impl Surfnet {
212    /// Start a surfnet with default settings (offline, transaction-mode blocks, 10 SOL payer).
213    pub async fn start() -> SurfnetResult<Self> {
214        SurfnetBuilder::default().start().await
215    }
216
217    /// Create a builder for custom configuration.
218    pub fn builder() -> SurfnetBuilder {
219        SurfnetBuilder::default()
220    }
221
222    /// The HTTP RPC URL (e.g. `http://127.0.0.1:12345`).
223    pub fn rpc_url(&self) -> &str {
224        &self.rpc_url
225    }
226
227    /// The WebSocket URL (e.g. `ws://127.0.0.1:12346`).
228    pub fn ws_url(&self) -> &str {
229        &self.ws_url
230    }
231
232    /// Create a new [`RpcClient`] connected to this surfnet.
233    pub fn rpc_client(&self) -> RpcClient {
234        RpcClient::new(&self.rpc_url)
235    }
236
237    /// The pre-funded payer keypair.
238    pub fn payer(&self) -> &Keypair {
239        &self.payer
240    }
241
242    /// Access cheatcode helpers for direct state manipulation.
243    pub fn cheatcodes(&self) -> Cheatcodes<'_> {
244        Cheatcodes::new(&self.rpc_url)
245    }
246
247    /// Get a reference to the simnet events receiver for observing runtime events.
248    pub fn events(&self) -> &Receiver<SimnetEvent> {
249        &self.simnet_events_rx
250    }
251
252    /// Send a command to the simnet runtime.
253    pub fn send_command(&self, command: SimnetCommand) -> SurfnetResult<()> {
254        self.simnet_commands_tx
255            .send(command)
256            .map_err(|e| SurfnetError::Runtime(format!("failed to send command: {e}")))
257    }
258
259    /// The unique instance ID for this surfnet.
260    pub fn instance_id(&self) -> &str {
261        &self.instance_id
262    }
263}
264
265impl Drop for Surfnet {
266    fn drop(&mut self) {
267        let _ = self.simnet_commands_tx.send(SimnetCommand::Terminate(None));
268    }
269}
270
271fn get_free_port() -> SurfnetResult<u16> {
272    let listener = TcpListener::bind("127.0.0.1:0")
273        .map_err(|e| SurfnetError::PortAllocation(e.to_string()))?;
274    let port = listener
275        .local_addr()
276        .map_err(|e| SurfnetError::PortAllocation(e.to_string()))?
277        .port();
278    drop(listener);
279    Ok(port)
280}
281
282fn wait_for_ready(events_rx: &Receiver<SimnetEvent>) -> SurfnetResult<()> {
283    loop {
284        match events_rx.recv() {
285            Ok(SimnetEvent::Ready(_)) => return Ok(()),
286            Ok(SimnetEvent::Aborted(err)) => return Err(SurfnetError::Aborted(err)),
287            Ok(SimnetEvent::Shutdown) => {
288                return Err(SurfnetError::Aborted(
289                    "surfnet shut down during startup".into(),
290                ));
291            }
292            Ok(_) => continue,
293            Err(e) => {
294                return Err(SurfnetError::Startup(format!(
295                    "events channel closed unexpectedly: {e}"
296                )));
297            }
298        }
299    }
300}
301
302fn wait_for_startup_airdrops(
303    rpc_url: &str,
304    addresses: &[Pubkey],
305    expected_lamports: u64,
306) -> SurfnetResult<()> {
307    let rpc_client = RpcClient::new(rpc_url.to_string());
308    let deadline = Instant::now() + Duration::from_secs(5);
309    let mut last_error = None;
310    let mut last_balances = vec![];
311
312    while Instant::now() < deadline {
313        last_balances.clear();
314        let mut all_match = true;
315
316        for address in addresses {
317            match rpc_client.get_balance_with_commitment(address, CommitmentConfig::processed()) {
318                Ok(response) => {
319                    last_balances.push((address.to_string(), response.value));
320                    if response.value != expected_lamports {
321                        all_match = false;
322                    }
323                }
324                Err(err) => {
325                    last_error = Some(err.to_string());
326                    all_match = false;
327                    break;
328                }
329            }
330        }
331
332        if all_match {
333            return Ok(());
334        }
335
336        sleep(Duration::from_millis(25));
337    }
338
339    let balance_summary = if last_balances.is_empty() {
340        "no balances observed".to_string()
341    } else {
342        last_balances
343            .iter()
344            .map(|(address, balance)| format!("{address}={balance}"))
345            .collect::<Vec<_>>()
346            .join(", ")
347    };
348
349    Err(SurfnetError::Startup(format!(
350        "startup balances not visible over RPC within timeout (expected {expected_lamports}); last balances: {balance_summary}; last error: {}",
351        last_error.unwrap_or_else(|| "none".to_string())
352    )))
353}