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