sui_indexer_alt_reader/
ledger_grpc_reader.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::collections::HashMap;
5use std::hash::Hash;
6use std::time::Duration;
7
8use anyhow::Context;
9use async_graphql::dataloader::DataLoader;
10use async_graphql::dataloader::Loader;
11use futures::future::try_join_all;
12use prometheus::Registry;
13use prost_types::FieldMask;
14use sui_rpc::Client;
15use sui_rpc::field::FieldMaskUtil;
16use sui_rpc::proto::proto_to_timestamp_ms;
17use sui_rpc::proto::sui::rpc::v2 as grpc;
18use sui_types::effects::TransactionEffects;
19use sui_types::event::Event;
20use sui_types::messages_checkpoint::CheckpointSummary;
21use sui_types::signature::GenericSignature;
22use sui_types::transaction::TransactionData;
23use tonic::transport::Uri;
24
25use crate::metrics::GrpcMetricsLayer;
26
27#[derive(clap::Args, Debug, Clone)]
28pub struct LedgerGrpcArgs {
29    /// Timeout for gRPC statements to the ledger service, in milliseconds.
30    #[arg(long)]
31    pub ledger_grpc_statement_timeout_ms: Option<u64>,
32
33    /// Maximum gRPC decoding message size for Ledger service responses, in bytes.
34    #[arg(long, default_value_t = 32 * 1024 * 1024)]
35    pub ledger_grpc_max_decoding_message_size: usize,
36}
37
38#[derive(Debug, Clone)]
39pub struct CheckpointedTransaction {
40    pub effects: Box<TransactionEffects>,
41    pub events: Option<Vec<Event>>,
42    pub transaction_data: Box<TransactionData>,
43    pub signatures: Vec<GenericSignature>,
44    pub timestamp_ms: Option<u64>,
45    pub cp_sequence_number: Option<u64>,
46    pub balance_changes: Vec<grpc::BalanceChange>,
47}
48
49/// A reader backed by gRPC LedgerService (sui-kv-rpc).
50///
51/// This connects to archival service that implements the same LedgerService gRPC interface
52/// as fullnode, but is backed by Bigtable for serving historical data.
53#[derive(Clone)]
54pub struct LedgerGrpcReader {
55    client: Client,
56    timeout: Option<Duration>,
57    max_batch_get_transactions: usize,
58    max_batch_get_objects: usize,
59}
60
61/// Maximum number of transaction digests `LedgerGrpcReader` will put in a single
62/// `BatchGetTransactions` call, matching the ledger gRPC/KV-RPC service's own hard cap.
63pub const MAX_BATCH_GET_TRANSACTIONS: usize = 200;
64
65/// Maximum number of object keys `LedgerGrpcReader` will put in a single `BatchGetObjects`
66/// call, matching the ledger gRPC/KV-RPC service's own hard cap.
67pub const MAX_BATCH_GET_OBJECTS: usize = 1000;
68
69/// Implemented by `LedgerGrpcReader` for each key type whose `Loader::load` needs to stay under
70/// the ledger service's batch-size limit — `DataLoader`'s own `max_batch_size` is only a dispatch
71/// trigger, not a hard cap on how many keys reach a single `Loader::load` call. `load_chunk`
72/// supplies the raw, single-chunk fetch and `chunk_size` supplies the limit; `load_chunked`'s
73/// default body splits `keys` into chunks, dispatches `load_chunk` concurrently per chunk, and
74/// merges the results.
75///
76/// `pub`, not `pub(crate)`: it's referenced through `LedgerGrpcReader`'s blanket `Loader<K>` impl
77/// below (`type Value = <Self as ChunkedLoader<K>>::Value`), and that impl is part of
78/// `LedgerGrpcReader`'s public interface since both the type and `Loader` are public.
79#[async_trait::async_trait]
80pub trait ChunkedLoader<K>
81where
82    K: Send + Sync + Hash + Eq + Clone + 'static,
83{
84    type Value: Send + Sync + Clone + 'static;
85    type Error: Send + Sync + Clone + 'static;
86
87    fn chunk_size(&self) -> usize;
88
89    async fn load_chunk(&self, keys: &[K]) -> Result<HashMap<K, Self::Value>, Self::Error>;
90
91    async fn load_chunked(&self, keys: &[K]) -> Result<HashMap<K, Self::Value>, Self::Error>
92    where
93        Self: Sync,
94    {
95        let limit = self.chunk_size();
96
97        let mut results = HashMap::new();
98        for batch in try_join_all(keys.chunks(limit).map(|chunk| self.load_chunk(chunk))).await? {
99            results.extend(batch);
100        }
101        Ok(results)
102    }
103}
104
105/// Covers every key type `LedgerGrpcReader` implements [`ChunkedLoader`] for. Coherent because
106/// `Self` (`LedgerGrpcReader`) is a concrete local type — only `K` is generic — unlike a bare
107/// `impl<K, T: ChunkedLoader<K>> Loader<K> for T`, which the orphan rules reject since neither
108/// `Loader` (foreign) nor `T` (an uncovered generic) is local.
109#[async_trait::async_trait]
110impl<K> Loader<K> for LedgerGrpcReader
111where
112    K: Send + Sync + Hash + Eq + Clone + 'static,
113    Self: ChunkedLoader<K>,
114{
115    type Value = <Self as ChunkedLoader<K>>::Value;
116    type Error = <Self as ChunkedLoader<K>>::Error;
117
118    async fn load(&self, keys: &[K]) -> Result<HashMap<K, Self::Value>, Self::Error> {
119        self.load_chunked(keys).await
120    }
121}
122
123impl LedgerGrpcArgs {
124    pub fn new(
125        statement_timeout_ms: Option<u64>,
126        max_decoding_message_size: Option<usize>,
127    ) -> Self {
128        let defaults = Self::default();
129        Self {
130            ledger_grpc_statement_timeout_ms: statement_timeout_ms,
131            ledger_grpc_max_decoding_message_size: max_decoding_message_size
132                .unwrap_or(defaults.ledger_grpc_max_decoding_message_size),
133        }
134    }
135
136    pub fn statement_timeout(&self) -> Option<std::time::Duration> {
137        self.ledger_grpc_statement_timeout_ms
138            .map(Duration::from_millis)
139    }
140}
141
142impl CheckpointedTransaction {
143    /// Read mask selecting everything needed to construct a `CheckpointedTransaction` from the gRPC
144    /// `ExecutedTransaction` proto.
145    pub fn read_mask() -> FieldMask {
146        FieldMask::from_paths([
147            "transaction.bcs",
148            "effects.bcs",
149            "events.bcs",
150            "signatures.bcs",
151            "checkpoint",
152            "timestamp",
153            "balance_changes",
154        ])
155    }
156}
157
158impl LedgerGrpcReader {
159    pub async fn new(
160        uri: Uri,
161        args: LedgerGrpcArgs,
162        prefix: Option<&str>,
163        registry: &Registry,
164        max_batch_get_transactions: usize,
165        max_batch_get_objects: usize,
166    ) -> anyhow::Result<Self> {
167        let timeout = args.statement_timeout();
168        let mut client = Client::new(uri)?
169            .with_max_decoding_message_size(args.ledger_grpc_max_decoding_message_size)
170            .request_layer(GrpcMetricsLayer::new(
171                prefix.unwrap_or("ledger_grpc"),
172                registry,
173            ));
174
175        if let Some(timeout) = timeout {
176            client = client.with_response_headers_timeout(timeout);
177        }
178
179        Ok(Self {
180            client,
181            timeout,
182            max_batch_get_transactions,
183            max_batch_get_objects,
184        })
185    }
186
187    pub(crate) fn as_data_loader(&self) -> DataLoader<Self> {
188        DataLoader::new(self.clone(), tokio::spawn)
189    }
190
191    pub(crate) fn max_batch_get_transactions(&self) -> usize {
192        self.max_batch_get_transactions
193    }
194
195    pub(crate) fn max_batch_get_objects(&self) -> usize {
196        self.max_batch_get_objects
197    }
198
199    pub async fn checkpoint_watermark(&self) -> anyhow::Result<CheckpointSummary> {
200        use grpc::GetCheckpointRequest;
201        use prost_types::FieldMask;
202        use sui_rpc::field::FieldMaskUtil;
203
204        let request =
205            GetCheckpointRequest::default().with_read_mask(FieldMask::from_paths(["summary.bcs"]));
206
207        let response = self.get_checkpoint(request).await?;
208
209        let checkpoint = response.checkpoint.context("No checkpoint returned")?;
210
211        checkpoint
212            .summary
213            .as_ref()
214            .and_then(|s| s.bcs.as_ref())
215            .context("Missing summary.bcs")?
216            .deserialize()
217            .context("Failed to deserialize checkpoint summary")
218    }
219
220    /// Resolve a checkpoint digest to its sequence number via the ledger service. Returns `None`
221    /// if no checkpoint with that digest is known.
222    pub async fn checkpoint_seq_by_digest(
223        &self,
224        digest: sui_types::digests::CheckpointDigest,
225    ) -> anyhow::Result<Option<u64>> {
226        use grpc::GetCheckpointRequest;
227        use prost_types::FieldMask;
228        use sui_rpc::field::FieldMaskUtil;
229
230        let sdk_digest = sui_sdk_types::Digest::new(digest.inner().to_owned());
231        let request = GetCheckpointRequest::by_digest(&sdk_digest)
232            .with_read_mask(FieldMask::from_paths(["sequence_number"]));
233
234        match self.get_checkpoint(request).await {
235            Ok(response) => {
236                let checkpoint = response.checkpoint.context("No checkpoint returned")?;
237                Ok(checkpoint.sequence_number)
238            }
239            Err(status) if status.code() == tonic::Code::NotFound => Ok(None),
240            Err(e) => Err(anyhow::anyhow!(e)),
241        }
242    }
243
244    // Public wrapper methods for gRPC calls with metrics instrumentation
245
246    pub async fn get_checkpoint(
247        &self,
248        request: grpc::GetCheckpointRequest,
249    ) -> Result<grpc::GetCheckpointResponse, tonic::Status> {
250        self.client
251            .clone()
252            .ledger_client()
253            .get_checkpoint(self.request(request))
254            .await
255            .map(|r| r.into_inner())
256    }
257
258    pub async fn batch_get_transactions(
259        &self,
260        request: grpc::BatchGetTransactionsRequest,
261    ) -> Result<grpc::BatchGetTransactionsResponse, tonic::Status> {
262        self.client
263            .clone()
264            .ledger_client()
265            .batch_get_transactions(self.request(request))
266            .await
267            .map(|r| r.into_inner())
268    }
269
270    pub async fn batch_get_objects(
271        &self,
272        request: grpc::BatchGetObjectsRequest,
273    ) -> Result<grpc::BatchGetObjectsResponse, tonic::Status> {
274        self.client
275            .clone()
276            .ledger_client()
277            .batch_get_objects(self.request(request))
278            .await
279            .map(|r| r.into_inner())
280    }
281
282    pub async fn get_transaction(
283        &self,
284        request: grpc::GetTransactionRequest,
285    ) -> Result<grpc::GetTransactionResponse, tonic::Status> {
286        self.client
287            .clone()
288            .ledger_client()
289            .get_transaction(self.request(request))
290            .await
291            .map(|r| r.into_inner())
292    }
293
294    /// Create a gRPC request, optionally with the grpc-timeout header if configured.
295    fn request<T>(&self, input: T) -> tonic::Request<T> {
296        let mut request = tonic::Request::new(input);
297        if let Some(timeout) = self.timeout {
298            request.set_timeout(timeout);
299        }
300        request
301    }
302}
303
304impl TryFrom<&grpc::ExecutedTransaction> for CheckpointedTransaction {
305    type Error = anyhow::Error;
306
307    fn try_from(executed: &grpc::ExecutedTransaction) -> anyhow::Result<Self> {
308        let full_tx: sui_types::full_checkpoint_content::ExecutedTransaction = executed
309            .try_into()
310            .context("Failed to convert ExecutedTransaction from proto")?;
311
312        let timestamp_ms = executed
313            .timestamp
314            .map(proto_to_timestamp_ms)
315            .transpose()
316            .with_context(|| format!("Failed to parse timestamp {:?}", executed.timestamp))?;
317
318        Ok(Self {
319            effects: Box::new(full_tx.effects),
320            events: full_tx.events.map(|events| events.data),
321            transaction_data: Box::new(full_tx.transaction),
322            signatures: full_tx.signatures,
323            timestamp_ms,
324            cp_sequence_number: executed.checkpoint,
325            balance_changes: executed.balance_changes.clone(),
326        })
327    }
328}
329
330impl Default for LedgerGrpcArgs {
331    fn default() -> Self {
332        Self {
333            ledger_grpc_statement_timeout_ms: None,
334            ledger_grpc_max_decoding_message_size: 32 * 1024 * 1024,
335        }
336    }
337}
338
339#[cfg(test)]
340pub(crate) mod test_support {
341    use std::net::SocketAddr;
342    use std::sync::Arc;
343    use std::sync::Mutex;
344    use std::time::Duration;
345
346    use prometheus::Registry;
347    use sui_rpc::proto::sui::rpc::v2::BatchGetObjectsRequest;
348    use sui_rpc::proto::sui::rpc::v2::BatchGetObjectsResponse;
349    use sui_rpc::proto::sui::rpc::v2::BatchGetTransactionsRequest;
350    use sui_rpc::proto::sui::rpc::v2::BatchGetTransactionsResponse;
351    use sui_rpc::proto::sui::rpc::v2::GetObjectRequest;
352    use sui_rpc::proto::sui::rpc::v2::ledger_service_server::LedgerService;
353    use sui_rpc::proto::sui::rpc::v2::ledger_service_server::LedgerServiceServer;
354    use tokio::net::TcpListener;
355    use tokio::task::JoinHandle;
356    use tokio_stream::wrappers::TcpListenerStream;
357    use tonic::Request;
358    use tonic::Response;
359    use tonic::Status;
360
361    use super::LedgerGrpcArgs;
362    use super::LedgerGrpcReader;
363    use super::MAX_BATCH_GET_OBJECTS;
364    use super::MAX_BATCH_GET_TRANSACTIONS;
365
366    /// Starts a [`MockLedgerServer`] and constructs a [`LedgerGrpcReader`]
367    /// pointed at it. Shared by every loader's chunking test.
368    pub(crate) async fn mock_reader() -> (LedgerGrpcReader, MockLedgerServer, JoinHandle<()>) {
369        let mock = MockLedgerServer::new();
370        let (addr, server) = mock.start().await.expect("start mock ledger service");
371        let reader = LedgerGrpcReader::new(
372            format!("http://{addr}").parse().unwrap(),
373            LedgerGrpcArgs::default(),
374            None,
375            &Registry::new(),
376            MAX_BATCH_GET_TRANSACTIONS,
377            MAX_BATCH_GET_OBJECTS,
378        )
379        .await
380        .expect("construct LedgerGrpcReader");
381        (reader, mock, server)
382    }
383
384    /// Asserts that `batches` (the digest lists recorded by
385    /// [`MockLedgerServer`] across however many gRPC calls a loader made) is
386    /// chunked to `limit`, and that their union reconstructs `expected`
387    /// exactly, with none lost or duplicated.
388    pub(crate) fn assert_chunked(batches: Vec<Vec<String>>, limit: usize, expected: &[String]) {
389        assert_eq!(batches.len(), expected.len().div_ceil(limit));
390        assert!(batches.iter().all(|batch| batch.len() <= limit));
391
392        let mut requested: Vec<String> = batches.into_iter().flatten().collect();
393        requested.sort();
394        let mut expected = expected.to_vec();
395        expected.sort();
396        assert_eq!(requested, expected);
397    }
398
399    /// A minimal in-process `LedgerService` that records the shape of every
400    /// `BatchGetTransactions`/`BatchGetObjects` call it receives and responds
401    /// with an empty result set. Every other RPC falls back to the generated
402    /// trait's default `unimplemented` behavior, which is all the chunking
403    /// tests that use this need.
404    #[derive(Clone, Default)]
405    pub(crate) struct MockLedgerServer {
406        transaction_batches: Arc<Mutex<Vec<Vec<String>>>>,
407        object_batches: Arc<Mutex<Vec<Vec<GetObjectRequest>>>>,
408    }
409
410    impl MockLedgerServer {
411        pub(crate) fn new() -> Self {
412            Self::default()
413        }
414
415        pub(crate) async fn start(&self) -> anyhow::Result<(SocketAddr, JoinHandle<()>)> {
416            let listener = TcpListener::bind("127.0.0.1:0").await?;
417            let addr = listener.local_addr()?;
418            let mock = self.clone();
419            let handle = tokio::spawn(async move {
420                let incoming = TcpListenerStream::new(listener);
421                tonic::transport::Server::builder()
422                    .add_service(LedgerServiceServer::new(mock))
423                    .serve_with_incoming(incoming)
424                    .await
425                    .ok();
426            });
427            tokio::time::sleep(Duration::from_millis(10)).await;
428            Ok((addr, handle))
429        }
430
431        pub(crate) fn transaction_batches(&self) -> Vec<Vec<String>> {
432            self.transaction_batches.lock().unwrap().clone()
433        }
434
435        pub(crate) fn object_batches(&self) -> Vec<Vec<GetObjectRequest>> {
436            self.object_batches.lock().unwrap().clone()
437        }
438    }
439
440    #[tonic::async_trait]
441    impl LedgerService for MockLedgerServer {
442        async fn batch_get_transactions(
443            &self,
444            request: Request<BatchGetTransactionsRequest>,
445        ) -> Result<Response<BatchGetTransactionsResponse>, Status> {
446            self.transaction_batches
447                .lock()
448                .unwrap()
449                .push(request.into_inner().digests);
450            Ok(Response::new(BatchGetTransactionsResponse::default()))
451        }
452
453        async fn batch_get_objects(
454            &self,
455            request: Request<BatchGetObjectsRequest>,
456        ) -> Result<Response<BatchGetObjectsResponse>, Status> {
457            self.object_batches
458                .lock()
459                .unwrap()
460                .push(request.into_inner().requests);
461            Ok(Response::new(BatchGetObjectsResponse::default()))
462        }
463    }
464}
sui_indexer_alt_reader/ledger_grpc_reader.rs

sui_indexer_alt_reader/
ledger_grpc_reader.rs
