mz_compute_client/protocol/

command.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 commands.

use std::time::Duration;

use mz_cluster_client::client::TryIntoProtocolNonce;
use mz_cluster_client::params::GrpcClientParameters;
use mz_compute_types::dataflows::RenderDataflowDescription;
use mz_dyncfg::ConfigUpdates;
use mz_expr::RowSetFinishing;
use mz_ore::tracing::OpenTelemetryContext;
use mz_persist_types::PersistLocation;
use mz_repr::{GlobalId, RelationDesc, Row, Timestamp};
use mz_storage_types::controller::CollectionMetadata;
use mz_tracing::params::TracingParameters;
use serde::{Deserialize, Serialize};
use timely::progress::frontier::Antichain;
use uuid::Uuid;

use crate::logging::LoggingConfig;

/// Compute protocol commands, sent by the compute controller to replicas.
///
/// Command sequences sent by the compute controller must be valid according to the [Protocol
/// Stages].
///
/// [Protocol Stages]: super#protocol-stages
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum ComputeCommand {
    /// `Hello` is the first command sent to a replica after a connection was established. It
    /// provides the replica with meta information about the connection.
    ///
    /// This command is special in that it is broadcast to all workers of a multi-worker replica.
    /// All subsequent commands, except `UpdateConfiguration`, are only sent to the first worker,
    /// which then distributes them to the other workers using a dataflow.
    Hello {
        /// A nonce unique to the current iteration of the compute protocol.
        ///
        /// The nonce allows identifying different iterations of the compute protocol. When the
        /// compute controller connects to a replica, it must send a nonce that is different from
        /// all nonces it sent to the same replica on previous connections. Multi-worker replicas
        /// should use the nonce to ensure that their individual workers agree on which protocol
        /// iteration they are in.
        nonce: Uuid,
    },

    /// `CreateInstance` must be sent after `Hello` to complete the [Creation Stage] of the compute
    /// protocol. Unlike `Hello`, it is only sent to the first worker of the replica, and then
    /// distributed through the timely runtime. `CreateInstance` instructs the replica to
    /// initialize its state to a point where it is ready to start maintaining dataflows.
    ///
    /// Upon receiving a `CreateInstance` command, the replica must further initialize logging
    /// dataflows according to the given [`LoggingConfig`].
    ///
    /// [Creation Stage]: super#creation-stage
    CreateInstance(Box<InstanceConfig>),

    /// `InitializationComplete` informs the replica about the end of the [Initialization Stage].
    /// Upon receiving this command, the replica should perform a reconciliation process, to ensure
    /// its dataflow state matches the state requested by the computation commands it received
    /// previously. The replica must now start sending responses to commands received previously,
    /// if it opted to defer them during the [Initialization Stage].
    ///
    /// [Initialization Stage]: super#initialization-stage
    InitializationComplete,

    /// `UpdateConfiguration` instructs the replica to update its configuration, according to the
    /// given [`ComputeParameters`].
    ///
    /// This command is special in that, like `Hello`, it is broadcast to all workers of the
    /// replica. However, unlike `Hello`, it is ignored by all workers except the first one, which
    /// distributes the command to the other workers through the timely runtime.
    /// `UpdateConfiguration` commands are broadcast only to allow the intermediary parts of the
    /// networking fabric to observe them and learn of configuration updates.
    ///
    /// Parameter updates transmitted through this command must be applied by the replica as soon
    /// as it receives the command, and they must be applied globally to all replica state, even
    /// dataflows and pending peeks that were created before the parameter update. This property
    /// allows the replica to hoist `UpdateConfiguration` commands during reconciliation.
    ///
    /// Configuration parameters that should not be applied globally, but only to specific
    /// dataflows or peeks, should be added to the [`RenderDataflowDescription`] or [`Peek`] types,
    /// rather than as [`ComputeParameters`].
    UpdateConfiguration(Box<ComputeParameters>),

    /// `CreateDataflow` instructs the replica to create a dataflow according to the given
    /// [`RenderDataflowDescription`].
    ///
    /// The [`RenderDataflowDescription`] must have the following properties:
    ///
    ///   * Dataflow imports are valid:
    ///     * Imported storage collections specified in [`source_imports`] exist and are readable by
    ///       the compute replica.
    ///     * Imported indexes specified in [`index_imports`] have been created on the replica
    ///       previously, by previous `CreateDataflow` commands.
    ///   * Dataflow imports are readable at the specified [`as_of`]. In other words: The `since`s of
    ///     imported collections are not beyond the dataflow [`as_of`].
    ///   * Dataflow exports have unique IDs, i.e., the IDs of exports from dataflows a replica is
    ///     instructed to create do not repeat (within a single protocol iteration).
    ///   * The dataflow objects defined in [`objects_to_build`] are topologically ordered according
    ///     to the dependency relation.
    ///
    /// A dataflow description that violates any of the above properties can cause the replica to
    /// exhibit undefined behavior, such as panicking or production of incorrect results. A replica
    /// should prefer panicking over producing incorrect results.
    ///
    /// After receiving a `CreateDataflow` command, if the created dataflow exports indexes or
    /// storage sinks, the replica must produce [`Frontiers`] responses that report the
    /// advancement of the frontiers of these compute collections.
    ///
    /// After receiving a `CreateDataflow` command, if the created dataflow exports subscribes, the
    /// replica must produce [`SubscribeResponse`]s that report the progress and results of the
    /// subscribes.
    ///
    /// After receiving a `CreateDataflow` command, if the created dataflow exports copy-to sinks,
    /// the replica must produce [`CopyToResponse`]s that report the results and completion of the
    /// copy-to sinks.
    ///
    /// After receiving a `CreateDataflow` command, if the created dataflow exports a persist sink
    /// (materialized view), the dataflow must not write any data to the sink until the replica
    /// receives an `AllowWrites` command for the corresponding collection.
    ///
    /// The replica may create the dataflow in a suspended state and defer starting the computation
    /// until it receives a corresponding `Schedule` command. Thus, to ensure dataflow execution,
    /// the compute controller should eventually send a `Schedule` command for each sent
    /// `CreateDataflow` command.
    ///
    /// [`objects_to_build`]: RenderDataflowDescription::objects_to_build
    /// [`source_imports`]: RenderDataflowDescription::source_imports
    /// [`index_imports`]: RenderDataflowDescription::index_imports
    /// [`as_of`]: RenderDataflowDescription::as_of
    /// [`Frontiers`]: super::response::ComputeResponse::Frontiers
    /// [`SubscribeResponse`]: super::response::ComputeResponse::SubscribeResponse
    /// [`CopyToResponse`]: super::response::ComputeResponse::CopyToResponse
    CreateDataflow(Box<RenderDataflowDescription<CollectionMetadata>>),

    /// `Schedule` allows the replica to start computation for a compute collection.
    ///
    /// It is invalid to send a `Schedule` command that references a collection that was not
    /// created by a corresponding `CreateDataflow` command before. Doing so may cause the replica
    /// to exhibit undefined behavior.
    ///
    /// It is also invalid to send a `Schedule` command that references a collection that has,
    /// through an `AllowCompaction` command, been allowed to compact to the empty frontier before.
    Schedule(GlobalId),

    /// `AllowWrites` allows the replica to start writing external state for a compute collection.
    ///
    /// It is invalid to send an `AllowWrites` command that references a compute collection
    /// that was not created by a corresponding `CreateDataflow` command before. Doing so may cause
    /// the replica to exhibit undefined behavior.
    ///
    /// It is also invalid to send an `AllowWrites` command for a compute collection that has
    /// been allowed to compact to the empty frontier through an `AllowCompaction` command before.
    ///
    /// NOTE: This command only allows writes to persist. Other external state, such as copy-to
    /// sinks, do not need an explicit permission to write.
    ///
    /// NOTE: We don't have a protocol in place that allows writes only after a
    /// certain, controller-determined, timestamp. Such a protocol would allow
    /// tighter control and could allow the instance to avoid work. However, it
    /// is more work to put in place the logic for that so we leave it as future
    /// work for now.
    ///
    /// Note: The `AllowWrites` command is per collection, but a dataflow could host multiple
    /// collections. We're accepting this impedance mismatch for now.
    AllowWrites(GlobalId),

    /// `AllowCompaction` informs the replica about the relaxation of external read capabilities on
    /// a compute collection exported by one of the replica's dataflows.
    ///
    /// The command names a collection and provides a frontier after which accumulations must be
    /// correct. The replica gains the liberty of compacting the corresponding maintained trace up
    /// through that frontier.
    ///
    /// It is invalid to send an `AllowCompaction` command that references a compute collection
    /// that was not created by a corresponding `CreateDataflow` command before. Doing so may cause
    /// the replica to exhibit undefined behavior.
    ///
    /// The `AllowCompaction` command only informs about external read requirements, not internal
    /// ones. The replica is responsible for ensuring that internal requirements are fulfilled at
    /// all times, so local dataflow inputs are not compacted beyond times at which they are still
    /// being read from.
    ///
    /// The read frontiers transmitted through `AllowCompaction`s may be beyond the corresponding
    /// collections' current `upper` frontiers. This signals that external readers are not
    /// interested in times up to the specified new read frontiers. Consequently, an empty read
    /// frontier signals that external readers are not interested in updates from the corresponding
    /// collection ever again, so the collection is not required anymore.
    ///
    /// Sending an `AllowCompaction` command with the empty frontier is the canonical way to drop
    /// compute collections.
    ///
    /// A replica that receives an `AllowCompaction` command with the empty frontier must
    /// eventually respond with [`Frontiers`] responses reporting empty frontiers for the
    /// same collection. ([#16271])
    ///
    /// [`Frontiers`]: super::response::ComputeResponse::Frontiers
    /// [#16271]: https://github.com/MaterializeInc/database-issues/issues/4699
    AllowCompaction {
        /// TODO(database-issues#7533): Add documentation.
        id: GlobalId,
        /// TODO(database-issues#7533): Add documentation.
        frontier: Antichain<Timestamp>,
    },

    /// `Peek` instructs the replica to perform a peek on a collection: either an index or a
    /// Persist-backed collection.
    ///
    /// The [`Peek`] description must have the following properties:
    ///
    ///   * If targeting an index, it has previously been created by a corresponding `CreateDataflow`
    ///     command. (If targeting a persist collection, that collection should exist.)
    ///   * The [`Peek::uuid`] is unique, i.e., the UUIDs of peeks a replica gets instructed to
    ///     perform do not repeat (within a single protocol iteration).
    ///
    /// A [`Peek`] description that violates any of the above properties can cause the replica to
    /// exhibit undefined behavior.
    ///
    /// Specifying a [`Peek::timestamp`] that is less than the target index's `since` frontier does
    /// not provoke undefined behavior. Instead, the replica must produce a [`PeekResponse::Error`]
    /// in response.
    ///
    /// After receiving a `Peek` command, the replica must eventually produce a single
    /// [`PeekResponse`]:
    ///
    ///    * For peeks that were not cancelled: [`Rows`], [`Stashed`], or [`Error`].
    ///    * For peeks that were cancelled: any response above, or [`Canceled`].
    ///
    /// [`PeekResponse`]: super::response::PeekResponse
    /// [`PeekResponse::Error`]: super::response::PeekResponse::Error
    /// [`Rows`]: super::response::PeekResponse::Rows
    /// [`Stashed`]: super::response::PeekResponse::Stashed
    /// [`Error`]: super::response::PeekResponse::Error
    /// [`Canceled`]: super::response::PeekResponse::Canceled
    Peek(Box<Peek>),

    /// `CancelPeek` instructs the replica to cancel the identified pending peek.
    ///
    /// It is invalid to send a `CancelPeek` command that references a peek that was not created
    /// by a corresponding `Peek` command before. Doing so may cause the replica to exhibit
    /// undefined behavior.
    ///
    /// If a replica cancels a peek in response to a `CancelPeek` command, it must respond with a
    /// [`PeekResponse::Canceled`]. The replica may also decide to fulfill the peek instead and
    /// return a different [`PeekResponse`], or it may already have returned a response to the
    /// specified peek. In these cases it must *not* return another [`PeekResponse`].
    ///
    /// [`PeekResponse`]: super::response::PeekResponse
    /// [`PeekResponse::Canceled`]: super::response::PeekResponse::Canceled
    CancelPeek {
        /// The identifier of the peek request to cancel.
        ///
        /// This Value must match a [`Peek::uuid`] value transmitted in a previous `Peek` command.
        uuid: Uuid,
    },
}

/// Configuration for a replica, passed with the `CreateInstance`. Replicas should halt
/// if the controller attempt to reconcile them with different values
/// for anything in this struct.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct InstanceConfig {
    /// Specification of introspection logging.
    pub logging: LoggingConfig,
    /// The offset relative to the replica startup at which it should expire. None disables feature.
    pub expiration_offset: Option<Duration>,
    /// The persist location where we can stash large peek results.
    pub peek_stash_persist_location: PersistLocation,
    /// Whether arrangements created by this replica use dictionary compression.
    ///
    /// Captured from `enable_arrangement_dictionary_compression_alpha` when the replica is created and
    /// held fixed for the replica's lifetime, so flipping the flag only affects replicas created
    /// afterwards rather than retroactively changing arrangements across the environment.
    pub arrangement_dictionary_compression: bool,
    /// A snapshot of the controller's dynamic configuration at replica creation, including any
    /// replica-scoped overrides.
    ///
    /// The replica applies this to its worker configuration before running create-time setup, so
    /// that work performed while handling `CreateInstance` (for example rendering introspection
    /// dataflows) observes controller-synced values rather than dyncfg defaults. The values are
    /// also delivered through the subsequent `UpdateConfiguration`, which applies globally and
    /// keeps the configuration current for the replica's lifetime. An empty snapshot leaves the
    /// worker at its defaults until that first `UpdateConfiguration`.
    pub initial_config: ConfigUpdates,
}

impl InstanceConfig {
    /// Check if the configuration is compatible with another configuration. This is true iff the
    /// logging configuration is equivalent, and the other configuration (non-strictly) strengthens
    /// the expiration offset.
    ///
    /// We consider a stricter offset compatible, which allows us to strengthen the value without
    /// forcing replica restarts. However, it also means that replicas will only pick up the new
    /// value after a restart.
    ///
    /// Dictionary compression is intentionally excluded from this check: it is captured at replica
    /// creation, so a change is always compatible and a running replica keeps the value it was
    /// created with (a new value is only picked up by replicas created afterwards).
    ///
    /// The initial config snapshot is likewise excluded: it carries dyncfg values that apply
    /// globally and are kept current through `UpdateConfiguration`, so a difference across
    /// reconnects is expected and does not require a restart.
    pub fn compatible_with(&self, other: &InstanceConfig) -> bool {
        // Destructure to protect against adding fields in the future.
        let InstanceConfig {
            logging: self_logging,
            expiration_offset: self_offset,
            peek_stash_persist_location: self_peek_stash_persist_location,
            // Captured at replica creation; intentionally not part of compatibility (see above).
            arrangement_dictionary_compression: _,
            // Globally-applied dyncfg snapshot, intentionally not part of compatibility (see above).
            initial_config: _,
        } = self;
        let InstanceConfig {
            logging: other_logging,
            expiration_offset: other_offset,
            peek_stash_persist_location: other_peek_stash_persist_location,
            arrangement_dictionary_compression: _,
            initial_config: _,
        } = other;

        // Logging is compatible if exactly the same.
        let logging_compatible = self_logging == other_logging;

        // The offsets are compatible of other_offset is less than or equal to self_offset, i.e., it
        // is a smaller offset and strengthens the offset.
        let self_offset = Antichain::from_iter(*self_offset);
        let other_offset = Antichain::from_iter(*other_offset);
        let offset_compatible = timely::PartialOrder::less_equal(&other_offset, &self_offset);

        let persist_location_compatible =
            self_peek_stash_persist_location == other_peek_stash_persist_location;

        logging_compatible && offset_compatible && persist_location_compatible
    }
}

/// Compute instance configuration parameters.
///
/// Parameters can be set (`Some`) or unset (`None`).
/// Unset parameters should be interpreted to mean "use the previous value".
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct ComputeParameters {
    /// An optional arbitrary string that describes the class of the workload
    /// this compute instance is running (e.g., `production` or `staging`).
    ///
    /// When `Some(x)`, a `workload_class=x` label is applied to all metrics
    /// exported by the metrics registry associated with the compute instance.
    pub workload_class: Option<Option<String>>,
    /// The maximum allowed size in bytes for results of peeks and subscribes.
    ///
    /// Peeks and subscribes that would return results larger than this maximum return the
    /// respective error responses instead:
    ///   * [`PeekResponse::Rows`] is replaced by [`PeekResponse::Error`].
    ///   * The [`SubscribeBatch::updates`] field is populated with an [`Err`] value.
    ///
    /// [`PeekResponse::Rows`]: super::response::PeekResponse::Rows
    /// [`PeekResponse::Error`]: super::response::PeekResponse::Error
    /// [`SubscribeBatch::updates`]: super::response::SubscribeBatch::updates
    pub max_result_size: Option<u64>,
    /// Tracing configuration.
    pub tracing: TracingParameters,
    /// gRPC client configuration.
    pub grpc_client: GrpcClientParameters,

    /// Config updates for components migrated to `mz_dyncfg`.
    pub dyncfg_updates: ConfigUpdates,
}

impl ComputeParameters {
    /// Update the parameter values with the set ones from `other`.
    pub fn update(&mut self, other: ComputeParameters) {
        let ComputeParameters {
            workload_class,
            max_result_size,
            tracing,
            grpc_client,
            dyncfg_updates,
        } = other;

        if workload_class.is_some() {
            self.workload_class = workload_class;
        }
        if max_result_size.is_some() {
            self.max_result_size = max_result_size;
        }

        self.tracing.update(tracing);
        self.grpc_client.update(grpc_client);

        self.dyncfg_updates.extend(dyncfg_updates);
    }

    /// Return whether all parameters are unset.
    pub fn all_unset(&self) -> bool {
        *self == Self::default()
    }
}

/// Metadata specific to the peek variant.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum PeekTarget {
    /// This peek is against an index. Since this should be held in memory on
    /// the target cluster, no additional coordinates are necessary.
    Index {
        /// The id of the (possibly transient) index.
        id: GlobalId,
    },
    /// This peek is against a Persist collection.
    Persist {
        /// The id of the backing Persist collection.
        id: GlobalId,
        /// The identifying metadata of the Persist shard.
        metadata: CollectionMetadata,
    },
}

impl PeekTarget {
    /// Returns the ID of the peeked collection.
    pub fn id(&self) -> GlobalId {
        match self {
            Self::Index { id } => *id,
            Self::Persist { id, .. } => *id,
        }
    }
}

/// Peek a collection, either in an arrangement or Persist.
///
/// This request elicits data from the worker, by naming the
/// collection and some actions to apply to the results before
/// returning them.
///
/// The `timestamp` member must be valid for the arrangement that
/// is referenced by `id`. This means that `AllowCompaction` for
/// this arrangement should not pass `timestamp` before this command.
/// Subsequent commands may arbitrarily compact the arrangements;
/// the dataflow runners are responsible for ensuring that they can
/// correctly answer the `Peek`.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct Peek {
    /// Target-specific metadata.
    pub target: PeekTarget,
    /// The relation description for the rows returned by this peek, before
    /// applying the [RowSetFinishing] but _after_ applying the given
    /// `map_filter_project`.
    pub result_desc: RelationDesc,
    /// If `Some`, then look up only the given keys from the collection (instead of a full scan).
    /// The vector is never empty.
    pub literal_constraints: Option<Vec<Row>>,
    /// The identifier of this peek request.
    ///
    /// Used in responses and cancellation requests.
    pub uuid: Uuid,
    /// The logical timestamp at which the collection is queried.
    pub timestamp: Timestamp,
    /// Actions to apply to the result set before returning them.
    pub finishing: RowSetFinishing,
    /// Linear operation to apply in-line on each result.
    pub map_filter_project: mz_expr::SafeMfpPlan,
    /// An `OpenTelemetryContext` to forward trace information along
    /// to the compute worker to allow associating traces between
    /// the compute controller and the compute worker.
    pub otel_ctx: OpenTelemetryContext,
}

impl TryIntoProtocolNonce for ComputeCommand {
    fn try_into_protocol_nonce(self) -> Result<Uuid, Self> {
        match self {
            ComputeCommand::Hello { nonce } => Ok(nonce),
            cmd => Err(cmd),
        }
    }
}

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

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