feat: implement StreamManager for gRPC subscription management

magicblock-chainlink/src/remote_account_provider/chain_laser_actor/mock.rs

@@ -0,0 +1,140 @@
+use std::sync::{Arc, Mutex};
+
+use async_trait::async_trait;
+use helius_laserstream::{grpc::SubscribeRequest, LaserstreamError};
+use tokio::sync::mpsc;
+use tokio_stream::wrappers::UnboundedReceiverStream;
+
+use super::{LaserResult, StreamFactory};
+use crate::remote_account_provider::{
+    chain_laser_actor::{LaserStreamWithHandle, StreamHandle},
+    RemoteAccountProviderResult,
+};
+
+/// A test mock that captures subscription requests and allows driving
+/// streams programmatically.
+#[derive(Clone)]
+pub struct MockStreamFactory {
+    /// Every `SubscribeRequest` passed to `subscribe()` is recorded
+    /// here so tests can assert on filter contents, commitment levels,
+    /// etc.
+    captured_requests: Arc<Mutex<Vec<SubscribeRequest>>>,
+
+    /// Requests sent through a `MockStreamHandle::write()` call are
+    /// recorded here so tests can verify handle-driven updates.
+    handle_requests: Arc<Mutex<Vec<SubscribeRequest>>>,
+
+    /// A sender that the test uses to push `LaserResult` items into
+    /// the streams returned by `subscribe()`.
+    /// Each call to `subscribe()` creates a new mpsc channel; the rx
+    /// side becomes the returned stream, and the tx side is stored
+    /// here so the test can drive updates.
+    stream_senders: Arc<Mutex<Vec<Arc<mpsc::UnboundedSender<LaserResult>>>>>,
+}
+
+impl MockStreamFactory {
+    /// Create a new mock stream factory
+    pub fn new() -> Self {
+        Self {
+            captured_requests: Arc::new(Mutex::new(Vec::new())),
+            handle_requests: Arc::new(Mutex::new(Vec::new())),
+            stream_senders: Arc::new(Mutex::new(Vec::new())),
+        }
+    }
+
+    /// Get the captured subscription requests (from `subscribe()`)
+    pub fn captured_requests(&self) -> Vec<SubscribeRequest> {
+        self.captured_requests.lock().unwrap().clone()
+    }
+
+    /// Get the requests sent through stream handles (from
+    /// `handle.write()`)
+    pub fn handle_requests(&self) -> Vec<SubscribeRequest> {
+        self.handle_requests.lock().unwrap().clone()
+    }
+
+    /// Push an error update to a specific stream
+    pub fn push_error_to_stream(&self, idx: usize, error: LaserstreamError) {
+        let senders = self.stream_senders.lock().unwrap();
+        if let Some(sender) = senders.get(idx) {
+            let _ = sender.send(Err(error));
+        }
+    }
+
+    /// Push an update to a specific stream by index
+    pub fn push_update_to_stream(&self, idx: usize, update: LaserResult) {
+        let senders = self.stream_senders.lock().unwrap();
+        if let Some(sender) = senders.get(idx) {
+            let _ = sender.send(update);
+        }
+    }
+
+    /// Get the number of active streams
+    pub fn active_stream_count(&self) -> usize {
+        self.stream_senders.lock().unwrap().len()
+    }
+
+    /// Close a specific stream by index
+    pub fn close_stream(&self, idx: usize) {
+        let mut senders = self.stream_senders.lock().unwrap();
+        if idx < senders.len() {
+            senders.remove(idx);
+        }
+    }
+}
+
+impl Default for MockStreamFactory {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// Mock handle that records write requests and drains them into the
+/// shared `handle_requests` vec on the factory.
+#[derive(Clone)]
+pub struct MockStreamHandle {
+    handle_requests: Arc<Mutex<Vec<SubscribeRequest>>>,
+}
+
+#[async_trait]
+impl StreamHandle for MockStreamHandle {
+    async fn write(
+        &self,
+        request: SubscribeRequest,
+    ) -> Result<(), LaserstreamError> {
+        self.handle_requests.lock().unwrap().push(request);
+        Ok(())
+    }
+}
+
+#[async_trait]
+impl StreamFactory<MockStreamHandle> for MockStreamFactory {
+    async fn subscribe(
+        &self,
+        request: SubscribeRequest,
+    ) -> RemoteAccountProviderResult<LaserStreamWithHandle<MockStreamHandle>>
+    {
+        // Record the initial subscribe request
+        self.captured_requests.lock().unwrap().push(request.clone());
+
+        // Create a channel for driving LaserResult items into the
+        // stream
+        let (stream_tx, stream_rx) = mpsc::unbounded_channel::<LaserResult>();
+        let stream = Box::pin(UnboundedReceiverStream::new(stream_rx));
+
+        let stream_tx = Arc::new(stream_tx);
+        self.stream_senders.lock().unwrap().push(stream_tx);
+
+        // The handle shares the factory's handle_requests vec so
+        // every write is visible to tests immediately.
+        let handle = MockStreamHandle {
+            handle_requests: Arc::clone(&self.handle_requests),
+        };
+
+        // Write the actual request to the handle (mirroring
+        // production behaviour of sending it over the network).
+        handle.write(request).await.unwrap();
+
+        Ok(LaserStreamWithHandle { stream, handle })
+    }
+}

magicblock-chainlink/src/remote_account_provider/chain_laser_actor/mod.rs

@@ -0,0 +1,72 @@
+use std::{collections::HashSet, pin::Pin, sync::Arc};
+
+use futures_util::Stream;
+use helius_laserstream::{
+    grpc::{SubscribeRequest, SubscribeUpdate},
+    LaserstreamError,
+};
+use parking_lot::RwLock;
+use solana_pubkey::Pubkey;
+use tokio::time::Duration;
+
+pub use self::{
+    actor::{ChainLaserActor, Slots},
+    stream_factory::{
+        LaserStreamWithHandle, StreamFactory, StreamFactoryImpl, StreamHandle,
+        StreamHandleImpl,
+    },
+    stream_manager::{StreamManager, StreamManagerConfig, StreamUpdateSource},
+};
+use crate::remote_account_provider::{
+    RemoteAccountProviderError, RemoteAccountProviderResult,
+};
+
+pub type SharedSubscriptions = Arc<RwLock<HashSet<Pubkey>>>;
+
+mod actor;
+#[cfg(test)]
+mod mock;
+mod stream_factory;
+mod stream_manager;
+
+/// Retry a `handle.write(request)` call with linear backoff.
+///
+/// Tries up to `MAX_RETRIES` (5) times with 50 ms × attempt
+/// backoff. Returns the original error after all retries are
+/// exhausted.
+pub(crate) async fn write_with_retry<S: StreamHandle>(
+    handle: &S,
+    task: &str,
+    request: SubscribeRequest,
+) -> RemoteAccountProviderResult<()> {
+    const MAX_RETRIES: usize = 5;
+    let mut retries = MAX_RETRIES;
+    let initial_retries = retries;
+
+    loop {
+        match handle.write(request.clone()).await {
+            Ok(()) => return Ok(()),
+            Err(err) => {
+                if retries > 0 {
+                    retries -= 1;
+                    let backoff_ms = 50u64 * (initial_retries - retries) as u64;
+                    tokio::time::sleep(Duration::from_millis(backoff_ms)).await;
+                    continue;
+                }
+                return Err(
+                    RemoteAccountProviderError::GrpcSubscriptionUpdateFailed(
+                        task.to_string(),
+                        MAX_RETRIES,
+                        format!("{err} ({err:?})"),
+                    ),
+                );
+            }
+        }
+    }
+}
+
+/// Result of a laser stream operation
+pub type LaserResult = Result<SubscribeUpdate, LaserstreamError>;
+
+/// A laser stream of subscription updates
+pub type LaserStream = Pin<Box<dyn Stream<Item = LaserResult> + Send>>;

magicblock-chainlink/src/remote_account_provider/chain_laser_actor/stream_factory.rs

@@ -0,0 +1,101 @@
+use async_trait::async_trait;
+use helius_laserstream::{
+    grpc::SubscribeRequest, LaserstreamError,
+    StreamHandle as HeliusStreamHandle,
+};
+
+use super::LaserStream;
+use crate::remote_account_provider::RemoteAccountProviderResult;
+
+/// A trait to represent the [HeliusStreamHandle].
+/// This is needed since we cannot create the helius one since
+/// [helius_laserstream::StreamHandle::write_tx] is private and there is no constructor.
+#[async_trait]
+pub trait StreamHandle {
+    /// Send a new subscription request to update the active subscription.
+    async fn write(
+        &self,
+        request: SubscribeRequest,
+    ) -> Result<(), LaserstreamError>;
+}
+
+pub struct StreamHandleImpl {
+    pub handle: HeliusStreamHandle,
+}
+
+#[async_trait]
+impl StreamHandle for StreamHandleImpl {
+    async fn write(
+        &self,
+        request: SubscribeRequest,
+    ) -> Result<(), LaserstreamError> {
+        // This async operation gets forwarded to the underlying subscription sender of the laser
+        // client and completes after the given item has been fully processed into the sink,
+        // including flushing.
+        // The assumption is that at that point it has been processed on the receiver end and the
+        // subscription is updated.
+        // See: https://github.com/helius-labs/laserstream-sdk/blob/v0.2.2/rust/src/client.rs#L196-L201
+        self.handle.write(request).await
+    }
+}
+
+/// Abstraction over stream creation for testability
+#[async_trait]
+pub trait StreamFactory<S: StreamHandle>: Send + Sync + 'static {
+    /// Create a stream for the given subscription request.
+    async fn subscribe(
+        &self,
+        request: SubscribeRequest,
+    ) -> RemoteAccountProviderResult<LaserStreamWithHandle<S>>;
+}
+
+pub struct LaserStreamWithHandle<S: StreamHandle> {
+    pub(crate) stream: LaserStream,
+    pub(crate) handle: S,
+}
+
+/// Production stream factory that wraps helius client subscribe
+pub struct StreamFactoryImpl {
+    config: helius_laserstream::LaserstreamConfig,
+}
+
+impl StreamFactoryImpl {
+    pub fn new(config: helius_laserstream::LaserstreamConfig) -> Self {
+        Self { config }
+    }
+}
+
+#[async_trait]
+impl StreamFactory<StreamHandleImpl> for StreamFactoryImpl {
+    /// This implementation creates the underlying gRPC stream with the request (which
+    /// returns immediately) and then writes the same request to the handle so it is sent
+    /// over the network before returning asynchronously.
+    async fn subscribe(
+        &self,
+        request: SubscribeRequest,
+    ) -> RemoteAccountProviderResult<LaserStreamWithHandle<StreamHandleImpl>>
+    {
+        // NOTE: this call returns immediately yielding subscription errors and account updates
+        // via the stream, thus the subscription has not been received yet upstream
+        // NOTE: we need to use the same request as otherwise there is a potential race condition
+        // where our `write` below completes before this call and the request is overwritten
+        // with an empty one.
+        // Inside helius_laserstream::client::subscribe the request is cloned for that reason to
+        // be able to attempt the request multiple times during reconnect attempts.
+        // Given our requests contain a max of about 2_000 pubkeys, an extra clone here is a small
+        // price to pay to avoid this race condition.
+        let (stream, handle) = helius_laserstream::client::subscribe(
+            self.config.clone(),
+            request.clone(),
+        );
+        let handle = StreamHandleImpl { handle };
+        // Write to the handle and await it which at least guarantees that it has
+        // been sent over the network, even though there is still no guarantee it has been
+        // processed and that the subscription became active immediately
+        super::write_with_retry(&handle, "subscribe", request).await?;
+        Ok(LaserStreamWithHandle {
+            stream: Box::pin(stream),
+            handle,
+        })
+    }
+}