sui_indexer_alt_reader/

consistent_reader.rs

// Copyright (c) Mysten Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

use std::future::Future;
use std::sync::Arc;
use std::time::Duration;

use anyhow::Context;
use anyhow::anyhow;
use prometheus::Registry;
use sui_indexer_alt_consistent_api::proto::rpc::consistent::v1alpha::consistent_service_client::ConsistentServiceClient;
use sui_types::base_types::ObjectDigest;
use sui_types::base_types::ObjectID;
use sui_types::base_types::ObjectRef;
use sui_types::base_types::SequenceNumber;
use tonic::transport::Channel;
use tracing::instrument;
use url::Url;

pub use sui_indexer_alt_consistent_api::proto::rpc::consistent::v1alpha as proto;

use crate::metrics::ConsistentReaderMetrics;

#[derive(clap::Args, Debug, Clone, Default)]
pub struct ConsistentReaderArgs {
    /// URL of the consistent store gRPC service
    #[arg(long)]
    pub consistent_store_url: Option<Url>,

    /// Time spent waiting for a request to complete, in milliseconds
    #[arg(long)]
    pub consistent_store_statement_timeout_ms: Option<u64>,
}

/// A reader backed by the consistent store gRPC service.
#[derive(Clone)]
pub struct ConsistentReader {
    client: Option<Client>,
    timeout: Option<Duration>,
    metrics: Arc<ConsistentReaderMetrics>,
}

/// Response from a paginated query.
pub struct Page<T> {
    pub results: Vec<Edge<T>>,
    pub has_previous_page: bool,
    pub has_next_page: bool,
}

/// An individual element from a paginated query, includes the value and its token (pagination
/// cursor).
pub struct Edge<T> {
    pub token: Vec<u8>,
    pub value: T,
}

type Client = ConsistentServiceClient<Channel>;

#[derive(thiserror::Error, Debug)]
pub enum Error {
    #[error(transparent)]
    Internal(#[from] anyhow::Error),

    #[error("{}", .0.message())]
    OutOfRange(#[source] tonic::Status),

    #[error("Consistent store client not configured")]
    NotConfigured,
}

impl ConsistentReaderArgs {
    pub fn statement_timeout(&self) -> Option<Duration> {
        self.consistent_store_statement_timeout_ms
            .map(Duration::from_millis)
    }
}

impl ConsistentReader {
    pub async fn new(
        prefix: Option<&str>,
        args: ConsistentReaderArgs,
        registry: &Registry,
    ) -> Result<Self, Error> {
        let client = if let Some(url) = &args.consistent_store_url {
            let mut endpoint = Channel::from_shared(url.to_string())
                .context("Failed to create channel for gRPC endpoint")?;

            if let Some(timeout) = args.statement_timeout() {
                endpoint = endpoint.timeout(timeout);
            }

            let channel = endpoint.connect_lazy();

            Some(ConsistentServiceClient::new(channel))
        } else {
            None
        };

        let timeout = args.statement_timeout();
        let metrics = ConsistentReaderMetrics::new(prefix, registry);

        Ok(Self {
            client,
            timeout,
            metrics,
        })
    }

    /// Get the consistent store's watermarks, as of the given `checkpoint`.
    #[instrument(skip(self), level = "debug")]
    pub async fn available_range(
        &self,
        checkpoint: u64,
    ) -> Result<proto::AvailableRangeResponse, Error> {
        self.request(
            "available_range",
            Some(checkpoint),
            |mut client, request| async move { client.available_range(request).await },
            proto::AvailableRangeRequest {},
        )
        .await
    }

    /// Batch get balances for multiple (owner, coin_type) pairs at a given checkpoint.
    #[instrument(skip(self), level = "debug")]
    pub async fn batch_get_balances(
        &self,
        checkpoint: u64,
        requests: Vec<(String, String)>,
    ) -> Result<Vec<proto::Balance>, Error> {
        let response = self
            .request(
                "batch_get_balances",
                Some(checkpoint),
                |mut client, request| async move { client.batch_get_balances(request).await },
                proto::BatchGetBalancesRequest {
                    requests: requests
                        .into_iter()
                        .map(|(owner, coin_type)| proto::GetBalanceRequest {
                            owner: Some(owner),
                            coin_type: Some(coin_type),
                        })
                        .collect(),
                },
            )
            .await?;

        let mut results = vec![];
        for balance in response.balances {
            results.push(balance);
        }

        Ok(results)
    }

    /// Get the balance for a single (owner, coin_type) pair at a given checkpoint.
    #[instrument(skip(self), level = "debug")]
    pub async fn get_balance(
        &self,
        checkpoint: Option<u64>,
        address: String,
        coin_type: String,
    ) -> Result<proto::Balance, Error> {
        self.request(
            "get_balance",
            checkpoint,
            |mut client, request| async move { client.get_balance(request).await },
            proto::GetBalanceRequest {
                owner: Some(address),
                coin_type: Some(coin_type),
            },
        )
        .await
    }

    /// Paginate coin balances for `address`, at checkpoint `checkpoint`.
    #[instrument(skip(self), level = "debug")]
    pub async fn list_balances(
        &self,
        checkpoint: Option<u64>,
        address: String,
        page_size: Option<u32>,
        after_token: Option<Vec<u8>>,
        before_token: Option<Vec<u8>>,
        is_from_front: bool,
    ) -> Result<Page<proto::Balance>, Error> {
        let response = self
            .request(
                "list_balances",
                checkpoint,
                |mut client, request| async move { client.list_balances(request).await },
                proto::ListBalancesRequest {
                    owner: Some(address),
                    page_size,
                    after_token: after_token.map(Into::into),
                    before_token: before_token.map(Into::into),
                    end: if is_from_front {
                        Some(proto::End::Front.into())
                    } else {
                        Some(proto::End::Back.into())
                    },
                },
            )
            .await?;

        let has_next_page = response.has_next_page();
        let has_previous_page = response.has_previous_page();

        let results = response
            .balances
            .into_iter()
            .map(|b| Edge {
                token: b.page_token.clone().unwrap_or_default().into(),
                value: b,
            })
            .collect();

        Ok(Page {
            results,
            has_next_page,
            has_previous_page,
        })
    }

    /// Paginate live objects with type filter `object_type`, at checkpoint `checkpoint`.
    #[instrument(skip(self), level = "debug")]
    pub async fn list_objects_by_type(
        &self,
        checkpoint: Option<u64>,
        object_type: String,
        page_size: Option<u32>,
        after_token: Option<Vec<u8>>,
        before_token: Option<Vec<u8>>,
        is_from_front: bool,
    ) -> Result<Page<ObjectRef>, Error> {
        let response = self
            .request(
                "list_objects_by_type",
                checkpoint,
                |mut client, request| async move { client.list_objects_by_type(request).await },
                proto::ListObjectsByTypeRequest {
                    object_type: Some(object_type),
                    page_size,
                    after_token: after_token.map(Into::into),
                    before_token: before_token.map(Into::into),
                    end: if is_from_front {
                        Some(proto::End::Front.into())
                    } else {
                        Some(proto::End::Back.into())
                    },
                },
            )
            .await?;

        let has_next_page = response.has_next_page();
        let has_previous_page = response.has_previous_page();

        let results = response
            .objects
            .into_iter()
            .map(TryFrom::try_from)
            .collect::<Result<Vec<_>, _>>()?;

        Ok(Page {
            results,
            has_next_page,
            has_previous_page,
        })
    }

    /// Paginate live objects at `checkpoint`, with owner described by `kind` and `address`, and an
    /// optional, `object_type`. If `object_type` is prefaced with `!`, objects of the type prefix
    /// will be excluded from the page of results.
    #[instrument(skip(self), level = "debug")]
    pub async fn list_owned_objects(
        &self,
        checkpoint: Option<u64>,
        kind: proto::owner::OwnerKind,
        address: Option<String>,
        object_type: Option<String>,
        page_size: Option<u32>,
        after_token: Option<Vec<u8>>,
        before_token: Option<Vec<u8>>,
        is_from_front: bool,
    ) -> Result<Page<ObjectRef>, Error> {
        let response = self
            .request(
                "list_owned_objects",
                checkpoint,
                |mut client, request| async move { client.list_owned_objects(request).await },
                proto::ListOwnedObjectsRequest {
                    owner: Some(proto::Owner {
                        kind: Some(kind.into()),
                        address,
                    }),
                    object_type,
                    page_size,
                    after_token: after_token.map(Into::into),
                    before_token: before_token.map(Into::into),
                    end: if is_from_front {
                        Some(proto::End::Front.into())
                    } else {
                        Some(proto::End::Back.into())
                    },
                },
            )
            .await?;

        let has_next_page = response.has_next_page();
        let has_previous_page = response.has_previous_page();

        let results = response
            .objects
            .into_iter()
            .map(TryFrom::try_from)
            .collect::<Result<Vec<_>, _>>()?;

        Ok(Page {
            results,
            has_next_page,
            has_previous_page,
        })
    }

    async fn request<I, O, Fut, F>(
        &self,
        method: &str,
        checkpoint: Option<u64>,
        response: F,
        input: I,
    ) -> Result<O, Error>
    where
        F: FnOnce(Client, tonic::Request<I>) -> Fut,
        Fut: Future<Output = Result<tonic::Response<O>, tonic::Status>>,
    {
        let Some(client) = self.client.clone() else {
            return Err(Error::NotConfigured);
        };

        self.metrics
            .requests_received
            .with_label_values(&[method])
            .inc();

        let _timer = self
            .metrics
            .latency
            .with_label_values(&[method])
            .start_timer();

        let mut request = tonic::Request::new(input);

        if let Some(timeout) = self.timeout {
            request.set_timeout(timeout);
        }

        if let Some(checkpoint) = checkpoint {
            request.metadata_mut().insert(
                proto::CHECKPOINT_HEIGHT_METADATA,
                checkpoint
                    .to_string()
                    .parse()
                    .with_context(|| format!("Invalid checkpoint {checkpoint}"))?,
            );
        }

        let response = response(client, request)
            .await
            .map(|r| r.into_inner())
            .map_err(Into::into);

        if response.is_ok() {
            self.metrics
                .requests_succeeded
                .with_label_values(&[method])
                .inc();
        } else {
            self.metrics
                .requests_failed
                .with_label_values(&[method])
                .inc()
        }

        response
    }
}

impl TryFrom<proto::Object> for Edge<ObjectRef> {
    type Error = Error;

    fn try_from(proto: proto::Object) -> Result<Self, Error> {
        let object_id: ObjectID = proto
            .object_id
            .context("object ID missing")?
            .parse()
            .context("invalid object ID")?;

        let digest: ObjectDigest = proto
            .digest
            .context("digest missing")?
            .parse()
            .context("invalid digest")?;

        let version: SequenceNumber = proto.version.context("version missing")?.into();
        let token: Vec<u8> = proto.page_token.unwrap_or_default().into();

        Ok(Edge {
            token,
            value: (object_id, version, digest),
        })
    }
}

impl From<tonic::Status> for Error {
    fn from(status: tonic::Status) -> Self {
        match status.code() {
            tonic::Code::OutOfRange => Error::OutOfRange(status),
            _ => Error::Internal(anyhow!(status.code()).context(status.message().to_string())),
        }
    }
}