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::task::{Context, Poll};
20use std::thread::Thread;
21use std::time::{Duration, Instant};
22
23use anyhow::Error;
24use mz_cluster::client::{ClusterClient, ClusterSpec, GuestClusterClient};
25use mz_cluster_client::client::TimelyConfig;
26use mz_compute_client::protocol::command::ComputeCommand;
27use mz_compute_client::protocol::history::ComputeCommandHistory;
28use mz_compute_client::protocol::response::ComputeResponse;
29use mz_ore::halt;
30use mz_ore::metrics::MetricsRegistry;
31use mz_ore::now::NowFn;
32use mz_ore::tracing::TracingHandle;
33use mz_persist_client::cache::PersistClientCache;
34use mz_rocksdb::config::SharedWriteBufferManager;
35use mz_service::client::GenericClient;
36use mz_storage::internal_control::{InternalCommandSender, InternalStorageCommand};
37use mz_storage::metrics::StorageMetrics;
38use mz_storage::storage_state::{StorageInstanceContext, StorageState, Worker as StorageWorker};
39use mz_storage_client::client::{StorageCommand, StorageResponse};
40use mz_storage_operators::persist::SharedBatches;
41use mz_storage_types::connections::ConnectionContext;
42use mz_txn_wal::operator::TxnsContext;
43use timely::progress::Antichain;
44use timely::worker::Worker as TimelyWorker;
45use tokio::sync::mpsc;
46use tokio::sync::mpsc::error::SendError;
47use tracing::{info, trace, warn};
48use uuid::Uuid;
49
50use crate::command_channel::{self, StorageLaneInput, UnifiedCommand};
51use crate::compute_state::{
52 ActiveComputeState, ComputeState, PeekPermits, PendingPeek, ReportedFrontier,
53};
54use crate::metrics::{ComputeMetrics, WorkerMetrics};
55
56/// Caller-provided configuration for compute.
57#[derive(Clone, Debug)]
58pub struct ComputeInstanceContext {
59 /// A directory that can be used for scratch work.
60 pub scratch_directory: Option<PathBuf>,
61 /// Whether to set core affinity for Timely workers.
62 pub worker_core_affinity: bool,
63 /// Context required to connect to an external sink from compute,
64 /// like the `CopyToS3OneshotSink` compute sink.
65 pub connection_context: ConnectionContext,
66}
67
68/// Which of a process's compute runtimes a given runtime is.
69///
70/// A clusterd process runs a single `Solo` runtime by default. When an interactive runtime is
71/// configured, the process instead runs a `Maintenance` and an `Interactive` runtime side by side.
72/// The named roles share per-process resources (persist cache, metrics registry, log spans). The
73/// role distinguishes them so that only the globals-owning runtime runs the non-idempotent
74/// process-global initializers, and so metric series and log spans do not collide.
75///
76/// `Solo` exists so the single-runtime default stays behaviorally identical to a deployment without
77/// a second runtime: no `role` metric label, and it owns the process globals just as the sole
78/// runtime always has.
79#[derive(Clone, Copy, Debug, PartialEq, Eq)]
80pub enum ComputeRuntimeRole {
81 /// The sole runtime of a single-runtime process. Owns index maintenance and the process-global
82 /// initializers.
83 Solo,
84 /// The maintenance runtime of a two-runtime process. Owns index maintenance and the
85 /// process-global initializers.
86 Maintenance,
87 /// The interactive runtime of a two-runtime process. Shares the process globals owned by
88 /// maintenance and serves reads.
89 ///
90 /// Test-only until the interactive runtime exists to construct it. It is present because the
91 /// `role` label's entire purpose is that two named roles register into one process registry
92 /// without colliding, and nothing else can express that: `Solo` registers the same metric names
93 /// with no `role` label, so prometheus rejects it alongside a named role for differing label
94 /// dimensions rather than treating it as a second series. Verifying non-collision therefore
95 /// needs a second *named* role.
96 ///
97 /// TODO: drop the `cfg` when the interactive runtime lands and constructs this.
98 #[cfg(test)]
99 Interactive,
100}
101
102impl ComputeRuntimeRole {
103 /// The `role` metric/log label for this role, or `None` for `Solo`.
104 ///
105 /// `Solo` omits the label so a single-runtime deployment registers exactly as it did before a
106 /// second runtime existed, keeping exact-match dashboards and alerts unchanged.
107 pub fn label(self) -> Option<&'static str> {
108 match self {
109 ComputeRuntimeRole::Solo => None,
110 ComputeRuntimeRole::Maintenance => Some("maintenance"),
111 #[cfg(test)]
112 ComputeRuntimeRole::Interactive => Some("interactive"),
113 }
114 }
115
116 /// Whether this role runs the non-idempotent, process-global initializers.
117 ///
118 /// `Solo` and `Maintenance` run them. An interactive runtime shares the same process and
119 /// inherits the globals maintenance installs, so re-running them would either double-apply a
120 /// non-idempotent effect or race maintenance.
121 ///
122 /// NOTE: every role a release build can construct owns the globals, so this is constantly true
123 /// outside tests. The distinction becomes load-bearing when the interactive runtime lands.
124 pub fn owns_process_globals(self) -> bool {
125 matches!(
126 self,
127 ComputeRuntimeRole::Solo | ComputeRuntimeRole::Maintenance
128 )
129 }
130}
131
132/// Configures the server with compute-specific metrics.
133#[derive(Clone)]
134struct Config {
135 /// `persist` client cache.
136 pub persist_clients: Arc<PersistClientCache>,
137 /// Shared persist batch state.
138 pub persist_batches: SharedBatches,
139 /// Context necessary for rendering txn-wal operators.
140 pub txns_ctx: TxnsContext,
141 /// A process-global handle to tracing configuration.
142 pub tracing_handle: Arc<TracingHandle>,
143 /// Metrics exposed by compute replicas.
144 pub metrics: ComputeMetrics,
145 /// Other configuration for compute.
146 pub context: ComputeInstanceContext,
147 /// The process-global metrics registry.
148 pub metrics_registry: MetricsRegistry,
149 /// The number of timely workers per process.
150 pub workers_per_process: usize,
151 /// Bounds how many offloaded peek walks run at once, shared by every worker this server runs.
152 ///
153 /// NOTE: per compute runtime, not global. A process running a maintenance and an interactive
154 /// runtime calls `serve` twice and admits the bound once per call.
155 pub peek_permits: Arc<PeekPermits>,
156 /// Configuration for hosting storage objects on this cluster, if enabled.
157 pub storage_guest: Option<Arc<StorageGuestConfig>>,
158}
159
160/// A per-worker channel delivering storage client connections.
161type StorageClientRx = mpsc::UnboundedReceiver<(
162 Uuid,
163 mpsc::UnboundedReceiver<StorageCommand>,
164 mpsc::UnboundedSender<StorageResponse>,
165)>;
166
167/// Configuration for hosting storage objects on the compute cluster.
168pub struct StorageGuestConfig {
169 /// Per-worker channels delivering storage client connections, indexed by local worker index.
170 client_rxs: Mutex<Vec<Option<StorageClientRx>>>,
171 /// Metrics for storage objects.
172 metrics: StorageMetrics,
173 /// Function to get wall time now.
174 now: NowFn,
175 /// Configuration for source and sink connections.
176 connection_context: ConnectionContext,
177 /// Other configuration for storage instances.
178 instance_context: StorageInstanceContext,
179 /// Shared rocksdb write buffer manager.
180 shared_rocksdb_write_buffer_manager: SharedWriteBufferManager,
181}
182
183/// Initiates a timely dataflow computation, processing compute commands.
184pub async fn serve(
185 timely_config: TimelyConfig,
186 role: ComputeRuntimeRole,
187 metrics_registry: &MetricsRegistry,
188 persist_clients: Arc<PersistClientCache>,
189 txns_ctx: TxnsContext,
190 tracing_handle: Arc<TracingHandle>,
191 context: ComputeInstanceContext,
192) -> Result<impl Fn() -> ComputeServerClient + use<>, Error> {
193 let workers_per_process = timely_config.workers;
194 let config = Config {
195 persist_batches: SharedBatches::new(),
196 persist_clients,
197 txns_ctx,
198 tracing_handle,
199 metrics: ComputeMetrics::register_with(metrics_registry, role),
200 context,
201 metrics_registry: metrics_registry.clone(),
202 workers_per_process,
203 peek_permits: Arc::new(PeekPermits::new(workers_per_process)),
204 storage_guest: None,
205 };
206
207 let (_worker_threads, client_builder) = serve_inner(config, timely_config).await?;
208 Ok(client_builder)
209}
210
211/// Initiates a timely dataflow computation that processes compute commands and additionally hosts
212/// storage objects, processing storage commands received over a separate client connection.
213///
214/// Returns client builders for both the compute and the storage side.
215pub async fn serve_unified(
216 timely_config: TimelyConfig,
217 role: ComputeRuntimeRole,
218 metrics_registry: &MetricsRegistry,
219 persist_clients: Arc<PersistClientCache>,
220 txns_ctx: TxnsContext,
221 tracing_handle: Arc<TracingHandle>,
222 context: ComputeInstanceContext,
223 now: NowFn,
224 storage_connection_context: ConnectionContext,
225 storage_instance_context: StorageInstanceContext,
226) -> Result<
227 (
228 impl Fn() -> ComputeServerClient + use<>,
229 impl Fn() -> GuestClusterClient<StorageCommand, StorageResponse> + use<>,
230 ),
231 Error,
232> {
233 let workers_per_process = timely_config.workers;
234
235 // Per-worker channels over which storage client connections are delivered.
236 let mut storage_client_txs = Vec::new();
237 let mut storage_client_rxs = Vec::new();
238 for _ in 0..workers_per_process {
239 let (tx, rx) = mpsc::unbounded_channel();
240 storage_client_txs.push(tx);
241 storage_client_rxs.push(Some(rx));
242 }
243
244 let storage_guest = StorageGuestConfig {
245 client_rxs: Mutex::new(storage_client_rxs),
246 metrics: StorageMetrics::register_with(metrics_registry),
247 now,
248 connection_context: storage_connection_context,
249 instance_context: storage_instance_context,
250 shared_rocksdb_write_buffer_manager: Default::default(),
251 };
252
253 let config = Config {
254 persist_batches: SharedBatches::new(),
255 persist_clients,
256 txns_ctx,
257 tracing_handle,
258 metrics: ComputeMetrics::register_with(metrics_registry, role),
259 context,
260 metrics_registry: metrics_registry.clone(),
261 workers_per_process,
262 peek_permits: Arc::new(PeekPermits::new(workers_per_process)),
263 storage_guest: Some(Arc::new(storage_guest)),
264 };
265
266 let (worker_threads, compute_client_builder) = serve_inner(config, timely_config).await?;
267
268 let storage_client_txs = Arc::new(storage_client_txs);
269 let storage_client_builder =
270 move || GuestClusterClient::new(Arc::clone(&storage_client_txs), worker_threads.clone());
271
272 Ok((compute_client_builder, storage_client_builder))
273}
274
275/// Builds the Timely cluster for the given config and returns its worker threads along with a
276/// builder for compute clients to it.
277async fn serve_inner(
278 config: Config,
279 timely_config: TimelyConfig,
280) -> Result<(Vec<Thread>, impl Fn() -> ComputeServerClient + use<>), Error> {
281 mz_timely_util::column_pager::metrics::register(
282 &config.metrics_registry,
283 mz_timely_util::column_pager::tiered_policy(),
284 );
285 mz_timely_util::pool_config::metrics::register(&config.metrics_registry);
286 mz_cluster::client::register_exert_policy_metrics(&config.metrics_registry);
287
288 let tokio_executor = tokio::runtime::Handle::current();
289
290 let timely_container = config.build_cluster(timely_config, tokio_executor).await?;
291 let worker_threads = timely_container.worker_threads();
292 let timely_container = Arc::new(Mutex::new(timely_container));
293
294 let client_builder =
295 move || ComputeServerClient(ClusterClient::new(Arc::clone(&timely_container)));
296
297 Ok((worker_threads, client_builder))
298}
299
300/// A client to the Timely cluster of a compute server.
301#[derive(Debug)]
302pub struct ComputeServerClient(ClusterClient<Config>);
303
304impl GenericClient<ComputeCommand, ComputeResponse> for ComputeServerClient {
305 fn send(&mut self, cmd: ComputeCommand) -> impl Future<Output = Result<(), Error>> + Send {
306 self.0.send(cmd)
307 }
308
309 fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Result<Option<ComputeResponse>, Error>> {
310 self.0.poll_recv(cx)
311 }
312}
313
314/// Error type returned on connection nonce changes.
315///
316/// A nonce change informs workers that subsequent commands come a from a new client connection
317/// and therefore require reconciliation.
318struct NonceChange(Uuid);
319
320/// Endpoint used by workers to receive compute commands.
321///
322/// Observes nonce changes in the command stream and converts them into receive errors.
323struct CommandReceiver {
324 /// The channel supplying commands.
325 inner: command_channel::Receiver,
326 /// The ID of the Timely worker.
327 worker_id: usize,
328 /// The nonce identifying the current cluster protocol incarnation.
329 nonce: Option<Uuid>,
330 /// A stash to enable peeking the next command, used in `try_recv`.
331 stashed_command: Option<ComputeCommand>,
332}
333
334impl CommandReceiver {
335 fn new(inner: command_channel::Receiver, worker_id: usize) -> Self {
336 Self {
337 inner,
338 worker_id,
339 nonce: None,
340 stashed_command: None,
341 }
342 }
343
344 /// Receive the next pending command, if any.
345 ///
346 /// If the next compute command has a different nonce, this method instead returns an `Err`
347 /// containing the new nonce.
348 fn try_recv(&mut self) -> Result<Option<WorkerCommand>, NonceChange> {
349 if let Some(command) = self.stashed_command.take() {
350 return Ok(Some(WorkerCommand::Compute(command)));
351 }
352 let Some(message) = self.inner.try_recv() else {
353 return Ok(None);
354 };
355
356 let (command, nonce) = match message {
357 UnifiedCommand::Compute(command, nonce) => (command, nonce),
358 UnifiedCommand::Storage(command) => {
359 trace!(
360 worker = self.worker_id,
361 ?command,
362 "received storage command"
363 );
364 return Ok(Some(WorkerCommand::Storage(command)));
365 }
366 };
367
368 trace!(worker = self.worker_id, %nonce, ?command, "received command");
369
370 if Some(nonce) == self.nonce {
371 Ok(Some(WorkerCommand::Compute(command)))
372 } else {
373 self.nonce = Some(nonce);
374 self.stashed_command = Some(command);
375 Err(NonceChange(nonce))
376 }
377 }
378}
379
380/// A command dispatched to the worker from the command channel.
381enum WorkerCommand {
382 /// A compute command.
383 Compute(ComputeCommand),
384 /// A storage-internal command, to be dispatched to the storage guest.
385 Storage(InternalStorageCommand),
386}
387
388/// Endpoint used by workers to send sending compute responses.
389///
390/// Tags responses with the current nonce, allowing receivers to filter out responses intended for
391/// previous client connections.
392pub(crate) struct ResponseSender {
393 /// The channel consuming responses.
394 inner: mpsc::UnboundedSender<(ComputeResponse, Uuid)>,
395 /// The ID of the Timely worker.
396 worker_id: usize,
397 /// The nonce identifying the current cluster protocol incarnation.
398 nonce: Option<Uuid>,
399}
400
401impl ResponseSender {
402 /// `pub(crate)` rather than private so the peek tests can build the sender a worker holds.
403 pub(crate) fn new(
404 inner: mpsc::UnboundedSender<(ComputeResponse, Uuid)>,
405 worker_id: usize,
406 ) -> Self {
407 Self {
408 inner,
409 worker_id,
410 nonce: None,
411 }
412 }
413
414 /// Set the cluster protocol nonce.
415 pub(crate) fn set_nonce(&mut self, nonce: Uuid) {
416 self.nonce = Some(nonce);
417 }
418
419 /// Send a compute response.
420 pub fn send(&self, response: ComputeResponse) -> Result<(), SendError<ComputeResponse>> {
421 let nonce = self.nonce.expect("nonce must be initialized");
422
423 trace!(worker = self.worker_id, %nonce, ?response, "sending response");
424 self.inner
425 .send((response, nonce))
426 .map_err(|SendError((resp, _))| SendError(resp))
427 }
428}
429
430/// State maintained for each worker thread.
431///
432/// Much of this state can be viewed as local variables for the worker thread,
433/// holding state that persists across function calls.
434struct Worker<'w> {
435 /// The underlying Timely worker.
436 timely_worker: &'w mut TimelyWorker,
437 /// The channel over which commands are received.
438 command_rx: CommandReceiver,
439 /// The channel over which responses are sent.
440 response_tx: ResponseSender,
441 compute_state: Option<ComputeState>,
442 /// Compute metrics.
443 metrics: WorkerMetrics,
444 /// A process-global cache of (blob_uri, consensus_uri) -> PersistClient.
445 /// This is intentionally shared between workers
446 persist_clients: Arc<PersistClientCache>,
447 persist_batches: SharedBatches,
448 /// Context necessary for rendering txn-wal operators.
449 txns_ctx: TxnsContext,
450 /// A process-global handle to tracing configuration.
451 tracing_handle: Arc<TracingHandle>,
452 context: ComputeInstanceContext,
453 /// The process-global metrics registry.
454 metrics_registry: MetricsRegistry,
455 /// The number of timely workers per process.
456 workers_per_process: usize,
457 /// Bounds how many offloaded peek walks run at once, shared by the workers of one `serve`
458 /// call rather than by the process.
459 peek_permits: Arc<PeekPermits>,
460 /// The hosted storage guest, if any.
461 storage: Option<StorageGuest>,
462}
463
464/// Per-worker state for hosting storage objects on the compute cluster.
465struct StorageGuest {
466 /// Channel delivering new storage client connections.
467 client_rx: StorageClientRx,
468 /// The current storage client connection, if any.
469 conn: Option<StorageConn>,
470 /// The hosted storage worker state.
471 storage_state: StorageState,
472 /// The last time storage maintenance ran.
473 last_maintenance: Instant,
474 /// The last time storage statistics were reported.
475 last_stats_time: Instant,
476}
477
478impl StorageGuest {
479 /// The longest the worker may park before the guest's next periodic duty (frontier reporting
480 /// or statistics collection) comes due, or `None` when no duty is pending.
481 ///
482 /// Mirrors the parking of storage's own server loop: the maintenance and statistics intervals
483 /// bound the park. A maintenance deadline in the past does not bound it, because maintenance
484 /// runs on the next wakeup anyway. The initial zero maintenance interval would otherwise turn
485 /// every park into a spin.
486 fn park_cap(&self) -> Option<Duration> {
487 // Periodic duties run only on a reconciled connection. Without one there is no deadline
488 // to meet, and connection and command arrivals unpark the worker.
489 let conn_serving = self
490 .conn
491 .as_ref()
492 .is_some_and(|conn| conn.reconcile_buf.is_none());
493 if !conn_serving {
494 return None;
495 }
496
497 let maintenance_interval = self.storage_state.server_maintenance_interval;
498 let stats_interval = self
499 .storage_state
500 .storage_configuration
501 .parameters
502 .statistics_collection_interval;
503
504 let next_maintenance =
505 (self.last_maintenance + maintenance_interval).checked_duration_since(Instant::now());
506 let next_stats = stats_interval.saturating_sub(self.last_stats_time.elapsed());
507 match next_maintenance {
508 Some(maintenance) => Some(maintenance.min(next_stats)),
509 None => Some(next_stats),
510 }
511 }
512}
513
514/// A storage client connection.
515struct StorageConn {
516 /// The channel over which storage commands are received.
517 command_rx: mpsc::UnboundedReceiver<StorageCommand>,
518 /// The channel over which storage responses are sent.
519 response_tx: mpsc::UnboundedSender<StorageResponse>,
520 /// Commands buffered for reconciliation, until `InitializationComplete` is received.
521 /// `None` once the connection is reconciled and serving.
522 reconcile_buf: Option<Vec<StorageCommand>>,
523}
524
525impl ClusterSpec for Config {
526 type Command = ComputeCommand;
527 type Response = ComputeResponse;
528
529 const NAME: &str = "compute";
530
531 fn run_worker(
532 &self,
533 timely_worker: &mut TimelyWorker,
534 client_rx: mpsc::UnboundedReceiver<(
535 Uuid,
536 mpsc::UnboundedReceiver<ComputeCommand>,
537 mpsc::UnboundedSender<ComputeResponse>,
538 )>,
539 ) {
540 if self.context.worker_core_affinity {
541 set_core_affinity(timely_worker.index());
542 }
543
544 let worker_id = timely_worker.index();
545 let metrics = self.metrics.for_worker(worker_id);
546
547 let local_index = worker_id % self.workers_per_process;
548
549 // Prepare the storage guest's inputs to the command channel, so
550 // storage-internal commands are sequenced through the same lane as compute commands.
551 let mut storage_lane_input = None;
552 let guest_setup = self.storage_guest.as_ref().map(|cfg| {
553 let storage_client_rx = cfg.client_rxs.lock().expect("poisoned")[local_index]
554 .take()
555 .expect("each worker takes its storage client_rx exactly once");
556 let (internal_tx, internal_rx) = std::sync::mpsc::channel();
557 let activator_slot = Rc::new(RefCell::new(None));
558 storage_lane_input = Some(StorageLaneInput {
559 rx: internal_rx,
560 activator_slot: Rc::clone(&activator_slot),
561 });
562 let internal_cmd_tx = InternalCommandSender::from_parts(internal_tx, activator_slot);
563 (Arc::clone(cfg), storage_client_rx, internal_cmd_tx)
564 });
565
566 // Create the command channel that broadcasts commands from worker 0 to other workers. We
567 // reuse this channel between client connections, to avoid bugs where different workers end
568 // up creating incompatible sides of the channel dataflow after reconnects.
569 // See database-issues#8964.
570 let (cmd_tx, cmd_rx) = command_channel::render(timely_worker, storage_lane_input);
571 let (resp_tx, resp_rx) = mpsc::unbounded_channel();
572
573 spawn_channel_adapter(client_rx, cmd_tx, resp_rx, worker_id);
574
575 // Create the storage guest state.
576 let storage = guest_setup.map(|(cfg, storage_client_rx, internal_cmd_tx)| {
577 let storage_state = StorageState::new_guest(
578 timely_worker.index(),
579 timely_worker.peers(),
580 internal_cmd_tx,
581 // The host dispatches internal commands from the unified command channel, so
582 // the guest reads no receiver of its own.
583 None,
584 cfg.metrics.clone(),
585 cfg.now.clone(),
586 cfg.connection_context.clone(),
587 cfg.instance_context.clone(),
588 Arc::clone(&self.persist_clients),
589 self.txns_ctx.clone(),
590 Arc::clone(&self.tracing_handle),
591 cfg.shared_rocksdb_write_buffer_manager.clone(),
592 );
593
594 StorageGuest {
595 client_rx: storage_client_rx,
596 conn: None,
597 storage_state,
598 last_maintenance: Instant::now(),
599 last_stats_time: Instant::now(),
600 }
601 });
602
603 Worker {
604 timely_worker,
605 command_rx: CommandReceiver::new(cmd_rx, worker_id),
606 response_tx: ResponseSender::new(resp_tx, worker_id),
607 metrics,
608 context: self.context.clone(),
609 persist_clients: Arc::clone(&self.persist_clients),
610 persist_batches: self.persist_batches.clone(),
611 txns_ctx: self.txns_ctx.clone(),
612 compute_state: None,
613 tracing_handle: Arc::clone(&self.tracing_handle),
614 metrics_registry: self.metrics_registry.clone(),
615 workers_per_process: self.workers_per_process,
616 peek_permits: Arc::clone(&self.peek_permits),
617 storage,
618 }
619 .run()
620 }
621}
622
623/// Set the current thread's core affinity, based on the given `worker_id`.
624#[cfg(not(target_os = "macos"))]
625fn set_core_affinity(worker_id: usize) {
626 use tracing::error;
627
628 let Some(mut core_ids) = core_affinity::get_core_ids() else {
629 error!(worker_id, "unable to get core IDs for setting affinity");
630 return;
631 };
632
633 // The `get_core_ids` docs don't say anything about a guaranteed order of the returned Vec,
634 // so sort it just to be safe.
635 core_ids.sort_unstable_by_key(|i| i.id);
636
637 // On multi-process replicas `worker_id` might be greater than the number of available cores.
638 // However, we assume that we always have at least as many cores as there are local workers.
639 // Violating this assumption is safe but might lead to degraded performance due to skew in core
640 // utilization.
641 let idx = worker_id % core_ids.len();
642 let core_id = core_ids[idx];
643
644 if core_affinity::set_for_current(core_id) {
645 info!(
646 worker_id,
647 core_id = core_id.id,
648 "set core affinity for worker"
649 );
650 } else {
651 error!(
652 worker_id,
653 core_id = core_id.id,
654 "failed to set core affinity for worker"
655 )
656 }
657}
658
659/// Set the current thread's core affinity, based on the given `worker_id`.
660#[cfg(target_os = "macos")]
661fn set_core_affinity(_worker_id: usize) {
662 // Setting core affinity is known to not work on Apple Silicon:
663 // https://github.com/Elzair/core_affinity_rs/issues/22
664 info!("setting core affinity is not supported on macOS");
665}
666
667impl<'w> Worker<'w> {
668 /// Runs a compute worker.
669 pub fn run(&mut self) {
670 // The command receiver is initialized without an nonce, so receiving the first command
671 // always triggers a nonce change.
672 let NonceChange(nonce) = self.recv_command().expect_err("change to first nonce");
673 self.set_nonce(nonce);
674
675 loop {
676 let Err(NonceChange(nonce)) = self.run_client();
677 self.set_nonce(nonce);
678 }
679 }
680
681 fn set_nonce(&mut self, nonce: Uuid) {
682 self.response_tx.set_nonce(nonce);
683 }
684
685 /// Handles commands for a client connection, returns when the nonce changes.
686 fn run_client(&mut self) -> Result<Infallible, NonceChange> {
687 self.reconcile()?;
688
689 // The last time we did periodic maintenance.
690 let mut last_maintenance = Instant::now();
691
692 // Commence normal operation.
693 loop {
694 // Get the maintenance interval, default to zero if we don't have a compute state.
695 let maintenance_interval = self
696 .compute_state
697 .as_ref()
698 .map_or(Duration::ZERO, |state| state.server_maintenance_interval);
699
700 let now = Instant::now();
701 // Determine if we need to perform maintenance, which is true if `maintenance_interval`
702 // time has passed since the last maintenance.
703 let sleep_duration;
704 if now >= last_maintenance + maintenance_interval {
705 last_maintenance = now;
706 sleep_duration = None;
707
708 // Report frontier information back the coordinator.
709 if let Some(mut compute_state) = self.activate_compute() {
710 compute_state.compute_state.traces.maintenance();
711 compute_state.report_frontiers();
712 compute_state.report_metrics();
713 compute_state.check_expiration();
714 }
715
716 self.metrics.record_shared_row_metrics();
717 } else {
718 // We didn't perform maintenance, sleep until the next maintenance interval.
719 let next_maintenance = last_maintenance + maintenance_interval;
720 sleep_duration = Some(next_maintenance.saturating_duration_since(now))
721 };
722
723 // Do not sleep while a peek waits for its turn. Only the sweep below gives it one,
724 // and nothing else leaves an activation behind to end the park.
725 let sleep_duration = match &self.compute_state {
726 Some(state) if state.peeks_awaiting_turn() => Some(Duration::ZERO),
727 _ => sleep_duration,
728 };
729
730 // With a storage guest, cap the park duration so the guest's periodic duties run on
731 // time.
732 let sleep_duration = match self.storage.as_ref().and_then(StorageGuest::park_cap) {
733 Some(cap) => Some(sleep_duration.map_or(cap, |d| d.min(cap))),
734 None => sleep_duration,
735 };
736
737 // Step the timely worker, recording the time taken.
738 let timer = self.metrics.timely_step_duration_seconds.start_timer();
739 if self.storage_guest_busy() {
740 self.timely_worker.step();
741 } else {
742 self.timely_worker.step_or_park(sleep_duration);
743 }
744 timer.observe_duration();
745
746 self.handle_pending_commands()?;
747
748 self.process_storage_guest();
749
750 if let Some(mut compute_state) = self.activate_compute() {
751 compute_state.process_peeks();
752 compute_state.process_subscribes();
753 compute_state.process_copy_tos();
754 }
755 }
756 }
757
758 fn handle_pending_commands(&mut self) -> Result<(), NonceChange> {
759 while let Some(cmd) = self.command_rx.try_recv()? {
760 match cmd {
761 WorkerCommand::Compute(cmd) => self.handle_command(cmd),
762 WorkerCommand::Storage(cmd) => self.handle_storage_internal_command(cmd),
763 }
764 }
765 Ok(())
766 }
767
768 /// Whether the storage guest has pending work that forbids parking.
769 ///
770 /// It is critical that we allow Timely to park iff there are no pending commands or async
771 /// worker responses, since those are delivered by other threads that only unpark us once, at
772 /// send time.
773 fn storage_guest_busy(&self) -> bool {
774 self.storage.as_ref().is_some_and(|guest| {
775 !guest.client_rx.is_empty()
776 || (!self.storage_commands_held()
777 && guest
778 .conn
779 .as_ref()
780 .is_some_and(|conn| !conn.command_rx.is_empty()))
781 || !guest.storage_state.async_worker.is_empty()
782 })
783 }
784
785 /// Whether the storage guest must leave external storage commands unread.
786 ///
787 /// Commands are held until compute's `CreateInstance` has initialized logging. Held commands
788 /// do not keep the worker busy. Handling `CreateInstance` releases them.
789 fn storage_commands_held(&self) -> bool {
790 // Storage commands render storage dataflows, and the loggers only observe operators
791 // created after their initialization, so a storage dataflow rendered earlier never
792 // appears in introspection. Nothing orders the storage and compute controller
793 // connections, so the storage controller can connect first. Every storage dataflow
794 // stems from an external storage command, so holding those back suffices.
795 self.compute_state.is_none()
796 }
797
798 /// Dispatch a storage-internal command from the command channel to
799 /// the storage guest. This is where all storage dataflow rendering happens.
800 fn handle_storage_internal_command(&mut self, cmd: InternalStorageCommand) {
801 let mut guest = self
802 .storage
803 .take()
804 .expect("the command channel carries storage commands only when a guest is hosted");
805
806 let mut worker = StorageWorker {
807 timely_worker: &mut *self.timely_worker,
808 client_rx: guest.client_rx,
809 storage_state: guest.storage_state,
810 };
811 worker.handle_internal_storage_command(cmd);
812
813 let StorageWorker {
814 timely_worker: _,
815 client_rx,
816 storage_state,
817 } = worker;
818 guest.client_rx = client_rx;
819 guest.storage_state = storage_state;
820 self.storage = Some(guest);
821 }
822
823 /// Process the storage guest's per-iteration duties: accept client
824 /// connections, handle external storage commands (buffering for reconciliation until
825 /// `InitializationComplete`), forward async worker responses, and report frontiers, dropped
826 /// collections, status updates, and statistics.
827 fn process_storage_guest(&mut self) {
828 let Some(mut guest) = self.storage.take() else {
829 return;
830 };
831
832 // Accept new client connections, replacing any current one. Every new connection starts
833 // with a reconciliation.
834 loop {
835 use tokio::sync::mpsc::error::TryRecvError;
836 match guest.client_rx.try_recv() {
837 Ok((_nonce, command_rx, response_tx)) => {
838 guest.conn = Some(StorageConn {
839 command_rx,
840 response_tx,
841 reconcile_buf: Some(Vec::new()),
842 });
843 }
844 Err(TryRecvError::Empty | TryRecvError::Disconnected) => break,
845 }
846 }
847
848 let commands_held = self.storage_commands_held();
849 let mut worker = StorageWorker {
850 timely_worker: &mut *self.timely_worker,
851 client_rx: guest.client_rx,
852 storage_state: guest.storage_state,
853 };
854
855 // Handle responses from the async worker. Only worker 0 does async processing, so only
856 // worker 0 ever receives any. This must run with or without a client connection:
857 // storage-internal commands keep flowing through the command channel while the
858 // controller is away and can issue async work, and an undrained response queue keeps
859 // `storage_guest_busy` true, turning every park into a spin.
860 while let Ok(response) = worker.storage_state.async_worker.try_recv() {
861 worker.handle_async_worker_response(response);
862 }
863
864 let conn = if commands_held {
865 None
866 } else {
867 guest.conn.take()
868 };
869 let Some(mut conn) = conn else {
870 let StorageWorker {
871 timely_worker: _,
872 client_rx,
873 storage_state,
874 } = worker;
875 guest.client_rx = client_rx;
876 guest.storage_state = storage_state;
877 self.storage = Some(guest);
878 return;
879 };
880
881 // Drain external storage commands.
882 let mut disconnected = false;
883 loop {
884 use tokio::sync::mpsc::error::TryRecvError;
885 match conn.command_rx.try_recv() {
886 Ok(StorageCommand::InitializationComplete) if conn.reconcile_buf.is_some() => {
887 let commands = conn.reconcile_buf.take().expect("checked above");
888 worker.reconcile_commands(commands);
889 }
890 Ok(cmd) => match &mut conn.reconcile_buf {
891 Some(buf) => buf.push(cmd),
892 None => worker.storage_state.handle_storage_command(cmd),
893 },
894 Err(TryRecvError::Empty) => break,
895 Err(TryRecvError::Disconnected) => {
896 disconnected = true;
897 break;
898 }
899 }
900 }
901
902 // Response-producing duties run only on a reconciled connection.
903 if conn.reconcile_buf.is_none() {
904 let maintenance_interval = worker.storage_state.server_maintenance_interval;
905 let now = Instant::now();
906 if now >= guest.last_maintenance + maintenance_interval {
907 guest.last_maintenance = now;
908 worker.report_frontier_progress(&conn.response_tx);
909 }
910
911 for id in std::mem::take(&mut worker.storage_state.dropped_ids) {
912 worker.send_storage_response(&conn.response_tx, StorageResponse::DroppedId(id));
913 }
914
915 worker.process_oneshot_ingestions(&conn.response_tx);
916 worker.report_status_updates(&conn.response_tx);
917
918 let stats_interval = worker
919 .storage_state
920 .storage_configuration
921 .parameters
922 .statistics_collection_interval;
923 if guest.last_stats_time.elapsed() >= stats_interval {
924 worker.report_storage_statistics(&conn.response_tx);
925 guest.last_stats_time = Instant::now();
926 }
927 }
928
929 let StorageWorker {
930 timely_worker: _,
931 client_rx,
932 storage_state,
933 } = worker;
934 guest.client_rx = client_rx;
935 guest.storage_state = storage_state;
936 guest.conn = (!disconnected).then_some(conn);
937 self.storage = Some(guest);
938 }
939
940 fn handle_command(&mut self, cmd: ComputeCommand) {
941 if matches!(&cmd, ComputeCommand::CreateInstance(_)) {
942 self.compute_state = Some(ComputeState::new(
943 Arc::clone(&self.persist_clients),
944 self.persist_batches.clone(),
945 self.txns_ctx.clone(),
946 self.metrics.clone(),
947 Arc::clone(&self.tracing_handle),
948 self.context.clone(),
949 self.metrics_registry.clone(),
950 self.workers_per_process,
951 Arc::clone(&self.peek_permits),
952 ));
953 }
954 self.activate_compute().unwrap().handle_compute_command(cmd);
955 }
956
957 fn activate_compute(&mut self) -> Option<ActiveComputeState<'_>> {
958 if let Some(compute_state) = &mut self.compute_state {
959 Some(ActiveComputeState {
960 timely_worker: &mut *self.timely_worker,
961 compute_state,
962 response_tx: &mut self.response_tx,
963 })
964 } else {
965 None
966 }
967 }
968
969 /// Receive the next compute command.
970 ///
971 /// This method blocks if no command is currently available, but takes care to step the Timely
972 /// worker while doing so.
973 fn recv_command(&mut self) -> Result<ComputeCommand, NonceChange> {
974 loop {
975 if let Some(cmd) = self.command_rx.try_recv()? {
976 match cmd {
977 WorkerCommand::Compute(cmd) => return Ok(cmd),
978 // Storage-internal commands are dispatched even while
979 // waiting for compute commands (e.g. during compute reconciliation), so
980 // storage dataflow construction keeps its lane position on all workers.
981 WorkerCommand::Storage(cmd) => {
982 self.handle_storage_internal_command(cmd);
983 continue;
984 }
985 }
986 }
987
988 // Keep serving the storage guest while blocked on compute
989 // commands, and avoid unbounded parks that would stall its periodic duties.
990 self.process_storage_guest();
991 let park_cap = self.storage.as_ref().and_then(StorageGuest::park_cap);
992
993 let start = Instant::now();
994 if self.storage_guest_busy() {
995 self.timely_worker.step();
996 } else if let Some(cap) = park_cap {
997 self.timely_worker.step_or_park(Some(cap));
998 } else {
999 self.timely_worker.step_or_park(None);
1000 }
1001 self.metrics
1002 .timely_step_duration_seconds
1003 .observe(start.elapsed().as_secs_f64());
1004 }
1005 }
1006
1007 /// Extract commands until `InitializationComplete`, and make the worker reflect those commands.
1008 ///
1009 /// This method is meant to be a function of the commands received thus far (as recorded in the
1010 /// compute state command history) and the new commands from `command_rx`. It should not be a
1011 /// function of other characteristics, like whether the worker has managed to respond to a peek
1012 /// or not. Some effort goes in to narrowing our view to only the existing commands we can be sure
1013 /// are live at all other workers.
1014 ///
1015 /// The methodology here is to drain `command_rx` until an `InitializationComplete`, at which point
1016 /// the prior commands are "reconciled" in. Reconciliation takes each goal dataflow and looks for an
1017 /// existing "compatible" dataflow (per `compatible()`) it can repurpose, with some additional tests
1018 /// to be sure that we can cut over from one to the other (no additional compaction, no tails/sinks).
1019 /// With any connections established, old orphaned dataflows are allow to compact away, and any new
1020 /// dataflows are created from scratch. "Kept" dataflows are allowed to compact up to any new `as_of`.
1021 ///
1022 /// Some additional tidying happens, cleaning up pending peeks, reported frontiers, and creating a new
1023 /// subscribe response buffer. We will need to be vigilant with future modifications to `ComputeState` to
1024 /// line up changes there with clean resets here.
1025 fn reconcile(&mut self) -> Result<(), NonceChange> {
1026 // To initialize the connection, we want to drain all commands until we receive a
1027 // `ComputeCommand::InitializationComplete` command to form a target command state.
1028 let mut new_commands = Vec::new();
1029 loop {
1030 match self.recv_command()? {
1031 ComputeCommand::InitializationComplete => break,
1032 command => new_commands.push(command),
1033 }
1034 }
1035
1036 // Commands we will need to apply before entering normal service.
1037 // These commands may include dropping existing dataflows, compacting existing dataflows,
1038 // and creating new dataflows, in addition to standard peek and compaction commands.
1039 // The result should be the same as if dropping all dataflows and running `new_commands`.
1040 let mut todo_commands = Vec::new();
1041 // We only have a compute history if we are in an initialized state
1042 // (i.e. after a `CreateInstance`).
1043 // If this is not the case, just copy `new_commands` into `todo_commands`.
1044 if let Some(compute_state) = &mut self.compute_state {
1045 // Reduce the installed commands.
1046 // Importantly, act as if all peeks may have been retired (as we cannot know otherwise).
1047 compute_state.command_history.discard_peeks();
1048 compute_state.command_history.reduce();
1049
1050 // At this point, we need to sort out which of the *certainly installed* dataflows are
1051 // suitable replacements for the requested dataflows. A dataflow is "certainly installed"
1052 // as of a frontier if its compaction allows it to go no further. We ignore peeks for this
1053 // reasoning, as we cannot be certain that peeks still exist at any other worker.
1054
1055 // Having reduced our installed command history retaining no peeks (above), we should be able
1056 // to use track down installed dataflows we can use as surrogates for requested dataflows (which
1057 // have retained all of their peeks, creating a more demanding `as_of` requirement).
1058 // NB: installed dataflows may still be allowed to further compact, and we should double check
1059 // this before being too confident. It should be rare without peeks, but could happen with e.g.
1060 // multiple outputs of a dataflow.
1061
1062 // The values with which a prior `CreateInstance` was called, if it was.
1063 let mut old_instance_config = None;
1064 // Index dataflows by `export_ids().collect()`, as this is a precondition for their compatibility.
1065 let mut old_dataflows = BTreeMap::default();
1066 // Maintain allowed compaction, in case installed identifiers may have been allowed to compact.
1067 let mut old_frontiers = BTreeMap::default();
1068 for command in compute_state.command_history.iter() {
1069 match command {
1070 ComputeCommand::CreateInstance(config) => {
1071 old_instance_config = Some(config);
1072 }
1073 ComputeCommand::CreateDataflow(dataflow) => {
1074 let export_ids = dataflow.export_ids().collect::<BTreeSet<_>>();
1075 old_dataflows.insert(export_ids, dataflow);
1076 }
1077 ComputeCommand::AllowCompaction { id, frontier } => {
1078 old_frontiers.insert(id, frontier);
1079 }
1080 _ => {
1081 // Nothing to do in these cases.
1082 }
1083 }
1084 }
1085
1086 // Compaction commands that can be applied to existing dataflows.
1087 let mut old_compaction = BTreeMap::default();
1088 // Exported identifiers from dataflows we retain.
1089 let mut retain_ids = BTreeSet::default();
1090
1091 // Traverse new commands, sorting out what remediation we can do.
1092 for command in new_commands.iter() {
1093 match command {
1094 ComputeCommand::CreateDataflow(dataflow) => {
1095 // Attempt to find an existing match for the dataflow.
1096 let as_of = dataflow.as_of.as_ref().unwrap();
1097 let export_ids = dataflow.export_ids().collect::<BTreeSet<_>>();
1098
1099 if let Some(old_dataflow) = old_dataflows.get(&export_ids) {
1100 let compatible = old_dataflow.compatible_with(dataflow);
1101 let uncompacted = !export_ids
1102 .iter()
1103 .flat_map(|id| old_frontiers.get(id))
1104 .any(|frontier| {
1105 !timely::PartialOrder::less_equal(
1106 *frontier,
1107 dataflow.as_of.as_ref().unwrap(),
1108 )
1109 });
1110
1111 // We cannot reconcile subscribe and copy-to sinks at the moment,
1112 // because the response buffer is shared, and to a first approximation
1113 // must be completely reformed.
1114 let subscribe_free = dataflow.subscribe_ids().next().is_none();
1115 let copy_to_free = dataflow.copy_to_ids().next().is_none();
1116
1117 // If we have replaced any dependency of this dataflow, we need to
1118 // replace this dataflow, to make it use the replacement.
1119 let dependencies_retained = dataflow
1120 .imported_index_ids()
1121 .all(|id| retain_ids.contains(&id));
1122
1123 if compatible
1124 && uncompacted
1125 && subscribe_free
1126 && copy_to_free
1127 && dependencies_retained
1128 {
1129 // Match found; remove the match from the deletion queue,
1130 // and compact its outputs to the dataflow's `as_of`.
1131 old_dataflows.remove(&export_ids);
1132 for id in export_ids.iter() {
1133 old_compaction.insert(*id, as_of.clone());
1134 }
1135 retain_ids.extend(export_ids);
1136 } else {
1137 warn!(
1138 ?export_ids,
1139 ?compatible,
1140 ?uncompacted,
1141 ?subscribe_free,
1142 ?copy_to_free,
1143 ?dependencies_retained,
1144 old_as_of = ?old_dataflow.as_of,
1145 new_as_of = ?as_of,
1146 "dataflow reconciliation failed",
1147 );
1148
1149 // Dump the full dataflow plans if they are incompatible, to
1150 // simplify debugging hard-to-reproduce reconciliation failures.
1151 if !compatible {
1152 warn!(
1153 old = ?old_dataflow,
1154 new = ?dataflow,
1155 "incompatible dataflows in reconciliation",
1156 );
1157 }
1158
1159 todo_commands
1160 .push(ComputeCommand::CreateDataflow(dataflow.clone()));
1161 }
1162
1163 compute_state.metrics.record_dataflow_reconciliation(
1164 compatible,
1165 uncompacted,
1166 subscribe_free,
1167 copy_to_free,
1168 dependencies_retained,
1169 );
1170 } else {
1171 todo_commands.push(ComputeCommand::CreateDataflow(dataflow.clone()));
1172 }
1173 }
1174 ComputeCommand::CreateInstance(config) => {
1175 // Cluster creation should not be performed again!
1176 if old_instance_config.map_or(false, |old| !old.compatible_with(config)) {
1177 halt!(
1178 "new instance configuration not compatible with existing instance configuration:\n{:?}\nvs\n{:?}",
1179 config,
1180 old_instance_config,
1181 );
1182 }
1183 }
1184 // All other commands we apply as requested.
1185 command => {
1186 todo_commands.push(command.clone());
1187 }
1188 }
1189 }
1190
1191 // Issue compaction commands first to reclaim resources.
1192 for (_, dataflow) in old_dataflows.iter() {
1193 for id in dataflow.export_ids() {
1194 // We want to drop anything that has not yet been dropped,
1195 // and nothing that has already been dropped.
1196 if old_frontiers.get(&id) != Some(&&Antichain::new()) {
1197 old_compaction.insert(id, Antichain::new());
1198 }
1199 }
1200 }
1201 for (&id, frontier) in &old_compaction {
1202 let frontier = frontier.clone();
1203 todo_commands.insert(0, ComputeCommand::AllowCompaction { id, frontier });
1204 }
1205
1206 // Clean up worker-local state.
1207 //
1208 // Various aspects of `ComputeState` need to be either uninstalled, or return to a blank slate.
1209 // All dropped dataflows should clean up after themselves, as we plan to install new dataflows
1210 // re-using the same identifiers.
1211 // All re-used dataflows should roll back any believed communicated information (e.g. frontiers)
1212 // so that they recommunicate that information as if from scratch.
1213
1214 // Remove all peeks, whether they have started or are still awaiting a turn.
1215 let queued = std::mem::take(&mut compute_state.queued_peeks);
1216 let pending = std::mem::take(&mut compute_state.pending_peeks);
1217 for peek in queued.into_iter().map(PendingPeek::Index).chain(pending) {
1218 // Log dropping the peek request.
1219 if let Some(logger) = compute_state.compute_logger.as_mut() {
1220 logger.log(&peek.as_log_event(false));
1221 }
1222 }
1223
1224 for (&id, collection) in compute_state.collections.iter_mut() {
1225 // Adjust reported frontiers:
1226 // * For dataflows we continue to use, reset to ensure we report something not
1227 // before the new `as_of` next.
1228 // * For dataflows we drop, set to the empty frontier, to ensure we don't report
1229 // anything for them.
1230 let retained = retain_ids.contains(&id);
1231 let compaction = old_compaction.remove(&id);
1232 let new_reported_frontier = match (retained, compaction) {
1233 (true, Some(new_as_of)) => ReportedFrontier::NotReported { lower: new_as_of },
1234 (true, None) => {
1235 unreachable!("retained dataflows are compacted to the new as_of")
1236 }
1237 (false, Some(new_frontier)) => {
1238 assert!(new_frontier.is_empty());
1239 ReportedFrontier::Reported(new_frontier)
1240 }
1241 (false, None) => {
1242 // Logging dataflows are implicitly retained and don't have a new as_of.
1243 // Reset them to the minimal frontier.
1244 ReportedFrontier::new()
1245 }
1246 };
1247
1248 collection.reset_reported_frontiers(new_reported_frontier);
1249
1250 // Sink tokens should be retained for retained dataflows, and dropped for dropped
1251 // dataflows.
1252 //
1253 // Dropping the tokens of active subscribe and copy-tos makes them place
1254 // `DroppedAt` responses into the respective response buffer. We drop those buffers
1255 // in the next step, which ensures that we don't send out `DroppedAt` responses for
1256 // subscribe/copy-tos dropped during reconciliation.
1257 if !retained {
1258 collection.sink_token = None;
1259 }
1260 }
1261
1262 // We must drop the response buffers as they are global across all subscribe/copy-tos.
1263 // If they were broken out by `GlobalId` then we could drop only the response buffers
1264 // of dataflows we drop.
1265 compute_state.subscribe_response_buffer = Rc::new(RefCell::new(Vec::new()));
1266 compute_state.copy_to_response_buffer = Rc::new(RefCell::new(Vec::new()));
1267
1268 // The controller expects the logging collections to be readable from the minimum time
1269 // initially. We cannot recreate the logging arrangements without restarting the
1270 // instance, but we can pad the compacted times with empty data. Doing so is sound
1271 // because logging collections from different replica incarnations are considered
1272 // distinct TVCs, so the controller doesn't expect any historical consistency from
1273 // these collections when it reconnects to a replica.
1274 //
1275 // TODO(database-issues#8152): Consider resolving this with controller-side reconciliation instead.
1276 if let Some(config) = old_instance_config {
1277 for id in config.logging.index_logs.values() {
1278 let trace = compute_state
1279 .traces
1280 .remove(id)
1281 .expect("logging trace exists");
1282 let padded = trace.into_padded();
1283 compute_state.traces.set(*id, padded);
1284 }
1285 }
1286 } else {
1287 todo_commands.clone_from(&new_commands);
1288 }
1289
1290 // Execute the commands to bring us to `new_commands`.
1291 for command in todo_commands.into_iter() {
1292 self.handle_command(command);
1293 }
1294
1295 // Overwrite `self.command_history` to reflect `new_commands`.
1296 // It is possible that there still isn't a compute state yet.
1297 if let Some(compute_state) = &mut self.compute_state {
1298 let mut command_history = ComputeCommandHistory::new(self.metrics.for_history());
1299 for command in new_commands.iter() {
1300 command_history.push(command.clone());
1301 }
1302 compute_state.command_history = command_history;
1303 }
1304 Ok(())
1305 }
1306}
1307
1308/// Spawn a task to bridge between [`ClusterClient`] and [`Worker`] channels.
1309///
1310/// The [`Worker`] expects a pair of persistent channels, with punctuation marking reconnects,
1311/// while the [`ClusterClient`] provides a new pair of channels on each reconnect.
1312fn spawn_channel_adapter(
1313 mut client_rx: mpsc::UnboundedReceiver<(
1314 Uuid,
1315 mpsc::UnboundedReceiver<ComputeCommand>,
1316 mpsc::UnboundedSender<ComputeResponse>,
1317 )>,
1318 command_tx: command_channel::Sender,
1319 mut response_rx: mpsc::UnboundedReceiver<(ComputeResponse, Uuid)>,
1320 worker_id: usize,
1321) {
1322 mz_ore::task::spawn(
1323 || format!("compute-channel-adapter-{worker_id}"),
1324 async move {
1325 // To make workers aware of the individual client connections, we tag forwarded
1326 // commands with the client nonce. Additionally, we use the nonce to filter out
1327 // responses with a different nonce, which are intended for different client
1328 // connections.
1329 //
1330 // It's possible that we receive responses with nonces from the past but also from the
1331 // future: Worker 0 might have received a new nonce before us and broadcasted it to our
1332 // Timely cluster. When we receive a response with a future nonce, we need to wait with
1333 // forwarding it until we have received the same nonce from a client connection.
1334 //
1335 // Nonces are not ordered so we don't know whether a response nonce is from the past or
1336 // the future. We thus assume that every response with an unknown nonce might be from
1337 // the future and stash them all. Every time we reconnect, we immediately send all
1338 // stashed responses with a matching nonce. Every time we receive a new response with a
1339 // nonce that matches our current one, we can discard the entire response stash as we
1340 // know that all stashed responses must be from the past.
1341 let mut stashed_responses = BTreeMap::<Uuid, Vec<ComputeResponse>>::new();
1342
1343 while let Some((nonce, mut command_rx, response_tx)) = client_rx.recv().await {
1344 // Send stashed responses for this client.
1345 if let Some(resps) = stashed_responses.remove(&nonce) {
1346 for resp in resps {
1347 let _ = response_tx.send(resp);
1348 }
1349 }
1350
1351 // Wait for a new response while forwarding received commands.
1352 let mut serve_rx_channels = async || loop {
1353 tokio::select! {
1354 msg = command_rx.recv() => match msg {
1355 Some(cmd) => command_tx.send((cmd, nonce)),
1356 None => return Err(()),
1357 },
1358 msg = response_rx.recv() => {
1359 return Ok(msg.expect("worker connected"));
1360 }
1361 }
1362 };
1363
1364 // Serve this connection until we see any of the channels disconnect.
1365 loop {
1366 let Ok((resp, resp_nonce)) = serve_rx_channels().await else {
1367 break;
1368 };
1369
1370 if resp_nonce == nonce {
1371 // Response for the current connection; forward it.
1372 stashed_responses.clear();
1373 if response_tx.send(resp).is_err() {
1374 break;
1375 }
1376 } else {
1377 // Response for a past or future connection; stash it.
1378 let stash = stashed_responses.entry(resp_nonce).or_default();
1379 stash.push(resp);
1380 }
1381 }
1382 }
1383 },
1384 );
1385}