Skip to main content

mz_storage/
storage_state.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//! Worker-local state for storage timely instances.
7//!
8//! One instance of a [`Worker`], along with its contained [`StorageState`], is
9//! part of an ensemble of storage workers that all run inside the same timely
10//! cluster. We call this worker a _storage worker_ to disambiguate it from
11//! other kinds of workers, potentially other components that might be sharing
12//! the same timely cluster.
13//!
14//! ## Controller and internal communication
15//!
16//! A worker receives _external_ [`StorageCommands`](StorageCommand) from the
17//! storage controller, via a channel. Storage workers also share an _internal_
18//! control/command fabric ([`internal_control`]). Internal commands go through
19//! a sequencer dataflow that ensures that all workers receive all commands in
20//! the same consistent order.
21//!
22//! We need to make sure that commands that cause dataflows to be rendered are
23//! processed in the same consistent order across all workers because timely
24//! requires this. To achieve this, we make sure that only internal commands can
25//! cause dataflows to be rendered. External commands (from the controller)
26//! cause internal commands to be broadcast (by only one worker), to get
27//! dataflows rendered.
28//!
29//! The internal command fabric is also used to broadcast messages from a local
30//! operator/worker to all workers. For example, when we need to tear down and
31//! restart a dataflow on all workers when an error is encountered.
32//!
33//! ## Async Storage Worker
34//!
35//! The storage worker has a companion [`AsyncStorageWorker`] that must be used
36//! when running code that requires `async`. This is needed because a timely
37//! main loop cannot run `async` code.
38//!
39//! ## Example flow of commands for `RunIngestion`
40//!
41//! With external commands, internal commands, and the async worker,
42//! understanding where and how commands from the controller are realized can
43//! get complicated. We will follow the complete flow for `RunIngestion`, as an
44//! example:
45//!
46//! 1. Worker receives a [`StorageCommand::RunIngestion`] command from the
47//!    controller.
48//! 2. This command is processed in [`StorageState::handle_storage_command`].
49//!    This step cannot render dataflows, because it does not have access to the
50//!    timely worker. It will only set up state that stays over the whole
51//!    lifetime of the source, such as the `reported_frontier`. Putting in place
52//!    this reported frontier will enable frontier reporting for that source. We
53//!    will not start reporting when we only see an internal command for
54//!    rendering a dataflow, which can "overtake" the external `RunIngestion`
55//!    command.
56//! 3. During processing of that command, we call
57//!    [`AsyncStorageWorker::update_ingestion_frontiers`], which causes a command to
58//!    be sent to the async worker.
59//! 4. We eventually get a response from the async worker:
60//!    [`AsyncStorageWorkerResponse::IngestionFrontiersUpdated`].
61//! 5. This response is handled in [`Worker::handle_async_worker_response`].
62//! 6. Handling that response causes a
63//!    [`InternalStorageCommand::CreateIngestionDataflow`] to be broadcast to
64//!    all workers via the internal command fabric.
65//! 7. This message will be processed (on each worker) in
66//!    [`Worker::handle_internal_storage_command`]. This is what will cause the
67//!    required dataflow to be rendered on all workers.
68//!
69//! The process described above assumes that the `RunIngestion` is _not_ an
70//! update, i.e. it is in response to a `CREATE SOURCE`-like statement.
71//!
72//! The primary distinction when handling a `RunIngestion` that represents an
73//! update, is that it might fill out new internal state in the mid-level
74//! clients on the way toward being run.
75
76use std::cell::RefCell;
77use std::collections::{BTreeMap, BTreeSet, VecDeque};
78use std::path::PathBuf;
79use std::rc::Rc;
80use std::sync::Arc;
81use std::thread;
82use std::time::Duration;
83
84use fail::fail_point;
85use mz_ore::now::NowFn;
86use mz_ore::soft_assert_or_log;
87use mz_ore::tracing::TracingHandle;
88use mz_persist_client::batch::ProtoBatch;
89use mz_persist_client::cache::PersistClientCache;
90use mz_persist_client::operators::shard_source::ErrorHandler;
91use mz_repr::{GlobalId, Timestamp};
92use mz_rocksdb::config::SharedWriteBufferManager;
93use mz_storage_client::client::{
94    RunIngestionCommand, StatusUpdate, StorageCommand, StorageResponse,
95};
96use mz_storage_types::AlterCompatible;
97use mz_storage_types::configuration::StorageConfiguration;
98use mz_storage_types::connections::ConnectionContext;
99use mz_storage_types::controller::CollectionMetadata;
100use mz_storage_types::dyncfgs::STORAGE_SERVER_MAINTENANCE_INTERVAL;
101use mz_storage_types::oneshot_sources::OneshotIngestionDescription;
102use mz_storage_types::sinks::StorageSinkDesc;
103use mz_storage_types::sources::IngestionDescription;
104use mz_timely_util::builder_async::PressOnDropButton;
105use mz_txn_wal::operator::TxnsContext;
106use timely::order::PartialOrder;
107use timely::progress::Timestamp as _;
108use timely::progress::frontier::Antichain;
109use timely::worker::Worker as TimelyWorker;
110use tokio::sync::mpsc::error::TryRecvError;
111use tokio::sync::{mpsc, watch};
112use tokio::time::Instant;
113use tracing::{debug, info, warn};
114use uuid::Uuid;
115
116use crate::internal_control::{
117    self, DataflowParameters, InternalCommandReceiver, InternalCommandSender,
118    InternalStorageCommand,
119};
120use crate::metrics::StorageMetrics;
121use crate::statistics::{AggregatedStatistics, SinkStatistics, SourceStatistics};
122use crate::storage_state::async_storage_worker::{AsyncStorageWorker, AsyncStorageWorkerResponse};
123
124pub mod async_storage_worker;
125
126type CommandReceiver = mpsc::UnboundedReceiver<StorageCommand>;
127type ResponseSender = mpsc::UnboundedSender<StorageResponse>;
128
129/// State maintained for each worker thread.
130///
131/// Much of this state can be viewed as local variables for the worker thread,
132/// holding state that persists across function calls.
133pub struct Worker<'w> {
134    /// The underlying Timely worker.
135    ///
136    /// NOTE: This is `pub` for testing.
137    pub timely_worker: &'w mut TimelyWorker,
138    /// The channel over which communication handles for newly connected clients
139    /// are delivered.
140    pub client_rx: mpsc::UnboundedReceiver<(Uuid, CommandReceiver, ResponseSender)>,
141    /// The state associated with collection ingress and egress.
142    pub storage_state: StorageState,
143}
144
145impl<'w> Worker<'w> {
146    /// Creates new `Worker` state from the given components.
147    pub fn new(
148        timely_worker: &'w mut TimelyWorker,
149        client_rx: mpsc::UnboundedReceiver<(Uuid, CommandReceiver, ResponseSender)>,
150        metrics: StorageMetrics,
151        now: NowFn,
152        connection_context: ConnectionContext,
153        instance_context: StorageInstanceContext,
154        persist_clients: Arc<PersistClientCache>,
155        txns_ctx: TxnsContext,
156        tracing_handle: Arc<TracingHandle>,
157        shared_rocksdb_write_buffer_manager: SharedWriteBufferManager,
158    ) -> Self {
159        // It is very important that we only create the internal control
160        // flow/command sequencer once because a) the worker state is re-used
161        // when a new client connects and b) dataflows that have already been
162        // rendered into the timely worker are reused as well.
163        //
164        // If we created a new sequencer every time we get a new client (likely
165        // because the controller re-started and re-connected), dataflows that
166        // were rendered before would still hold a handle to the old sequencer
167        // but we would not read their commands anymore.
168        let (internal_cmd_tx, internal_cmd_rx) =
169            internal_control::setup_command_sequencer(timely_worker);
170
171        let storage_state = StorageState::new_guest(
172            timely_worker.index(),
173            timely_worker.peers(),
174            internal_cmd_tx,
175            Some(internal_cmd_rx),
176            metrics,
177            now,
178            connection_context,
179            instance_context,
180            persist_clients,
181            txns_ctx,
182            tracing_handle,
183            shared_rocksdb_write_buffer_manager,
184        );
185
186        // TODO(aljoscha): We might want `async_worker` and `internal_cmd_tx` to
187        // be fields of `Worker` instead of `StorageState`, but at least for the
188        // command flow sources and sinks need access to that. We can refactor
189        // this once we have a clearer boundary between what sources/sinks need
190        // and the full "power" of the internal command flow, which should stay
191        // internal to the worker/not be exposed to source/sink implementations.
192        Self {
193            timely_worker,
194            client_rx,
195            storage_state,
196        }
197    }
198}
199
200impl StorageState {
201    /// Creates per-worker storage state, for hosting on any Timely worker.
202    ///
203    /// The caller provides the internal command channel endpoints: the native storage worker wires
204    /// them to the sequencer dataflow, a foreign host wires the sender to its own sequencing
205    /// channel and passes no receiver, since it dispatches internal commands itself.
206    /// Must be called on the hosting worker's thread, because the async worker unparks the
207    /// creating thread.
208    pub fn new_guest(
209        timely_worker_index: usize,
210        timely_worker_peers: usize,
211        internal_cmd_tx: InternalCommandSender,
212        internal_cmd_rx: Option<InternalCommandReceiver>,
213        metrics: StorageMetrics,
214        now: NowFn,
215        connection_context: ConnectionContext,
216        instance_context: StorageInstanceContext,
217        persist_clients: Arc<PersistClientCache>,
218        txns_ctx: TxnsContext,
219        tracing_handle: Arc<TracingHandle>,
220        shared_rocksdb_write_buffer_manager: SharedWriteBufferManager,
221    ) -> Self {
222        let storage_configuration =
223            StorageConfiguration::new(connection_context, mz_dyncfgs::all_dyncfgs());
224
225        // We always initialize as read_only=true. Only when we're explicitly
226        // allowed do we switch to doing writes.
227        let (read_only_tx, read_only_rx) = watch::channel(true);
228
229        // Similar to the internal command sequencer, it is very important that
230        // we only create the async worker once because a) the worker state is
231        // re-used when a new client connects and b) commands that have already
232        // been sent and might yield a response will be lost if a new iteration
233        // of `run_client` creates a new async worker.
234        //
235        // If we created a new async worker every time we get a new client
236        // (likely because the controller re-started and re-connected), we can
237        // get into an inconsistent state where we think that a dataflow has
238        // been rendered, for example because there is an entry in
239        // `StorageState::ingestions`, while there is not yet a dataflow. This
240        // happens because the dataflow only gets rendered once we get a
241        // response from the async worker and send off an internal command.
242        //
243        // The core idea is that both the sequencer and the async worker are
244        // part of the per-worker state, and must be treated as such, meaning
245        // they must survive between invocations of `run_client`.
246
247        // TODO(aljoscha): This thread unparking business seems brittle, but that's
248        // also how the command channel works currently. We can wrap it inside a
249        // struct that holds both a channel and a `Thread`, but I don't
250        // think that would help too much.
251        let async_worker = async_storage_worker::AsyncStorageWorker::new(
252            thread::current(),
253            Arc::clone(&persist_clients),
254        );
255        let cluster_memory_limit = instance_context.cluster_memory_limit;
256
257        let storage_state = StorageState {
258            source_uppers: BTreeMap::new(),
259            source_tokens: BTreeMap::new(),
260            metrics,
261            reported_frontiers: BTreeMap::new(),
262            ingestions: BTreeMap::new(),
263            exports: BTreeMap::new(),
264            oneshot_ingestions: BTreeMap::new(),
265            now,
266            timely_worker_index,
267            timely_worker_peers,
268            instance_context,
269            persist_clients,
270            txns_ctx,
271            sink_tokens: BTreeMap::new(),
272            sink_write_frontiers: BTreeMap::new(),
273            dropped_ids: Vec::new(),
274            aggregated_statistics: AggregatedStatistics::new(
275                timely_worker_index,
276                timely_worker_peers,
277            ),
278            shared_status_updates: Default::default(),
279            latest_status_updates: Default::default(),
280            initial_status_reported: Default::default(),
281            internal_cmd_tx,
282            internal_cmd_rx,
283            read_only_tx,
284            read_only_rx,
285            async_worker,
286            storage_configuration,
287            dataflow_parameters: DataflowParameters::new(
288                shared_rocksdb_write_buffer_manager,
289                cluster_memory_limit,
290            ),
291            tracing_handle,
292            server_maintenance_interval: Duration::ZERO,
293        };
294
295        storage_state
296    }
297}
298
299/// Worker-local state related to the ingress or egress of collections of data.
300pub struct StorageState {
301    /// The highest observed upper frontier for collection.
302    ///
303    /// This is shared among all source instances, so that they can jointly advance the
304    /// frontier even as other instances are created and dropped. Ideally, the Storage
305    /// module would eventually provide one source of truth on this rather than multiple,
306    /// and we should aim for that but are not there yet.
307    pub source_uppers: BTreeMap<GlobalId, Rc<RefCell<Antichain<mz_repr::Timestamp>>>>,
308    /// Handles to created sources, keyed by ID
309    /// NB: The type of the tokens must not be changed to something other than `PressOnDropButton`
310    /// to prevent usage of custom shutdown tokens that are tricky to get right.
311    pub source_tokens: BTreeMap<GlobalId, Vec<PressOnDropButton>>,
312    /// Metrics for storage objects.
313    pub metrics: StorageMetrics,
314    /// Tracks the conditional write frontiers we have reported.
315    pub reported_frontiers: BTreeMap<GlobalId, Antichain<Timestamp>>,
316    /// Descriptions of each installed ingestion.
317    pub ingestions: BTreeMap<GlobalId, IngestionDescription<CollectionMetadata>>,
318    /// Descriptions of each installed export.
319    pub exports: BTreeMap<GlobalId, StorageSinkDesc<CollectionMetadata, mz_repr::Timestamp>>,
320    /// Descriptions of oneshot ingestions that are currently running.
321    pub oneshot_ingestions: BTreeMap<uuid::Uuid, OneshotIngestionDescription<ProtoBatch>>,
322    /// Undocumented
323    pub now: NowFn,
324    /// Index of the associated timely dataflow worker.
325    pub timely_worker_index: usize,
326    /// Peers in the associated timely dataflow worker.
327    pub timely_worker_peers: usize,
328    /// Other configuration for sources and sinks.
329    pub instance_context: StorageInstanceContext,
330    /// A process-global cache of (blob_uri, consensus_uri) -> PersistClient.
331    /// This is intentionally shared between workers
332    pub persist_clients: Arc<PersistClientCache>,
333    /// Context necessary for rendering txn-wal operators.
334    pub txns_ctx: TxnsContext,
335    /// Tokens that should be dropped when a dataflow is dropped to clean up
336    /// associated state.
337    /// NB: The type of the tokens must not be changed to something other than `PressOnDropButton`
338    /// to prevent usage of custom shutdown tokens that are tricky to get right.
339    pub sink_tokens: BTreeMap<GlobalId, Vec<PressOnDropButton>>,
340    /// Frontier of sink writes (all subsequent writes will be at times at or
341    /// equal to this frontier)
342    pub sink_write_frontiers: BTreeMap<GlobalId, Rc<RefCell<Antichain<Timestamp>>>>,
343    /// Collection ids that have been dropped but not yet reported as dropped
344    pub dropped_ids: Vec<GlobalId>,
345
346    /// Statistics for sources and sinks.
347    pub aggregated_statistics: AggregatedStatistics,
348
349    /// A place shared with running dataflows, so that health operators, can
350    /// report status updates back to us.
351    ///
352    /// **NOTE**: Operators that append to this collection should take care to only add new
353    /// status updates if the status of the ingestion/export in question has _changed_.
354    pub shared_status_updates: Rc<RefCell<Vec<StatusUpdate>>>,
355
356    /// The latest status update for each object.
357    pub latest_status_updates: BTreeMap<GlobalId, StatusUpdate>,
358
359    /// Whether we have reported the initial status after connecting to a new client.
360    /// This is reset to false when a new client connects.
361    pub initial_status_reported: bool,
362
363    /// Sender for cluster-internal storage commands. These can be sent from
364    /// within workers/operators and will be distributed to all workers. For
365    /// example, for shutting down an entire dataflow from within a
366    /// operator/worker.
367    pub internal_cmd_tx: InternalCommandSender,
368    /// Receiver for cluster-internal storage commands. `None` when the state is hosted outside
369    /// the storage server, whose host dispatches internal commands itself.
370    pub internal_cmd_rx: Option<InternalCommandReceiver>,
371
372    /// When this replica/cluster is in read-only mode it must not affect any
373    /// changes to external state. This flag can only be changed by a
374    /// [StorageCommand::AllowWrites].
375    ///
376    /// Everything running on this replica/cluster must obey this flag. At the
377    /// time of writing, nothing currently looks at this flag.
378    /// TODO(benesch): fix this.
379    ///
380    /// NOTE: In the future, we might want a more complicated flag, for example
381    /// something that tells us after which timestamp we are allowed to write.
382    /// In this first version we are keeping things as simple as possible!
383    pub read_only_rx: watch::Receiver<bool>,
384
385    /// Send-side for read-only state.
386    pub read_only_tx: watch::Sender<bool>,
387
388    /// Async worker companion, used for running code that requires async, which
389    /// the timely main loop cannot do.
390    pub async_worker: AsyncStorageWorker<mz_repr::Timestamp>,
391
392    /// Configuration for source and sink connections.
393    pub storage_configuration: StorageConfiguration,
394    /// Dynamically configurable parameters that control how dataflows are rendered.
395    /// NOTE(guswynn): we should consider moving these into `storage_configuration`.
396    pub dataflow_parameters: DataflowParameters,
397
398    /// A process-global handle to tracing configuration.
399    pub tracing_handle: Arc<TracingHandle>,
400
401    /// Interval at which to perform server maintenance tasks. Set to a zero interval to
402    /// perform maintenance with every `step_or_park` invocation.
403    pub server_maintenance_interval: Duration,
404}
405
406impl StorageState {
407    /// Return an error handler that triggers a suspend and restart of the corresponding storage
408    /// dataflow.
409    pub fn error_handler(&self, context: &'static str, id: GlobalId) -> ErrorHandler {
410        let tx = self.internal_cmd_tx.clone();
411        ErrorHandler::signal(move |e| {
412            tx.send(InternalStorageCommand::SuspendAndRestart {
413                id,
414                reason: format!("{context}: {e:#}"),
415            })
416        })
417    }
418}
419
420/// Extra context for a storage instance.
421/// This is extra information that is used when rendering source
422/// and sinks that is not tied to the source/connection configuration itself.
423#[derive(Clone)]
424pub struct StorageInstanceContext {
425    /// A directory that can be used for scratch work.
426    pub scratch_directory: Option<PathBuf>,
427    /// The memory limit of the materialize cluster replica. This will
428    /// be used to calculate and configure the maximum inflight bytes for backpressure
429    pub cluster_memory_limit: Option<usize>,
430    /// Whether the replica can swap.
431    pub swap_available: bool,
432}
433
434impl StorageInstanceContext {
435    /// Build a new `StorageInstanceContext`.
436    pub fn new(
437        scratch_directory: Option<PathBuf>,
438        cluster_memory_limit: Option<usize>,
439        swap_available: bool,
440    ) -> Self {
441        Self {
442            scratch_directory,
443            cluster_memory_limit,
444            swap_available,
445        }
446    }
447
448    /// Returns a `rocksdb::Env` for a new RocksDB instance.
449    ///
450    /// With a scratch directory this is the default `Env`, which stores data
451    /// on the host filesystem. Without one, RocksDB runs in memory, and every
452    /// call returns a fresh in-memory `Env`. State written through an `Env`
453    /// is only reachable through that same `Env`, so a per-instance `Env`
454    /// isolates instances from each other and from previous incarnations of
455    /// themselves. Background threads are process-wide either way, both
456    /// variants delegate them to the default `Env`.
457    pub fn rocksdb_env(&self) -> Result<rocksdb::Env, rocksdb::Error> {
458        if self.scratch_directory.is_some() {
459            rocksdb::Env::new()
460        } else {
461            rocksdb::Env::mem_env()
462        }
463    }
464}
465
466impl<'w> Worker<'w> {
467    /// Waits for client connections and runs them to completion.
468    pub fn run(&mut self) {
469        while let Some((_nonce, rx, tx)) = self.client_rx.blocking_recv() {
470            self.run_client(rx, tx);
471        }
472    }
473
474    /// Runs this (timely) storage worker until the given `command_rx` is
475    /// disconnected.
476    ///
477    /// See the [module documentation](crate::storage_state) for this
478    /// workers responsibilities, how it communicates with the other workers and
479    /// how commands flow from the controller and through the workers.
480    fn run_client(&mut self, mut command_rx: CommandReceiver, response_tx: ResponseSender) {
481        // At this point, all workers are still reading from the command flow.
482        if self.reconcile(&mut command_rx).is_err() {
483            return;
484        }
485
486        // The last time we reported statistics.
487        let mut last_stats_time = Instant::now();
488
489        // The last time we did periodic maintenance.
490        let mut last_maintenance = std::time::Instant::now();
491
492        let mut disconnected = false;
493        while !disconnected {
494            let config = &self.storage_state.storage_configuration;
495            let stats_interval = config.parameters.statistics_collection_interval;
496
497            let maintenance_interval = self.storage_state.server_maintenance_interval;
498
499            let now = std::time::Instant::now();
500            // Determine if we need to perform maintenance, which is true if `maintenance_interval`
501            // time has passed since the last maintenance.
502            let sleep_duration;
503            if now >= last_maintenance + maintenance_interval {
504                last_maintenance = now;
505                sleep_duration = None;
506
507                self.report_frontier_progress(&response_tx);
508            } else {
509                // We didn't perform maintenance, sleep until the next maintenance interval.
510                let next_maintenance = last_maintenance + maintenance_interval;
511                sleep_duration = Some(next_maintenance.saturating_duration_since(now))
512            }
513
514            // Ask Timely to execute a unit of work.
515            //
516            // If there are no pending commands or responses from the async
517            // worker, we ask Timely to park the thread if there's nothing to
518            // do. We rely on another thread unparking us when there's new work
519            // to be done, e.g., when sending a command or when new Kafka
520            // messages have arrived.
521            //
522            // It is critical that we allow Timely to park iff there are no
523            // pending commands or responses. The command may have already been
524            // consumed by the call to `client_rx.recv`. See:
525            // https://github.com/MaterializeInc/materialize/pull/13973#issuecomment-1200312212
526            if command_rx.is_empty() && self.storage_state.async_worker.is_empty() {
527                // Make sure we wake up again to report any pending statistics updates.
528                let mut park_duration = stats_interval.saturating_sub(last_stats_time.elapsed());
529                if let Some(sleep_duration) = sleep_duration {
530                    park_duration = std::cmp::min(sleep_duration, park_duration);
531                }
532                self.timely_worker.step_or_park(Some(park_duration));
533            } else {
534                self.timely_worker.step();
535            }
536
537            // Rerport any dropped ids
538            for id in std::mem::take(&mut self.storage_state.dropped_ids) {
539                self.send_storage_response(&response_tx, StorageResponse::DroppedId(id));
540            }
541
542            self.process_oneshot_ingestions(&response_tx);
543
544            self.report_status_updates(&response_tx);
545
546            if last_stats_time.elapsed() >= stats_interval {
547                self.report_storage_statistics(&response_tx);
548                last_stats_time = Instant::now();
549            }
550
551            // Handle any received commands.
552            loop {
553                match command_rx.try_recv() {
554                    Ok(cmd) => self.storage_state.handle_storage_command(cmd),
555                    Err(TryRecvError::Empty) => break,
556                    Err(TryRecvError::Disconnected) => {
557                        disconnected = true;
558                        break;
559                    }
560                }
561            }
562
563            // Handle responses from the async worker.
564            while let Ok(response) = self.storage_state.async_worker.try_recv() {
565                self.handle_async_worker_response(response);
566            }
567
568            // Handle any received commands.
569            while let Some(command) = self
570                .storage_state
571                .internal_cmd_rx
572                .as_ref()
573                .expect("storage server always wires a receiver")
574                .try_recv()
575            {
576                self.handle_internal_storage_command(command);
577            }
578        }
579    }
580
581    /// Entry point for applying a response from the async storage worker.
582    pub fn handle_async_worker_response(
583        &self,
584        async_response: AsyncStorageWorkerResponse<mz_repr::Timestamp>,
585    ) {
586        // NOTE: If we want to share the load of async processing we
587        // have to change `handle_storage_command` and change this
588        // assert.
589        assert_eq!(
590            self.timely_worker.index(),
591            0,
592            "only worker #0 is doing async processing"
593        );
594        match async_response {
595            AsyncStorageWorkerResponse::IngestionFrontiersUpdated {
596                id,
597                ingestion_description,
598                as_of,
599                resume_uppers,
600                source_resume_uppers,
601            } => {
602                self.storage_state.internal_cmd_tx.send(
603                    InternalStorageCommand::CreateIngestionDataflow {
604                        id,
605                        ingestion_description,
606                        as_of,
607                        resume_uppers,
608                        source_resume_uppers,
609                    },
610                );
611            }
612            AsyncStorageWorkerResponse::ExportFrontiersUpdated { id, description } => {
613                self.storage_state
614                    .internal_cmd_tx
615                    .send(InternalStorageCommand::RunSinkDataflow(id, description));
616            }
617            AsyncStorageWorkerResponse::DropDataflow(id) => {
618                self.storage_state
619                    .internal_cmd_tx
620                    .send(InternalStorageCommand::DropDataflow(vec![id]));
621            }
622        }
623    }
624
625    /// Entry point for applying an internal storage command.
626    pub fn handle_internal_storage_command(&mut self, internal_cmd: InternalStorageCommand) {
627        match internal_cmd {
628            InternalStorageCommand::SuspendAndRestart { id, reason } => {
629                info!(
630                    "worker {}/{} initiating suspend-and-restart for {id} because of: {reason}",
631                    self.timely_worker.index(),
632                    self.timely_worker.peers(),
633                );
634
635                let maybe_ingestion = self.storage_state.ingestions.get(&id).cloned();
636                if let Some(ingestion_description) = maybe_ingestion {
637                    // Yank the token of the previously existing source dataflow.Note that this
638                    // token also includes any source exports/subsources.
639                    let maybe_token = self.storage_state.source_tokens.remove(&id);
640                    if maybe_token.is_none() {
641                        // Something has dropped the source. Make sure we don't
642                        // accidentally re-create it.
643                        return;
644                    }
645
646                    // This needs to be done by one worker, which will
647                    // broadcasts a `CreateIngestionDataflow` command to all
648                    // workers based on the response that contains the
649                    // resumption upper.
650                    //
651                    // Doing this separately on each worker could lead to
652                    // differing resume_uppers which might lead to all kinds of
653                    // mayhem.
654                    //
655                    // TODO(aljoscha): If we ever become worried that this is
656                    // putting undue pressure on worker 0 we can pick the
657                    // designated worker for a source/sink based on `id.hash()`.
658                    if self.timely_worker.index() == 0 {
659                        for (id, _) in ingestion_description.source_exports.iter() {
660                            self.storage_state
661                                .aggregated_statistics
662                                .advance_global_epoch(*id);
663                        }
664                        self.storage_state
665                            .async_worker
666                            .update_ingestion_frontiers(id, ingestion_description);
667                    }
668
669                    // Continue with other commands.
670                    return;
671                }
672
673                let maybe_sink = self.storage_state.exports.get(&id).cloned();
674                if let Some(sink_description) = maybe_sink {
675                    // Yank the token of the previously existing sink
676                    // dataflow.
677                    let maybe_token = self.storage_state.sink_tokens.remove(&id);
678
679                    if maybe_token.is_none() {
680                        // Something has dropped the sink. Make sure we don't
681                        // accidentally re-create it.
682                        return;
683                    }
684
685                    // This needs to be broadcast by one worker and go through
686                    // the internal command fabric, to ensure consistent
687                    // ordering of dataflow rendering across all workers.
688                    if self.timely_worker.index() == 0 {
689                        self.storage_state
690                            .aggregated_statistics
691                            .advance_global_epoch(id);
692                        self.storage_state
693                            .async_worker
694                            .update_sink_frontiers(id, sink_description);
695                    }
696
697                    // Continue with other commands.
698                    return;
699                }
700
701                if !self
702                    .storage_state
703                    .ingestions
704                    .values()
705                    .any(|v| v.source_exports.contains_key(&id))
706                {
707                    // Our current approach to dropping a source results in a race between shard
708                    // finalization (which happens in the controller) and dataflow shutdown (which
709                    // happens in clusterd). If a source is created and dropped fast enough -or the
710                    // two commands get sufficiently delayed- then it's possible to receive a
711                    // SuspendAndRestart command for an unknown source. We cannot assert that this
712                    // never happens but we log an error here to track how often this happens.
713                    warn!(
714                        "got InternalStorageCommand::SuspendAndRestart for something that is not a source or sink: {id}"
715                    );
716                }
717            }
718            InternalStorageCommand::CreateIngestionDataflow {
719                id: ingestion_id,
720                mut ingestion_description,
721                as_of,
722                mut resume_uppers,
723                mut source_resume_uppers,
724            } => {
725                info!(
726                    ?as_of,
727                    ?resume_uppers,
728                    "worker {}/{} trying to (re-)start ingestion {ingestion_id}",
729                    self.timely_worker.index(),
730                    self.timely_worker.peers(),
731                );
732
733                // We initialize statistics before we prune finished exports. We
734                // still want to export statistics for these, plus the rendering
735                // machinery will get confused if there are not at least
736                // statistics for the "main" source.
737                for (export_id, export) in ingestion_description.source_exports.iter() {
738                    let resume_upper = resume_uppers[export_id].clone();
739                    self.storage_state.aggregated_statistics.initialize_source(
740                        *export_id,
741                        ingestion_id,
742                        resume_upper.clone(),
743                        as_of.clone(),
744                        || {
745                            SourceStatistics::new(
746                                *export_id,
747                                self.storage_state.timely_worker_index,
748                                &self.storage_state.metrics.source_statistics,
749                                ingestion_id,
750                                &export.storage_metadata.data_shard,
751                                export.data_config.envelope.clone(),
752                                resume_upper,
753                                as_of.clone(),
754                            )
755                        },
756                    );
757                }
758
759                let finished_exports: BTreeSet<GlobalId> = resume_uppers
760                    .iter()
761                    .filter(|(_, frontier)| frontier.is_empty())
762                    .map(|(id, _)| *id)
763                    .collect();
764
765                resume_uppers.retain(|id, _| !finished_exports.contains(id));
766                source_resume_uppers.retain(|id, _| !finished_exports.contains(id));
767                ingestion_description
768                    .source_exports
769                    .retain(|id, _| !finished_exports.contains(id));
770
771                for id in ingestion_description.collection_ids() {
772                    // If there is already a shared upper, we re-use it, to make
773                    // sure that parties that are already using the shared upper
774                    // can continue doing so.
775                    let source_upper = self
776                        .storage_state
777                        .source_uppers
778                        .entry(id.clone())
779                        .or_insert_with(|| {
780                            Rc::new(RefCell::new(Antichain::from_elem(Timestamp::minimum())))
781                        });
782
783                    let mut source_upper = source_upper.borrow_mut();
784                    if !source_upper.is_empty() {
785                        source_upper.clear();
786                        source_upper.insert(mz_repr::Timestamp::minimum());
787                    }
788                }
789
790                // If all subsources of the source are finished, we can skip rendering entirely.
791                // Also, if `as_of` is empty, the dataflow has been finalized, so we can skip it as
792                // well.
793                //
794                // TODO(guswynn|petrosagg): this is a bit hacky, and is a consequence of storage state
795                // management being a bit of a mess. we should clean this up and remove weird if
796                // statements like this.
797                if resume_uppers.values().all(|frontier| frontier.is_empty()) || as_of.is_empty() {
798                    info!(
799                        ?resume_uppers,
800                        ?as_of,
801                        "worker {}/{} skipping building ingestion dataflow \
802                        for {ingestion_id} because the ingestion is finished",
803                        self.timely_worker.index(),
804                        self.timely_worker.peers(),
805                    );
806                    return;
807                }
808
809                crate::render::build_ingestion_dataflow(
810                    self.timely_worker,
811                    &mut self.storage_state,
812                    ingestion_id,
813                    ingestion_description,
814                    as_of,
815                    resume_uppers,
816                    source_resume_uppers,
817                );
818            }
819            InternalStorageCommand::RunOneshotIngestion {
820                ingestion_id,
821                collection_id,
822                collection_meta,
823                request,
824            } => {
825                crate::render::build_oneshot_ingestion_dataflow(
826                    self.timely_worker,
827                    &mut self.storage_state,
828                    ingestion_id,
829                    collection_id,
830                    collection_meta,
831                    request,
832                );
833            }
834            InternalStorageCommand::RunSinkDataflow(sink_id, sink_description) => {
835                info!(
836                    "worker {}/{} trying to (re-)start sink {sink_id}",
837                    self.timely_worker.index(),
838                    self.timely_worker.peers(),
839                );
840
841                {
842                    // If there is already a shared write frontier, we re-use it, to
843                    // make sure that parties that are already using the shared
844                    // frontier can continue doing so.
845                    let sink_write_frontier = self
846                        .storage_state
847                        .sink_write_frontiers
848                        .entry(sink_id.clone())
849                        .or_insert_with(|| Rc::new(RefCell::new(Antichain::new())));
850
851                    let mut sink_write_frontier = sink_write_frontier.borrow_mut();
852                    sink_write_frontier.clear();
853                    sink_write_frontier.insert(mz_repr::Timestamp::minimum());
854                }
855                self.storage_state
856                    .aggregated_statistics
857                    .initialize_sink(sink_id, || {
858                        SinkStatistics::new(
859                            sink_id,
860                            self.storage_state.timely_worker_index,
861                            &self.storage_state.metrics.sink_statistics,
862                        )
863                    });
864
865                crate::render::build_export_dataflow(
866                    self.timely_worker,
867                    &mut self.storage_state,
868                    sink_id,
869                    sink_description,
870                );
871            }
872            InternalStorageCommand::DropDataflow(ids) => {
873                for id in &ids {
874                    // Clean up per-source / per-sink state.
875                    self.storage_state.source_uppers.remove(id);
876                    self.storage_state.source_tokens.remove(id);
877
878                    self.storage_state.sink_tokens.remove(id);
879                    self.storage_state.sink_write_frontiers.remove(id);
880
881                    self.storage_state.aggregated_statistics.deinitialize(*id);
882                }
883            }
884            InternalStorageCommand::UpdateConfiguration { storage_parameters } => {
885                self.storage_state
886                    .dataflow_parameters
887                    .update(storage_parameters.clone());
888                self.storage_state
889                    .storage_configuration
890                    .update(storage_parameters);
891
892                // Clear out the updates as we no longer forward them to anyone else to process.
893                // We clone `StorageState::storage_configuration` many times during rendering
894                // and want to avoid cloning these unused updates.
895                self.storage_state
896                    .storage_configuration
897                    .parameters
898                    .dyncfg_updates = Default::default();
899
900                // Remember the maintenance interval locally to avoid reading it from the config set on
901                // every server iteration.
902                self.storage_state.server_maintenance_interval =
903                    STORAGE_SERVER_MAINTENANCE_INTERVAL
904                        .get(self.storage_state.storage_configuration.config_set());
905
906                // Apply storage's upsert spill flag to both stash flavors'
907                // mechanisms: the storage chunk spill gate (chunked flavor)
908                // and the storage-owned column pager (paged flavor). The
909                // buffer pool, the pager pool, and their budgets are the
910                // shared ones configured by compute's `apply_worker_config`
911                // (compute and storage run in the same process).
912                //
913                // The flag is replica-scoped: the storage controller merges
914                // per-replica overrides into the `UpdateConfiguration`
915                // commands it sends, so reading this worker's `ConfigSet`
916                // here observes them.
917                {
918                    use mz_storage_types::dyncfgs::ENABLE_UPSERT_PAGED_SPILL;
919
920                    let enabled = ENABLE_UPSERT_PAGED_SPILL
921                        .get(self.storage_state.storage_configuration.config_set());
922                    debug!(
923                        worker = self.timely_worker.index(),
924                        enabled, "upsert stash spill: applying gate",
925                    );
926                    crate::upsert::upsert_stash_spill::set_enabled(enabled);
927                    crate::upsert::upsert_stash_pager::set_enabled(enabled);
928                }
929            }
930            InternalStorageCommand::StatisticsUpdate { sources, sinks } => self
931                .storage_state
932                .aggregated_statistics
933                .ingest(sources, sinks),
934        }
935    }
936
937    /// Emit information about write frontier progress, along with information that should
938    /// be made durable for this to be the case.
939    ///
940    /// The write frontier progress is "conditional" in that it is not until the information is made
941    /// durable that the data are emitted to downstream workers, and indeed they should not rely on
942    /// the completeness of what they hear until the information is made durable.
943    ///
944    /// Specifically, this sends information about new timestamp bindings created by dataflow workers,
945    /// with the understanding if that if made durable (and ack'd back to the workers) the source will
946    /// in fact progress with this write frontier.
947    pub fn report_frontier_progress(&mut self, response_tx: &ResponseSender) {
948        let mut new_uppers = Vec::new();
949
950        // Check if any observed frontier should advance the reported frontiers.
951        for (id, frontier) in self
952            .storage_state
953            .source_uppers
954            .iter()
955            .chain(self.storage_state.sink_write_frontiers.iter())
956        {
957            let Some(reported_frontier) = self.storage_state.reported_frontiers.get_mut(id) else {
958                // Frontier reporting has not yet been started for this object.
959                // Potentially because this timely worker has not yet seen the
960                // `CreateSources` command.
961                continue;
962            };
963
964            let observed_frontier = frontier.borrow();
965
966            // Only do a thing if it *advances* the frontier, not just *changes* the frontier.
967            // This is protection against `frontier` lagging behind what we have conditionally reported.
968            if PartialOrder::less_than(reported_frontier, &observed_frontier) {
969                new_uppers.push((*id, observed_frontier.clone()));
970                reported_frontier.clone_from(&observed_frontier);
971            }
972        }
973
974        for (id, upper) in new_uppers {
975            self.send_storage_response(response_tx, StorageResponse::FrontierUpper(id, upper));
976        }
977    }
978
979    /// Pumps latest status updates from the buffer shared with operators and
980    /// reports any updates that need reporting.
981    pub fn report_status_updates(&mut self, response_tx: &ResponseSender) {
982        // If we haven't done the initial status report, report all current statuses
983        if !self.storage_state.initial_status_reported {
984            // We pull initially reported status updates to "now", so that they
985            // sort as the latest update in internal status collections. This
986            // makes it so that a newly bootstrapped envd can append status
987            // updates to internal status collections that report an accurate
988            // view as of the time when they came up.
989            let now_ts = mz_ore::now::to_datetime((self.storage_state.now)());
990            let status_updates = self
991                .storage_state
992                .latest_status_updates
993                .values()
994                .cloned()
995                .map(|mut update| {
996                    update.timestamp = now_ts.clone();
997                    update
998                });
999            for update in status_updates {
1000                self.send_storage_response(response_tx, StorageResponse::StatusUpdate(update));
1001            }
1002            self.storage_state.initial_status_reported = true;
1003        }
1004
1005        // Pump updates into our state and stage them for reporting.
1006        for shared_update in self.storage_state.shared_status_updates.take() {
1007            self.send_storage_response(
1008                response_tx,
1009                StorageResponse::StatusUpdate(shared_update.clone()),
1010            );
1011
1012            self.storage_state
1013                .latest_status_updates
1014                .insert(shared_update.id, shared_update);
1015        }
1016    }
1017
1018    /// Report source statistics back to the controller.
1019    pub fn report_storage_statistics(&mut self, response_tx: &ResponseSender) {
1020        let (sources, sinks) = self.storage_state.aggregated_statistics.emit_local();
1021        if !sources.is_empty() || !sinks.is_empty() {
1022            self.storage_state
1023                .internal_cmd_tx
1024                .send(InternalStorageCommand::StatisticsUpdate { sources, sinks })
1025        }
1026
1027        let (sources, sinks) = self.storage_state.aggregated_statistics.snapshot();
1028        if !sources.is_empty() || !sinks.is_empty() {
1029            self.send_storage_response(
1030                response_tx,
1031                StorageResponse::StatisticsUpdates(sources, sinks),
1032            );
1033        }
1034    }
1035
1036    /// Send a response to the coordinator.
1037    pub fn send_storage_response(&self, response_tx: &ResponseSender, response: StorageResponse) {
1038        // Ignore send errors because the coordinator is free to ignore our
1039        // responses. This happens during shutdown.
1040        let _ = response_tx.send(response);
1041    }
1042
1043    /// Forward completed oneshot ingestion results to the coordinator.
1044    pub fn process_oneshot_ingestions(&mut self, response_tx: &ResponseSender) {
1045        for (ingestion_id, ingestion_state) in &mut self.storage_state.oneshot_ingestions {
1046            loop {
1047                match ingestion_state.results.try_recv() {
1048                    Ok(result) => {
1049                        let response = match result {
1050                            Ok(maybe_batch) => maybe_batch.into_iter().map(Result::Ok).collect(),
1051                            Err(err) => vec![Err(err)],
1052                        };
1053                        let staged_batches = BTreeMap::from([(*ingestion_id, response)]);
1054                        let _ = response_tx.send(StorageResponse::StagedBatches(staged_batches));
1055                    }
1056                    Err(TryRecvError::Empty) => {
1057                        break;
1058                    }
1059                    Err(TryRecvError::Disconnected) => {
1060                        break;
1061                    }
1062                }
1063            }
1064        }
1065    }
1066
1067    /// Extract commands until `InitializationComplete`, and make the worker
1068    /// reflect those commands. If the worker can not be made to reflect the
1069    /// commands, return an error.
1070    fn reconcile(&mut self, command_rx: &mut CommandReceiver) -> Result<(), ()> {
1071        // To initialize the connection, we want to drain all commands until we
1072        // receive a `StorageCommand::InitializationComplete` command to form a
1073        // target command state.
1074        let mut commands = vec![];
1075        loop {
1076            match command_rx.blocking_recv().ok_or(())? {
1077                StorageCommand::InitializationComplete => break,
1078                command => commands.push(command),
1079            }
1080        }
1081
1082        self.reconcile_commands(commands);
1083        Ok(())
1084    }
1085
1086    /// Reconciles the worker state with the given target command state,
1087    /// which the caller has drained from a new client connection up to (exclusive) the
1088    /// `InitializationComplete` marker.
1089    pub fn reconcile_commands(&mut self, mut commands: Vec<StorageCommand>) {
1090        let worker_id = self.timely_worker.index();
1091
1092        // Track which frontiers this envd expects; we will also set their
1093        // initial timestamp to the minimum timestamp to reset them as we don't
1094        // know what frontiers the new envd expects.
1095        let mut expected_objects = BTreeSet::new();
1096
1097        let mut drop_commands = BTreeSet::new();
1098        let mut running_ingestion_descriptions = self.storage_state.ingestions.clone();
1099        let mut running_exports_descriptions = self.storage_state.exports.clone();
1100
1101        let mut create_oneshot_ingestions: BTreeSet<Uuid> = BTreeSet::new();
1102        let mut cancel_oneshot_ingestions: BTreeSet<Uuid> = BTreeSet::new();
1103
1104        for command in &mut commands {
1105            match command {
1106                StorageCommand::Hello { .. } => {
1107                    panic!("Hello must be captured before")
1108                }
1109                StorageCommand::AllowCompaction(id, since) => {
1110                    info!(%worker_id, ?id, ?since, "reconcile: received AllowCompaction command");
1111
1112                    // collect all "drop commands". These are `AllowCompaction`
1113                    // commands that compact to the empty since. Then, later, we make sure
1114                    // we retain only those `Create*` commands that are not dropped. We
1115                    // assume that the `AllowCompaction` command is ordered after the
1116                    // `Create*` commands but don't assert that.
1117                    // WIP: Should we assert?
1118                    if since.is_empty() {
1119                        drop_commands.insert(*id);
1120                    }
1121                }
1122                StorageCommand::RunIngestion(ingestion) => {
1123                    info!(%worker_id, ?ingestion, "reconcile: received RunIngestion command");
1124
1125                    // Ensure that ingestions are forward-rolling alter compatible.
1126                    let prev = running_ingestion_descriptions
1127                        .insert(ingestion.id, ingestion.description.clone());
1128
1129                    if let Some(prev_ingest) = prev {
1130                        // If the new ingestion is not exactly equal to the currently running
1131                        // ingestion, we must either track that we need to synthesize an update
1132                        // command to change the ingestion, or panic.
1133                        prev_ingest
1134                            .alter_compatible(ingestion.id, &ingestion.description)
1135                            .expect("only alter compatible ingestions permitted");
1136                    }
1137                }
1138                StorageCommand::RunSink(export) => {
1139                    info!(%worker_id, ?export, "reconcile: received RunSink command");
1140
1141                    // Ensure that exports are forward-rolling alter compatible.
1142                    let prev =
1143                        running_exports_descriptions.insert(export.id, export.description.clone());
1144
1145                    if let Some(prev_export) = prev {
1146                        prev_export
1147                            .alter_compatible(export.id, &export.description)
1148                            .expect("only alter compatible exports permitted");
1149                    }
1150                }
1151                StorageCommand::RunOneshotIngestion(ingestion) => {
1152                    info!(
1153                        %worker_id,
1154                        ingestion_id = %ingestion.ingestion_id,
1155                        collection_id = %ingestion.collection_id,
1156                        "reconcile: received RunOneshotIngestion command",
1157                    );
1158                    create_oneshot_ingestions.insert(ingestion.ingestion_id);
1159                }
1160                StorageCommand::CancelOneshotIngestion(uuid) => {
1161                    info!(%worker_id, %uuid, "reconcile: received CancelOneshotIngestion command");
1162                    cancel_oneshot_ingestions.insert(*uuid);
1163                }
1164                StorageCommand::InitializationComplete
1165                | StorageCommand::AllowWrites
1166                | StorageCommand::UpdateConfiguration(_) => (),
1167            }
1168        }
1169
1170        let mut seen_most_recent_definition = BTreeSet::new();
1171
1172        // We iterate over this backward to ensure that we keep only the most recent ingestion
1173        // description.
1174        let mut filtered_commands = VecDeque::new();
1175        for mut command in commands.into_iter().rev() {
1176            let mut should_keep = true;
1177            match &mut command {
1178                StorageCommand::Hello { .. } => {
1179                    panic!("Hello must be captured before")
1180                }
1181                StorageCommand::RunIngestion(ingestion) => {
1182                    // Subsources can be dropped independently of their
1183                    // primary source, so we evaluate them in a separate
1184                    // loop.
1185                    for export_id in ingestion
1186                        .description
1187                        .source_exports
1188                        .keys()
1189                        .filter(|export_id| **export_id != ingestion.id)
1190                    {
1191                        if drop_commands.remove(export_id) {
1192                            info!(%worker_id, %export_id, "reconcile: dropping subsource");
1193                            self.storage_state.dropped_ids.push(*export_id);
1194                        }
1195                    }
1196
1197                    if drop_commands.remove(&ingestion.id)
1198                        || self.storage_state.dropped_ids.contains(&ingestion.id)
1199                    {
1200                        info!(%worker_id, %ingestion.id, "reconcile: dropping ingestion");
1201
1202                        // If an ingestion is dropped, so too must all of
1203                        // its subsources (i.e. ingestion exports, as well
1204                        // as its progress subsource).
1205                        for id in ingestion.description.collection_ids() {
1206                            drop_commands.remove(&id);
1207                            self.storage_state.dropped_ids.push(id);
1208                        }
1209                        should_keep = false;
1210                    } else {
1211                        let most_recent_defintion =
1212                            seen_most_recent_definition.insert(ingestion.id);
1213
1214                        if most_recent_defintion {
1215                            // If this is the most recent definition, this
1216                            // is what we will be running when
1217                            // reconciliation completes. This definition
1218                            // must not include any dropped subsources.
1219                            ingestion.description.source_exports.retain(|export_id, _| {
1220                                !self.storage_state.dropped_ids.contains(export_id)
1221                            });
1222
1223                            // After clearing any dropped subsources, we can
1224                            // state that we expect all of these to exist.
1225                            expected_objects.extend(ingestion.description.collection_ids());
1226                        }
1227
1228                        let running_ingestion = self.storage_state.ingestions.get(&ingestion.id);
1229
1230                        // We keep only:
1231                        // - The most recent version of the ingestion, which
1232                        //   is why these commands are run in reverse.
1233                        // - Ingestions whose descriptions are not exactly
1234                        //   those that are currently running.
1235                        should_keep = most_recent_defintion
1236                            && running_ingestion != Some(&ingestion.description)
1237                    }
1238                }
1239                StorageCommand::RunSink(export) => {
1240                    if drop_commands.remove(&export.id)
1241                        // If there were multiple `RunSink` in the command
1242                        // stream, we want to ensure none of them are
1243                        // retained.
1244                        || self.storage_state.dropped_ids.contains(&export.id)
1245                    {
1246                        info!(%worker_id, %export.id, "reconcile: dropping sink");
1247
1248                        // Make sure that we report back that the ID was
1249                        // dropped.
1250                        self.storage_state.dropped_ids.push(export.id);
1251
1252                        should_keep = false
1253                    } else {
1254                        expected_objects.insert(export.id);
1255
1256                        let running_sink = self.storage_state.exports.get(&export.id);
1257
1258                        // We keep only:
1259                        // - The most recent version of the sink, which
1260                        //   is why these commands are run in reverse.
1261                        // - Sinks whose descriptions are not exactly
1262                        //   those that are currently running.
1263                        should_keep = seen_most_recent_definition.insert(export.id)
1264                            && running_sink != Some(&export.description);
1265                    }
1266                }
1267                StorageCommand::RunOneshotIngestion(ingestion) => {
1268                    let already_running = self
1269                        .storage_state
1270                        .oneshot_ingestions
1271                        .contains_key(&ingestion.ingestion_id);
1272                    let was_canceled = cancel_oneshot_ingestions.contains(&ingestion.ingestion_id);
1273
1274                    should_keep = !already_running && !was_canceled;
1275                }
1276                StorageCommand::CancelOneshotIngestion(ingestion_id) => {
1277                    let already_running = self
1278                        .storage_state
1279                        .oneshot_ingestions
1280                        .contains_key(ingestion_id);
1281                    should_keep = already_running;
1282                }
1283                StorageCommand::InitializationComplete
1284                | StorageCommand::AllowWrites
1285                | StorageCommand::UpdateConfiguration(_)
1286                | StorageCommand::AllowCompaction(_, _) => (),
1287            }
1288            if should_keep {
1289                filtered_commands.push_front(command);
1290            }
1291        }
1292        let commands = filtered_commands;
1293
1294        // Make sure all the "drop commands" matched up with a source or sink.
1295        // This is also what the regular handler logic for `AllowCompaction`
1296        // would do.
1297        soft_assert_or_log!(
1298            drop_commands.is_empty(),
1299            "AllowCompaction commands for non-existent IDs {:?}",
1300            drop_commands
1301        );
1302
1303        // Determine the ID of all objects we did _not_ see; these are
1304        // considered stale.
1305        let stale_objects = self
1306            .storage_state
1307            .ingestions
1308            .values()
1309            .map(|i| i.collection_ids())
1310            .flatten()
1311            .chain(self.storage_state.exports.keys().copied())
1312            // Objects are considered stale if we did not see them re-created.
1313            .filter(|id| !expected_objects.contains(id))
1314            .collect::<Vec<_>>();
1315        let stale_oneshot_ingestions = self
1316            .storage_state
1317            .oneshot_ingestions
1318            .keys()
1319            .filter(|ingestion_id| {
1320                let to_create = create_oneshot_ingestions.contains(ingestion_id);
1321                let to_drop = cancel_oneshot_ingestions.contains(ingestion_id);
1322                mz_ore::soft_assert_or_log!(
1323                    !(!to_create && to_drop),
1324                    "attempting to drop oneshot source {ingestion_id} that is not expected to be created during reconciliation"
1325                );
1326                !to_create && !to_drop
1327            })
1328            .copied()
1329            .collect::<Vec<_>>();
1330
1331        info!(
1332            %worker_id, ?expected_objects, ?stale_objects, ?stale_oneshot_ingestions,
1333            "reconcile: modifing storage state to match expected objects",
1334        );
1335
1336        for id in stale_objects {
1337            self.storage_state.drop_collection(id);
1338        }
1339        for id in stale_oneshot_ingestions {
1340            self.storage_state.drop_oneshot_ingestion(id);
1341        }
1342
1343        // Do not report dropping any objects that do not belong to expected
1344        // objects.
1345        self.storage_state
1346            .dropped_ids
1347            .retain(|id| expected_objects.contains(id));
1348
1349        // Do not report any frontiers that do not belong to expected objects.
1350        // Note that this set of objects can differ from the set of sources and
1351        // sinks.
1352        self.storage_state
1353            .reported_frontiers
1354            .retain(|id, _| expected_objects.contains(id));
1355
1356        // Reset the reported frontiers for the remaining objects.
1357        for (_, frontier) in &mut self.storage_state.reported_frontiers {
1358            *frontier = Antichain::from_elem(<_>::minimum());
1359        }
1360
1361        // Reset the initial status reported flag when a new client connects
1362        self.storage_state.initial_status_reported = false;
1363
1364        // Execute the modified commands.
1365        for command in commands {
1366            self.storage_state.handle_storage_command(command);
1367        }
1368    }
1369}
1370
1371impl StorageState {
1372    /// Entry point for applying a storage command.
1373    ///
1374    /// NOTE: This does not have access to the timely worker and therefore
1375    /// cannot render dataflows. For dataflow rendering, this needs to either
1376    /// send asynchronous command to the `async_worker` or internal
1377    /// commands to the `internal_cmd_tx`.
1378    pub fn handle_storage_command(&mut self, cmd: StorageCommand) {
1379        match cmd {
1380            StorageCommand::Hello { .. } => panic!("Hello must be captured before"),
1381            StorageCommand::InitializationComplete => (),
1382            StorageCommand::AllowWrites => {
1383                self.read_only_tx
1384                    .send(false)
1385                    .expect("we're holding one other end");
1386                self.persist_clients.cfg().enable_compaction();
1387            }
1388            StorageCommand::UpdateConfiguration(params) => {
1389                // These can be done from all workers safely.
1390                debug!("Applying configuration update: {params:?}");
1391
1392                // We serialize the dyncfg updates in StorageParameters, but configure
1393                // persist separately.
1394                self.persist_clients
1395                    .cfg()
1396                    .apply_from(&params.dyncfg_updates);
1397
1398                params.tracing.apply(self.tracing_handle.as_ref());
1399
1400                if let Some(log_filter) = &params.tracing.log_filter {
1401                    self.storage_configuration
1402                        .connection_context
1403                        .librdkafka_log_level =
1404                        mz_ore::tracing::crate_level(&log_filter.clone().into(), "librdkafka");
1405                }
1406
1407                // This needs to be broadcast by one worker and go through
1408                // the internal command fabric, to ensure consistent
1409                // ordering of dataflow rendering across all workers.
1410                if self.timely_worker_index == 0 {
1411                    self.internal_cmd_tx
1412                        .send(InternalStorageCommand::UpdateConfiguration {
1413                            storage_parameters: *params,
1414                        })
1415                }
1416            }
1417            StorageCommand::RunIngestion(ingestion) => {
1418                let RunIngestionCommand { id, description } = *ingestion;
1419
1420                // Remember the ingestion description to facilitate possible
1421                // reconciliation later.
1422                self.ingestions.insert(id, description.clone());
1423
1424                // Initialize shared frontier reporting.
1425                for id in description.collection_ids() {
1426                    self.reported_frontiers
1427                        .entry(id)
1428                        .or_insert_with(|| Antichain::from_elem(mz_repr::Timestamp::minimum()));
1429                }
1430
1431                // This needs to be done by one worker, which will broadcasts a
1432                // `CreateIngestionDataflow` command to all workers based on the response that
1433                // contains the resumption upper.
1434                //
1435                // Doing this separately on each worker could lead to differing resume_uppers
1436                // which might lead to all kinds of mayhem.
1437                //
1438                // n.b. the ingestion on each worker uses the description from worker 0––not the
1439                // ingestion in the local storage state. This is something we might have
1440                // interest in fixing in the future, e.g. materialize#19907
1441                if self.timely_worker_index == 0 {
1442                    self.async_worker
1443                        .update_ingestion_frontiers(id, description);
1444                }
1445            }
1446            StorageCommand::RunOneshotIngestion(oneshot) => {
1447                if self.timely_worker_index == 0 {
1448                    self.internal_cmd_tx
1449                        .send(InternalStorageCommand::RunOneshotIngestion {
1450                            ingestion_id: oneshot.ingestion_id,
1451                            collection_id: oneshot.collection_id,
1452                            collection_meta: oneshot.collection_meta,
1453                            request: oneshot.request,
1454                        });
1455                }
1456            }
1457            StorageCommand::CancelOneshotIngestion(id) => {
1458                self.drop_oneshot_ingestion(id);
1459            }
1460            StorageCommand::RunSink(export) => {
1461                // Remember the sink description to facilitate possible
1462                // reconciliation later.
1463                let prev = self.exports.insert(export.id, export.description.clone());
1464
1465                // New sink, add state.
1466                if prev.is_none() {
1467                    self.reported_frontiers.insert(
1468                        export.id,
1469                        Antichain::from_elem(mz_repr::Timestamp::minimum()),
1470                    );
1471                }
1472
1473                // This needs to be broadcast by one worker and go through the internal command
1474                // fabric, to ensure consistent ordering of dataflow rendering across all
1475                // workers.
1476                if self.timely_worker_index == 0 {
1477                    self.internal_cmd_tx
1478                        .send(InternalStorageCommand::RunSinkDataflow(
1479                            export.id,
1480                            export.description,
1481                        ));
1482                }
1483            }
1484            StorageCommand::AllowCompaction(id, frontier) => {
1485                soft_assert_or_log!(
1486                    self.exports.contains_key(&id) || self.reported_frontiers.contains_key(&id),
1487                    "AllowCompaction command for non-existent {id}"
1488                );
1489
1490                if frontier.is_empty() {
1491                    // Indicates that we may drop `id`, as there are no more valid times to read.
1492                    self.drop_collection(id);
1493                }
1494            }
1495        }
1496    }
1497
1498    /// Drop the identified storage collection from the storage state.
1499    fn drop_collection(&mut self, id: GlobalId) {
1500        fail_point!("crash_on_drop");
1501
1502        self.ingestions.remove(&id);
1503        self.exports.remove(&id);
1504
1505        let _ = self.latest_status_updates.remove(&id);
1506
1507        // This will stop reporting of frontiers.
1508        //
1509        // If this object still has its frontiers reported, we will notify the
1510        // client envd of the drop.
1511        if self.reported_frontiers.remove(&id).is_some() {
1512            // The only actions left are internal cleanup, so we can commit to
1513            // the client that these objects have been dropped.
1514            //
1515            // This must be done now rather than in response to `DropDataflow`,
1516            // otherwise we introduce the possibility of a timing issue where:
1517            // - We remove all tracking state from the storage state and send
1518            //   `DropDataflow` (i.e. this block).
1519            // - While waiting to process that command, we reconcile with a new
1520            //   envd. That envd has already committed to its catalog that this
1521            //   object no longer exists.
1522            // - We process the `DropDataflow` command, and identify that this
1523            //   object has been dropped.
1524            // - The next time `dropped_ids` is processed, we send a response
1525            //   that this ID has been dropped, but the upstream state has no
1526            //   record of that object having ever existed.
1527            self.dropped_ids.push(id);
1528        }
1529
1530        // Send through async worker for correct ordering with RunIngestion, and
1531        // dropping the dataflow is done on async worker response.
1532        if self.timely_worker_index == 0 {
1533            self.async_worker.drop_dataflow(id);
1534        }
1535    }
1536
1537    /// Drop the identified oneshot ingestion from the storage state.
1538    fn drop_oneshot_ingestion(&mut self, ingestion_id: uuid::Uuid) {
1539        let prev = self.oneshot_ingestions.remove(&ingestion_id);
1540        info!(%ingestion_id, existed = %prev.is_some(), "dropping oneshot ingestion");
1541    }
1542}