sui_indexer_alt_framework/ingestion/
store_client.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::sync::Arc;
5
6use anyhow::Context;
7use bytes::Bytes;
8use object_store::Error;
9use object_store::ObjectStore;
10use object_store::ObjectStoreExt;
11use object_store::RetryConfig;
12use object_store::path::Path as ObjectPath;
13use prometheus::IntCounter;
14use serde::de::DeserializeOwned;
15use sui_types::digests::ChainIdentifier;
16
17use crate::ingestion::decode;
18use crate::ingestion::ingestion_client::CheckpointError;
19use crate::ingestion::ingestion_client::CheckpointResult;
20use crate::ingestion::ingestion_client::IngestionClientTrait;
21use crate::types::full_checkpoint_content::Checkpoint;
22
23/// Watermark of the `checkpoint_blob` pipeline (`sui-checkpoint-blob-indexer`), which writes the
24/// `{checkpoint}.binpb.zst` files this client reads, at the path `sui-indexer-alt-object-store`
25/// stores it.
26pub(crate) const WATERMARK_PATH: &str = "_metadata/watermarks/checkpoint_blob.json";
27
28pub struct StoreIngestionClient {
29    store: Arc<dyn ObjectStore>,
30    /// Counter incremented (in the [`IngestionClientTrait`] impl) by the size in bytes of each
31    /// fetched checkpoint payload. `None` for callers that only use this client for one-shot
32    /// metadata fetches (e.g. `end_of_epoch_checkpoints`) and don't need a metric.
33    total_ingested_bytes: Option<IntCounter>,
34}
35
36#[derive(serde::Deserialize, serde::Serialize)]
37pub(crate) struct ObjectStoreWatermark {
38    /// `None` until the writer commits its first checkpoint.
39    pub checkpoint_hi_inclusive: Option<u64>,
40}
41
42impl StoreIngestionClient {
43    pub fn new(store: Arc<dyn ObjectStore>, total_ingested_bytes: Option<IntCounter>) -> Self {
44        Self {
45            store,
46            total_ingested_bytes,
47        }
48    }
49
50    /// Fetch metadata mapping epoch IDs to the sequence numbers of their last checkpoints.
51    /// The response is a JSON-encoded array of checkpoint sequence numbers.
52    pub async fn end_of_epoch_checkpoints<T: DeserializeOwned>(&self) -> anyhow::Result<T> {
53        let bytes = self.bytes(ObjectPath::from("epochs.json")).await?;
54        Ok(serde_json::from_slice(&bytes)?)
55    }
56
57    /// Fetch and decode checkpoint data by sequence number.
58    pub async fn checkpoint(&self, checkpoint: u64) -> anyhow::Result<Checkpoint> {
59        let bytes = self.checkpoint_bytes(checkpoint).await?;
60        // zstd decompress + prost decode + proto -> Checkpoint is multi-ms of CPU
61        // work; offload to the blocking pool so it doesn't stall the reactor.
62        let decoded = tokio::task::spawn_blocking(move || decode::checkpoint(&bytes))
63            .await
64            .context("decode task panicked")??;
65        Ok(decoded)
66    }
67
68    async fn checkpoint_bytes(&self, checkpoint: u64) -> object_store::Result<Bytes> {
69        self.bytes(ObjectPath::from(format!("{checkpoint}.binpb.zst")))
70            .await
71    }
72
73    async fn bytes(&self, path: ObjectPath) -> object_store::Result<Bytes> {
74        let result = self.store.get(&path).await?;
75        result.bytes().await
76    }
77
78    async fn watermark_checkpoint_hi_inclusive(&self) -> anyhow::Result<Option<u64>> {
79        let bytes = match self.bytes(ObjectPath::from(WATERMARK_PATH)).await {
80            Ok(bytes) => bytes,
81            Err(Error::NotFound { .. }) => return Ok(None),
82            Err(e) => return Err(e).context(format!("error reading {WATERMARK_PATH}")),
83        };
84
85        let watermark: ObjectStoreWatermark =
86            serde_json::from_slice(&bytes).context(format!("error parsing {WATERMARK_PATH}"))?;
87
88        Ok(watermark.checkpoint_hi_inclusive)
89    }
90}
91
92#[async_trait::async_trait]
93impl IngestionClientTrait for StoreIngestionClient {
94    async fn chain_id(&self) -> anyhow::Result<ChainIdentifier> {
95        let checkpoint = self.checkpoint(0).await?;
96        Ok((*checkpoint.summary.digest()).into())
97    }
98
99    /// Fetch a checkpoint from the remote store.
100    ///
101    /// Transient errors include:
102    ///
103    /// - failures to issue a request, (network errors, redirect issues, etc)
104    /// - request timeouts,
105    /// - rate limiting,
106    /// - server errors (5xx),
107    /// - issues getting a full response.
108    async fn checkpoint(&self, checkpoint: u64) -> CheckpointResult {
109        let bytes = self
110            .checkpoint_bytes(checkpoint)
111            .await
112            .map_err(|e| match e {
113                Error::NotFound { .. } => CheckpointError::NotFound,
114                e => CheckpointError::Fetch(e.into()),
115            })?;
116
117        if let Some(counter) = &self.total_ingested_bytes {
118            counter.inc_by(bytes.len() as u64);
119        }
120        // zstd decompress + prost decode + proto -> Checkpoint is multi-ms of CPU
121        // work; offload to the blocking pool so it doesn't stall the reactor.
122        tokio::task::spawn_blocking(move || decode::checkpoint(&bytes))
123            .await
124            .map_err(|e| CheckpointError::Fetch(anyhow::anyhow!("decode task panicked: {e}")))?
125            .map_err(CheckpointError::Decode)
126    }
127
128    async fn latest_checkpoint_number(&self) -> anyhow::Result<u64> {
129        self.watermark_checkpoint_hi_inclusive()
130            .await
131            .map(|cp| cp.unwrap_or(0))
132    }
133}
134
135/// Disable object_store's internal retries so that transient errors (429s, 5xx) propagate
136/// immediately to the framework's own retry logic.
137pub(super) fn retry_config() -> RetryConfig {
138    RetryConfig {
139        max_retries: 0,
140        ..Default::default()
141    }
142}
143
144#[cfg(test)]
145pub(crate) mod tests {
146    use axum::http::StatusCode;
147    use object_store::ClientOptions;
148    use object_store::http::HttpBuilder;
149    use std::sync::Mutex;
150    use std::sync::atomic::AtomicUsize;
151    use std::sync::atomic::Ordering;
152    use std::time::Duration;
153    use wiremock::Mock;
154    use wiremock::MockServer;
155    use wiremock::Request;
156    use wiremock::Respond;
157    use wiremock::ResponseTemplate;
158    use wiremock::matchers::method;
159    use wiremock::matchers::path;
160    use wiremock::matchers::path_regex;
161
162    use crate::ingestion::error::Error;
163    use crate::ingestion::ingestion_client::IngestionClient;
164    use crate::ingestion::test_utils::test_checkpoint_data;
165    use crate::metrics::tests::test_ingestion_metrics;
166
167    use super::*;
168
169    pub(crate) async fn respond_with(server: &MockServer, response: impl Respond + 'static) {
170        Mock::given(method("GET"))
171            .and(path_regex(r"/\d+\.binpb\.zst"))
172            .respond_with(response)
173            .mount(server)
174            .await;
175    }
176
177    /// Mount a high-priority mock for checkpoint 0 used by `StoreIngestionClient::chain_id()`.
178    pub(crate) async fn respond_with_chain_id(server: &MockServer) {
179        Mock::given(method("GET"))
180            .and(path("/0.binpb.zst"))
181            .respond_with(status(StatusCode::OK).set_body_bytes(test_checkpoint_data(0)))
182            .with_priority(1)
183            .mount(server)
184            .await;
185    }
186
187    /// Returns the expected chain_id produced by `StoreIngestionClient::chain_id()` when
188    /// `respond_with_chain_id` is mounted.
189    pub(crate) fn expected_chain_id() -> ChainIdentifier {
190        let bytes = test_checkpoint_data(0);
191        let checkpoint = decode::checkpoint(&bytes).unwrap();
192        (*checkpoint.summary.digest()).into()
193    }
194
195    pub(crate) fn status(code: StatusCode) -> ResponseTemplate {
196        ResponseTemplate::new(code.as_u16())
197    }
198
199    fn remote_test_client(uri: String) -> IngestionClient {
200        let store = HttpBuilder::new()
201            .with_url(uri)
202            .with_client_options(ClientOptions::default().with_allow_http(true))
203            .build()
204            .map(Arc::new)
205            .unwrap();
206        IngestionClient::with_store(store, test_ingestion_metrics()).unwrap()
207    }
208
209    async fn test_latest_checkpoint_number(watermark: ResponseTemplate) -> anyhow::Result<u64> {
210        let server = MockServer::start().await;
211
212        Mock::given(method("GET"))
213            .and(path(WATERMARK_PATH))
214            .respond_with(watermark)
215            .mount(&server)
216            .await;
217
218        let store = HttpBuilder::new()
219            .with_url(server.uri())
220            .with_client_options(ClientOptions::default().with_allow_http(true))
221            .build()
222            .map(Arc::new)
223            .unwrap();
224        let client = StoreIngestionClient::new(store, None);
225
226        IngestionClientTrait::latest_checkpoint_number(&client).await
227    }
228
229    #[tokio::test]
230    async fn test_latest_checkpoint_no_watermark() {
231        assert_eq!(
232            test_latest_checkpoint_number(status(StatusCode::NOT_FOUND))
233                .await
234                .unwrap(),
235            0
236        )
237    }
238
239    #[tokio::test]
240    async fn test_latest_checkpoint_corrupt_watermark() {
241        assert!(
242            test_latest_checkpoint_number(status(StatusCode::OK).set_body_string("<"))
243                .await
244                .is_err()
245        )
246    }
247
248    #[tokio::test]
249    async fn test_latest_checkpoint_from_watermark() {
250        let body = serde_json::json!({"checkpoint_hi_inclusive": 1}).to_string();
251        assert_eq!(
252            test_latest_checkpoint_number(status(StatusCode::OK).set_body_string(body))
253                .await
254                .unwrap(),
255            1
256        )
257    }
258
259    /// The watermark this client reads is the one `sui-indexer-alt-object-store` writes for the
260    /// pipeline that produces checkpoint files.
261    #[tokio::test]
262    async fn test_latest_checkpoint_from_object_store_indexer_watermark() {
263        use sui_indexer_alt_framework_store_traits::CommitterWatermark;
264        use sui_indexer_alt_framework_store_traits::Connection;
265        use sui_indexer_alt_framework_store_traits::Store;
266
267        let memory: Arc<dyn ObjectStore> = Arc::new(object_store::memory::InMemory::new());
268        let indexer_store = sui_indexer_alt_object_store::ObjectStore::new(memory.clone());
269        let mut conn = indexer_store.connect().await.unwrap();
270        let client = StoreIngestionClient::new(memory, None);
271
272        conn.init_watermark("checkpoint_blob", None).await.unwrap();
273        assert_eq!(
274            IngestionClientTrait::latest_checkpoint_number(&client)
275                .await
276                .unwrap(),
277            0
278        );
279
280        conn.set_committer_watermark(
281            "checkpoint_blob",
282            CommitterWatermark {
283                epoch_hi_inclusive: 3,
284                checkpoint_hi_inclusive: 1234,
285                tx_hi: 5678,
286                timestamp_ms_hi_inclusive: 9,
287            },
288        )
289        .await
290        .unwrap();
291        assert_eq!(
292            IngestionClientTrait::latest_checkpoint_number(&client)
293                .await
294                .unwrap(),
295            1234
296        );
297    }
298
299    #[tokio::test]
300    async fn fail_on_not_found() {
301        let server = MockServer::start().await;
302        respond_with(&server, status(StatusCode::NOT_FOUND)).await;
303
304        let client = remote_test_client(server.uri());
305        let error = client.checkpoint(42).await.unwrap_err();
306
307        assert!(matches!(error, Error::NotFound(42)));
308    }
309
310    /// Assume that failures to send the request to the remote store are due to temporary
311    /// connectivity issues, and retry them.
312    #[tokio::test]
313    async fn retry_on_request_error() {
314        let server = MockServer::start().await;
315
316        let times: Mutex<u64> = Mutex::new(0);
317        respond_with(&server, move |r: &Request| {
318            let mut times = times.lock().unwrap();
319            *times += 1;
320            match (*times, r.url.path()) {
321                // The first request will trigger a redirect to 999999.binpb.zst no matter what
322                // the original request was for -- triggering a request error.
323                (1, _) => status(StatusCode::MOVED_PERMANENTLY)
324                    .append_header("Location", "/999999.binpb.zst"),
325
326                // Set-up checkpoint 999999 as an infinite redirect loop.
327                (_, "/999999.binpb.zst") => {
328                    status(StatusCode::MOVED_PERMANENTLY).append_header("Location", r.url.as_str())
329                }
330
331                // Subsequently, requests will succeed.
332                _ => status(StatusCode::OK).set_body_bytes(test_checkpoint_data(42)),
333            }
334        })
335        .await;
336        respond_with_chain_id(&server).await;
337
338        let client = remote_test_client(server.uri());
339        let envelope = client.checkpoint(42).await.unwrap();
340
341        assert_eq!(42, envelope.checkpoint.summary.sequence_number);
342        assert_eq!(envelope.chain_id, expected_chain_id());
343    }
344
345    /// Assume that certain errors will recover by themselves, and keep retrying with an
346    /// exponential back-off. These errors include: 5xx (server) errors, 408 (timeout), and 429
347    /// (rate limiting).
348    #[tokio::test]
349    async fn retry_on_transient_server_error() {
350        let server = MockServer::start().await;
351        let times: Mutex<u64> = Mutex::new(0);
352        respond_with(&server, move |_: &Request| {
353            let mut times = times.lock().unwrap();
354            *times += 1;
355            match *times {
356                1 => status(StatusCode::INTERNAL_SERVER_ERROR),
357                2 => status(StatusCode::REQUEST_TIMEOUT),
358                3 => status(StatusCode::TOO_MANY_REQUESTS),
359                _ => status(StatusCode::OK).set_body_bytes(test_checkpoint_data(42)),
360            }
361        })
362        .await;
363        respond_with_chain_id(&server).await;
364
365        let client = remote_test_client(server.uri());
366        let envelope = client.checkpoint(42).await.unwrap();
367
368        assert_eq!(42, envelope.checkpoint.summary.sequence_number);
369        assert_eq!(envelope.chain_id, expected_chain_id());
370    }
371
372    /// Treat deserialization failure as another kind of transient error -- all checkpoint data
373    /// that is fetched should be valid (deserializable as a `Checkpoint`).
374    #[tokio::test]
375    async fn retry_on_deserialization_error() {
376        let server = MockServer::start().await;
377        let times: Mutex<u64> = Mutex::new(0);
378        respond_with(&server, move |_: &Request| {
379            let mut times = times.lock().unwrap();
380            *times += 1;
381            if *times < 3 {
382                status(StatusCode::OK).set_body_bytes(vec![])
383            } else {
384                status(StatusCode::OK).set_body_bytes(test_checkpoint_data(42))
385            }
386        })
387        .await;
388        respond_with_chain_id(&server).await;
389
390        let client = remote_test_client(server.uri());
391        let envelope = client.checkpoint(42).await.unwrap();
392
393        assert_eq!(42, envelope.checkpoint.summary.sequence_number);
394        assert_eq!(envelope.chain_id, expected_chain_id());
395    }
396
397    /// Test that timeout errors are retried as transient errors.
398    /// The first request will timeout, the second will succeed.
399    #[tokio::test]
400    async fn retry_on_timeout() {
401        let server = MockServer::start().await;
402        let times: Arc<AtomicUsize> = Arc::new(AtomicUsize::new(0));
403        let times_clone = times.clone();
404
405        // First request will delay longer than timeout, second will succeed immediately
406        respond_with(&server, move |_: &Request| {
407            match times_clone.fetch_add(1, Ordering::Relaxed) {
408                0 => {
409                    // Delay longer than our test timeout (2 seconds)
410                    std::thread::sleep(Duration::from_secs(4));
411                    status(StatusCode::OK).set_body_bytes(test_checkpoint_data(42))
412                }
413                _ => {
414                    // Respond immediately on retry attempts
415                    status(StatusCode::OK).set_body_bytes(test_checkpoint_data(42))
416                }
417            }
418        })
419        .await;
420        respond_with_chain_id(&server).await;
421
422        let options = ClientOptions::default()
423            .with_allow_http(true)
424            .with_timeout(Duration::from_secs(2));
425        let store = HttpBuilder::new()
426            .with_url(server.uri())
427            .with_client_options(options)
428            .build()
429            .map(Arc::new)
430            .unwrap();
431        let ingestion_client =
432            IngestionClient::with_store(store, test_ingestion_metrics()).unwrap();
433
434        // This should timeout once, then succeed on retry
435        let envelope = ingestion_client.checkpoint(42).await.unwrap();
436        assert_eq!(42, envelope.checkpoint.summary.sequence_number);
437        assert_eq!(envelope.chain_id, expected_chain_id());
438
439        // Verify that the server received exactly 2 requests (1 timeout + 1 successful retry)
440        // The chain_id request for checkpoint 0 is handled by a separate mock.
441        let final_count = times.load(Ordering::Relaxed);
442        assert_eq!(final_count, 2);
443    }
444}
sui_indexer_alt_framework/ingestion/store_client.rs

sui_indexer_alt_framework/ingestion/
store_client.rs
