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