sui_indexer_alt_framework/ingestion/
streaming_client.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::sync::Arc;
5use std::time::Duration;
6
7use anyhow::Context;
8use anyhow::anyhow;
9use async_trait::async_trait;
10use futures::StreamExt;
11use futures::stream::BoxStream;
12use sui_rpc::headers::X_SUI_CHAIN_ID;
13use sui_rpc::proto::sui::rpc::v2::SubscribeCheckpointsRequest;
14use sui_rpc::proto::sui::rpc::v2::subscription_service_client::SubscriptionServiceClient;
15use sui_types::digests::ChainIdentifier;
16use sui_types::messages_checkpoint::CheckpointDigest;
17use tokio_stream::adapters::Peekable;
18use tonic::Code;
19use tonic::Status;
20use tonic::transport::Endpoint;
21use tonic::transport::Uri;
22use tracing::warn;
23use url::Url;
24
25use crate::ingestion::MAX_GRPC_MESSAGE_SIZE_BYTES;
26use crate::ingestion::error::Error;
27use crate::ingestion::error::Result;
28use crate::ingestion::ingestion_client::IngestionClient;
29use crate::metrics::IngestionMetrics;
30use crate::types::full_checkpoint_content::Checkpoint;
31
32pub struct CheckpointStream {
33    pub stream: Peekable<BoxStream<'static, Result<Checkpoint>>>,
34    pub chain_id: ChainIdentifier,
35}
36
37/// Trait representing a client for streaming checkpoint data.
38#[async_trait]
39pub trait CheckpointStreamingClient {
40    /// Returns the CheckpointStream and chain id.
41    async fn connect(&self) -> Result<CheckpointStream>;
42
43    /// Returns the latest checkpoint number available from the streaming source.
44    async fn latest_checkpoint_number(&self) -> Result<u64> {
45        let mut stream = self.connect().await?;
46
47        match stream.stream.next().await {
48            Some(Ok(checkpoint)) => Ok(checkpoint.summary.sequence_number),
49            Some(Err(e)) => Err(e),
50            None => Err(Error::StreamingError(anyhow!("Stream ended unexpectedly"))),
51        }
52    }
53}
54
55#[derive(clap::Args, Clone, Debug, Default)]
56pub struct StreamingClientArgs {
57    /// gRPC endpoint for streaming checkpoints. If it also serves `LedgerService`, the checkpoints
58    /// between the indexer and the start of a newly connected stream are fetched from it, rather
59    /// than from the ingestion source.
60    #[clap(long, env)]
61    pub streaming_url: Option<Uri>,
62}
63
64/// gRPC-based implementation of the CheckpointStreamingClient trait.
65#[derive(Clone)]
66pub struct GrpcStreamingClient {
67    uri: Uri,
68    connection_timeout: Duration,
69    statement_timeout: Duration,
70}
71
72impl StreamingClientArgs {
73    /// The client that fills the gap in front of a newly connected stream, fetching checkpoints
74    /// from the streaming endpoint's `LedgerService` and reporting to `metrics`.
75    ///
76    /// `None` if there is no streaming endpoint, or if it refuses a ledger request (because it
77    /// does not serve `LedgerService`, or requires credentials), in which case stream gaps are
78    /// filled from the ingestion source. An endpoint that does not answer within
79    /// `connection_timeout` keeps its client, because the client is only used right after a
80    /// stream to that endpoint connects.
81    pub async fn gap_fill_client(
82        &self,
83        connection_timeout: Duration,
84        metrics: Arc<IngestionMetrics>,
85    ) -> Result<Option<IngestionClient>> {
86        let Some(uri) = &self.streaming_url else {
87            return Ok(None);
88        };
89
90        let url = Url::parse(&uri.to_string())
91            .with_context(|| format!("Invalid streaming URL for gap filling: {uri}"))
92            .map_err(Error::StreamingError)?;
93
94        let client = IngestionClient::with_grpc(url, None, None, metrics)?;
95        match tokio::time::timeout(connection_timeout, client.latest_checkpoint_number()).await {
96            Ok(Err(error)) if refuses_ledger_requests(&error) => {
97                warn!(
98                    %uri,
99                    "Streaming endpoint refused a ledger request, filling stream gaps from the \
100                     ingestion source instead: {error}"
101                );
102                Ok(None)
103            }
104            _ => Ok(Some(client)),
105        }
106    }
107}
108
109impl GrpcStreamingClient {
110    pub fn new(uri: Uri, connection_timeout: Duration, statement_timeout: Duration) -> Self {
111        Self {
112            uri,
113            connection_timeout,
114            statement_timeout,
115        }
116    }
117}
118
119#[async_trait]
120impl CheckpointStreamingClient for GrpcStreamingClient {
121    async fn connect(&self) -> Result<CheckpointStream> {
122        let endpoint = Endpoint::from(self.uri.clone())
123            .connect_timeout(self.connection_timeout)
124            .timeout(self.connection_timeout);
125
126        let mut client = SubscriptionServiceClient::connect(endpoint)
127            .await
128            .map_err(|err| Error::RpcClientError(Status::from_error(err.into())))?
129            .max_decoding_message_size(MAX_GRPC_MESSAGE_SIZE_BYTES);
130
131        let mut request = SubscribeCheckpointsRequest::default();
132        request.read_mask = Some(Checkpoint::proto_field_mask());
133
134        let response = client
135            .subscribe_checkpoints(request)
136            .await
137            .map_err(Error::RpcClientError)?;
138
139        let chain_id_value = response.metadata().get(X_SUI_CHAIN_ID).ok_or_else(|| {
140            Error::StreamingError(anyhow!("Chain ID not found in response metadata"))
141        })?;
142        let chain_id: ChainIdentifier = chain_id_value
143            .to_str()
144            .map_err(|e| Error::StreamingError(anyhow!("Chain ID is not valid ASCII: {e}")))?
145            .parse::<CheckpointDigest>()
146            .map_err(|e| Error::StreamingError(anyhow!("Chain ID parse error: {e}")))?
147            .into();
148
149        let stream = response
150            .into_inner()
151            .map(|result| async move {
152                match result {
153                    Ok(response) => {
154                        let checkpoint = response
155                            .checkpoint
156                            .context("Checkpoint data missing in response")
157                            .map_err(Error::StreamingError)?;
158                        // Proto -> Checkpoint conversion is multi-ms of CPU work;
159                        // offload to the blocking pool so it doesn't stall the reactor.
160                        // Combined with `.buffered(4)` below, up to 4 decodes can run
161                        // concurrently while new bytes keep flowing from gRPC.
162                        tokio::task::spawn_blocking(move || {
163                            Checkpoint::try_from(&checkpoint).context("Failed to parse checkpoint")
164                        })
165                        .await
166                        .map_err(|e| Error::StreamingError(anyhow!("decode task panicked: {e}")))?
167                        .map_err(Error::StreamingError)
168                    }
169                    Err(e) => Err(Error::RpcClientError(e)),
170                }
171            })
172            .buffered(4);
173        let stream = wrap_stream(stream, self.statement_timeout);
174
175        Ok(CheckpointStream { stream, chain_id })
176    }
177}
178
179/// Wraps a stream with a per-item timeout. Converts the resulting `Err(Elapsed)` into
180/// `Err(StreamingError)` if it occurs.
181fn wrap_stream(
182    stream: impl futures::Stream<Item = Result<Checkpoint>> + Send + 'static,
183    statement_timeout: Duration,
184) -> Peekable<BoxStream<'static, Result<Checkpoint>>> {
185    let stream = tokio_stream::StreamExt::timeout(stream, statement_timeout)
186        .map(move |result| match result {
187            Err(_elapsed) => Err(Error::StreamingError(anyhow!(
188                "Statement timeout after {statement_timeout:?}"
189            ))),
190            Ok(result) => result,
191        })
192        .boxed();
193    tokio_stream::StreamExt::peekable(stream)
194}
195
196/// Whether `error` (from a ledger request) means the endpoint will never serve ledger requests to
197/// this client, as opposed to failing transiently.
198fn refuses_ledger_requests(error: &anyhow::Error) -> bool {
199    error.downcast_ref::<Status>().is_some_and(|status| {
200        matches!(
201            status.code(),
202            Code::Unimplemented | Code::PermissionDenied | Code::Unauthenticated
203        )
204    })
205}
206
207#[cfg(test)]
208mod tests {
209    use std::net::SocketAddr;
210    use std::time::Duration;
211
212    use sui_rpc::proto::sui::rpc::v2::SubscribeCheckpointsRequest;
213    use sui_rpc::proto::sui::rpc::v2::SubscribeCheckpointsResponse;
214    use sui_rpc::proto::sui::rpc::v2::subscription_service_server::SubscriptionService;
215    use sui_rpc::proto::sui::rpc::v2::subscription_service_server::SubscriptionServiceServer;
216    use tonic::transport::Server;
217
218    use crate::metrics::tests::test_ingestion_metrics;
219
220    use super::*;
221
222    /// A gRPC server that accepts connections but never responds to
223    /// subscribe_checkpoints, simulating a stalled RPC handshake.
224    struct HangingSubscriptionService;
225
226    #[tonic::async_trait]
227    impl SubscriptionService for HangingSubscriptionService {
228        async fn subscribe_checkpoints(
229            &self,
230            _request: tonic::Request<SubscribeCheckpointsRequest>,
231        ) -> std::result::Result<
232            tonic::Response<
233                BoxStream<'static, std::result::Result<SubscribeCheckpointsResponse, Status>>,
234            >,
235            Status,
236        > {
237            futures::future::pending().await
238        }
239    }
240
241    #[tokio::test]
242    async fn subscribe_checkpoints_times_out_on_stalled_server() {
243        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
244        let addr: SocketAddr = listener.local_addr().unwrap();
245
246        tokio::spawn(async move {
247            let incoming = tokio_stream::wrappers::TcpListenerStream::new(listener);
248            Server::builder()
249                .add_service(SubscriptionServiceServer::new(HangingSubscriptionService))
250                .serve_with_incoming(incoming)
251                .await
252                .unwrap();
253        });
254
255        let timeout = Duration::from_millis(200);
256        let uri: Uri = format!("http://{addr}").parse().unwrap();
257        let client = GrpcStreamingClient::new(uri, timeout, timeout);
258
259        let start = std::time::Instant::now();
260        let result = client.connect().await;
261        let elapsed = start.elapsed();
262
263        assert!(result.is_err(), "expected timeout error");
264        assert!(
265            elapsed < Duration::from_secs(5),
266            "connect() took {elapsed:?}, should have timed out in ~200ms"
267        );
268    }
269
270    /// A streaming endpoint without `LedgerService` yields no gap-fill client, so stream gaps
271    /// are filled from the ingestion source.
272    #[tokio::test]
273    async fn no_gap_fill_client_without_ledger_service() {
274        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
275        let addr: SocketAddr = listener.local_addr().unwrap();
276
277        tokio::spawn(async move {
278            let incoming = tokio_stream::wrappers::TcpListenerStream::new(listener);
279            Server::builder()
280                .add_service(SubscriptionServiceServer::new(HangingSubscriptionService))
281                .serve_with_incoming(incoming)
282                .await
283                .unwrap();
284        });
285
286        let args = StreamingClientArgs {
287            streaming_url: Some(format!("http://{addr}").parse().unwrap()),
288        };
289        let client = args
290            .gap_fill_client(Duration::from_secs(5), test_ingestion_metrics())
291            .await
292            .unwrap();
293
294        assert!(client.is_none());
295    }
296
297    /// An unreachable streaming endpoint keeps its gap-fill client, because it is only used
298    /// right after a stream to that endpoint connects.
299    #[tokio::test]
300    async fn gap_fill_client_kept_when_endpoint_unreachable() {
301        // Bind and immediately release a port, so connections to it are refused.
302        let addr = tokio::net::TcpListener::bind("127.0.0.1:0")
303            .await
304            .unwrap()
305            .local_addr()
306            .unwrap();
307
308        let args = StreamingClientArgs {
309            streaming_url: Some(format!("http://{addr}").parse().unwrap()),
310        };
311        let client = args
312            .gap_fill_client(Duration::from_secs(5), test_ingestion_metrics())
313            .await
314            .unwrap();
315
316        assert!(client.is_some());
317    }
318}
319
320#[cfg(test)]
321pub mod test_utils {
322    use std::pin::Pin;
323    use std::sync::Arc;
324    use std::sync::Mutex;
325    use std::sync::atomic::AtomicUsize;
326    use std::sync::atomic::Ordering;
327    use std::time::Duration;
328    use std::time::Instant;
329
330    use futures::Stream;
331
332    use crate::types::test_checkpoint_data_builder::TestCheckpointBuilder;
333
334    use super::*;
335
336    enum StreamAction {
337        Checkpoint(u64),
338        Error,
339        Timeout {
340            deadline: Option<Instant>,
341            duration: Duration,
342        },
343    }
344
345    struct MockStreamState {
346        actions: Arc<Mutex<Vec<StreamAction>>>,
347    }
348
349    impl Stream for MockStreamState {
350        type Item = Result<Checkpoint>;
351
352        fn poll_next(
353            self: Pin<&mut Self>,
354            _cx: &mut std::task::Context<'_>,
355        ) -> std::task::Poll<Option<Self::Item>> {
356            let mut actions = self.actions.lock().unwrap();
357            if actions.is_empty() {
358                return std::task::Poll::Ready(None);
359            }
360
361            match &actions[0] {
362                StreamAction::Checkpoint(seq) => {
363                    let seq = *seq;
364                    actions.remove(0);
365                    let mut builder = TestCheckpointBuilder::new(seq);
366                    std::task::Poll::Ready(Some(Ok(builder.build_checkpoint())))
367                }
368                StreamAction::Error => {
369                    actions.remove(0);
370                    std::task::Poll::Ready(Some(Err(Error::StreamingError(anyhow::anyhow!(
371                        "Mock streaming error"
372                    )))))
373                }
374                StreamAction::Timeout { deadline, duration } => match deadline {
375                    None => {
376                        let deadline = Instant::now() + *duration;
377                        actions[0] = StreamAction::Timeout {
378                            deadline: Some(deadline),
379                            duration: *duration,
380                        };
381                        std::task::Poll::Pending
382                    }
383                    Some(deadline_instant) => {
384                        if Instant::now() >= *deadline_instant {
385                            actions.remove(0);
386                            drop(actions);
387                            self.poll_next(_cx)
388                        } else {
389                            std::task::Poll::Pending
390                        }
391                    }
392                },
393            }
394        }
395    }
396
397    /// Mock streaming client for testing with predefined checkpoints.
398    pub struct MockStreamingClient {
399        actions: Arc<Mutex<Vec<StreamAction>>>,
400        connection_failures_remaining: AtomicUsize,
401        connection_timeouts_remaining: AtomicUsize,
402        /// How long mock timeout actions hang (must be > statement_timeout for timeouts to fire).
403        timeout_duration: Duration,
404        /// Statement timeout applied to the stream wrapper.
405        statement_timeout: Duration,
406    }
407
408    impl MockStreamingClient {
409        pub fn mock_chain_id() -> ChainIdentifier {
410            CheckpointDigest::new([1; 32]).into()
411        }
412
413        pub fn new<I>(checkpoint_range: I, timeout_duration: Option<Duration>) -> Self
414        where
415            I: IntoIterator<Item = u64>,
416        {
417            let timeout_duration = timeout_duration.unwrap_or(Duration::from_secs(5));
418            Self {
419                actions: Arc::new(Mutex::new(
420                    checkpoint_range
421                        .into_iter()
422                        .map(StreamAction::Checkpoint)
423                        .collect(),
424                )),
425                connection_failures_remaining: AtomicUsize::new(0),
426                connection_timeouts_remaining: AtomicUsize::new(0),
427                statement_timeout: timeout_duration / 2,
428                timeout_duration,
429            }
430        }
431
432        /// Make `connect` fail for the next N calls
433        pub fn fail_connection_times(mut self, times: usize) -> Self {
434            self.connection_failures_remaining = AtomicUsize::new(times);
435            self
436        }
437
438        /// Make `connect` timeout for the next N calls
439        pub fn fail_connection_with_timeout(mut self, times: usize) -> Self {
440            self.connection_timeouts_remaining = AtomicUsize::new(times);
441            self
442        }
443
444        /// Insert an error at the back of the queue.
445        pub fn insert_error(&mut self) {
446            self.actions.lock().unwrap().push(StreamAction::Error);
447        }
448
449        /// Insert a timeout at the back of the queue (causes poll_next to return Pending).
450        pub fn insert_timeout(&mut self) {
451            self.insert_timeout_with_duration(self.timeout_duration)
452        }
453
454        /// Insert a timeout with custom duration.
455        pub fn insert_timeout_with_duration(&mut self, duration: Duration) {
456            self.actions.lock().unwrap().push(StreamAction::Timeout {
457                deadline: None,
458                duration,
459            });
460        }
461
462        /// Insert a checkpoint at the back of the queue.
463        pub fn insert_checkpoint(&mut self, sequence_number: u64) {
464            self.insert_checkpoint_range([sequence_number])
465        }
466
467        pub fn insert_checkpoint_range<I>(&mut self, checkpoint_range: I)
468        where
469            I: IntoIterator<Item = u64>,
470        {
471            let mut actions = self.actions.lock().unwrap();
472            for sequence_number in checkpoint_range {
473                actions.push(StreamAction::Checkpoint(sequence_number));
474            }
475        }
476    }
477
478    #[async_trait]
479    impl CheckpointStreamingClient for MockStreamingClient {
480        async fn connect(&self) -> Result<CheckpointStream> {
481            if self.connection_timeouts_remaining.load(Ordering::Relaxed) > 0 {
482                self.connection_timeouts_remaining
483                    .fetch_sub(1, Ordering::Relaxed);
484                // Simulate a connection timeout
485                tokio::time::sleep(self.timeout_duration).await;
486                return Err(Error::StreamingError(anyhow::anyhow!(
487                    "Mock connection timeout"
488                )));
489            }
490            if self.connection_failures_remaining.load(Ordering::Relaxed) > 0 {
491                self.connection_failures_remaining
492                    .fetch_sub(1, Ordering::Relaxed);
493                return Err(Error::StreamingError(anyhow::anyhow!(
494                    "Mock connection failure"
495                )));
496            }
497            let stream_state = MockStreamState {
498                actions: Arc::clone(&self.actions),
499            };
500            Ok(CheckpointStream {
501                stream: wrap_stream(stream_state, self.statement_timeout),
502                chain_id: Self::mock_chain_id(),
503            })
504        }
505    }
506}
sui_indexer_alt_framework/ingestion/streaming_client.rs

sui_indexer_alt_framework/ingestion/
streaming_client.rs
