mz_compute_client/protocol/

response.rs

// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0.

//! Compute protocol responses.

use std::fmt;

use bytesize::ByteSize;
use mz_expr::EvalError;
use mz_expr::row::RowCollection;
use mz_ore::cast::CastFrom;
use mz_ore::tracing::OpenTelemetryContext;
use mz_persist_client::batch::ProtoBatch;
use mz_persist_types::ShardId;
use mz_repr::{GlobalId, RelationDesc, Timestamp, UpdateCollection};
use mz_storage_types::errors::DataflowError;
use serde::{Deserialize, Serialize};
use timely::progress::frontier::Antichain;
use uuid::Uuid;

/// Compute protocol responses, sent by replicas to the compute controller.
///
/// Replicas send `ComputeResponse`s in response to [`ComputeCommand`]s they previously received
/// from the compute controller.
///
/// [`ComputeCommand`]: super::command::ComputeCommand
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum ComputeResponse {
    /// `Frontiers` announces the advancement of the various frontiers of the specified compute
    /// collection.
    ///
    /// Replicas must send `Frontiers` responses for compute collections that are indexes or
    /// storage sinks. Replicas must not send `Frontiers` responses for subscribes and copy-tos
    /// ([#16274]).
    ///
    /// Replicas must never report regressing frontiers. Specifically:
    ///
    ///   * The first frontier of any kind reported for a collection must not be less than that
    ///     collection's initial `as_of` frontier.
    ///   * Subsequent reported frontiers for a collection must not be less than any frontier of
    ///     the same kind reported previously for the same collection.
    ///
    /// Replicas must send `Frontiers` responses that report each frontier kind to have advanced to
    /// the empty frontier in response to an [`AllowCompaction` command] that allows compaction of
    /// the collection to to the empty frontier, unless the frontier has previously advanced to the
    /// empty frontier as part of the regular dataflow computation. ([#16271])
    ///
    /// Once a frontier was reported to have been advanced to the empty frontier, the replica must
    /// not send further `Frontiers` responses with non-`None` values for that frontier kind.
    ///
    /// The replica must not send `Frontiers` responses for collections that have not
    /// been created previously by a [`CreateDataflow` command] or by a [`CreateInstance`
    /// command].
    ///
    /// [`AllowCompaction` command]: super::command::ComputeCommand::AllowCompaction
    /// [`CreateDataflow` command]: super::command::ComputeCommand::CreateDataflow
    /// [`CreateInstance` command]: super::command::ComputeCommand::CreateInstance
    /// [#16271]: https://github.com/MaterializeInc/database-issues/issues/4699
    /// [#16274]: https://github.com/MaterializeInc/database-issues/issues/4701
    Frontiers(GlobalId, FrontiersResponse),

    /// `PeekResponse` reports the result of a previous [`Peek` command]. The peek is identified by
    /// a `Uuid` that matches the command's [`Peek::uuid`].
    ///
    /// The replica must send exactly one `PeekResponse` for every [`Peek` command] it received.
    ///
    /// If the replica did not receive a [`CancelPeek` command] for a peek, it must not send a
    /// [`Canceled`] response for that peek. If the replica did receive a [`CancelPeek` command]
    /// for a peek, it may send any of the three [`PeekResponse`] variants.
    ///
    /// The replica must not send `PeekResponse`s for peek IDs that were not previously specified
    /// in a [`Peek` command].
    ///
    /// [`Peek` command]: super::command::ComputeCommand::Peek
    /// [`CancelPeek` command]: super::command::ComputeCommand::CancelPeek
    /// [`Peek::uuid`]: super::command::Peek::uuid
    /// [`Canceled`]: PeekResponse::Canceled
    PeekResponse(Uuid, PeekResponse, OpenTelemetryContext),

    /// `SubscribeResponse` reports the results emitted by an active subscribe over some time
    /// interval.
    ///
    /// For each subscribe that was installed by a previous [`CreateDataflow` command], the
    /// replica must emit [`Batch`] responses that cover the entire time interval from the
    /// minimum time until the subscribe advances to the empty frontier or is
    /// dropped. The time intervals of consecutive [`Batch`]es must be increasing, contiguous,
    /// non-overlapping, and non-empty. All updates transmitted in a batch must be consolidated and
    /// have times within that batch’s time interval. All updates' times must be greater than or
    /// equal to `as_of`. The `upper` of the first [`Batch`] of a subscribe must not be less than
    /// that subscribe's initial `as_of` frontier.
    ///
    /// The replica must send [`DroppedAt`] responses if the subscribe was dropped in response to
    /// an [`AllowCompaction` command] that advanced its read frontier to the empty frontier. The
    /// [`DroppedAt`] frontier must be the upper frontier of the last emitted batch.
    ///
    /// The replica must not send a [`DroppedAt`] response if the subscribe’s upper frontier
    /// (reported by [`Batch`] responses) has advanced to the empty frontier (e.g. because its
    /// inputs advanced to the empty frontier).
    ///
    /// Once a subscribe was reported to have advanced to the empty frontier, or has been dropped:
    ///
    ///   * It must no longer read from its inputs.
    ///   * The replica must not send further `SubscribeResponse`s for that subscribe.
    ///
    /// The replica must not send `SubscribeResponse`s for subscribes that have not been
    /// created previously by a [`CreateDataflow` command].
    ///
    /// [`Batch`]: SubscribeResponse::Batch
    /// [`DroppedAt`]: SubscribeResponse::DroppedAt
    /// [`CreateDataflow` command]: super::command::ComputeCommand::CreateDataflow
    /// [`AllowCompaction` command]: super::command::ComputeCommand::AllowCompaction
    SubscribeResponse(GlobalId, SubscribeResponse),

    /// `CopyToResponse` reports the completion of an S3-oneshot sink.
    ///
    /// The replica must send exactly one `CopyToResponse` for every S3-oneshot sink previously
    /// created by a [`CreateDataflow` command].
    ///
    /// The replica must not send `CopyToResponse`s for S3-oneshot sinks that were not previously
    /// created by a [`CreateDataflow` command].
    ///
    /// [`CreateDataflow` command]: super::command::ComputeCommand::CreateDataflow
    CopyToResponse(GlobalId, CopyToResponse),

    /// `Status` reports status updates from replicas to the controller.
    ///
    /// `Status` responses are a way for replicas to stream back introspection data that the
    /// controller can then announce to its clients. They have no effect on the lifecycles of
    /// compute collections. Correct operation of the Compute layer must not rely on `Status`
    /// responses being sent or received.
    ///
    /// `Status` responses that are specific to collections must only be sent for collections that
    /// (a) have previously been created by a [`CreateDataflow` command] and (b) have not yet
    /// been reported to have advanced to the empty frontier.
    ///
    /// [`CreateDataflow` command]: super::command::ComputeCommand::CreateDataflow
    Status(StatusResponse),
}

/// A response reporting advancement of frontiers of a compute collection.
///
/// All contained frontier fields are optional. `None` values imply that the respective frontier
/// has not advanced and the previously reported value is still current.
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct FrontiersResponse {
    /// The collection's new write frontier, if any.
    ///
    /// Upon receiving an updated `write_frontier`, the controller may assume that the contents of the
    /// collection are sealed for all times less than that frontier. Once it has reported the
    /// `write_frontier` as the empty frontier, the replica must no longer change the contents of the
    /// collection.
    pub write_frontier: Option<Antichain<Timestamp>>,
    /// The collection's new input frontier, if any.
    ///
    /// Upon receiving an updated `input_frontier`, the controller may assume that the replica has
    /// finished reading from the collection’s inputs up to that frontier. Once it has reported the
    /// `input_frontier` as the empty frontier, the replica must no longer read from the
    /// collection's inputs.
    pub input_frontier: Option<Antichain<Timestamp>>,
    /// The collection's new output frontier, if any.
    ///
    /// Upon receiving an updated `output_frontier`, the controller may assume that the replica
    /// has finished processing the collection's input up to that frontier.
    ///
    /// The `output_frontier` is often equal to the `write_frontier`, but not always. Some
    /// collections can jump their write frontiers ahead of the times they have finished
    /// processing, causing the `output_frontier` to lag behind the `write_frontier`. Collections
    /// writing materialized views do so in two cases:
    ///
    ///  * `REFRESH` MVs jump their write frontier ahead to the next refresh time.
    ///  * In a multi-replica cluster, slower replicas observe and report the write frontier of the
    ///    fastest replica, by witnessing advancements of the target persist shard's `upper`.
    pub output_frontier: Option<Antichain<Timestamp>>,
}

impl FrontiersResponse {
    /// Returns whether there are any contained updates.
    pub fn has_updates(&self) -> bool {
        self.write_frontier.is_some()
            || self.input_frontier.is_some()
            || self.output_frontier.is_some()
    }
}

/// The response from a `Peek`.
///
/// Note that each `Peek` expects to generate exactly one `PeekResponse`, i.e.
/// we expect a 1:1 contract between `Peek` and `PeekResponse`.
///
/// Encoded with bincode, which identifies a variant by its position, so both ends of a connection
/// have to agree on this declaration. They do: the CTP handshake refuses a peer whose version
/// differs from its own (`mz_service::transport`), so a replica never speaks to a controller that
/// declares a different shape, and the encoding is free to change with the type.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum PeekResponse {
    /// Returned rows of a successful peek.
    Rows(Vec<RowCollection>),
    /// Results of the peek were stashed in persist batches.
    Stashed(Box<StashedPeekResponse>),
    /// Error of an unsuccessful peek.
    Error(PeekError),
    /// The peek was canceled.
    Canceled,
}

impl PeekResponse {
    /// Return the size of row bytes stored inline in this response.
    pub fn inline_byte_len(&self) -> usize {
        match self {
            Self::Rows(rows) => rows.iter().map(|r| r.byte_len()).sum(),
            Self::Stashed(stashed) => stashed.inline_rows.iter().map(|r| r.byte_len()).sum(),
            Self::Error(_) | Self::Canceled => 0,
        }
    }
}

/// The error of an unsuccessful peek.
///
/// The variant decides what the user sees: a `Dataflow` error keeps the structure the dataflow
/// produced and so gets the same SQLSTATE constant folding would have assigned, a
/// `RowIterationLimitExceeded` names a limit the user can raise, and an `Unstructured` error is
/// reported as an internal error. Prefer the structured variants whenever the source has one.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum PeekError {
    /// An error produced while executing the dataflow, for example evaluating an expression over
    /// a collection.
    Dataflow(Box<DataflowError>),
    /// An error from the peek machinery itself, with no structured form.
    Unstructured(String),
    /// A worker examined more rows than `compute_peek_row_iteration_limit` allows.
    RowIterationLimitExceeded {
        /// The limit that was in effect, in rows.
        limit: usize,
    },
    /// A worker accumulated more answer bytes than `max_result_size` allows.
    ResultExceedsMaxSize {
        /// The ceiling that was in effect, in bytes.
        max_result_size: usize,
    },
}

impl PeekError {
    /// Constructs an unstructured peek error.
    pub fn unstructured(message: impl Into<String>) -> Self {
        Self::Unstructured(message.into())
    }
}

impl fmt::Display for PeekError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Dataflow(error) => error.fmt(f),
            Self::Unstructured(error) => f.write_str(error),
            Self::RowIterationLimitExceeded { limit } => write!(
                f,
                "query exceeded the configured row iteration limit of {limit} rows"
            ),
            Self::ResultExceedsMaxSize { max_result_size } => write!(
                f,
                "result exceeds max size of {}",
                ByteSize::b(u64::cast_from(*max_result_size))
            ),
        }
    }
}

impl From<DataflowError> for PeekError {
    fn from(error: DataflowError) -> Self {
        Self::Dataflow(Box::new(error))
    }
}

impl From<EvalError> for PeekError {
    fn from(error: EvalError) -> Self {
        Self::Dataflow(Box::new(error.into()))
    }
}

/// Response from a peek whose results have been stashed into persist.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct StashedPeekResponse {
    /// The number of rows stored in response batches. This is the sum of the
    /// diff values of the contained rows.
    ///
    /// This does _NOT_ include rows in `inline_rows`.
    pub num_rows_batches: u64,
    /// The sum of the encoded sizes of all batches in this response.
    pub encoded_size_bytes: usize,
    /// [RelationDesc] for the rows in these stashed batches of results.
    pub relation_desc: RelationDesc,
    /// The [ShardId] under which result batches have been stashed.
    pub shard_id: ShardId,
    /// Batches of Rows, must be combined with responses from other workers and
    /// consolidated before sending back via a client.
    pub batches: Vec<ProtoBatch>,
    /// Rows that have not been uploaded to the stash, because their total size
    /// did not go above the threshold for using the peek stash.
    ///
    /// We will have a mix of stashed responses and inline responses because the
    /// result sizes across different workers can and will vary.
    pub inline_rows: Vec<RowCollection>,
}

impl StashedPeekResponse {
    /// Total count of [mz_repr::Row]s represented by this collection, considering a
    /// possible `OFFSET` and `LIMIT`.
    pub fn num_rows(&self, offset: usize, limit: Option<usize>) -> usize {
        let num_stashed_rows: usize = usize::cast_from(self.num_rows_batches);
        let num_inline_rows: usize = self.inline_rows.iter().map(|r| r.count()).sum();
        RowCollection::offset_limit(num_stashed_rows + num_inline_rows, offset, limit)
    }

    /// The size in bytes of the encoded rows in this result.
    pub fn size_bytes(&self) -> usize {
        let inline_size: usize = self.inline_rows.iter().map(|r| r.byte_len()).sum();

        self.encoded_size_bytes + inline_size
    }
}

/// Various responses that can be communicated after a COPY TO command.
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
pub enum CopyToResponse {
    /// Returned number of rows for a successful COPY TO.
    RowCount(u64),
    /// Error of an unsuccessful COPY TO.
    Error(String),
    /// The COPY TO sink dataflow was dropped.
    Dropped,
}

/// Various responses that can be communicated about the progress of a SUBSCRIBE command.
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
pub enum SubscribeResponse {
    /// A batch of updates over a non-empty interval of time.
    Batch(SubscribeBatch),
    /// The SUBSCRIBE dataflow was dropped, leaving updates from this frontier onward unspecified.
    DroppedAt(Antichain<Timestamp>),
}

impl SubscribeResponse {
    /// Converts `self` to an error if a maximum size is exceeded.
    pub fn to_error_if_exceeds(&mut self, max_result_size: usize) {
        if let SubscribeResponse::Batch(batch) = self {
            batch.to_error_if_exceeds(max_result_size);
        }
    }
}

/// A batch of updates for the interval `[lower, upper)`.
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
pub struct SubscribeBatch {
    /// The lower frontier of the batch of updates.
    pub lower: Antichain<Timestamp>,
    /// The upper frontier of the batch of updates.
    pub upper: Antichain<Timestamp>,
    /// All updates greater than `lower` and not greater than `upper`.
    ///
    /// Each of the update collections is sorted first by time, and then by the ordering specified
    /// by the order-by of the subscribe. There is no implied ordering between different collections
    /// of updates in the [Vec].
    ///
    /// It's typical to include just a single collection as part of the clusterd's response, but we
    /// may aggregate different batches together later as we combine results from different workers.
    ///
    /// An `Err` variant can be used to indicate e.g. that the size of the updates exceeds internal limits.
    pub updates: Result<Vec<UpdateCollection>, String>,
}

impl SubscribeBatch {
    /// Converts `self` to an error if a maximum size is exceeded.
    fn to_error_if_exceeds(&mut self, max_result_size: usize) {
        if let Ok(updates) = &self.updates {
            let total_size: usize = updates.iter().map(|updates| updates.byte_len()).sum();
            if total_size > max_result_size {
                self.updates = Err(format!(
                    "result exceeds max size of {}",
                    ByteSize::b(u64::cast_from(max_result_size))
                ));
            }
        }
    }
}

/// Status updates replicas can report to the controller.
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
pub enum StatusResponse {
    /// No status responses are implemented currently, but we're leaving the infrastructure around
    /// in anticipation of materialize#31246.
    Placeholder,
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Test to ensure the size of the `ComputeResponse` enum doesn't regress.
    #[mz_ore::test]
    fn test_compute_response_size() {
        assert_eq!(std::mem::size_of::<ComputeResponse>(), 112);
    }
}