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