sui_indexer_alt_framework/ingestion/
streaming_client.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::time::Duration;
5
6use anyhow::Context;
7use anyhow::anyhow;
8use async_trait::async_trait;
9use futures::StreamExt;
10use futures::stream::BoxStream;
11use sui_rpc::headers::X_SUI_CHAIN_ID;
12use sui_rpc::proto::sui::rpc::v2::SubscribeCheckpointsRequest;
13use sui_rpc::proto::sui::rpc::v2::subscription_service_client::SubscriptionServiceClient;
14use sui_types::digests::ChainIdentifier;
15use sui_types::messages_checkpoint::CheckpointDigest;
16use tokio_stream::adapters::Peekable;
17use tonic::Status;
18use tonic::transport::Endpoint;
19use tonic::transport::Uri;
20
21use crate::ingestion::MAX_GRPC_MESSAGE_SIZE_BYTES;
22use crate::ingestion::error::Error;
23use crate::ingestion::error::Result;
24use crate::types::full_checkpoint_content::Checkpoint;
25
26pub struct CheckpointStream {
27    pub stream: Peekable<BoxStream<'static, Result<Checkpoint>>>,
28    pub chain_id: ChainIdentifier,
29}
30
31/// Trait representing a client for streaming checkpoint data.
32#[async_trait]
33pub trait CheckpointStreamingClient {
34    /// Returns the CheckpointStream and chain id.
35    async fn connect(&self) -> Result<CheckpointStream>;
36
37    /// Returns the latest checkpoint number available from the streaming source.
38    async fn latest_checkpoint_number(&self) -> Result<u64> {
39        let mut stream = self.connect().await?;
40
41        match stream.stream.next().await {
42            Some(Ok(checkpoint)) => Ok(checkpoint.summary.sequence_number),
43            Some(Err(e)) => Err(e),
44            None => Err(Error::StreamingError(anyhow!("Stream ended unexpectedly"))),
45        }
46    }
47}
48
49#[derive(clap::Args, Clone, Debug, Default)]
50pub struct StreamingClientArgs {
51    /// gRPC endpoint for streaming checkpoints
52    #[clap(long, env)]
53    pub streaming_url: Option<Uri>,
54}
55
56/// gRPC-based implementation of the CheckpointStreamingClient trait.
57#[derive(Clone)]
58pub struct GrpcStreamingClient {
59    uri: Uri,
60    connection_timeout: Duration,
61    statement_timeout: Duration,
62}
63
64impl GrpcStreamingClient {
65    pub fn new(uri: Uri, connection_timeout: Duration, statement_timeout: Duration) -> Self {
66        Self {
67            uri,
68            connection_timeout,
69            statement_timeout,
70        }
71    }
72}
73
74#[async_trait]
75impl CheckpointStreamingClient for GrpcStreamingClient {
76    async fn connect(&self) -> Result<CheckpointStream> {
77        let endpoint = Endpoint::from(self.uri.clone())
78            .connect_timeout(self.connection_timeout)
79            .timeout(self.connection_timeout);
80
81        let mut client = SubscriptionServiceClient::connect(endpoint)
82            .await
83            .map_err(|err| Error::RpcClientError(Status::from_error(err.into())))?
84            .max_decoding_message_size(MAX_GRPC_MESSAGE_SIZE_BYTES);
85
86        let mut request = SubscribeCheckpointsRequest::default();
87        request.read_mask = Some(Checkpoint::proto_field_mask());
88
89        let response = client
90            .subscribe_checkpoints(request)
91            .await
92            .map_err(Error::RpcClientError)?;
93
94        let chain_id_value = response.metadata().get(X_SUI_CHAIN_ID).ok_or_else(|| {
95            Error::StreamingError(anyhow!("Chain ID not found in response metadata"))
96        })?;
97        let chain_id: ChainIdentifier = chain_id_value
98            .to_str()
99            .map_err(|e| Error::StreamingError(anyhow!("Chain ID is not valid ASCII: {e}")))?
100            .parse::<CheckpointDigest>()
101            .map_err(|e| Error::StreamingError(anyhow!("Chain ID parse error: {e}")))?
102            .into();
103
104        let stream = response
105            .into_inner()
106            .map(|result| async move {
107                match result {
108                    Ok(response) => {
109                        let checkpoint = response
110                            .checkpoint
111                            .context("Checkpoint data missing in response")
112                            .map_err(Error::StreamingError)?;
113                        // Proto -> Checkpoint conversion is multi-ms of CPU work;
114                        // offload to the blocking pool so it doesn't stall the reactor.
115                        // Combined with `.buffered(4)` below, up to 4 decodes can run
116                        // concurrently while new bytes keep flowing from gRPC.
117                        tokio::task::spawn_blocking(move || {
118                            Checkpoint::try_from(&checkpoint).context("Failed to parse checkpoint")
119                        })
120                        .await
121                        .map_err(|e| Error::StreamingError(anyhow!("decode task panicked: {e}")))?
122                        .map_err(Error::StreamingError)
123                    }
124                    Err(e) => Err(Error::RpcClientError(e)),
125                }
126            })
127            .buffered(4);
128        let stream = wrap_stream(stream, self.statement_timeout);
129
130        Ok(CheckpointStream { stream, chain_id })
131    }
132}
133
134/// Wraps a stream with a per-item timeout. Converts the resulting `Err(Elapsed)` into
135/// `Err(StreamingError)` if it occurs.
136fn wrap_stream(
137    stream: impl futures::Stream<Item = Result<Checkpoint>> + Send + 'static,
138    statement_timeout: Duration,
139) -> Peekable<BoxStream<'static, Result<Checkpoint>>> {
140    let stream = tokio_stream::StreamExt::timeout(stream, statement_timeout)
141        .map(move |result| match result {
142            Err(_elapsed) => Err(Error::StreamingError(anyhow!(
143                "Statement timeout after {statement_timeout:?}"
144            ))),
145            Ok(result) => result,
146        })
147        .boxed();
148    tokio_stream::StreamExt::peekable(stream)
149}
150
151#[cfg(test)]
152mod tests {
153    use std::net::SocketAddr;
154    use std::time::Duration;
155
156    use sui_rpc::proto::sui::rpc::v2::SubscribeCheckpointsRequest;
157    use sui_rpc::proto::sui::rpc::v2::SubscribeCheckpointsResponse;
158    use sui_rpc::proto::sui::rpc::v2::subscription_service_server::SubscriptionService;
159    use sui_rpc::proto::sui::rpc::v2::subscription_service_server::SubscriptionServiceServer;
160    use tonic::transport::Server;
161
162    use super::*;
163
164    /// A gRPC server that accepts connections but never responds to
165    /// subscribe_checkpoints, simulating a stalled RPC handshake.
166    struct HangingSubscriptionService;
167
168    #[tonic::async_trait]
169    impl SubscriptionService for HangingSubscriptionService {
170        async fn subscribe_checkpoints(
171            &self,
172            _request: tonic::Request<SubscribeCheckpointsRequest>,
173        ) -> std::result::Result<
174            tonic::Response<
175                BoxStream<'static, std::result::Result<SubscribeCheckpointsResponse, Status>>,
176            >,
177            Status,
178        > {
179            futures::future::pending().await
180        }
181    }
182
183    #[tokio::test]
184    async fn subscribe_checkpoints_times_out_on_stalled_server() {
185        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
186        let addr: SocketAddr = listener.local_addr().unwrap();
187
188        tokio::spawn(async move {
189            let incoming = tokio_stream::wrappers::TcpListenerStream::new(listener);
190            Server::builder()
191                .add_service(SubscriptionServiceServer::new(HangingSubscriptionService))
192                .serve_with_incoming(incoming)
193                .await
194                .unwrap();
195        });
196
197        let timeout = Duration::from_millis(200);
198        let uri: Uri = format!("http://{addr}").parse().unwrap();
199        let client = GrpcStreamingClient::new(uri, timeout, timeout);
200
201        let start = std::time::Instant::now();
202        let result = client.connect().await;
203        let elapsed = start.elapsed();
204
205        assert!(result.is_err(), "expected timeout error");
206        assert!(
207            elapsed < Duration::from_secs(5),
208            "connect() took {elapsed:?}, should have timed out in ~200ms"
209        );
210    }
211}
212
213#[cfg(test)]
214pub mod test_utils {
215    use std::pin::Pin;
216    use std::sync::Arc;
217    use std::sync::Mutex;
218    use std::sync::atomic::AtomicUsize;
219    use std::sync::atomic::Ordering;
220    use std::time::Duration;
221    use std::time::Instant;
222
223    use futures::Stream;
224
225    use crate::types::test_checkpoint_data_builder::TestCheckpointBuilder;
226
227    use super::*;
228
229    enum StreamAction {
230        Checkpoint(u64),
231        Error,
232        Timeout {
233            deadline: Option<Instant>,
234            duration: Duration,
235        },
236    }
237
238    struct MockStreamState {
239        actions: Arc<Mutex<Vec<StreamAction>>>,
240    }
241
242    impl Stream for MockStreamState {
243        type Item = Result<Checkpoint>;
244
245        fn poll_next(
246            self: Pin<&mut Self>,
247            _cx: &mut std::task::Context<'_>,
248        ) -> std::task::Poll<Option<Self::Item>> {
249            let mut actions = self.actions.lock().unwrap();
250            if actions.is_empty() {
251                return std::task::Poll::Ready(None);
252            }
253
254            match &actions[0] {
255                StreamAction::Checkpoint(seq) => {
256                    let seq = *seq;
257                    actions.remove(0);
258                    let mut builder = TestCheckpointBuilder::new(seq);
259                    std::task::Poll::Ready(Some(Ok(builder.build_checkpoint())))
260                }
261                StreamAction::Error => {
262                    actions.remove(0);
263                    std::task::Poll::Ready(Some(Err(Error::StreamingError(anyhow::anyhow!(
264                        "Mock streaming error"
265                    )))))
266                }
267                StreamAction::Timeout { deadline, duration } => match deadline {
268                    None => {
269                        let deadline = Instant::now() + *duration;
270                        actions[0] = StreamAction::Timeout {
271                            deadline: Some(deadline),
272                            duration: *duration,
273                        };
274                        std::task::Poll::Pending
275                    }
276                    Some(deadline_instant) => {
277                        if Instant::now() >= *deadline_instant {
278                            actions.remove(0);
279                            drop(actions);
280                            self.poll_next(_cx)
281                        } else {
282                            std::task::Poll::Pending
283                        }
284                    }
285                },
286            }
287        }
288    }
289
290    /// Mock streaming client for testing with predefined checkpoints.
291    pub struct MockStreamingClient {
292        actions: Arc<Mutex<Vec<StreamAction>>>,
293        connection_failures_remaining: AtomicUsize,
294        connection_timeouts_remaining: AtomicUsize,
295        /// How long mock timeout actions hang (must be > statement_timeout for timeouts to fire).
296        timeout_duration: Duration,
297        /// Statement timeout applied to the stream wrapper.
298        statement_timeout: Duration,
299    }
300
301    impl MockStreamingClient {
302        pub fn mock_chain_id() -> ChainIdentifier {
303            CheckpointDigest::new([1; 32]).into()
304        }
305
306        pub fn new<I>(checkpoint_range: I, timeout_duration: Option<Duration>) -> Self
307        where
308            I: IntoIterator<Item = u64>,
309        {
310            let timeout_duration = timeout_duration.unwrap_or(Duration::from_secs(5));
311            Self {
312                actions: Arc::new(Mutex::new(
313                    checkpoint_range
314                        .into_iter()
315                        .map(StreamAction::Checkpoint)
316                        .collect(),
317                )),
318                connection_failures_remaining: AtomicUsize::new(0),
319                connection_timeouts_remaining: AtomicUsize::new(0),
320                statement_timeout: timeout_duration / 2,
321                timeout_duration,
322            }
323        }
324
325        /// Make `connect` fail for the next N calls
326        pub fn fail_connection_times(mut self, times: usize) -> Self {
327            self.connection_failures_remaining = AtomicUsize::new(times);
328            self
329        }
330
331        /// Make `connect` timeout for the next N calls
332        pub fn fail_connection_with_timeout(mut self, times: usize) -> Self {
333            self.connection_timeouts_remaining = AtomicUsize::new(times);
334            self
335        }
336
337        /// Insert an error at the back of the queue.
338        pub fn insert_error(&mut self) {
339            self.actions.lock().unwrap().push(StreamAction::Error);
340        }
341
342        /// Insert a timeout at the back of the queue (causes poll_next to return Pending).
343        pub fn insert_timeout(&mut self) {
344            self.insert_timeout_with_duration(self.timeout_duration)
345        }
346
347        /// Insert a timeout with custom duration.
348        pub fn insert_timeout_with_duration(&mut self, duration: Duration) {
349            self.actions.lock().unwrap().push(StreamAction::Timeout {
350                deadline: None,
351                duration,
352            });
353        }
354
355        /// Insert a checkpoint at the back of the queue.
356        pub fn insert_checkpoint(&mut self, sequence_number: u64) {
357            self.insert_checkpoint_range([sequence_number])
358        }
359
360        pub fn insert_checkpoint_range<I>(&mut self, checkpoint_range: I)
361        where
362            I: IntoIterator<Item = u64>,
363        {
364            let mut actions = self.actions.lock().unwrap();
365            for sequence_number in checkpoint_range {
366                actions.push(StreamAction::Checkpoint(sequence_number));
367            }
368        }
369    }
370
371    #[async_trait]
372    impl CheckpointStreamingClient for MockStreamingClient {
373        async fn connect(&self) -> Result<CheckpointStream> {
374            if self.connection_timeouts_remaining.load(Ordering::Relaxed) > 0 {
375                self.connection_timeouts_remaining
376                    .fetch_sub(1, Ordering::Relaxed);
377                // Simulate a connection timeout
378                tokio::time::sleep(self.timeout_duration).await;
379                return Err(Error::StreamingError(anyhow::anyhow!(
380                    "Mock connection timeout"
381                )));
382            }
383            if self.connection_failures_remaining.load(Ordering::Relaxed) > 0 {
384                self.connection_failures_remaining
385                    .fetch_sub(1, Ordering::Relaxed);
386                return Err(Error::StreamingError(anyhow::anyhow!(
387                    "Mock connection failure"
388                )));
389            }
390            let stream_state = MockStreamState {
391                actions: Arc::clone(&self.actions),
392            };
393            Ok(CheckpointStream {
394                stream: wrap_stream(stream_state, self.statement_timeout),
395                chain_id: Self::mock_chain_id(),
396            })
397        }
398    }
399}
sui_indexer_alt_framework/ingestion/streaming_client.rs

sui_indexer_alt_framework/ingestion/
streaming_client.rs
