Skip to main content

mz_compute/
server.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! An interactive dataflow server.
11
12use std::cell::RefCell;
13use std::collections::{BTreeMap, BTreeSet};
14use std::convert::Infallible;
15use std::fmt::Debug;
16use std::path::PathBuf;
17use std::rc::Rc;
18use std::sync::{Arc, Mutex};
19use std::time::{Duration, Instant};
20
21use anyhow::Error;
22use mz_cluster::client::{ClusterClient, ClusterSpec};
23use mz_cluster_client::client::TimelyConfig;
24use mz_compute_client::protocol::command::ComputeCommand;
25use mz_compute_client::protocol::history::ComputeCommandHistory;
26use mz_compute_client::protocol::response::ComputeResponse;
27use mz_compute_client::service::ComputeClient;
28use mz_ore::halt;
29use mz_ore::metrics::MetricsRegistry;
30use mz_ore::tracing::TracingHandle;
31use mz_persist_client::cache::PersistClientCache;
32use mz_storage_types::connections::ConnectionContext;
33use mz_timely_util::capture::EventLink;
34use mz_txn_wal::operator::TxnsContext;
35use timely::logging::TimelyEvent;
36use timely::progress::Antichain;
37use timely::worker::Worker as TimelyWorker;
38use tokio::sync::mpsc;
39use tokio::sync::mpsc::error::SendError;
40use tracing::{info, trace, warn};
41use uuid::Uuid;
42
43use crate::command_channel;
44use crate::compute_state::{ActiveComputeState, ComputeState, ReportedFrontier};
45use crate::metrics::{ComputeMetrics, WorkerMetrics};
46
47/// Caller-provided configuration for compute.
48#[derive(Clone, Debug)]
49pub struct ComputeInstanceContext {
50    /// A directory that can be used for scratch work.
51    pub scratch_directory: Option<PathBuf>,
52    /// Whether to set core affinity for Timely workers.
53    pub worker_core_affinity: bool,
54    /// Context required to connect to an external sink from compute,
55    /// like the `CopyToS3OneshotSink` compute sink.
56    pub connection_context: ConnectionContext,
57}
58
59/// Which of a process's compute runtimes a given runtime is.
60///
61/// A clusterd process runs a single `Solo` runtime by default. When an interactive runtime is
62/// configured, the process instead runs a `Maintenance` and an `Interactive` runtime side by side.
63/// The named roles share per-process resources (persist cache, metrics registry, log spans). The
64/// role distinguishes them so that only the globals-owning runtime runs the non-idempotent
65/// process-global initializers, and so metric series and log spans do not collide.
66///
67/// `Solo` exists so the single-runtime default stays behaviorally identical to a deployment without
68/// a second runtime: no `role` metric label, and it owns the process globals just as the sole
69/// runtime always has.
70#[derive(Clone, Copy, Debug, PartialEq, Eq)]
71pub enum ComputeRuntimeRole {
72    /// The sole runtime of a single-runtime process. Owns index maintenance and the process-global
73    /// initializers.
74    Solo,
75    /// The maintenance runtime of a two-runtime process. Owns index maintenance and the
76    /// process-global initializers.
77    Maintenance,
78    /// The interactive runtime of a two-runtime process. Shares the process globals owned by
79    /// maintenance and serves reads.
80    ///
81    /// Test-only until the interactive runtime exists to construct it. It is present because the
82    /// `role` label's entire purpose is that two named roles register into one process registry
83    /// without colliding, and nothing else can express that: `Solo` registers the same metric names
84    /// with no `role` label, so prometheus rejects it alongside a named role for differing label
85    /// dimensions rather than treating it as a second series. Verifying non-collision therefore
86    /// needs a second *named* role.
87    ///
88    /// TODO: drop the `cfg` when the interactive runtime lands and constructs this.
89    #[cfg(test)]
90    Interactive,
91}
92
93impl ComputeRuntimeRole {
94    /// The `role` metric/log label for this role, or `None` for `Solo`.
95    ///
96    /// `Solo` omits the label so a single-runtime deployment registers exactly as it did before a
97    /// second runtime existed, keeping exact-match dashboards and alerts unchanged.
98    pub fn label(self) -> Option<&'static str> {
99        match self {
100            ComputeRuntimeRole::Solo => None,
101            ComputeRuntimeRole::Maintenance => Some("maintenance"),
102            #[cfg(test)]
103            ComputeRuntimeRole::Interactive => Some("interactive"),
104        }
105    }
106
107    /// Whether this role runs the non-idempotent, process-global initializers.
108    ///
109    /// `Solo` and `Maintenance` run them. An interactive runtime shares the same process and
110    /// inherits the globals maintenance installs, so re-running them would either double-apply a
111    /// non-idempotent effect or race maintenance.
112    ///
113    /// NOTE: every role a release build can construct owns the globals, so this is constantly true
114    /// outside tests. The distinction becomes load-bearing when the interactive runtime lands.
115    pub fn owns_process_globals(self) -> bool {
116        matches!(
117            self,
118            ComputeRuntimeRole::Solo | ComputeRuntimeRole::Maintenance
119        )
120    }
121}
122
123/// Type alias for the storage timely log reader.
124pub(crate) type StorageTimelyLogReader =
125    Arc<EventLink<mz_repr::Timestamp, Vec<(Duration, TimelyEvent)>>>;
126
127/// Configures the server with compute-specific metrics.
128#[derive(Clone)]
129struct Config {
130    /// `persist` client cache.
131    pub persist_clients: Arc<PersistClientCache>,
132    /// Context necessary for rendering txn-wal operators.
133    pub txns_ctx: TxnsContext,
134    /// A process-global handle to tracing configuration.
135    pub tracing_handle: Arc<TracingHandle>,
136    /// Metrics exposed by compute replicas.
137    pub metrics: ComputeMetrics,
138    /// Other configuration for compute.
139    pub context: ComputeInstanceContext,
140    /// The process-global metrics registry.
141    pub metrics_registry: MetricsRegistry,
142    /// The number of timely workers per process.
143    pub workers_per_process: usize,
144    /// A reader for each storage worker in this process.
145    pub storage_log_readers: Arc<Mutex<Vec<Option<StorageTimelyLogReader>>>>,
146}
147
148/// Initiates a timely dataflow computation, processing compute commands.
149pub async fn serve(
150    timely_config: TimelyConfig,
151    role: ComputeRuntimeRole,
152    metrics_registry: &MetricsRegistry,
153    persist_clients: Arc<PersistClientCache>,
154    txns_ctx: TxnsContext,
155    tracing_handle: Arc<TracingHandle>,
156    context: ComputeInstanceContext,
157    storage_log_readers: Vec<StorageTimelyLogReader>,
158) -> Result<impl Fn() -> Box<dyn ComputeClient> + use<>, Error> {
159    let workers_per_process = timely_config.workers;
160    // Normalize the log-reader vec to exactly one slot per local worker. Empty
161    // input means logging is disabled; pad with `None` so index-based access is
162    // always in bounds.
163    let storage_log_readers = if storage_log_readers.is_empty() {
164        (0..workers_per_process).map(|_| None).collect()
165    } else {
166        assert_eq!(storage_log_readers.len(), workers_per_process);
167        storage_log_readers.into_iter().map(Some).collect()
168    };
169    mz_timely_util::column_pager::metrics::register(
170        metrics_registry,
171        mz_timely_util::column_pager::tiered_policy(),
172    );
173    mz_timely_util::pool_config::metrics::register(metrics_registry);
174
175    let config = Config {
176        persist_clients,
177        txns_ctx,
178        tracing_handle,
179        metrics: ComputeMetrics::register_with(metrics_registry, role),
180        context,
181        metrics_registry: metrics_registry.clone(),
182        workers_per_process,
183        storage_log_readers: Arc::new(Mutex::new(storage_log_readers)),
184    };
185    let tokio_executor = tokio::runtime::Handle::current();
186
187    let timely_container = config.build_cluster(timely_config, tokio_executor).await?;
188    let timely_container = Arc::new(Mutex::new(timely_container));
189
190    let client_builder = move || {
191        let client = ClusterClient::new(Arc::clone(&timely_container));
192        let client: Box<dyn ComputeClient> = Box::new(client);
193        client
194    };
195
196    Ok(client_builder)
197}
198
199/// Error type returned on connection nonce changes.
200///
201/// A nonce change informs workers that subsequent commands come a from a new client connection
202/// and therefore require reconciliation.
203struct NonceChange(Uuid);
204
205/// Endpoint used by workers to receive compute commands.
206///
207/// Observes nonce changes in the command stream and converts them into receive errors.
208struct CommandReceiver {
209    /// The channel supplying commands.
210    inner: command_channel::Receiver,
211    /// The ID of the Timely worker.
212    worker_id: usize,
213    /// The nonce identifying the current cluster protocol incarnation.
214    nonce: Option<Uuid>,
215    /// A stash to enable peeking the next command, used in `try_recv`.
216    stashed_command: Option<ComputeCommand>,
217}
218
219impl CommandReceiver {
220    fn new(inner: command_channel::Receiver, worker_id: usize) -> Self {
221        Self {
222            inner,
223            worker_id,
224            nonce: None,
225            stashed_command: None,
226        }
227    }
228
229    /// Receive the next pending command, if any.
230    ///
231    /// If the next command has a different nonce, this method instead returns an `Err`
232    /// containing the new nonce.
233    fn try_recv(&mut self) -> Result<Option<ComputeCommand>, NonceChange> {
234        if let Some(command) = self.stashed_command.take() {
235            return Ok(Some(command));
236        }
237        let Some((command, nonce)) = self.inner.try_recv() else {
238            return Ok(None);
239        };
240
241        trace!(worker = self.worker_id, %nonce, ?command, "received command");
242
243        if Some(nonce) == self.nonce {
244            Ok(Some(command))
245        } else {
246            self.nonce = Some(nonce);
247            self.stashed_command = Some(command);
248            Err(NonceChange(nonce))
249        }
250    }
251}
252
253/// Endpoint used by workers to send sending compute responses.
254///
255/// Tags responses with the current nonce, allowing receivers to filter out responses intended for
256/// previous client connections.
257pub(crate) struct ResponseSender {
258    /// The channel consuming responses.
259    inner: mpsc::UnboundedSender<(ComputeResponse, Uuid)>,
260    /// The ID of the Timely worker.
261    worker_id: usize,
262    /// The nonce identifying the current cluster protocol incarnation.
263    nonce: Option<Uuid>,
264}
265
266impl ResponseSender {
267    fn new(inner: mpsc::UnboundedSender<(ComputeResponse, Uuid)>, worker_id: usize) -> Self {
268        Self {
269            inner,
270            worker_id,
271            nonce: None,
272        }
273    }
274
275    /// Set the cluster protocol nonce.
276    fn set_nonce(&mut self, nonce: Uuid) {
277        self.nonce = Some(nonce);
278    }
279
280    /// Send a compute response.
281    pub fn send(&self, response: ComputeResponse) -> Result<(), SendError<ComputeResponse>> {
282        let nonce = self.nonce.expect("nonce must be initialized");
283
284        trace!(worker = self.worker_id, %nonce, ?response, "sending response");
285        self.inner
286            .send((response, nonce))
287            .map_err(|SendError((resp, _))| SendError(resp))
288    }
289}
290
291/// State maintained for each worker thread.
292///
293/// Much of this state can be viewed as local variables for the worker thread,
294/// holding state that persists across function calls.
295struct Worker<'w> {
296    /// The underlying Timely worker.
297    timely_worker: &'w mut TimelyWorker,
298    /// The channel over which commands are received.
299    command_rx: CommandReceiver,
300    /// The channel over which responses are sent.
301    response_tx: ResponseSender,
302    compute_state: Option<ComputeState>,
303    /// Compute metrics.
304    metrics: WorkerMetrics,
305    /// A process-global cache of (blob_uri, consensus_uri) -> PersistClient.
306    /// This is intentionally shared between workers
307    persist_clients: Arc<PersistClientCache>,
308    /// Context necessary for rendering txn-wal operators.
309    txns_ctx: TxnsContext,
310    /// A process-global handle to tracing configuration.
311    tracing_handle: Arc<TracingHandle>,
312    context: ComputeInstanceContext,
313    /// The process-global metrics registry.
314    metrics_registry: MetricsRegistry,
315    /// The number of timely workers per process.
316    workers_per_process: usize,
317    /// Reader for storage timely logging events.
318    storage_log_reader: Option<StorageTimelyLogReader>,
319}
320
321impl ClusterSpec for Config {
322    type Command = ComputeCommand;
323    type Response = ComputeResponse;
324
325    const NAME: &str = "compute";
326
327    fn run_worker(
328        &self,
329        timely_worker: &mut TimelyWorker,
330        client_rx: mpsc::UnboundedReceiver<(
331            Uuid,
332            mpsc::UnboundedReceiver<ComputeCommand>,
333            mpsc::UnboundedSender<ComputeResponse>,
334        )>,
335    ) {
336        if self.context.worker_core_affinity {
337            set_core_affinity(timely_worker.index());
338        }
339
340        let worker_id = timely_worker.index();
341        let metrics = self.metrics.for_worker(worker_id);
342
343        // Take this worker's storage log reader, indexed by local worker index
344        // so compute worker x matches storage worker x.
345        let local_index = worker_id % self.workers_per_process;
346        let storage_log_reader = self.storage_log_readers.lock().unwrap()[local_index].take();
347
348        // Create the command channel that broadcasts commands from worker 0 to other workers. We
349        // reuse this channel between client connections, to avoid bugs where different workers end
350        // up creating incompatible sides of the channel dataflow after reconnects.
351        // See database-issues#8964.
352        let (cmd_tx, cmd_rx) = command_channel::render(timely_worker);
353        let (resp_tx, resp_rx) = mpsc::unbounded_channel();
354
355        spawn_channel_adapter(client_rx, cmd_tx, resp_rx, worker_id);
356
357        Worker {
358            timely_worker,
359            command_rx: CommandReceiver::new(cmd_rx, worker_id),
360            response_tx: ResponseSender::new(resp_tx, worker_id),
361            metrics,
362            context: self.context.clone(),
363            persist_clients: Arc::clone(&self.persist_clients),
364            txns_ctx: self.txns_ctx.clone(),
365            compute_state: None,
366            tracing_handle: Arc::clone(&self.tracing_handle),
367            metrics_registry: self.metrics_registry.clone(),
368            workers_per_process: self.workers_per_process,
369            storage_log_reader,
370        }
371        .run()
372    }
373}
374
375/// Set the current thread's core affinity, based on the given `worker_id`.
376#[cfg(not(target_os = "macos"))]
377fn set_core_affinity(worker_id: usize) {
378    use tracing::error;
379
380    let Some(mut core_ids) = core_affinity::get_core_ids() else {
381        error!(worker_id, "unable to get core IDs for setting affinity");
382        return;
383    };
384
385    // The `get_core_ids` docs don't say anything about a guaranteed order of the returned Vec,
386    // so sort it just to be safe.
387    core_ids.sort_unstable_by_key(|i| i.id);
388
389    // On multi-process replicas `worker_id` might be greater than the number of available cores.
390    // However, we assume that we always have at least as many cores as there are local workers.
391    // Violating this assumption is safe but might lead to degraded performance due to skew in core
392    // utilization.
393    let idx = worker_id % core_ids.len();
394    let core_id = core_ids[idx];
395
396    if core_affinity::set_for_current(core_id) {
397        info!(
398            worker_id,
399            core_id = core_id.id,
400            "set core affinity for worker"
401        );
402    } else {
403        error!(
404            worker_id,
405            core_id = core_id.id,
406            "failed to set core affinity for worker"
407        )
408    }
409}
410
411/// Set the current thread's core affinity, based on the given `worker_id`.
412#[cfg(target_os = "macos")]
413fn set_core_affinity(_worker_id: usize) {
414    // Setting core affinity is known to not work on Apple Silicon:
415    // https://github.com/Elzair/core_affinity_rs/issues/22
416    info!("setting core affinity is not supported on macOS");
417}
418
419impl<'w> Worker<'w> {
420    /// Runs a compute worker.
421    pub fn run(&mut self) {
422        // The command receiver is initialized without an nonce, so receiving the first command
423        // always triggers a nonce change.
424        let NonceChange(nonce) = self.recv_command().expect_err("change to first nonce");
425        self.set_nonce(nonce);
426
427        loop {
428            let Err(NonceChange(nonce)) = self.run_client();
429            self.set_nonce(nonce);
430        }
431    }
432
433    fn set_nonce(&mut self, nonce: Uuid) {
434        self.response_tx.set_nonce(nonce);
435    }
436
437    /// Handles commands for a client connection, returns when the nonce changes.
438    fn run_client(&mut self) -> Result<Infallible, NonceChange> {
439        self.reconcile()?;
440
441        // The last time we did periodic maintenance.
442        let mut last_maintenance = Instant::now();
443
444        // Commence normal operation.
445        loop {
446            // Get the maintenance interval, default to zero if we don't have a compute state.
447            let maintenance_interval = self
448                .compute_state
449                .as_ref()
450                .map_or(Duration::ZERO, |state| state.server_maintenance_interval);
451
452            let now = Instant::now();
453            // Determine if we need to perform maintenance, which is true if `maintenance_interval`
454            // time has passed since the last maintenance.
455            let sleep_duration;
456            if now >= last_maintenance + maintenance_interval {
457                last_maintenance = now;
458                sleep_duration = None;
459
460                // Report frontier information back the coordinator.
461                if let Some(mut compute_state) = self.activate_compute() {
462                    compute_state.compute_state.traces.maintenance();
463                    compute_state.report_frontiers();
464                    compute_state.report_metrics();
465                    compute_state.check_expiration();
466                }
467
468                self.metrics.record_shared_row_metrics();
469            } else {
470                // We didn't perform maintenance, sleep until the next maintenance interval.
471                let next_maintenance = last_maintenance + maintenance_interval;
472                sleep_duration = Some(next_maintenance.saturating_duration_since(now))
473            };
474
475            // Step the timely worker, recording the time taken.
476            let timer = self.metrics.timely_step_duration_seconds.start_timer();
477            self.timely_worker.step_or_park(sleep_duration);
478            timer.observe_duration();
479
480            self.handle_pending_commands()?;
481
482            if let Some(mut compute_state) = self.activate_compute() {
483                compute_state.process_peeks();
484                compute_state.process_subscribes();
485                compute_state.process_copy_tos();
486            }
487        }
488    }
489
490    fn handle_pending_commands(&mut self) -> Result<(), NonceChange> {
491        while let Some(cmd) = self.command_rx.try_recv()? {
492            self.handle_command(cmd);
493        }
494        Ok(())
495    }
496
497    fn handle_command(&mut self, cmd: ComputeCommand) {
498        if matches!(&cmd, ComputeCommand::CreateInstance(_)) {
499            self.compute_state = Some(ComputeState::new(
500                Arc::clone(&self.persist_clients),
501                self.txns_ctx.clone(),
502                self.metrics.clone(),
503                Arc::clone(&self.tracing_handle),
504                self.context.clone(),
505                self.metrics_registry.clone(),
506                self.workers_per_process,
507                self.storage_log_reader.take(),
508            ));
509        }
510        self.activate_compute().unwrap().handle_compute_command(cmd);
511    }
512
513    fn activate_compute(&mut self) -> Option<ActiveComputeState<'_>> {
514        if let Some(compute_state) = &mut self.compute_state {
515            Some(ActiveComputeState {
516                timely_worker: &mut *self.timely_worker,
517                compute_state,
518                response_tx: &mut self.response_tx,
519            })
520        } else {
521            None
522        }
523    }
524
525    /// Receive the next compute command.
526    ///
527    /// This method blocks if no command is currently available, but takes care to step the Timely
528    /// worker while doing so.
529    fn recv_command(&mut self) -> Result<ComputeCommand, NonceChange> {
530        loop {
531            if let Some(cmd) = self.command_rx.try_recv()? {
532                return Ok(cmd);
533            }
534
535            let start = Instant::now();
536            self.timely_worker.step_or_park(None);
537            self.metrics
538                .timely_step_duration_seconds
539                .observe(start.elapsed().as_secs_f64());
540        }
541    }
542
543    /// Extract commands until `InitializationComplete`, and make the worker reflect those commands.
544    ///
545    /// This method is meant to be a function of the commands received thus far (as recorded in the
546    /// compute state command history) and the new commands from `command_rx`. It should not be a
547    /// function of other characteristics, like whether the worker has managed to respond to a peek
548    /// or not. Some effort goes in to narrowing our view to only the existing commands we can be sure
549    /// are live at all other workers.
550    ///
551    /// The methodology here is to drain `command_rx` until an `InitializationComplete`, at which point
552    /// the prior commands are "reconciled" in. Reconciliation takes each goal dataflow and looks for an
553    /// existing "compatible" dataflow (per `compatible()`) it can repurpose, with some additional tests
554    /// to be sure that we can cut over from one to the other (no additional compaction, no tails/sinks).
555    /// With any connections established, old orphaned dataflows are allow to compact away, and any new
556    /// dataflows are created from scratch. "Kept" dataflows are allowed to compact up to any new `as_of`.
557    ///
558    /// Some additional tidying happens, cleaning up pending peeks, reported frontiers, and creating a new
559    /// subscribe response buffer. We will need to be vigilant with future modifications to `ComputeState` to
560    /// line up changes there with clean resets here.
561    fn reconcile(&mut self) -> Result<(), NonceChange> {
562        // To initialize the connection, we want to drain all commands until we receive a
563        // `ComputeCommand::InitializationComplete` command to form a target command state.
564        let mut new_commands = Vec::new();
565        loop {
566            match self.recv_command()? {
567                ComputeCommand::InitializationComplete => break,
568                command => new_commands.push(command),
569            }
570        }
571
572        // Commands we will need to apply before entering normal service.
573        // These commands may include dropping existing dataflows, compacting existing dataflows,
574        // and creating new dataflows, in addition to standard peek and compaction commands.
575        // The result should be the same as if dropping all dataflows and running `new_commands`.
576        let mut todo_commands = Vec::new();
577        // We only have a compute history if we are in an initialized state
578        // (i.e. after a `CreateInstance`).
579        // If this is not the case, just copy `new_commands` into `todo_commands`.
580        if let Some(compute_state) = &mut self.compute_state {
581            // Reduce the installed commands.
582            // Importantly, act as if all peeks may have been retired (as we cannot know otherwise).
583            compute_state.command_history.discard_peeks();
584            compute_state.command_history.reduce();
585
586            // At this point, we need to sort out which of the *certainly installed* dataflows are
587            // suitable replacements for the requested dataflows. A dataflow is "certainly installed"
588            // as of a frontier if its compaction allows it to go no further. We ignore peeks for this
589            // reasoning, as we cannot be certain that peeks still exist at any other worker.
590
591            // Having reduced our installed command history retaining no peeks (above), we should be able
592            // to use track down installed dataflows we can use as surrogates for requested dataflows (which
593            // have retained all of their peeks, creating a more demanding `as_of` requirement).
594            // NB: installed dataflows may still be allowed to further compact, and we should double check
595            // this before being too confident. It should be rare without peeks, but could happen with e.g.
596            // multiple outputs of a dataflow.
597
598            // The values with which a prior `CreateInstance` was called, if it was.
599            let mut old_instance_config = None;
600            // Index dataflows by `export_ids().collect()`, as this is a precondition for their compatibility.
601            let mut old_dataflows = BTreeMap::default();
602            // Maintain allowed compaction, in case installed identifiers may have been allowed to compact.
603            let mut old_frontiers = BTreeMap::default();
604            for command in compute_state.command_history.iter() {
605                match command {
606                    ComputeCommand::CreateInstance(config) => {
607                        old_instance_config = Some(config);
608                    }
609                    ComputeCommand::CreateDataflow(dataflow) => {
610                        let export_ids = dataflow.export_ids().collect::<BTreeSet<_>>();
611                        old_dataflows.insert(export_ids, dataflow);
612                    }
613                    ComputeCommand::AllowCompaction { id, frontier } => {
614                        old_frontiers.insert(id, frontier);
615                    }
616                    _ => {
617                        // Nothing to do in these cases.
618                    }
619                }
620            }
621
622            // Compaction commands that can be applied to existing dataflows.
623            let mut old_compaction = BTreeMap::default();
624            // Exported identifiers from dataflows we retain.
625            let mut retain_ids = BTreeSet::default();
626
627            // Traverse new commands, sorting out what remediation we can do.
628            for command in new_commands.iter() {
629                match command {
630                    ComputeCommand::CreateDataflow(dataflow) => {
631                        // Attempt to find an existing match for the dataflow.
632                        let as_of = dataflow.as_of.as_ref().unwrap();
633                        let export_ids = dataflow.export_ids().collect::<BTreeSet<_>>();
634
635                        if let Some(old_dataflow) = old_dataflows.get(&export_ids) {
636                            let compatible = old_dataflow.compatible_with(dataflow);
637                            let uncompacted = !export_ids
638                                .iter()
639                                .flat_map(|id| old_frontiers.get(id))
640                                .any(|frontier| {
641                                    !timely::PartialOrder::less_equal(
642                                        *frontier,
643                                        dataflow.as_of.as_ref().unwrap(),
644                                    )
645                                });
646
647                            // We cannot reconcile subscribe and copy-to sinks at the moment,
648                            // because the response buffer is shared, and to a first approximation
649                            // must be completely reformed.
650                            let subscribe_free = dataflow.subscribe_ids().next().is_none();
651                            let copy_to_free = dataflow.copy_to_ids().next().is_none();
652
653                            // If we have replaced any dependency of this dataflow, we need to
654                            // replace this dataflow, to make it use the replacement.
655                            let dependencies_retained = dataflow
656                                .imported_index_ids()
657                                .all(|id| retain_ids.contains(&id));
658
659                            if compatible
660                                && uncompacted
661                                && subscribe_free
662                                && copy_to_free
663                                && dependencies_retained
664                            {
665                                // Match found; remove the match from the deletion queue,
666                                // and compact its outputs to the dataflow's `as_of`.
667                                old_dataflows.remove(&export_ids);
668                                for id in export_ids.iter() {
669                                    old_compaction.insert(*id, as_of.clone());
670                                }
671                                retain_ids.extend(export_ids);
672                            } else {
673                                warn!(
674                                    ?export_ids,
675                                    ?compatible,
676                                    ?uncompacted,
677                                    ?subscribe_free,
678                                    ?copy_to_free,
679                                    ?dependencies_retained,
680                                    old_as_of = ?old_dataflow.as_of,
681                                    new_as_of = ?as_of,
682                                    "dataflow reconciliation failed",
683                                );
684
685                                // Dump the full dataflow plans if they are incompatible, to
686                                // simplify debugging hard-to-reproduce reconciliation failures.
687                                if !compatible {
688                                    warn!(
689                                        old = ?old_dataflow,
690                                        new = ?dataflow,
691                                        "incompatible dataflows in reconciliation",
692                                    );
693                                }
694
695                                todo_commands
696                                    .push(ComputeCommand::CreateDataflow(dataflow.clone()));
697                            }
698
699                            compute_state.metrics.record_dataflow_reconciliation(
700                                compatible,
701                                uncompacted,
702                                subscribe_free,
703                                copy_to_free,
704                                dependencies_retained,
705                            );
706                        } else {
707                            todo_commands.push(ComputeCommand::CreateDataflow(dataflow.clone()));
708                        }
709                    }
710                    ComputeCommand::CreateInstance(config) => {
711                        // Cluster creation should not be performed again!
712                        if old_instance_config.map_or(false, |old| !old.compatible_with(config)) {
713                            halt!(
714                                "new instance configuration not compatible with existing instance configuration:\n{:?}\nvs\n{:?}",
715                                config,
716                                old_instance_config,
717                            );
718                        }
719                    }
720                    // All other commands we apply as requested.
721                    command => {
722                        todo_commands.push(command.clone());
723                    }
724                }
725            }
726
727            // Issue compaction commands first to reclaim resources.
728            for (_, dataflow) in old_dataflows.iter() {
729                for id in dataflow.export_ids() {
730                    // We want to drop anything that has not yet been dropped,
731                    // and nothing that has already been dropped.
732                    if old_frontiers.get(&id) != Some(&&Antichain::new()) {
733                        old_compaction.insert(id, Antichain::new());
734                    }
735                }
736            }
737            for (&id, frontier) in &old_compaction {
738                let frontier = frontier.clone();
739                todo_commands.insert(0, ComputeCommand::AllowCompaction { id, frontier });
740            }
741
742            // Clean up worker-local state.
743            //
744            // Various aspects of `ComputeState` need to be either uninstalled, or return to a blank slate.
745            // All dropped dataflows should clean up after themselves, as we plan to install new dataflows
746            // re-using the same identifiers.
747            // All re-used dataflows should roll back any believed communicated information (e.g. frontiers)
748            // so that they recommunicate that information as if from scratch.
749
750            // Remove all pending peeks.
751            for (_, peek) in std::mem::take(&mut compute_state.pending_peeks) {
752                // Log dropping the peek request.
753                if let Some(logger) = compute_state.compute_logger.as_mut() {
754                    logger.log(&peek.as_log_event(false));
755                }
756            }
757
758            for (&id, collection) in compute_state.collections.iter_mut() {
759                // Adjust reported frontiers:
760                //  * For dataflows we continue to use, reset to ensure we report something not
761                //    before the new `as_of` next.
762                //  * For dataflows we drop, set to the empty frontier, to ensure we don't report
763                //    anything for them.
764                let retained = retain_ids.contains(&id);
765                let compaction = old_compaction.remove(&id);
766                let new_reported_frontier = match (retained, compaction) {
767                    (true, Some(new_as_of)) => ReportedFrontier::NotReported { lower: new_as_of },
768                    (true, None) => {
769                        unreachable!("retained dataflows are compacted to the new as_of")
770                    }
771                    (false, Some(new_frontier)) => {
772                        assert!(new_frontier.is_empty());
773                        ReportedFrontier::Reported(new_frontier)
774                    }
775                    (false, None) => {
776                        // Logging dataflows are implicitly retained and don't have a new as_of.
777                        // Reset them to the minimal frontier.
778                        ReportedFrontier::new()
779                    }
780                };
781
782                collection.reset_reported_frontiers(new_reported_frontier);
783
784                // Sink tokens should be retained for retained dataflows, and dropped for dropped
785                // dataflows.
786                //
787                // Dropping the tokens of active subscribe and copy-tos makes them place
788                // `DroppedAt` responses into the respective response buffer. We drop those buffers
789                // in the next step, which ensures that we don't send out `DroppedAt` responses for
790                // subscribe/copy-tos dropped during reconciliation.
791                if !retained {
792                    collection.sink_token = None;
793                }
794            }
795
796            // We must drop the response buffers as they are global across all subscribe/copy-tos.
797            // If they were broken out by `GlobalId` then we could drop only the response buffers
798            // of dataflows we drop.
799            compute_state.subscribe_response_buffer = Rc::new(RefCell::new(Vec::new()));
800            compute_state.copy_to_response_buffer = Rc::new(RefCell::new(Vec::new()));
801
802            // The controller expects the logging collections to be readable from the minimum time
803            // initially. We cannot recreate the logging arrangements without restarting the
804            // instance, but we can pad the compacted times with empty data. Doing so is sound
805            // because logging collections from different replica incarnations are considered
806            // distinct TVCs, so the controller doesn't expect any historical consistency from
807            // these collections when it reconnects to a replica.
808            //
809            // TODO(database-issues#8152): Consider resolving this with controller-side reconciliation instead.
810            if let Some(config) = old_instance_config {
811                for id in config.logging.index_logs.values() {
812                    let trace = compute_state
813                        .traces
814                        .remove(id)
815                        .expect("logging trace exists");
816                    let padded = trace.into_padded();
817                    compute_state.traces.set(*id, padded);
818                }
819            }
820        } else {
821            todo_commands.clone_from(&new_commands);
822        }
823
824        // Execute the commands to bring us to `new_commands`.
825        for command in todo_commands.into_iter() {
826            self.handle_command(command);
827        }
828
829        // Overwrite `self.command_history` to reflect `new_commands`.
830        // It is possible that there still isn't a compute state yet.
831        if let Some(compute_state) = &mut self.compute_state {
832            let mut command_history = ComputeCommandHistory::new(self.metrics.for_history());
833            for command in new_commands.iter() {
834                command_history.push(command.clone());
835            }
836            compute_state.command_history = command_history;
837        }
838        Ok(())
839    }
840}
841
842/// Spawn a task to bridge between [`ClusterClient`] and [`Worker`] channels.
843///
844/// The [`Worker`] expects a pair of persistent channels, with punctuation marking reconnects,
845/// while the [`ClusterClient`] provides a new pair of channels on each reconnect.
846fn spawn_channel_adapter(
847    mut client_rx: mpsc::UnboundedReceiver<(
848        Uuid,
849        mpsc::UnboundedReceiver<ComputeCommand>,
850        mpsc::UnboundedSender<ComputeResponse>,
851    )>,
852    command_tx: command_channel::Sender,
853    mut response_rx: mpsc::UnboundedReceiver<(ComputeResponse, Uuid)>,
854    worker_id: usize,
855) {
856    mz_ore::task::spawn(
857        || format!("compute-channel-adapter-{worker_id}"),
858        async move {
859            // To make workers aware of the individual client connections, we tag forwarded
860            // commands with the client nonce. Additionally, we use the nonce to filter out
861            // responses with a different nonce, which are intended for different client
862            // connections.
863            //
864            // It's possible that we receive responses with nonces from the past but also from the
865            // future: Worker 0 might have received a new nonce before us and broadcasted it to our
866            // Timely cluster. When we receive a response with a future nonce, we need to wait with
867            // forwarding it until we have received the same nonce from a client connection.
868            //
869            // Nonces are not ordered so we don't know whether a response nonce is from the past or
870            // the future. We thus assume that every response with an unknown nonce might be from
871            // the future and stash them all. Every time we reconnect, we immediately send all
872            // stashed responses with a matching nonce. Every time we receive a new response with a
873            // nonce that matches our current one, we can discard the entire response stash as we
874            // know that all stashed responses must be from the past.
875            let mut stashed_responses = BTreeMap::<Uuid, Vec<ComputeResponse>>::new();
876
877            while let Some((nonce, mut command_rx, response_tx)) = client_rx.recv().await {
878                // Send stashed responses for this client.
879                if let Some(resps) = stashed_responses.remove(&nonce) {
880                    for resp in resps {
881                        let _ = response_tx.send(resp);
882                    }
883                }
884
885                // Wait for a new response while forwarding received commands.
886                let mut serve_rx_channels = async || loop {
887                    tokio::select! {
888                        msg = command_rx.recv() => match msg {
889                            Some(cmd) => command_tx.send((cmd, nonce)),
890                            None => return Err(()),
891                        },
892                        msg = response_rx.recv() => {
893                            return Ok(msg.expect("worker connected"));
894                        }
895                    }
896                };
897
898                // Serve this connection until we see any of the channels disconnect.
899                loop {
900                    let Ok((resp, resp_nonce)) = serve_rx_channels().await else {
901                        break;
902                    };
903
904                    if resp_nonce == nonce {
905                        // Response for the current connection; forward it.
906                        stashed_responses.clear();
907                        if response_tx.send(resp).is_err() {
908                            break;
909                        }
910                    } else {
911                        // Response for a past or future connection; stash it.
912                        let stash = stashed_responses.entry(resp_nonce).or_default();
913                        stash.push(resp);
914                    }
915                }
916            }
917        },
918    );
919}