compute: publish maintained indexes into the sharing registry

src/clusterd/src/lib.rs

@@ -22,6 +22,7 @@ use mz_build_info::{BuildInfo, build_info};
 use mz_cloud_resources::AwsExternalIdPrefix;
 use mz_cluster_client::client::TimelyConfig;
 use mz_compute::server::{ComputeInstanceContext, ComputeRuntimeRole};
+use mz_compute::sharing::ArrangementSharingRegistry;
 use mz_http_util::DynamicFilterTarget;
 use mz_orchestrator_tracing::{StaticTracingConfig, TracingCliArgs};
 use mz_ore::cli::{self, CliConfig};
@@ -473,11 +474,16 @@ async fn run(args: Args) -> Result<(), anyhow::Error> {
     );
 
     // Start compute server.
+    //
+    // The sharing registry is per process rather than per runtime: a reader on one runtime looks up
+    // the slot a publisher on another runtime filled, so both must hold the same registry.
+    let sharing_registry = ArrangementSharingRegistry::new();
     let compute_client_builder = mz_compute::server::serve(
         compute_timely_config,
         ComputeRuntimeRole::Solo,
         &metrics_registry,
         persist_clients,
+        sharing_registry,
         txns_ctx,
         tracing_handle,
         ComputeInstanceContext {

src/compute/src/compute_state.rs

@@ -75,7 +75,8 @@ use crate::logging::compute::{CollectionLogging, ComputeEvent, PeekEvent};
 use crate::logging::initialize::LoggingTraces;
 use crate::metrics::{CollectionMetrics, WorkerMetrics};
 use crate::render::{LinearJoinSpec, StartSignal};
-use crate::server::{ComputeInstanceContext, ResponseSender};
+use crate::server::{ComputeInstanceContext, ComputeRuntimeRole, ResponseSender};
+use crate::sharing::ArrangementSharingRegistry;
 
 mod peek_result_iterator;
 mod peek_stash;
@@ -189,6 +190,11 @@ pub struct ComputeState {
     /// A process-global cache of (blob_uri, consensus_uri) -> PersistClient.
     /// This is intentionally shared between workers.
     pub persist_clients: Arc<PersistClientCache>,
+    /// A per-process registry of published index arrangements.
+    ///
+    /// Intentionally shared between all workers of the process, each of which publishes into its own
+    /// worker-ordinal slot. `Clone` shares the same underlying map.
+    pub sharing_registry: ArrangementSharingRegistry,
     /// Context necessary for rendering txn-wal operators.
     pub txns_ctx: TxnsContext,
     /// History of commands received by this workers and all its peers.
@@ -247,12 +253,20 @@ pub struct ComputeState {
 
     /// The storage worker forwards its introspection logs to the compute worker.
     pub storage_log_reader: Option<crate::server::StorageTimelyLogReader>,
+
+    /// Which of the process's compute runtimes this state belongs to.
+    ///
+    /// Only the maintenance runtime runs the non-idempotent process-global initializers. The
+    /// interactive runtime shares the same process and inherits those globals.
+    role: ComputeRuntimeRole,
 }
 
 impl ComputeState {
     /// Construct a new `ComputeState`.
     pub fn new(
+        role: ComputeRuntimeRole,
         persist_clients: Arc<PersistClientCache>,
+        sharing_registry: ArrangementSharingRegistry,
         txns_ctx: TxnsContext,
         metrics: WorkerMetrics,
         tracing_handle: Arc<TracingHandle>,
@@ -273,6 +287,7 @@ impl ComputeState {
             peek_stash_persist_location: None,
             compute_logger: None,
             persist_clients,
+            sharing_registry,
             txns_ctx,
             command_history,
             max_result_size: u64::MAX,
@@ -288,9 +303,15 @@ impl ComputeState {
             init_system_time: mz_ore::now::SYSTEM_TIME(),
             replica_expiration: Antichain::default(),
             storage_log_reader,
+            role,
         }
     }
 
+    /// Which of the process's compute runtimes this state serves.
+    pub(crate) fn role(&self) -> ComputeRuntimeRole {
+        self.role
+    }
+
     /// Return a mutable reference to the identified collection.
     ///
     /// Panics if the collection doesn't exist.
@@ -955,6 +976,8 @@ impl<'a> ActiveComputeState<'a> {
             Rc::clone(&self.compute_state.worker_config),
             self.compute_state.workers_per_process,
             storage_log_reader,
+            self.compute_state.role(),
+            self.compute_state.sharing_registry.clone(),
         );
 
         let dataflow_index = Rc::new(dataflow_index);

src/compute/src/logging/initialize.rs

@@ -16,7 +16,7 @@ use differential_dataflow::logging::{DifferentialEvent, DifferentialEventBuilder
 use mz_compute_client::logging::{LogVariant, LoggingConfig};
 use mz_dyncfg::ConfigSet;
 use mz_ore::metrics::MetricsRegistry;
-use mz_repr::{Diff, Timestamp};
+use mz_repr::{Diff, GlobalId, Timestamp};
 use mz_storage_operators::persist_source::Subtime;
 use mz_timely_util::columnar::Column;
 use mz_timely_util::columnar::builder::ColumnBuilder;
@@ -35,7 +35,10 @@ use crate::extensions::arrange::{KeyCollection, MzArrange};
 use crate::logging::compute::{ComputeEvent, ComputeEventBuilder};
 use crate::logging::{BatchLogger, EventQueue, SharedLoggingState};
 use crate::render::errors::DataflowErrorSer;
-use crate::typedefs::{ErrBatcher, ErrBuilder};
+use crate::server::ComputeRuntimeRole;
+use crate::shared_trace::PublishArrangement;
+use crate::sharing::ArrangementSharingRegistry;
+use crate::typedefs::{ErrAgent, ErrBatcher, ErrBuilder, RowRowAgent};
 
 /// Initialize logging dataflows.
 ///
@@ -48,6 +51,8 @@ pub fn initialize(
     worker_config: Rc<ConfigSet>,
     workers_per_process: usize,
     storage_log_reader: Option<crate::server::StorageTimelyLogReader>,
+    role: ComputeRuntimeRole,
+    sharing_registry: ArrangementSharingRegistry,
 ) -> LoggingTraces {
     let interval_ms = std::cmp::max(1, config.interval.as_millis());
 
@@ -74,6 +79,8 @@ pub fn initialize(
         worker_config,
         workers_per_process,
         storage_log_reader,
+        role,
+        sharing_registry,
     };
 
     // Depending on whether we should log the creation of the logging dataflows, we register the
@@ -114,6 +121,11 @@ struct LoggingContext<'a> {
     workers_per_process: usize,
     /// Optional reader for storage timely logging events.
     storage_log_reader: Option<crate::server::StorageTimelyLogReader>,
+    /// This runtime's role. Only `Maintenance` publishes its logging indexes into the sharing
+    /// registry.
+    role: ComputeRuntimeRole,
+    /// The per-process registry maintenance publishes its logging indexes into.
+    sharing_registry: ArrangementSharingRegistry,
 }
 
 pub(crate) struct LoggingTraces {
@@ -206,6 +218,20 @@ impl LoggingContext<'_> {
             let traces = collections
                 .into_iter()
                 .map(|(log, collection)| {
+                    // Publish maintenance's logging index into the sharing registry so the
+                    // interactive runtime serves introspection peeks from it. Gated on the
+                    // Maintenance role inside the helper, so this is a no-op (adds no operators) on
+                    // Interactive and Solo.
+                    if let Some(&id) = self.config.index_logs.get(&log) {
+                        publish_logging_index(
+                            self.role,
+                            &self.sharing_registry,
+                            &scope,
+                            id,
+                            &collection.trace,
+                            &errs,
+                        );
+                    }
                     let bundle = TraceBundle::new(collection.trace, errs.clone())
                         .with_drop(collection.token);
                     (log, bundle)
@@ -353,3 +379,135 @@ impl ExtractTimestamp for (Timestamp, Subtime) {
         self.0
     }
 }
+
+/// Publishes a maintenance logging index's `oks`/`errs` arrangements into the sharing registry so
+/// the interactive runtime serves introspection peeks from them.
+///
+/// Gated strictly on the `Maintenance` role. Interactive must not publish: it reads maintenance's
+/// slot, and its own (empty) copy would clobber it. Solo has no registry peer. The gate is
+/// deliberately stricter than `ComputeRuntimeRole::publishes`, which also admits Interactive.
+///
+/// The arrangements are re-imported from their trace handles into `scope`. The original arrange
+/// streams are consumed inside the per-log construction regions, so only the trace handles survive
+/// here, and `Arranged::publish` needs a live arrangement stream on this scope to attach its
+/// publisher operator.
+fn publish_logging_index(
+    role: ComputeRuntimeRole,
+    registry: &ArrangementSharingRegistry,
+    scope: &timely::dataflow::Scope<'_, Timestamp>,
+    id: GlobalId,
+    oks_trace: &RowRowAgent<Timestamp, Diff>,
+    errs_trace: &ErrAgent<Timestamp, Diff>,
+) {
+    if role != ComputeRuntimeRole::Maintenance {
+        return;
+    }
+
+    // Re-import the trace handles to obtain live arrangement streams `publish` can attach a
+    // publisher operator to. The publisher refreshes its published chain from the trace, the
+    // authoritative source, so the re-import replay only drives the publisher's wakeups.
+    let oks = oks_trace
+        .clone()
+        .import_named(scope.clone(), &format!("PublishLog({id})"));
+    let errs = errs_trace
+        .clone()
+        .import_named(scope.clone(), &format!("PublishLogErr({id})"));
+
+    // Adopt the registry's placeholder for `id` rather than publishing fresh and inserting: whichever
+    // side, this maintenance publish or an interactive import ahead of it, touches `id` first creates
+    // the slot, so filling it in place cannot overwrite a placeholder a reader has already imported.
+    //
+    // Both halves signal on seal. An introspection read whose result is an error (a division-by-zero
+    // surfacing in `mz_compute_error_counts_raw_unified`) carries its data on the errs stream, so an
+    // oks-only signal would leave it stuck.
+    let worker_index = scope.index();
+    let slot = registry.get_or_create_placeholder(id, worker_index, scope.peers());
+    let oks_registry = registry.clone();
+    PublishArrangement::adopt(&oks, &slot.oks, move || {
+        oks_registry.note_frontier(id, worker_index)
+    });
+    let errs_registry = registry.clone();
+    PublishArrangement::adopt(&errs, &slot.errs, move || {
+        errs_registry.note_frontier(id, worker_index)
+    });
+    // `get_or_create_placeholder` does not notify on create.
+    registry.notify(id, worker_index);
+}
+
+#[cfg(test)]
+mod tests {
+    use differential_dataflow::input::Input;
+    use mz_repr::{Diff, GlobalId, Row, Timestamp};
+    use mz_row_spine::{RowRowBatcher, RowRowBuilder};
+    use mz_timely_util::columnation::ColumnationChunker;
+
+    use crate::extensions::arrange::{KeyCollection, MzArrange};
+    use crate::render::errors::DataflowErrorSer;
+    use crate::server::ComputeRuntimeRole;
+    use crate::sharing::ArrangementSharingRegistry;
+    use crate::typedefs::{ErrBatcher, ErrBuilder, ErrSpine, RowRowSpine};
+
+    use super::publish_logging_index;
+
+    /// A logging/introspection index is a `RowRow` `oks` arrangement plus an (empty) `errs`
+    /// arrangement, published into the sharing registry only by the maintenance runtime. Interactive
+    /// and Solo must not publish: interactive reads maintenance's slot rather than clobbering it with
+    /// its own empty copy, and Solo has no registry peer.
+    ///
+    /// Builds real `RowRow`/`Err` arrangements (the exact types the logging path produces) and drives
+    /// [`publish_logging_index`] for each role, asserting only maintenance ends up published.
+    #[mz_ore::test]
+    fn maintenance_publishes_logging_index_others_do_not() {
+        for (role, expect_published) in [
+            (ComputeRuntimeRole::Maintenance, true),
+            (ComputeRuntimeRole::Interactive, false),
+            (ComputeRuntimeRole::Solo, false),
+        ] {
+            let id = GlobalId::System(1);
+            let registry = ArrangementSharingRegistry::new();
+            let registry_in = registry.clone();
+
+            timely::execute_directly(move |worker| {
+                worker.dataflow::<Timestamp, _, _>(|scope| {
+                    let (mut oks_input, oks_collection) =
+                        scope.new_collection::<(Row, Row), Diff>();
+                    let oks = oks_collection.mz_arrange::<
+                        ColumnationChunker<_>,
+                        RowRowBatcher<_, _>,
+                        RowRowBuilder<_, _>,
+                        RowRowSpine<_, _>,
+                    >("test log oks");
+
+                    let (mut errs_input, errs_collection) =
+                        scope.new_collection::<DataflowErrorSer, Diff>();
+                    let errs = KeyCollection::from(errs_collection).mz_arrange::<
+                        ColumnationChunker<_>,
+                        ErrBatcher<_, _>,
+                        ErrBuilder<_, _>,
+                        ErrSpine<_, _>,
+                    >("test log errs");
+
+                    publish_logging_index(
+                        role,
+                        &registry_in,
+                        &scope.clone(),
+                        id,
+                        &oks.trace,
+                        &errs.trace,
+                    );
+
+                    oks_input.advance_to(Timestamp::from(1_u64));
+                    oks_input.flush();
+                    errs_input.advance_to(Timestamp::from(1_u64));
+                    errs_input.flush();
+                });
+            });
+
+            assert_eq!(
+                registry.handles(&id, 0).is_some(),
+                expect_published,
+                "role {role:?} publication mismatch"
+            );
+        }
+    }
+}