mz_adapter/coord/sequencer/inner/
subscribe.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10use futures::StreamExt;
11use mz_adapter_types::connection::ConnectionId;
12use mz_adapter_types::dyncfgs::SUBSCRIBE_MAX_BUFFERED_BYTES;
13use mz_cluster_client::ReplicaId;
14use mz_compute_types::ComputeInstanceId;
15use mz_compute_types::dataflows::LirDataflowDescription;
16use mz_ore::collections::CollectionExt;
17use mz_ore::instrument;
18use mz_repr::GlobalId;
19use mz_sql::plan;
20use std::collections::BTreeSet;
21use std::sync::{Arc, Mutex};
22use timely::PartialOrder;
23use tokio::sync::mpsc;
24use tokio_stream::wrappers::UnboundedReceiverStream;
25use uuid::Uuid;
26
27use crate::active_compute_sink::{
28    ActiveComputeSink, ActiveComputeSinkRetireReason, ActiveSubscribe, ActiveSubscribeOwner,
29    SubscribeBacklogAccounting,
30};
31use crate::command::ExecuteResponse;
32use crate::coord::Coordinator;
33use crate::coord::appends::BuiltinTableAppendNotify;
34use crate::coord::peek::PeekResponseUnary;
35use crate::error::AdapterError;
36use crate::statement_logging::StatementLoggingId;
37use crate::{ExecuteContextGuard, ReadHolds};
38
39impl Coordinator {
40    #[instrument]
41    pub(crate) async fn implement_subscribe(
42        &mut self,
43        statement_logging_id: Option<StatementLoggingId>,
44        df_desc: LirDataflowDescription,
45        dependency_ids: BTreeSet<GlobalId>,
46        cluster_id: ComputeInstanceId,
47        replica_id: Option<ReplicaId>,
48        conn_id: ConnectionId,
49        session_uuid: Uuid,
50        read_holds: ReadHolds,
51        plan: plan::SubscribePlan,
52    ) -> Result<(ExecuteResponse, BuiltinTableAppendNotify), AdapterError> {
53        let sink_id = df_desc.sink_id();
54
55        let (tx, rx) = mpsc::unbounded_channel::<PeekResponseUnary>();
56        let backlog_accounting = Arc::new(Mutex::new(SubscribeBacklogAccounting::default()));
57        let max_buffered_bytes =
58            SUBSCRIBE_MAX_BUFFERED_BYTES.get(self.catalog().system_config().dyncfgs());
59        let active_subscribe = ActiveSubscribe {
60            owner: ActiveSubscribeOwner::Session {
61                conn_id: conn_id.clone(),
62                session_uuid,
63            },
64            channel: tx,
65            backlog_accounting: Arc::clone(&backlog_accounting),
66            max_buffered_bytes,
67            emit_progress: plan.emit_progress,
68            as_of: df_desc
69                .as_of
70                .as_ref()
71                .and_then(|t| t.as_option())
72                .copied()
73                .expect("set to Some in an earlier stage"),
74            arity: df_desc
75                .sink_exports
76                .values()
77                .into_element()
78                .from_desc
79                .arity(),
80            cluster_id,
81            depends_on: dependency_ids,
82            start_time: self.now(),
83            output: plan.output,
84            internal: false,
85        };
86        active_subscribe.initialize();
87
88        // A subscribe whose `up_to` is not beyond its `as_of` covers an empty range and
89        // produces no data beyond the progress row `initialize` already emitted. It must
90        // not install a sink: the sink only finishes once its input frontier reaches
91        // `up_to`, which a live input never does at `MAX`, so `AS OF MAX UP TO MAX` would
92        // never complete. Retiring `active_subscribe` drops its channel sender, which
93        // completes the subscribe for the client.
94        let is_empty = df_desc
95            .as_of
96            .as_ref()
97            .is_some_and(|as_of| PartialOrder::less_equal(&df_desc.until, as_of));
98        let write_notify: BuiltinTableAppendNotify = if is_empty {
99            active_subscribe.retire(ActiveComputeSinkRetireReason::Finished);
100            // No `mz_subscriptions` bookkeeping was deferred.
101            Box::pin(std::future::ready(()))
102        } else {
103            // Register bookkeeping for the new SUBSCRIBE and ship its dataflow. The
104            // `mz_subscriptions` write is deferred to a group commit (see
105            // `add_active_compute_sink`) rather than committed inline, so it does not block
106            // the coordinator loop on a timestamp-oracle round trip. We hand the notify back
107            // so the caller can wait before returning the `SUBSCRIBE` response to the
108            // subscribing session.
109            let write_notify = self
110                .add_active_compute_sink(sink_id, ActiveComputeSink::Subscribe(active_subscribe));
111
112            // Ship the dataflow, handling errors gracefully. With the frontend subscribe
113            // sequencing, a dependency can be dropped between sequencing (on the session
114            // task) and here. The read holds acquired during sequencing don't prevent that:
115            // they hold back compaction, not drops.
116            if let Err(e) = self
117                .try_ship_dataflow(df_desc, cluster_id, replica_id)
118                .await
119            {
120                // Clean up the active compute sink that was added above, since the dataflow
121                // was never created. If we don't do this, the sink_id remains in
122                // `drop_sinks` but no collection exists in the compute controller, causing
123                // a panic when the connection terminates. This also retracts the deferred
124                // `mz_subscriptions` write, so `write_notify` can be dropped.
125                self.remove_active_compute_sink(sink_id).await;
126                return Err(
127                    AdapterError::concurrent_dependency_drop_from_dataflow_creation_error(e),
128                );
129            }
130            write_notify
131        };
132
133        // Explicitly drop read holds, just to make it obvious what's happening.
134        drop(read_holds);
135
136        // Wrap the receiver so draining a message releases its footprint from the
137        // shared accounting. FIFO delivery keeps the queue aligned with the
138        // channel, so popping the oldest footprint matches the message just
139        // drained. This keeps the accounting equal to the currently buffered
140        // depth, which the coordinator watches to bound this subscribe.
141        let rx = UnboundedReceiverStream::new(rx).map(move |response| {
142            backlog_accounting
143                .lock()
144                .expect("subscribe backlog accounting poisoned")
145                .pop();
146            response
147        });
148        let resp = ExecuteResponse::Subscribing {
149            rx: Box::new(rx),
150            ctx_extra: ExecuteContextGuard::new(statement_logging_id, self.internal_cmd_tx.clone()),
151            instance_id: cluster_id,
152        };
153        let resp = match plan.copy_to {
154            None => resp,
155            Some(format) => ExecuteResponse::CopyTo {
156                format,
157                resp: Box::new(resp),
158            },
159        };
160        Ok((resp, write_notify))
161    }
162}
mz_adapter/coord/sequencer/inner/subscribe.rs

mz_adapter/coord/sequencer/inner/
subscribe.rs
