sui_indexer_alt_framework/ingestion/
streaming_client.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::pin::Pin;
5use std::time::Duration;
6
7use anyhow::Context;
8use async_trait::async_trait;
9use futures::{Stream, StreamExt};
10use sui_rpc::proto::sui::rpc::v2::{
11    SubscribeCheckpointsRequest, subscription_service_client::SubscriptionServiceClient,
12};
13use sui_rpc_api::client::checkpoint_data_field_mask;
14use tonic::{
15    Status,
16    transport::{Endpoint, Uri},
17};
18
19use crate::ingestion::MAX_GRPC_MESSAGE_SIZE_BYTES;
20use crate::ingestion::error::{Error, Result};
21use crate::types::full_checkpoint_content::Checkpoint;
22
23/// Type alias for a stream of checkpoint data.
24pub type CheckpointStream = Pin<Box<dyn Stream<Item = Result<Checkpoint>> + Send>>;
25
26/// Trait representing a client for streaming checkpoint data.
27#[async_trait]
28pub trait CheckpointStreamingClient {
29    async fn connect(&mut self) -> Result<CheckpointStream>;
30}
31
32#[derive(clap::Args, Clone, Debug, Default)]
33pub struct StreamingClientArgs {
34    /// gRPC endpoint for streaming checkpoints
35    #[clap(long, env)]
36    pub streaming_url: Option<Uri>,
37}
38
39/// gRPC-based implementation of the CheckpointStreamingClient trait.
40pub struct GrpcStreamingClient {
41    uri: Uri,
42    connection_timeout: Duration,
43}
44
45impl GrpcStreamingClient {
46    pub fn new(uri: Uri, connection_timeout: Duration) -> Self {
47        Self {
48            uri,
49            connection_timeout,
50        }
51    }
52}
53
54#[async_trait]
55impl CheckpointStreamingClient for GrpcStreamingClient {
56    async fn connect(&mut self) -> Result<CheckpointStream> {
57        let endpoint = Endpoint::from(self.uri.clone()).connect_timeout(self.connection_timeout);
58
59        let mut client = SubscriptionServiceClient::connect(endpoint)
60            .await
61            .map_err(|err| Error::RpcClientError(Status::from_error(err.into())))?
62            .max_decoding_message_size(MAX_GRPC_MESSAGE_SIZE_BYTES);
63
64        let mut request = SubscribeCheckpointsRequest::default();
65        request.read_mask = Some(checkpoint_data_field_mask());
66
67        let stream = client
68            .subscribe_checkpoints(request)
69            .await
70            .map_err(Error::RpcClientError)?
71            .into_inner();
72
73        let converted_stream = stream.map(|result| match result {
74            Ok(response) => response
75                .checkpoint
76                .context("Checkpoint data missing in response")
77                .and_then(|checkpoint| {
78                    sui_types::full_checkpoint_content::Checkpoint::try_from(&checkpoint)
79                        .context("Failed to parse checkpoint")
80                })
81                .map_err(Error::StreamingError),
82            Err(e) => Err(Error::RpcClientError(e)),
83        });
84
85        Ok(Box::pin(converted_stream))
86    }
87}
88
89#[cfg(test)]
90pub mod test_utils {
91    use super::*;
92    use crate::types::test_checkpoint_data_builder::TestCheckpointBuilder;
93    use std::sync::{Arc, Mutex};
94    use std::time::{Duration, Instant};
95
96    enum StreamAction {
97        Checkpoint(u64),
98        Error,
99        Timeout {
100            deadline: Option<Instant>,
101            duration: Duration,
102        },
103    }
104
105    struct MockStreamState {
106        actions: Arc<Mutex<Vec<StreamAction>>>,
107    }
108
109    impl Stream for MockStreamState {
110        type Item = Result<Checkpoint>;
111
112        fn poll_next(
113            self: Pin<&mut Self>,
114            _cx: &mut std::task::Context<'_>,
115        ) -> std::task::Poll<Option<Self::Item>> {
116            let mut actions = self.actions.lock().unwrap();
117            if actions.is_empty() {
118                return std::task::Poll::Ready(None);
119            }
120
121            match &actions[0] {
122                StreamAction::Checkpoint(seq) => {
123                    let seq = *seq;
124                    actions.remove(0);
125                    let mut builder = TestCheckpointBuilder::new(seq);
126                    std::task::Poll::Ready(Some(Ok(builder.build_checkpoint())))
127                }
128                StreamAction::Error => {
129                    actions.remove(0);
130                    std::task::Poll::Ready(Some(Err(Error::StreamingError(anyhow::anyhow!(
131                        "Mock streaming error"
132                    )))))
133                }
134                StreamAction::Timeout { deadline, duration } => match deadline {
135                    None => {
136                        let deadline = Instant::now() + *duration;
137                        actions[0] = StreamAction::Timeout {
138                            deadline: Some(deadline),
139                            duration: *duration,
140                        };
141                        std::task::Poll::Pending
142                    }
143                    Some(deadline_instant) => {
144                        if Instant::now() >= *deadline_instant {
145                            actions.remove(0);
146                            drop(actions);
147                            self.poll_next(_cx)
148                        } else {
149                            std::task::Poll::Pending
150                        }
151                    }
152                },
153            }
154        }
155    }
156
157    /// Mock streaming client for testing with predefined checkpoints.
158    pub struct MockStreamingClient {
159        actions: Arc<Mutex<Vec<StreamAction>>>,
160        connection_failures_remaining: usize,
161        connection_timeouts_remaining: usize,
162        timeout_duration: Duration,
163    }
164
165    impl MockStreamingClient {
166        pub fn new<I>(checkpoint_range: I, timeout_duration: Option<Duration>) -> Self
167        where
168            I: IntoIterator<Item = u64>,
169        {
170            Self {
171                actions: Arc::new(Mutex::new(
172                    checkpoint_range
173                        .into_iter()
174                        .map(StreamAction::Checkpoint)
175                        .collect(),
176                )),
177                connection_failures_remaining: 0,
178                connection_timeouts_remaining: 0,
179                timeout_duration: timeout_duration.unwrap_or(Duration::from_secs(5)),
180            }
181        }
182
183        /// Make `connect` fail for the next N calls
184        pub fn fail_connection_times(mut self, times: usize) -> Self {
185            self.connection_failures_remaining = times;
186            self
187        }
188
189        /// Make `connect` timeout for the next N calls
190        pub fn fail_connection_with_timeout(mut self, times: usize) -> Self {
191            self.connection_timeouts_remaining = times;
192            self
193        }
194
195        /// Insert an error at the back of the queue.
196        pub fn insert_error(&mut self) {
197            self.actions.lock().unwrap().push(StreamAction::Error);
198        }
199
200        /// Insert a timeout at the back of the queue (causes poll_next to return Pending).
201        pub fn insert_timeout(&mut self) {
202            self.insert_timeout_with_duration(self.timeout_duration)
203        }
204
205        /// Insert a timeout with custom duration.
206        pub fn insert_timeout_with_duration(&mut self, duration: Duration) {
207            self.actions.lock().unwrap().push(StreamAction::Timeout {
208                deadline: None,
209                duration,
210            });
211        }
212
213        /// Insert a checkpoint at the back of the queue.
214        pub fn insert_checkpoint(&mut self, sequence_number: u64) {
215            self.insert_checkpoint_range([sequence_number])
216        }
217
218        pub fn insert_checkpoint_range<I>(&mut self, checkpoint_range: I)
219        where
220            I: IntoIterator<Item = u64>,
221        {
222            let mut actions = self.actions.lock().unwrap();
223            for sequence_number in checkpoint_range {
224                actions.push(StreamAction::Checkpoint(sequence_number));
225            }
226        }
227    }
228
229    #[async_trait]
230    impl CheckpointStreamingClient for MockStreamingClient {
231        async fn connect(&mut self) -> Result<CheckpointStream> {
232            if self.connection_timeouts_remaining > 0 {
233                self.connection_timeouts_remaining -= 1;
234                // Simulate a connection timeout
235                tokio::time::sleep(self.timeout_duration).await;
236                return Err(Error::StreamingError(anyhow::anyhow!(
237                    "Mock connection timeout"
238                )));
239            }
240            if self.connection_failures_remaining > 0 {
241                self.connection_failures_remaining -= 1;
242                return Err(Error::StreamingError(anyhow::anyhow!(
243                    "Mock connection failure"
244                )));
245            }
246            let stream = MockStreamState {
247                actions: Arc::clone(&self.actions),
248            };
249
250            Ok(Box::pin(stream))
251        }
252    }
253}
sui_indexer_alt_framework/ingestion/streaming_client.rs

sui_indexer_alt_framework/ingestion/
streaming_client.rs
