1use std::time::Duration;
5
6use anyhow::Context;
7use anyhow::bail;
8use anyhow::ensure;
9use bytes::Bytes;
10use futures::Stream;
11use futures::StreamExt;
12use prometheus::Registry;
13use sui_rpc::Client;
14use sui_rpc::proto::sui::rpc::v2 as proto;
15use sui_rpc::proto::sui::rpc::v2::ExecutedTransaction;
16use tonic::transport::Uri;
17use tracing::warn;
18
19use crate::ledger_grpc_reader::LedgerGrpcArgs;
20use crate::metrics::GrpcMetricsLayer;
21
22#[derive(Clone)]
24pub struct AlphaLedgerGrpcReader {
25 client: Client,
26 timeout: Option<Duration>,
27}
28
29#[derive(Debug, Clone)]
31pub struct PageItem<T> {
32 pub payload: T,
33 pub cursor: Bytes,
34}
35
36#[derive(Debug, Clone)]
38pub struct StreamPage<T> {
39 pub items: Vec<PageItem<T>>,
41 first_wm_cursor: Option<Bytes>,
42 last_wm_cursor: Option<Bytes>,
43 pub end_reason: Option<proto::QueryEndReason>,
44}
45
46#[derive(Debug)]
47enum FrameKind<T> {
48 Frame {
49 payload: Option<T>,
50 cursor: Option<Bytes>,
51 end_reason: Option<proto::QueryEndReason>,
52 },
53 Unknown,
55}
56
57impl AlphaLedgerGrpcReader {
58 pub async fn new(
59 uri: Uri,
60 args: LedgerGrpcArgs,
61 prefix: Option<&str>,
62 registry: &Registry,
63 ) -> anyhow::Result<Self> {
64 let timeout = args.statement_timeout();
65 let mut client = Client::new(uri)?
66 .with_max_decoding_message_size(args.ledger_grpc_max_decoding_message_size)
67 .request_layer(GrpcMetricsLayer::new(
68 prefix.unwrap_or("ledger_grpc"),
69 registry,
70 ));
71
72 if let Some(timeout) = timeout {
73 client = client.with_response_headers_timeout(timeout);
74 }
75
76 Ok(Self { client, timeout })
77 }
78
79 pub async fn list_transactions(
80 &self,
81 request: proto::ListTransactionsRequest,
82 ) -> anyhow::Result<StreamPage<ExecutedTransaction>> {
83 let stream = self
84 .client
85 .clone()
86 .ledger_client()
87 .list_transactions(self.request(request))
88 .await
89 .context("ListTransactions stream open failed")?
90 .into_inner();
91
92 drain_list_stream("ListTransactions", stream).await
93 }
94
95 pub async fn list_events(
96 &self,
97 request: proto::ListEventsRequest,
98 ) -> anyhow::Result<StreamPage<proto::Event>> {
99 let stream = self
100 .client
101 .clone()
102 .ledger_client()
103 .list_events(self.request(request))
104 .await
105 .context("ListEvents stream open failed")?
106 .into_inner();
107
108 drain_list_stream("ListEvents", stream).await
109 }
110
111 pub async fn list_checkpoints(
112 &self,
113 request: proto::ListCheckpointsRequest,
114 ) -> anyhow::Result<StreamPage<proto::Checkpoint>> {
115 let stream = self
116 .client
117 .clone()
118 .ledger_client()
119 .list_checkpoints(self.request(request))
120 .await
121 .context("ListCheckpoints stream open failed")?
122 .into_inner();
123
124 drain_list_stream("ListCheckpoints", stream).await
125 }
126
127 fn request<T>(&self, input: T) -> tonic::Request<T> {
129 let mut request = tonic::Request::new(input);
130 if let Some(timeout) = self.timeout {
131 request.set_timeout(timeout);
132 }
133 request
134 }
135}
136
137impl<T> StreamPage<T> {
138 pub fn has_more(&self) -> bool {
146 use proto::QueryEndReason as R;
147 match self.end_reason {
148 None => true,
149 Some(R::Unknown | R::ItemLimit | R::ScanLimit) => true,
150 Some(R::LedgerTip | R::CheckpointBound) => false,
151 Some(R::CursorBound) => self.last_cursor().is_some(),
152 Some(_) => true,
155 }
156 }
157
158 pub fn first_cursor(&self) -> Option<&Bytes> {
161 self.first_wm_cursor
162 .as_ref()
163 .or_else(|| self.items.first().map(|item| &item.cursor))
164 }
165
166 pub fn last_cursor(&self) -> Option<&Bytes> {
169 self.last_wm_cursor
170 .as_ref()
171 .or_else(|| self.items.last().map(|item| &item.cursor))
172 }
173
174 #[cfg(feature = "testing")]
178 pub fn for_test(
179 items: Vec<PageItem<T>>,
180 first_wm_cursor: Option<Bytes>,
181 last_wm_cursor: Option<Bytes>,
182 end_reason: Option<proto::QueryEndReason>,
183 ) -> Self {
184 Self {
185 items,
186 first_wm_cursor,
187 last_wm_cursor,
188 end_reason,
189 }
190 }
191
192 fn apply(&mut self, frame: FrameKind<T>) -> bool {
196 let FrameKind::Frame {
197 payload,
198 cursor,
199 end_reason,
200 } = frame
201 else {
202 warn!("ignoring unrecognized frame");
203 return false;
204 };
205 match payload {
206 Some(payload) => {
207 let cursor = cursor.expect("TryFrom validated item cursor");
208 self.last_wm_cursor = None;
209 self.items.push(PageItem { payload, cursor });
210 }
211 None => {
212 if let Some(cursor) = cursor {
213 self.last_wm_cursor = Some(cursor.clone());
214 if self.items.is_empty() && self.first_wm_cursor.is_none() {
215 self.first_wm_cursor = Some(cursor);
216 }
217 }
218 }
219 }
220 if let Some(reason) = end_reason {
221 self.end_reason = Some(reason);
225 return true;
226 }
227 false
228 }
229}
230
231impl<T> Default for StreamPage<T> {
232 fn default() -> Self {
233 Self {
234 items: Vec::new(),
235 first_wm_cursor: None,
236 last_wm_cursor: None,
237 end_reason: None,
238 }
239 }
240}
241
242impl TryFrom<proto::ListTransactionsResponse> for FrameKind<ExecutedTransaction> {
243 type Error = anyhow::Error;
244
245 fn try_from(response: proto::ListTransactionsResponse) -> anyhow::Result<Self> {
246 classify_frame(response.transaction, response.watermark, response.end)
247 }
248}
249
250impl TryFrom<proto::ListEventsResponse> for FrameKind<proto::Event> {
251 type Error = anyhow::Error;
252
253 fn try_from(response: proto::ListEventsResponse) -> anyhow::Result<Self> {
254 classify_frame(response.event, response.watermark, response.end)
255 }
256}
257
258impl TryFrom<proto::ListCheckpointsResponse> for FrameKind<proto::Checkpoint> {
259 type Error = anyhow::Error;
260
261 fn try_from(response: proto::ListCheckpointsResponse) -> anyhow::Result<Self> {
262 classify_frame(response.checkpoint, response.watermark, response.end)
263 }
264}
265
266fn classify_frame<T>(
269 payload: Option<T>,
270 watermark: Option<proto::Watermark>,
271 end: Option<proto::QueryEnd>,
272) -> anyhow::Result<FrameKind<T>> {
273 let cursor = watermark.and_then(|w| w.cursor);
274 let end_reason = end.map(|e| e.reason());
275
276 if payload.is_none() && cursor.is_none() && end_reason.is_none() {
277 return Ok(FrameKind::Unknown);
278 }
279 if payload.is_some() && cursor.is_none() {
280 bail!("Item frame missing watermark.cursor");
281 }
282
283 Ok(FrameKind::Frame {
284 payload,
285 cursor,
286 end_reason,
287 })
288}
289
290async fn drain_list_stream<R, T, S>(
291 rpc_name: &'static str,
292 stream: S,
293) -> anyhow::Result<StreamPage<T>>
294where
295 R: TryInto<FrameKind<T>, Error = anyhow::Error>,
296 S: Stream<Item = Result<R, tonic::Status>>,
297{
298 futures::pin_mut!(stream);
299 let mut page = StreamPage::default();
300 while let Some(result) = stream.next().await {
301 match result {
302 Ok(response) => {
303 let frame = response
304 .try_into()
305 .with_context(|| format!("{rpc_name}: malformed frame"))?;
306 if page.apply(frame) {
308 break;
309 }
310 }
311 Err(status)
315 if matches!(
316 status.code(),
317 tonic::Code::DeadlineExceeded | tonic::Code::Cancelled
318 ) =>
319 {
320 break;
321 }
322 Err(status) => {
324 bail!(
325 "{rpc_name}: stream error {:?}: {}",
326 status.code(),
327 status.message()
328 );
329 }
330 }
331 }
332
333 ensure!(
337 !page.has_more() || page.last_cursor().is_some(),
338 "{rpc_name}: server reported more results but did not provide resume cursor — cannot continue",
339 );
340
341 Ok(page)
342}
343
344#[cfg(test)]
345mod tests {
346 use super::*;
347
348 fn item_response(cursor: &[u8]) -> proto::ListTransactionsResponse {
350 let mut watermark = proto::Watermark::default();
351 watermark.cursor = Some(Bytes::copy_from_slice(cursor));
352 let mut response = proto::ListTransactionsResponse::default();
353 response.transaction = Some(ExecutedTransaction::default());
354 response.watermark = Some(watermark);
355 response
356 }
357
358 fn watermark_response(cursor: &[u8]) -> proto::ListTransactionsResponse {
359 let mut watermark = proto::Watermark::default();
360 watermark.cursor = Some(Bytes::copy_from_slice(cursor));
361 let mut response = proto::ListTransactionsResponse::default();
362 response.watermark = Some(watermark);
363 response
364 }
365
366 fn end_response(reason: proto::QueryEndReason) -> proto::ListTransactionsResponse {
367 let mut end = proto::QueryEnd::default();
368 end.reason = Some(reason as i32);
369 let mut response = proto::ListTransactionsResponse::default();
370 response.end = Some(end);
371 response
372 }
373
374 fn end_response_with_cursor(
375 cursor: &[u8],
376 reason: proto::QueryEndReason,
377 ) -> proto::ListTransactionsResponse {
378 let mut response = end_response(reason);
379 let mut watermark = proto::Watermark::default();
380 watermark.cursor = Some(Bytes::copy_from_slice(cursor));
381 response.watermark = Some(watermark);
382 response
383 }
384
385 fn frame(r: proto::ListTransactionsResponse) -> FrameKind<ExecutedTransaction> {
386 r.try_into().expect("test fixture should be well-formed")
387 }
388
389 async fn drain_iter(
390 responses: Vec<Result<proto::ListTransactionsResponse, tonic::Status>>,
391 ) -> anyhow::Result<StreamPage<ExecutedTransaction>> {
392 drain_list_stream::<_, ExecutedTransaction, _>(
393 "ListTransactions",
394 futures::stream::iter(responses),
395 )
396 .await
397 }
398
399 #[test]
400 fn drains_items_tracking_latest_cursor_and_end_reason() {
401 let mut page: StreamPage<ExecutedTransaction> = StreamPage::default();
402 page.apply(frame(item_response(b"c1")));
403 page.apply(frame(watermark_response(b"w2")));
404 let mut last = item_response(b"c3");
405 last.end = end_response(proto::QueryEndReason::ItemLimit).end;
406 page.apply(frame(last));
407 assert_eq!(page.items.len(), 2);
408 assert_eq!(page.items[0].cursor.as_ref(), b"c1".as_ref());
411 assert_eq!(page.items[1].cursor.as_ref(), b"c3".as_ref());
412 assert_eq!(
413 page.first_cursor().map(|c| c.as_ref()),
414 Some(b"c1".as_ref())
415 );
416 assert_eq!(page.last_cursor().map(|c| c.as_ref()), Some(b"c3".as_ref()));
417 assert_eq!(page.end_reason, Some(proto::QueryEndReason::ItemLimit));
418 }
419
420 #[test]
421 fn standalone_watermark_advances_cursor_without_items() {
422 let mut page: StreamPage<ExecutedTransaction> = StreamPage::default();
423 page.apply(frame(watermark_response(b"w1")));
424 page.apply(frame(end_response_with_cursor(
425 b"w2",
426 proto::QueryEndReason::LedgerTip,
427 )));
428
429 assert!(page.items.is_empty());
430 assert_eq!(
431 page.first_cursor().map(|c| c.as_ref()),
432 Some(b"w1".as_ref())
433 );
434 assert_eq!(page.last_cursor().map(|c| c.as_ref()), Some(b"w2".as_ref()));
435 assert_eq!(page.end_reason, Some(proto::QueryEndReason::LedgerTip));
436 }
437
438 #[test]
439 fn apply_signals_stop_only_on_end_frame() {
440 let mut page: StreamPage<ExecutedTransaction> = StreamPage::default();
443 assert!(!page.apply(frame(item_response(b"c1"))));
444 assert!(!page.apply(frame(watermark_response(b"w1"))));
445 assert!(!page.apply(frame(proto::ListTransactionsResponse::default())));
447 assert!(page.apply(frame(end_response(proto::QueryEndReason::LedgerTip))));
448
449 let mut page: StreamPage<ExecutedTransaction> = StreamPage::default();
450 let mut item_end = item_response(b"c2");
451 item_end.end = end_response(proto::QueryEndReason::ItemLimit).end;
452 assert!(page.apply(frame(item_end)));
453 assert_eq!(page.items.len(), 1);
454 assert_eq!(page.items[0].cursor.as_ref(), b"c2".as_ref());
455 assert_eq!(page.end_reason, Some(proto::QueryEndReason::ItemLimit));
456 }
457
458 #[test]
459 fn first_wm_cursor_set_to_first_pre_item_watermark() {
460 let mut page: StreamPage<ExecutedTransaction> = StreamPage::default();
462 page.apply(frame(watermark_response(b"w1")));
463 page.apply(frame(item_response(b"c2")));
464 page.apply(frame(watermark_response(b"w3")));
465 page.apply(frame(item_response(b"c4")));
466
467 assert_eq!(
468 page.first_cursor().map(|c| c.as_ref()),
469 Some(b"w1".as_ref())
470 );
471 assert_eq!(page.last_cursor().map(|c| c.as_ref()), Some(b"c4".as_ref()));
472 }
473
474 #[test]
475 fn first_wm_cursor_not_set_after_items() {
476 let mut page: StreamPage<ExecutedTransaction> = StreamPage::default();
478 page.apply(frame(item_response(b"c2")));
479 page.apply(frame(watermark_response(b"w3")));
480 page.apply(frame(item_response(b"c4")));
481 page.apply(frame(watermark_response(b"w1")));
482
483 assert_eq!(
484 page.first_cursor().map(|c| c.as_ref()),
485 Some(b"c2".as_ref())
486 );
487 assert!(page.first_wm_cursor.is_none());
488 }
489
490 #[test]
491 fn trailing_watermark_advances_past_last_item() {
492 let mut page: StreamPage<ExecutedTransaction> = StreamPage::default();
493 page.apply(frame(item_response(b"c1")));
494 page.apply(frame(watermark_response(b"w2")));
495
496 assert_eq!(page.last_cursor().map(|c| c.as_ref()), Some(b"w2".as_ref()));
499 assert_eq!(
501 page.first_cursor().map(|c| c.as_ref()),
502 Some(b"c1".as_ref())
503 );
504 }
505
506 #[test]
507 fn has_more_true_when_truncated_or_timed_out() {
508 for reason in [
510 proto::QueryEndReason::ItemLimit,
511 proto::QueryEndReason::ScanLimit,
512 ] {
513 let mut page: StreamPage<ExecutedTransaction> = StreamPage::default();
514 page.apply(frame(end_response(reason)));
515 assert!(page.has_more(), "expected has_more for {reason:?}");
516 }
517
518 let page: StreamPage<ExecutedTransaction> = StreamPage::default();
522 assert!(page.has_more());
523 }
524
525 #[test]
526 fn has_more_false_on_authoritative_terminals() {
527 for reason in [
531 proto::QueryEndReason::CheckpointBound,
532 proto::QueryEndReason::LedgerTip,
533 proto::QueryEndReason::CursorBound,
534 ] {
535 let mut page: StreamPage<ExecutedTransaction> = StreamPage::default();
536 page.apply(frame(end_response(reason)));
537 assert!(!page.has_more(), "expected !has_more for {reason:?}");
538 }
539 }
540
541 #[test]
542 fn has_more_true_on_cursor_bound_with_tracked_cursor() {
543 let mut page: StreamPage<ExecutedTransaction> = StreamPage::default();
544 let mut response = item_response(b"c1");
545 response.end = end_response(proto::QueryEndReason::CursorBound).end;
546 page.apply(frame(response));
547
548 assert_eq!(page.last_cursor().map(|c| c.as_ref()), Some(b"c1".as_ref()));
549 assert!(
550 page.has_more(),
551 "CursorBound with tracked cursor should not be terminal"
552 );
553 }
554
555 #[test]
556 fn apply_end_with_unknown_reason_folds_to_unknown() {
557 let mut end = proto::QueryEnd::default();
558 end.reason = Some(i32::MAX);
559 let mut response = proto::ListTransactionsResponse::default();
560 response.end = Some(end);
561
562 let mut page: StreamPage<ExecutedTransaction> = StreamPage::default();
563 page.apply(frame(response));
564 assert_eq!(page.end_reason, Some(proto::QueryEndReason::Unknown));
565 }
566
567 #[test]
568 fn apply_unknown_frame_does_not_mutate_page() {
569 let response = proto::ListTransactionsResponse::default();
572
573 let mut page: StreamPage<ExecutedTransaction> = StreamPage::default();
574 page.apply(frame(response));
575 assert!(page.items.is_empty());
576 assert_eq!(page.first_cursor(), None);
577 assert_eq!(page.last_cursor(), None);
578 assert_eq!(page.end_reason, None);
579 }
580
581 #[test]
582 fn try_from_item_without_cursor_errors() {
583 let mut response = proto::ListTransactionsResponse::default();
587 response.transaction = Some(ExecutedTransaction::default());
588
589 let result: anyhow::Result<FrameKind<ExecutedTransaction>> = response.try_into();
590 let err = result.expect_err("missing item cursor should error");
591 assert!(
592 err.to_string()
593 .contains("Item frame missing watermark.cursor"),
594 "unexpected error: {err:#}"
595 );
596 }
597
598 #[tokio::test]
599 async fn drain_preserves_partial_progress_on_timeout() {
600 let page = drain_iter(vec![
601 Ok(item_response(b"c1")),
602 Ok(item_response(b"c2")),
603 Err(tonic::Status::cancelled("client channel timeout")),
604 ])
605 .await
606 .expect("partial progress should be preserved");
607
608 assert_eq!(page.items.len(), 2);
609 assert_eq!(page.last_cursor().map(|c| c.as_ref()), Some(b"c2".as_ref()));
610 assert_eq!(page.end_reason, None);
611 assert!(page.has_more());
612 }
613
614 #[tokio::test]
615 async fn drain_errors_on_zero_progress_timeout() {
616 drain_iter(vec![Err(tonic::Status::deadline_exceeded("server budget"))])
617 .await
618 .expect_err("zero-progress timeout should error");
619 }
620
621 #[tokio::test]
622 async fn drain_errors_on_zero_progress_half_close() {
623 drain_iter(vec![])
624 .await
625 .expect_err("zero-progress half-close should error");
626 }
627
628 #[tokio::test]
629 async fn drain_propagates_non_timeout_status() {
630 drain_iter(vec![
634 Ok(item_response(b"c1")),
635 Err(tonic::Status::internal("upstream blew up")),
636 ])
637 .await
638 .expect_err("non-timeout status should propagate as error");
639 }
640
641 #[tokio::test]
642 async fn drain_errors_on_malformed_item_frame() {
643 let mut malformed = proto::ListTransactionsResponse::default();
647 malformed.transaction = Some(ExecutedTransaction::default());
648
649 drain_iter(vec![Ok(item_response(b"c0")), Ok(malformed)])
650 .await
651 .expect_err("malformed Item frame should error the drain");
652 }
653
654 #[tokio::test]
655 async fn drain_returns_page_on_half_close_after_progress() {
656 let page = drain_iter(vec![Ok(item_response(b"c1"))])
659 .await
660 .expect("partial-progress half-close should succeed");
661
662 assert_eq!(page.items.len(), 1);
663 assert_eq!(page.last_cursor().map(|c| c.as_ref()), Some(b"c1".as_ref()));
664 assert_eq!(page.end_reason, None);
665 }
666}