mz_cluster/

client.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.

//! An interactive cluster server.

use std::fmt;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use std::task::{Context, Poll};
use std::thread::Thread;

use anyhow::{Error, anyhow};
use async_trait::async_trait;
use differential_dataflow::trace::ExertionLogic;
use mz_cluster_client::client::{TimelyConfig, TryIntoProtocolNonce};
use mz_ore::metric;
use mz_ore::metrics::{ComputedUIntGauge, MetricsRegistry};
use mz_service::client::{GenericClient, Partitionable, Partitioned};
use mz_service::local::LocalClient;
use timely::WorkerConfig;
use timely::communication::allocator::zero_copy::bytes_slab::BytesRefill;
use timely::communication::initialize::WorkerGuards;
use timely::execute::execute_from;
use timely::worker::Worker as TimelyWorker;
use tokio::runtime::Handle;
use tokio::sync::mpsc;
use tracing::{info, info_span};
use uuid::Uuid;

use crate::communication::initialize_networking;

type PartitionedClient<C, R> = Partitioned<LocalClient<C, R>, C, R>;

/// A client managing access to the local portion of a Timely cluster
pub struct ClusterClient<C>
where
    C: ClusterSpec,
    (C::Command, C::Response): Partitionable<C::Command, C::Response>,
{
    /// The actual client to talk to the cluster
    inner: Option<PartitionedClient<C::Command, C::Response>>,
    /// The running timely instance
    timely_container: Arc<Mutex<TimelyContainer<C>>>,
}

/// Metadata about timely workers in this process.
pub struct TimelyContainer<C: ClusterSpec> {
    /// Channels over which to send endpoints for wiring up a new Client
    client_txs: Vec<
        mpsc::UnboundedSender<(
            Uuid,
            mpsc::UnboundedReceiver<C::Command>,
            mpsc::UnboundedSender<C::Response>,
        )>,
    >,
    /// Thread guards that keep worker threads alive
    worker_guards: WorkerGuards<()>,
}

impl<C: ClusterSpec> TimelyContainer<C> {
    /// The threads of the Timely workers in this process.
    pub fn worker_threads(&self) -> Vec<Thread> {
        self.worker_guards
            .guards()
            .iter()
            .map(|h| h.thread().clone())
            .collect()
    }
}

/// A client to a secondary ("guest") command stream served by workers of
/// an existing Timely cluster. Like [`ClusterClient`], but the per-worker client channels are
/// provided externally instead of coming from a [`TimelyContainer`] built for this command type.
pub struct GuestClusterClient<Cmd, Resp>
where
    (Cmd, Resp): Partitionable<Cmd, Resp>,
{
    /// Per-worker channels over which to send endpoints for wiring up a new client.
    client_txs: Arc<
        Vec<
            mpsc::UnboundedSender<(
                Uuid,
                mpsc::UnboundedReceiver<Cmd>,
                mpsc::UnboundedSender<Resp>,
            )>,
        >,
    >,
    /// The worker threads, for unparking on send.
    worker_threads: Vec<Thread>,
    /// The actual client to talk to the cluster.
    inner: Option<PartitionedClient<Cmd, Resp>>,
}

impl<Cmd, Resp> GuestClusterClient<Cmd, Resp>
where
    Cmd: fmt::Debug + Send + TryIntoProtocolNonce,
    Resp: fmt::Debug + Send,
    (Cmd, Resp): Partitionable<Cmd, Resp>,
{
    /// Create a new `GuestClusterClient`.
    pub fn new(
        client_txs: Arc<
            Vec<
                mpsc::UnboundedSender<(
                    Uuid,
                    mpsc::UnboundedReceiver<Cmd>,
                    mpsc::UnboundedSender<Resp>,
                )>,
            >,
        >,
        worker_threads: Vec<Thread>,
    ) -> Self {
        Self {
            client_txs,
            worker_threads,
            inner: None,
        }
    }

    fn connect(&mut self, nonce: Uuid) {
        let mut command_txs = Vec::new();
        let mut response_rxs = Vec::new();
        for client_tx in self.client_txs.iter() {
            let (cmd_tx, cmd_rx) = mpsc::unbounded_channel();
            let (resp_tx, resp_rx) = mpsc::unbounded_channel();

            client_tx
                .send((nonce, cmd_rx, resp_tx))
                .expect("worker not dropped");

            command_txs.push(cmd_tx);
            response_rxs.push(resp_rx);
        }

        self.inner = Some(LocalClient::new_partitioned(
            response_rxs,
            command_txs,
            self.worker_threads.clone(),
        ));
    }
}

impl<Cmd, Resp> fmt::Debug for GuestClusterClient<Cmd, Resp>
where
    (Cmd, Resp): Partitionable<Cmd, Resp>,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("GuestClusterClient").finish_non_exhaustive()
    }
}

impl<Cmd, Resp> GenericClient<Cmd, Resp> for GuestClusterClient<Cmd, Resp>
where
    Cmd: fmt::Debug + Send + TryIntoProtocolNonce,
    Resp: fmt::Debug + Send,
    (Cmd, Resp): Partitionable<Cmd, Resp>,
{
    async fn send(&mut self, cmd: Cmd) -> Result<(), Error> {
        match cmd.try_into_protocol_nonce() {
            Ok(nonce) => {
                self.connect(nonce);
                Ok(())
            }
            Err(cmd) => self.inner.as_mut().expect("initialized").send(cmd).await,
        }
    }

    fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Result<Option<Resp>, Error>> {
        match self.inner.as_mut() {
            Some(client) => client.poll_recv(cx),
            None => Poll::Pending,
        }
    }
}

impl<C> ClusterClient<C>
where
    C: ClusterSpec,
    (C::Command, C::Response): Partitionable<C::Command, C::Response>,
{
    /// Create a new `ClusterClient`.
    pub fn new(timely_container: Arc<Mutex<TimelyContainer<C>>>) -> Self {
        Self {
            timely_container,
            inner: None,
        }
    }

    /// Connect to the Timely cluster with the given client nonce.
    fn connect(&mut self, nonce: Uuid) -> Result<(), Error> {
        let timely = self.timely_container.lock().expect("poisoned");

        let mut command_txs = Vec::new();
        let mut response_rxs = Vec::new();
        for client_tx in &timely.client_txs {
            let (cmd_tx, cmd_rx) = mpsc::unbounded_channel();
            let (resp_tx, resp_rx) = mpsc::unbounded_channel();

            client_tx
                .send((nonce, cmd_rx, resp_tx))
                .expect("worker not dropped");

            command_txs.push(cmd_tx);
            response_rxs.push(resp_rx);
        }

        self.inner = Some(LocalClient::new_partitioned(
            response_rxs,
            command_txs,
            timely.worker_threads(),
        ));
        Ok(())
    }
}

impl<C> fmt::Debug for ClusterClient<C>
where
    C: ClusterSpec,
    (C::Command, C::Response): Partitionable<C::Command, C::Response>,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ClusterClient")
            .field("inner", &self.inner)
            .finish_non_exhaustive()
    }
}

impl<C> GenericClient<C::Command, C::Response> for ClusterClient<C>
where
    C: ClusterSpec,
    (C::Command, C::Response): Partitionable<C::Command, C::Response>,
{
    async fn send(&mut self, cmd: C::Command) -> Result<(), Error> {
        // Changing this debug statement requires changing the replica-isolation test
        tracing::debug!("ClusterClient send={:?}", &cmd);

        match cmd.try_into_protocol_nonce() {
            Ok(nonce) => self.connect(nonce),
            Err(cmd) => self.inner.as_mut().expect("initialized").send(cmd).await,
        }
    }

    fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Result<Option<C::Response>, Error>> {
        match self.inner.as_mut() {
            Some(client) => client.poll_recv(cx),
            None => Poll::Pending,
        }
    }
}

/// Specification for a Timely cluster to which a [`ClusterClient`] connects.
///
/// This trait is used to make the [`ClusterClient`] generic over the compute and storage cluster
/// implementations.
#[async_trait]
pub trait ClusterSpec: Clone + Send + Sync + 'static {
    /// The cluster command type.
    type Command: fmt::Debug + Send + TryIntoProtocolNonce;
    /// The cluster response type.
    type Response: fmt::Debug + Send;

    /// The name of this cluster ("compute" or "storage").
    const NAME: &str;

    /// Run the given Timely worker.
    fn run_worker(
        &self,
        timely_worker: &mut TimelyWorker,
        client_rx: mpsc::UnboundedReceiver<(
            Uuid,
            mpsc::UnboundedReceiver<Self::Command>,
            mpsc::UnboundedSender<Self::Response>,
        )>,
    );

    /// Build a Timely cluster using the given config.
    async fn build_cluster(
        &self,
        config: TimelyConfig,
        tokio_executor: Handle,
    ) -> Result<TimelyContainer<Self>, Error> {
        info!("Building timely container with config {config:?}");
        let (client_txs, client_rxs): (Vec<_>, Vec<_>) = (0..config.workers)
            .map(|_| mpsc::unbounded_channel())
            .unzip();
        let client_rxs: Mutex<Vec<_>> = Mutex::new(client_rxs.into_iter().map(Some).collect());

        let refill = if config.enable_zero_copy_lgalloc {
            BytesRefill {
                logic: Arc::new(|size| Box::new(alloc::lgalloc_refill(size))),
                limit: config.zero_copy_limit,
            }
        } else {
            BytesRefill {
                logic: Arc::new(|size| Box::new(vec![0; size])),
                limit: config.zero_copy_limit,
            }
        };

        let (builders, other) = initialize_networking(
            config.workers,
            config.process,
            config.addresses.clone(),
            refill,
            config.enable_zero_copy,
        )
        .await?;

        let mut worker_config = WorkerConfig::default();

        // We set a custom exertion logic for proportionality > 0. A proportionality value of 0
        // means that no arrangement merge effort is exerted and merging occurs only in response to
        // updates.
        if config.arrangement_exert_proportionality > 0 {
            let merge_effort = Some(1000);

            // ExertionLogic defines a function to determine if a spine is sufficiently tidied.
            // Its arguments are an iterator over the index of a layer, the count of batches in the
            // layer and the length of batches at the layer. The iterator enumerates layers from the
            // largest to the smallest layer.

            let arc: ExertionLogic = Arc::new(move |layers| {
                EXERT_POLICY_CALLS.fetch_add(1, Ordering::Relaxed);
                let mut prop = config.arrangement_exert_proportionality;

                // Layers are ordered from largest to smallest.
                // Skip to the largest occupied layer.
                let layers = layers
                    .iter()
                    .copied()
                    .skip_while(|(_idx, count, _len)| *count == 0);

                let mut first = true;
                for (_idx, count, len) in layers {
                    if count > 1 {
                        // Found an in-progress merge that we should continue.
                        EXERT_POLICY_MERGE_GRANTS.fetch_add(1, Ordering::Relaxed);
                        return merge_effort;
                    }

                    if !first && prop > 0 && len > 0 {
                        // Found a non-empty batch within `arrangement_exert_proportionality` of
                        // the largest one.
                        EXERT_POLICY_CONSOLIDATION_GRANTS.fetch_add(1, Ordering::Relaxed);
                        return merge_effort;
                    }

                    first = false;
                    prop /= 2;
                }

                None
            });
            worker_config.set::<ExertionLogic>("differential/default_exert_logic".to_string(), arc);
        }

        let spec = self.clone();
        let worker_guards = execute_from(builders, other, worker_config, move |timely_worker| {
            let worker_idx = timely_worker.index();

            // Per worker tracing span, lets us identify Timely clusters and workers in the logs.
            let span = info_span!("timely", name = Self::NAME, worker_id = worker_idx);
            let _span_guard = span.enter();

            // Every Timely instance in this process names its threads `timely:work-N`, restarting
            // the index at 0, so storage and compute worker threads collide under the same OS
            // thread name. Rename to disambiguate them for profilers and `top -H`.
            mz_ore::process::set_current_thread_name(&format!("{}:{worker_idx}", Self::NAME));

            let _tokio_guard = tokio_executor.enter();
            let client_rx = client_rxs.lock().unwrap()[worker_idx % config.workers]
                .take()
                .unwrap();
            spec.run_worker(timely_worker, client_rx);
        })
        .map_err(|e| anyhow!(e))?;

        Ok(TimelyContainer {
            client_txs,
            worker_guards,
        })
    }
}

// Evaluations of the exertion policy installed above, process-wide across
// every timely runtime and trace implementation. A spine evaluates the policy
// both when exerted and after each insert, so evaluations and grants count
// decisions, not applications of effort. With
// `arrangement_exert_proportionality` at zero no policy is installed and the
// counts stay zero.
static EXERT_POLICY_CALLS: AtomicU64 = AtomicU64::new(0);
static EXERT_POLICY_MERGE_GRANTS: AtomicU64 = AtomicU64::new(0);
static EXERT_POLICY_CONSOLIDATION_GRANTS: AtomicU64 = AtomicU64::new(0);

/// Register gauges for the arrangement exertion policy's decisions.
///
/// The counters are process-wide, so only the first call registers them, and
/// they cover every timely runtime in the process, not only the caller's.
pub fn register_exert_policy_metrics(registry: &MetricsRegistry) {
    static REGISTERED: OnceLock<()> = OnceLock::new();
    REGISTERED.get_or_init(|| {
        let _: ComputedUIntGauge = registry.register_computed_gauge(
            metric!(name: "mz_arrangement_exert_policy_calls_total", help: "Arrangement exertion policy evaluations."),
            || EXERT_POLICY_CALLS.load(Ordering::Relaxed),
        );
        for (reason, grants) in [
            ("active_merge", &EXERT_POLICY_MERGE_GRANTS),
            ("consolidation", &EXERT_POLICY_CONSOLIDATION_GRANTS),
        ] {
            let _: ComputedUIntGauge = registry.register_computed_gauge(
                metric!(name: "mz_arrangement_exert_policy_grants_total", help: "Arrangement exertion policy evaluations that returned effort, by reason.", const_labels: {"reason" => reason}),
                move || grants.load(Ordering::Relaxed),
            );
        }
    });
}

mod alloc {
    /// A Timely communication refill function that uses lgalloc.
    ///
    /// Returns a handle to lgalloc'ed memory if lgalloc can handle
    /// the request, otherwise we fall back to a heap allocation. In either case, the handle must
    /// be dropped to free the memory.
    pub(crate) fn lgalloc_refill(size: usize) -> LgallocHandle {
        match lgalloc::allocate::<u8>(size) {
            Ok((pointer, capacity, handle)) => {
                let handle = Some(handle);
                LgallocHandle {
                    handle,
                    pointer,
                    capacity,
                }
            }
            Err(_) => {
                // Allocate memory
                let mut alloc = vec![0_u8; size];
                // Ensure that the length matches the capacity.
                alloc.shrink_to_fit();
                // Get a pointer to the allocated memory.
                let pointer = std::ptr::NonNull::new(alloc.as_mut_ptr()).unwrap();
                // Forget the vector to avoid dropping it. We'll free the memory in `drop`.
                std::mem::forget(alloc);
                LgallocHandle {
                    handle: None,
                    pointer,
                    capacity: size,
                }
            }
        }
    }

    /// A handle to memory allocated by lgalloc. This can either be memory allocated by lgalloc or
    /// memory allocated by Vec. If the handle is set, it's lgalloc memory. If the handle is None,
    /// it's a regular heap allocation.
    pub(crate) struct LgallocHandle {
        /// Lgalloc handle, set if the memory was allocated by lgalloc.
        handle: Option<lgalloc::Handle>,
        /// Pointer to the allocated memory. Always well-aligned, but can be dangling.
        pointer: std::ptr::NonNull<u8>,
        /// Capacity of the allocated memory in bytes.
        capacity: usize,
    }

    // SAFETY: `LgallocHandle` exclusively owns its allocation (either an
    // lgalloc handle or a leaked `Vec`), and the `NonNull<u8>` is just a pointer
    // into that owned memory. Transferring the handle transfers ownership, so it
    // is safe to send across threads. Required because timely's `BytesRefill`
    // produces `Box<dyn DerefMut<Target=[u8]> + Send>`.
    unsafe impl Send for LgallocHandle {}

    impl std::ops::Deref for LgallocHandle {
        type Target = [u8];
        #[inline(always)]
        fn deref(&self) -> &Self::Target {
            unsafe { std::slice::from_raw_parts(self.pointer.as_ptr(), self.capacity) }
        }
    }

    impl std::ops::DerefMut for LgallocHandle {
        #[inline(always)]
        fn deref_mut(&mut self) -> &mut Self::Target {
            unsafe { std::slice::from_raw_parts_mut(self.pointer.as_ptr(), self.capacity) }
        }
    }

    impl Drop for LgallocHandle {
        fn drop(&mut self) {
            // If we have a handle, it's lgalloc memory. Otherwise, it's a heap allocation.
            if let Some(handle) = self.handle.take() {
                lgalloc::deallocate(handle);
            } else {
                unsafe { Vec::from_raw_parts(self.pointer.as_ptr(), 0, self.capacity) };
            }
            // Update pointer and capacity such that we don't double-free.
            self.pointer = std::ptr::NonNull::dangling();
            self.capacity = 0;
        }
    }
}

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

    /// Registration is guarded process-wide, so under a test runner that
    /// shares the process another test may have registered first, into its
    /// own registry. Either way a second call must not panic, and this
    /// registry holds both families or neither.
    #[mz_ore::test]
    fn registering_twice_does_not_panic() {
        let registry = MetricsRegistry::new();
        register_exert_policy_metrics(&registry);
        register_exert_policy_metrics(&registry);
        let names: Vec<_> = registry
            .gather()
            .into_iter()
            .map(|family| family.name().to_string())
            .filter(|name| name.starts_with("mz_arrangement_exert_policy"))
            .collect();
        assert!(matches!(names.len(), 0 | 2), "{names:?}");
    }
}