neon/pageserver/page_api/src/model.rs

//! Structs representing the canonical page service API.
//!
//! These mirror the autogenerated Protobuf types. The differences are:
//!
//! - Types that are in fact required by the API are not Options. The protobuf "required"
//!   attribute is deprecated and 'prost' marks a lot of members as optional because of that.
//!   (See <https://github.com/tokio-rs/prost/issues/800> for a gripe on this)
//!
//! - Use more precise datatypes, e.g. Lsn and uints shorter than 32 bits.
//!
//! - Validate protocol invariants, via try_from() and try_into().

use bytes::Bytes;
use postgres_ffi::Oid;
use smallvec::SmallVec;
// TODO: split out Lsn, RelTag, SlruKind, Oid and other basic types to a separate crate, to avoid
// pulling in all of their other crate dependencies when building the client.
use utils::lsn::Lsn;

use crate::proto;

/// A protocol error. Typically returned via try_from() or try_into().
#[derive(thiserror::Error, Debug)]
pub enum ProtocolError {
    #[error("field '{0}' has invalid value '{1}'")]
    Invalid(&'static str, String),
    #[error("required field '{0}' is missing")]
    Missing(&'static str),
}

impl ProtocolError {
    /// Helper to generate a new ProtocolError::Invalid for the given field and value.
    pub fn invalid(field: &'static str, value: impl std::fmt::Debug) -> Self {
        Self::Invalid(field, format!("{value:?}"))
    }
}

impl From<ProtocolError> for tonic::Status {
    fn from(err: ProtocolError) -> Self {
        tonic::Status::invalid_argument(format!("{err}"))
    }
}

/// The LSN a request should read at.
#[derive(Clone, Copy, Debug)]
pub struct ReadLsn {
    /// The request's read LSN.
    pub request_lsn: Lsn,
    /// If given, the caller guarantees that the page has not been modified since this LSN. Must be
    /// smaller than or equal to request_lsn. This allows the Pageserver to serve an old page
    /// without waiting for the request LSN to arrive. Valid for all request types.
    ///
    /// It is undefined behaviour to make a request such that the page was, in fact, modified
    /// between request_lsn and not_modified_since_lsn. The Pageserver might detect it and return an
    /// error, or it might return the old page version or the new page version. Setting
    /// not_modified_since_lsn equal to request_lsn is always safe, but can lead to unnecessary
    /// waiting.
    pub not_modified_since_lsn: Option<Lsn>,
}

impl ReadLsn {
    /// Validates the ReadLsn.
    pub fn validate(&self) -> Result<(), ProtocolError> {
        if self.request_lsn == Lsn::INVALID {
            return Err(ProtocolError::invalid("request_lsn", self.request_lsn));
        }
        if self.not_modified_since_lsn > Some(self.request_lsn) {
            return Err(ProtocolError::invalid(
                "not_modified_since_lsn",
                self.not_modified_since_lsn,
            ));
        }
        Ok(())
    }
}

impl TryFrom<proto::ReadLsn> for ReadLsn {
    type Error = ProtocolError;

    fn try_from(pb: proto::ReadLsn) -> Result<Self, Self::Error> {
        let read_lsn = Self {
            request_lsn: Lsn(pb.request_lsn),
            not_modified_since_lsn: match pb.not_modified_since_lsn {
                0 => None,
                lsn => Some(Lsn(lsn)),
            },
        };
        read_lsn.validate()?;
        Ok(read_lsn)
    }
}

impl TryFrom<ReadLsn> for proto::ReadLsn {
    type Error = ProtocolError;

    fn try_from(read_lsn: ReadLsn) -> Result<Self, Self::Error> {
        read_lsn.validate()?;
        Ok(Self {
            request_lsn: read_lsn.request_lsn.0,
            not_modified_since_lsn: read_lsn.not_modified_since_lsn.unwrap_or_default().0,
        })
    }
}

impl From<&ReadLsn> for proto::ReadLsn {
    fn from(value: &ReadLsn) -> proto::ReadLsn {
        proto::ReadLsn {
            request_lsn: value.request_lsn.into(),
            not_modified_since_lsn: value.not_modified_since_lsn.unwrap_or_default().0,
        }
    }
}

// RelTag is defined in pageserver_api::reltag.
pub type RelTag = pageserver_api::reltag::RelTag;

impl TryFrom<proto::RelTag> for RelTag {
    type Error = ProtocolError;

    fn try_from(pb: proto::RelTag) -> Result<Self, Self::Error> {
        Ok(Self {
            spcnode: pb.spc_oid,
            dbnode: pb.db_oid,
            relnode: pb.rel_number,
            forknum: pb
                .fork_number
                .try_into()
                .map_err(|_| ProtocolError::invalid("fork_number", pb.fork_number))?,
        })
    }
}

impl From<RelTag> for proto::RelTag {
    fn from(rel_tag: RelTag) -> Self {
        Self {
            spc_oid: rel_tag.spcnode,
            db_oid: rel_tag.dbnode,
            rel_number: rel_tag.relnode,
            fork_number: rel_tag.forknum as u32,
        }
    }
}

impl From<&RelTag> for proto::RelTag {
    fn from(value: &RelTag) -> proto::RelTag {
        proto::RelTag {
            spc_oid: value.spcnode,
            db_oid: value.dbnode,
            rel_number: value.relnode,
            fork_number: value.forknum as u32,
        }
    }
}
/// Checks whether a relation exists, at the given LSN. Only valid on shard 0, other shards error.
#[derive(Clone, Copy, Debug)]
pub struct CheckRelExistsRequest {
    pub read_lsn: ReadLsn,
    pub rel: RelTag,
}

impl TryFrom<proto::CheckRelExistsRequest> for CheckRelExistsRequest {
    type Error = ProtocolError;

    fn try_from(pb: proto::CheckRelExistsRequest) -> Result<Self, Self::Error> {
        Ok(Self {
            read_lsn: pb
                .read_lsn
                .ok_or(ProtocolError::Missing("read_lsn"))?
                .try_into()?,
            rel: pb.rel.ok_or(ProtocolError::Missing("rel"))?.try_into()?,
        })
    }
}

impl From<&CheckRelExistsRequest> for proto::CheckRelExistsRequest {
    fn from(value: &CheckRelExistsRequest) -> proto::CheckRelExistsRequest {
        proto::CheckRelExistsRequest {
            read_lsn: Some((&value.read_lsn).into()),
            rel: Some((&value.rel).into()),
        }
    }
}
pub type CheckRelExistsResponse = bool;

impl From<proto::CheckRelExistsResponse> for CheckRelExistsResponse {
    fn from(pb: proto::CheckRelExistsResponse) -> Self {
        pb.exists
    }
}

impl From<CheckRelExistsResponse> for proto::CheckRelExistsResponse {
    fn from(exists: CheckRelExistsResponse) -> Self {
        Self { exists }
    }
}

/// Requests a base backup at a given LSN.
#[derive(Clone, Copy, Debug)]
pub struct GetBaseBackupRequest {
    /// The LSN to fetch a base backup at.
    pub read_lsn: ReadLsn,
    /// If true, logical replication slots will not be created.
    pub replica: bool,
}

impl TryFrom<proto::GetBaseBackupRequest> for GetBaseBackupRequest {
    type Error = ProtocolError;

    fn try_from(pb: proto::GetBaseBackupRequest) -> Result<Self, Self::Error> {
        Ok(Self {
            read_lsn: pb
                .read_lsn
                .ok_or(ProtocolError::Missing("read_lsn"))?
                .try_into()?,
            replica: pb.replica,
        })
    }
}

impl From<&GetBaseBackupRequest> for proto::GetBaseBackupRequest {
    fn from(value: &GetBaseBackupRequest) -> proto::GetBaseBackupRequest {
        proto::GetBaseBackupRequest {
            read_lsn: Some((&value.read_lsn).into()),
            replica: value.replica,
        }
    }
}

impl TryFrom<GetBaseBackupRequest> for proto::GetBaseBackupRequest {
    type Error = ProtocolError;

    fn try_from(request: GetBaseBackupRequest) -> Result<Self, Self::Error> {
        Ok(Self {
            read_lsn: Some(request.read_lsn.try_into()?),
            replica: request.replica,
        })
    }
}

pub type GetBaseBackupResponseChunk = Bytes;

impl TryFrom<proto::GetBaseBackupResponseChunk> for GetBaseBackupResponseChunk {
    type Error = ProtocolError;

    fn try_from(pb: proto::GetBaseBackupResponseChunk) -> Result<Self, Self::Error> {
        if pb.chunk.is_empty() {
            return Err(ProtocolError::Missing("chunk"));
        }
        Ok(pb.chunk)
    }
}

impl TryFrom<GetBaseBackupResponseChunk> for proto::GetBaseBackupResponseChunk {
    type Error = ProtocolError;

    fn try_from(chunk: GetBaseBackupResponseChunk) -> Result<Self, Self::Error> {
        if chunk.is_empty() {
            return Err(ProtocolError::Missing("chunk"));
        }
        Ok(Self { chunk })
    }
}

/// Requests the size of a database, as # of bytes. Only valid on shard 0, other shards will error.
#[derive(Clone, Copy, Debug)]
pub struct GetDbSizeRequest {
    pub read_lsn: ReadLsn,
    pub db_oid: Oid,
}

impl TryFrom<proto::GetDbSizeRequest> for GetDbSizeRequest {
    type Error = ProtocolError;

    fn try_from(pb: proto::GetDbSizeRequest) -> Result<Self, Self::Error> {
        Ok(Self {
            read_lsn: pb
                .read_lsn
                .ok_or(ProtocolError::Missing("read_lsn"))?
                .try_into()?,
            db_oid: pb.db_oid,
        })
    }
}

impl From<&GetDbSizeRequest> for proto::GetDbSizeRequest {
    fn from(value: &GetDbSizeRequest) -> proto::GetDbSizeRequest {
        proto::GetDbSizeRequest {
            read_lsn: Some((&value.read_lsn).into()),
            db_oid: value.db_oid,
        }
    }
}
impl TryFrom<GetDbSizeRequest> for proto::GetDbSizeRequest {
    type Error = ProtocolError;

    fn try_from(request: GetDbSizeRequest) -> Result<Self, Self::Error> {
        Ok(Self {
            read_lsn: Some(request.read_lsn.try_into()?),
            db_oid: request.db_oid,
        })
    }
}

pub type GetDbSizeResponse = u64;

impl From<proto::GetDbSizeResponse> for GetDbSizeResponse {
    fn from(pb: proto::GetDbSizeResponse) -> Self {
        pb.num_bytes
    }
}

impl From<GetDbSizeResponse> for proto::GetDbSizeResponse {
    fn from(num_bytes: GetDbSizeResponse) -> Self {
        Self { num_bytes }
    }
}

/// Requests one or more pages.
#[derive(Clone, Debug)]
pub struct GetPageRequest {
    /// A request ID. Will be included in the response. Should be unique for in-flight requests on
    /// the stream.
    pub request_id: RequestID,
    /// The request class.
    pub request_class: GetPageClass,
    /// The LSN to read at.
    pub read_lsn: ReadLsn,
    /// The relation to read from.
    pub rel: RelTag,
    /// Page numbers to read. Must belong to the remote shard.
    ///
    /// Multiple pages will be executed as a single batch by the Pageserver, amortizing layer access
    /// costs and parallelizing them. This may increase the latency of any individual request, but
    /// improves the overall latency and throughput of the batch as a whole.
    pub block_numbers: SmallVec<[u32; 1]>,
}

impl TryFrom<proto::GetPageRequest> for GetPageRequest {
    type Error = ProtocolError;

    fn try_from(pb: proto::GetPageRequest) -> Result<Self, Self::Error> {
        if pb.block_number.is_empty() {
            return Err(ProtocolError::Missing("block_number"));
        }
        Ok(Self {
            request_id: pb.request_id,
            request_class: pb.request_class.into(),
            read_lsn: pb
                .read_lsn
                .ok_or(ProtocolError::Missing("read_lsn"))?
                .try_into()?,
            rel: pb.rel.ok_or(ProtocolError::Missing("rel"))?.try_into()?,
            block_numbers: pb.block_number.into(),
        })
    }
}

impl From<&GetPageRequest> for proto::GetPageRequest {
    fn from(request: &GetPageRequest) -> proto::GetPageRequest {
        proto::GetPageRequest {
            request_id: request.request_id,
            request_class: request.request_class.into(),
            read_lsn: Some(request.read_lsn.try_into().unwrap()),
            rel: Some(request.rel.into()),
            block_number: request.block_numbers.clone().into_vec(),
        }
    }
}
impl TryFrom<GetPageRequest> for proto::GetPageRequest {
    type Error = ProtocolError;

    fn try_from(request: GetPageRequest) -> Result<Self, Self::Error> {
        if request.block_numbers.is_empty() {
            return Err(ProtocolError::Missing("block_number"));
        }
        Ok(Self {
            request_id: request.request_id,
            request_class: request.request_class.into(),
            read_lsn: Some(request.read_lsn.try_into()?),
            rel: Some(request.rel.into()),
            block_number: request.block_numbers.into_vec(),
        })
    }
}

/// A GetPage request ID.
pub type RequestID = u64;

/// A GetPage request class.
#[derive(Clone, Copy, Debug)]
pub enum GetPageClass {
    /// Unknown class. For backwards compatibility: used when an older client version sends a class
    /// that a newer server version has removed.
    Unknown,
    /// A normal request. This is the default.
    Normal,
    /// A prefetch request. NB: can only be classified on pg < 18.
    Prefetch,
    /// A background request (e.g. vacuum).
    Background,
}

impl From<proto::GetPageClass> for GetPageClass {
    fn from(pb: proto::GetPageClass) -> Self {
        match pb {
            proto::GetPageClass::Unknown => Self::Unknown,
            proto::GetPageClass::Normal => Self::Normal,
            proto::GetPageClass::Prefetch => Self::Prefetch,
            proto::GetPageClass::Background => Self::Background,
        }
    }
}

impl From<i32> for GetPageClass {
    fn from(class: i32) -> Self {
        proto::GetPageClass::try_from(class)
            .unwrap_or(proto::GetPageClass::Unknown)
            .into()
    }
}

impl From<GetPageClass> for proto::GetPageClass {
    fn from(class: GetPageClass) -> Self {
        match class {
            GetPageClass::Unknown => Self::Unknown,
            GetPageClass::Normal => Self::Normal,
            GetPageClass::Prefetch => Self::Prefetch,
            GetPageClass::Background => Self::Background,
        }
    }
}

impl From<GetPageClass> for i32 {
    fn from(class: GetPageClass) -> Self {
        proto::GetPageClass::from(class).into()
    }
}

/// A GetPage response.
///
/// A batch response will contain all of the requested pages. We could eagerly emit individual pages
/// as soon as they are ready, but on a readv() Postgres holds buffer pool locks on all pages in the
/// batch and we'll only return once the entire batch is ready, so no one can make use of the
/// individual pages.
#[derive(Clone, Debug)]
pub struct GetPageResponse {
    /// The original request's ID.
    pub request_id: RequestID,
    /// The response status code.
    pub status_code: GetPageStatusCode,
    /// A string describing the status, if any.
    pub reason: Option<String>,
    /// The 8KB page images, in the same order as the request. Empty if status != OK.
    pub page_images: SmallVec<[Bytes; 1]>,
}

impl From<proto::GetPageResponse> for GetPageResponse {
    fn from(pb: proto::GetPageResponse) -> Self {
        Self {
            request_id: pb.request_id,
            status_code: pb.status_code.into(),
            reason: Some(pb.reason).filter(|r| !r.is_empty()),
            page_images: pb.page_image.into(),
        }
    }
}

impl From<GetPageResponse> for proto::GetPageResponse {
    fn from(response: GetPageResponse) -> Self {
        Self {
            request_id: response.request_id,
            status_code: response.status_code.into(),
            reason: response.reason.unwrap_or_default(),
            page_image: response.page_images.into_vec(),
        }
    }
}

/// A GetPage response status code.
///
/// These are effectively equivalent to gRPC statuses. However, we use a bidirectional stream
/// (potentially shared by many backends), and a gRPC status response would terminate the stream so
/// we send GetPageResponse messages with these codes instead.
#[derive(Clone, Copy, Debug)]
pub enum GetPageStatusCode {
    /// Unknown status. For forwards compatibility: used when an older client version receives a new
    /// status code from a newer server version.
    Unknown,
    /// The request was successful.
    Ok,
    /// The page did not exist. The tenant/timeline/shard has already been validated during stream
    /// setup.
    NotFound,
    /// The request was invalid.
    InvalidRequest,
    /// The request failed due to an internal server error.
    InternalError,
    /// The tenant is rate limited. Slow down and retry later.
    SlowDown,
}

impl From<proto::GetPageStatusCode> for GetPageStatusCode {
    fn from(pb: proto::GetPageStatusCode) -> Self {
        match pb {
            proto::GetPageStatusCode::Unknown => Self::Unknown,
            proto::GetPageStatusCode::Ok => Self::Ok,
            proto::GetPageStatusCode::NotFound => Self::NotFound,
            proto::GetPageStatusCode::InvalidRequest => Self::InvalidRequest,
            proto::GetPageStatusCode::InternalError => Self::InternalError,
            proto::GetPageStatusCode::SlowDown => Self::SlowDown,
        }
    }
}

impl From<i32> for GetPageStatusCode {
    fn from(status_code: i32) -> Self {
        proto::GetPageStatusCode::try_from(status_code)
            .unwrap_or(proto::GetPageStatusCode::Unknown)
            .into()
    }
}

impl From<GetPageStatusCode> for proto::GetPageStatusCode {
    fn from(status_code: GetPageStatusCode) -> Self {
        match status_code {
            GetPageStatusCode::Unknown => Self::Unknown,
            GetPageStatusCode::Ok => Self::Ok,
            GetPageStatusCode::NotFound => Self::NotFound,
            GetPageStatusCode::InvalidRequest => Self::InvalidRequest,
            GetPageStatusCode::InternalError => Self::InternalError,
            GetPageStatusCode::SlowDown => Self::SlowDown,
        }
    }
}

impl From<GetPageStatusCode> for i32 {
    fn from(status_code: GetPageStatusCode) -> Self {
        proto::GetPageStatusCode::from(status_code).into()
    }
}

// Fetches the size of a relation at a given LSN, as # of blocks. Only valid on shard 0, other
// shards will error.
pub struct GetRelSizeRequest {
    pub read_lsn: ReadLsn,
    pub rel: RelTag,
}

impl TryFrom<proto::GetRelSizeRequest> for GetRelSizeRequest {
    type Error = ProtocolError;

    fn try_from(proto: proto::GetRelSizeRequest) -> Result<Self, Self::Error> {
        Ok(Self {
            read_lsn: proto
                .read_lsn
                .ok_or(ProtocolError::Missing("read_lsn"))?
                .try_into()?,
            rel: proto.rel.ok_or(ProtocolError::Missing("rel"))?.try_into()?,
        })
    }
}

impl From<&GetRelSizeRequest> for proto::GetRelSizeRequest {
    fn from(value: &GetRelSizeRequest) -> proto::GetRelSizeRequest {
        proto::GetRelSizeRequest {
            read_lsn: Some((&value.read_lsn).into()),
            rel: Some((&value.rel).into()),
        }
    }
}
impl TryFrom<GetRelSizeRequest> for proto::GetRelSizeRequest {
    type Error = ProtocolError;

    fn try_from(request: GetRelSizeRequest) -> Result<Self, Self::Error> {
        Ok(Self {
            read_lsn: Some(request.read_lsn.try_into()?),
            rel: Some(request.rel.into()),
        })
    }
}

pub type GetRelSizeResponse = u32;

impl From<proto::GetRelSizeResponse> for GetRelSizeResponse {
    fn from(proto: proto::GetRelSizeResponse) -> Self {
        proto.num_blocks
    }
}

impl From<GetRelSizeResponse> for proto::GetRelSizeResponse {
    fn from(num_blocks: GetRelSizeResponse) -> Self {
        Self { num_blocks }
    }
}

/// Requests an SLRU segment. Only valid on shard 0, other shards will error.
pub struct GetSlruSegmentRequest {
    pub read_lsn: ReadLsn,
    pub kind: SlruKind,
    pub segno: u32,
}

impl TryFrom<proto::GetSlruSegmentRequest> for GetSlruSegmentRequest {
    type Error = ProtocolError;

    fn try_from(pb: proto::GetSlruSegmentRequest) -> Result<Self, Self::Error> {
        Ok(Self {
            read_lsn: pb
                .read_lsn
                .ok_or(ProtocolError::Missing("read_lsn"))?
                .try_into()?,
            kind: u8::try_from(pb.kind)
                .ok()
                .and_then(SlruKind::from_repr)
                .ok_or_else(|| ProtocolError::invalid("slru_kind", pb.kind))?,
            segno: pb.segno,
        })
    }
}

impl TryFrom<GetSlruSegmentRequest> for proto::GetSlruSegmentRequest {
    type Error = ProtocolError;

    fn try_from(request: GetSlruSegmentRequest) -> Result<Self, Self::Error> {
        Ok(Self {
            read_lsn: Some(request.read_lsn.try_into()?),
            kind: request.kind as u32,
            segno: request.segno,
        })
    }
}

pub type GetSlruSegmentResponse = Bytes;

impl TryFrom<proto::GetSlruSegmentResponse> for GetSlruSegmentResponse {
    type Error = ProtocolError;

    fn try_from(pb: proto::GetSlruSegmentResponse) -> Result<Self, Self::Error> {
        if pb.segment.is_empty() {
            return Err(ProtocolError::Missing("segment"));
        }
        Ok(pb.segment)
    }
}

impl TryFrom<GetSlruSegmentResponse> for proto::GetSlruSegmentResponse {
    type Error = ProtocolError;

    fn try_from(segment: GetSlruSegmentResponse) -> Result<Self, Self::Error> {
        if segment.is_empty() {
            return Err(ProtocolError::Missing("segment"));
        }
        Ok(Self { segment })
    }
}

// SlruKind is defined in pageserver_api::reltag.
pub type SlruKind = pageserver_api::reltag::SlruKind;