helius_stream/
stream.rs

1use std::collections::HashMap;
2use std::net::TcpStream;
3
4use base64::Engine;
5use log::{debug, warn};
6use serde::Deserialize;
7use tungstenite::{connect, stream::MaybeTlsStream, Message, WebSocket};
8
9use crate::error::StreamError;
10use crate::health::{ReconnectPolicy, StreamHealth};
11use crate::types::{AccountUpdate, Pubkey, StreamConfig, StreamState};
12
13/// A resilient Helius WebSocket client.
14///
15/// Subscribes to Solana account updates and tracks health (gaps, staleness,
16/// reconnect attempts). Synchronous; if you need async, wrap calls in
17/// `tokio::task::spawn_blocking`.
18pub struct HeliusStream {
19    ws: WebSocket<MaybeTlsStream<TcpStream>>,
20    config: StreamConfig,
21    state: StreamState,
22    health: StreamHealth,
23    reconnect: ReconnectPolicy,
24    /// Maps server-assigned subscription_id → pubkey.
25    subs: HashMap<u64, Pubkey>,
26    /// Subscribe requests sent but not yet acknowledged.
27    pending_subs: Vec<(u64, Pubkey)>,
28    next_request_id: u64,
29}
30
31#[derive(Debug, Deserialize)]
32struct WsNotification {
33    method: Option<String>,
34    params: Option<WsParams>,
35}
36
37#[derive(Debug, Deserialize)]
38struct WsParams {
39    result: Option<WsResult>,
40    subscription: Option<u64>,
41}
42
43#[derive(Debug, Deserialize)]
44struct WsResult {
45    context: Option<WsContext>,
46    value: Option<WsAccountValue>,
47}
48
49#[derive(Debug, Deserialize)]
50struct WsContext {
51    slot: u64,
52}
53
54#[derive(Debug, Deserialize)]
55#[serde(rename_all = "camelCase")]
56struct WsAccountValue {
57    lamports: u64,
58    data: serde_json::Value,
59}
60
61impl HeliusStream {
62    /// Open a WebSocket connection to the configured endpoint.
63    pub fn connect(config: StreamConfig) -> Result<Self, StreamError> {
64        let url = config.ws_url();
65        let (ws, _) = connect(&url).map_err(|e| StreamError::Connect(e.to_string()))?;
66        Ok(Self {
67            ws,
68            config,
69            state: StreamState::Connected { since: std::time::SystemTime::now() },
70            health: StreamHealth::new(),
71            reconnect: ReconnectPolicy::new(),
72            subs: HashMap::new(),
73            pending_subs: Vec::new(),
74            next_request_id: 1,
75        })
76    }
77
78    /// Subscribe to account updates for a 32-byte pubkey.
79    pub fn subscribe_account(&mut self, pubkey: Pubkey) -> Result<(), StreamError> {
80        if self.subs.values().any(|pk| *pk == pubkey)
81            || self.pending_subs.iter().any(|(_, pk)| *pk == pubkey)
82        {
83            return Ok(());
84        }
85        let addr = bs58::encode(pubkey).into_string();
86        let id = self.next_request_id;
87        self.next_request_id += 1;
88        let req = serde_json::json!({
89            "jsonrpc": "2.0",
90            "id": id,
91            "method": "accountSubscribe",
92            "params": [&addr, {"encoding": "base64", "commitment": "processed"}]
93        });
94        self.ws
95            .send(Message::Text(req.to_string()))
96            .map_err(|e| StreamError::Send(e.to_string()))?;
97        self.pending_subs.push((id, pubkey));
98        debug!(target: "helius_stream", "subscribed account {}", &addr[..addr.len().min(8)]);
99        Ok(())
100    }
101
102    /// Convenience: subscribe by base58 string.
103    pub fn subscribe_account_b58(&mut self, addr: &str) -> Result<Pubkey, StreamError> {
104        let decoded = bs58::decode(addr)
105            .into_vec()
106            .map_err(|e| StreamError::InvalidPubkey(addr.into(), e.to_string()))?;
107        if decoded.len() != 32 {
108            return Err(StreamError::InvalidPubkey(
109                addr.into(),
110                format!("expected 32 bytes, got {}", decoded.len()),
111            ));
112        }
113        let mut pk = [0u8; 32];
114        pk.copy_from_slice(&decoded);
115        self.subscribe_account(pk)?;
116        Ok(pk)
117    }
118
119    /// Block until the next account update arrives, or return None on close.
120    pub fn next_update(&mut self) -> Option<AccountUpdate> {
121        loop {
122            match self.ws.read() {
123                Ok(Message::Text(text)) => {
124                    let val: serde_json::Value = match serde_json::from_str(&text) {
125                        Ok(v) => v,
126                        Err(_) => continue,
127                    };
128                    // subscribe-ack: {"result": <sub_id>, "id": <req_id>}
129                    if let (Some(sub_id), Some(req_id)) = (
130                        val.get("result").and_then(|v| v.as_u64()),
131                        val.get("id").and_then(|v| v.as_u64()),
132                    ) {
133                        if let Some(pos) = self.pending_subs.iter().position(|(id, _)| *id == req_id) {
134                            let (_, pk) = self.pending_subs.remove(pos);
135                            self.subs.insert(sub_id, pk);
136                            debug!(target: "helius_stream", "subscribe ack req_id={} sub_id={}", req_id, sub_id);
137                        }
138                        continue;
139                    }
140                    // accountNotification
141                    let notif: WsNotification = match serde_json::from_value(val) {
142                        Ok(n) => n,
143                        Err(_) => continue,
144                    };
145                    if notif.method.as_deref() != Some("accountNotification") {
146                        continue;
147                    }
148                    let params = match notif.params {
149                        Some(p) => p,
150                        None => continue,
151                    };
152                    let sub_id = params.subscription.unwrap_or(0);
153                    let pubkey = self.subs.get(&sub_id).copied().unwrap_or([0u8; 32]);
154                    let result = match params.result {
155                        Some(r) => r,
156                        None => continue,
157                    };
158                    let slot = result.context.map(|c| c.slot).unwrap_or(0);
159                    let (lamports, data) = match result.value {
160                        Some(v) => {
161                            let lam = v.lamports;
162                            let raw = v
163                                .data
164                                .as_array()
165                                .and_then(|a| a.first())
166                                .and_then(|s| s.as_str())
167                                .unwrap_or("");
168                            (lam, decode_base64(raw))
169                        }
170                        None => (0, vec![]),
171                    };
172
173                    let gap = self.health.record_update(slot);
174                    self.state = if gap {
175                        let gap_slots = slot.saturating_sub(self.health.last_slot());
176                        StreamState::Degraded { gap_slots }
177                    } else {
178                        StreamState::Connected { since: std::time::SystemTime::now() }
179                    };
180
181                    return Some(AccountUpdate {
182                        pubkey,
183                        lamports,
184                        data,
185                        slot,
186                        write_version: slot,
187                    });
188                }
189                Ok(Message::Ping(data)) => {
190                    let _ = self.ws.send(Message::Pong(data));
191                }
192                Ok(Message::Close(_)) => {
193                    warn!(target: "helius_stream", "connection closed by server");
194                    self.state = StreamState::Failed { reason: "server close".into() };
195                    return None;
196                }
197                Err(e) => {
198                    warn!(target: "helius_stream", "ws error: {e}");
199                    self.state = StreamState::Failed { reason: e.to_string() };
200                    return None;
201                }
202                _ => {}
203            }
204        }
205    }
206
207    /// Re-check freshness and update internal state if the stream went stale.
208    pub fn check_staleness(&mut self) {
209        if self.health.is_stale(self.config.staleness_threshold) {
210            let ms = self.health.last_update().elapsed().as_millis() as u64;
211            self.state = StreamState::Stale { stale_for_ms: ms };
212        }
213    }
214
215    /// True if state is `Connected` or mildly `Degraded` (≤2 slot gap).
216    /// Use this as a circuit breaker before acting on data.
217    pub fn is_safe_for_simulation(&self) -> bool {
218        match &self.state {
219            StreamState::Connected { .. } => true,
220            StreamState::Degraded { gap_slots } => *gap_slots <= 2,
221            _ => false,
222        }
223    }
224
225    pub fn state(&self) -> &StreamState { &self.state }
226    pub fn health(&self) -> &StreamHealth { &self.health }
227    pub fn reconnect_policy_mut(&mut self) -> &mut ReconnectPolicy { &mut self.reconnect }
228    pub fn subscription_count(&self) -> usize { self.subs.len() }
229    pub fn pending_count(&self) -> usize { self.pending_subs.len() }
230}
231
232fn decode_base64(s: &str) -> Vec<u8> {
233    base64::engine::general_purpose::STANDARD
234        .decode(s)
235        .unwrap_or_default()
236}
helius_stream/stream.rs

helius_stream/
stream.rs
