Merge pull request #33713 from ggevay/peek-frontend

New peek sequencing in Adapter Frontend

Author: ggevay

misc/python/materialize/checks/all_checks/statement_logging.py

@@ -17,9 +17,13 @@ class StatementLogging(Check):
     def initialize(self) -> Testdrive:
         return Testdrive(
             dedent(
+                # TODO(peek-seq): enable_frontend_peek_sequencing when it supports statement logging.
                 """
                 $ postgres-execute connection=postgres://mz_system@${testdrive.materialize-internal-sql-addr}
                 ALTER SYSTEM SET statement_logging_max_sample_rate TO 1.0
+
+                $[version>=2600100] postgres-execute connection=postgres://mz_system@${testdrive.materialize-internal-sql-addr}
+                ALTER SYSTEM SET enable_frontend_peek_sequencing = false;
                 """
             )
         )

misc/python/materialize/mzcompose/__init__.py

@@ -138,6 +138,10 @@ def get_variable_system_parameters(
     version: MzVersion,
     force_source_table_syntax: bool,
 ) -> list[VariableSystemParameter]:
+    """Note: Only the default is tested unless we explicitly select "System Parameters: Random" in trigger-ci.
+    These defaults are applied _after_ applying the settings from `get_minimal_system_parameters`.
+    """
+
     return [
         # -----
         # To reduce CRDB load as we are struggling with it in CI (values based on load test environment):
@@ -190,6 +194,11 @@ def get_variable_system_parameters(
             "true",
             ["true", "false"],
         ),
+        VariableSystemParameter(
+            "enable_frontend_peek_sequencing",
+            "false",
+            ["true", "false"],
+        ),
         VariableSystemParameter(
             "kafka_default_metadata_fetch_interval",
             "1s",

misc/python/materialize/parallel_workload/action.py

@@ -172,6 +172,7 @@ def errors_to_ignore(self, exe: Executor) -> list[str]:
                     "real-time source dropped before ingesting the upstream system's visible frontier",  # Expected, see https://buildkite.com/materialize/nightly/builds/9399#0191be17-1f4c-4321-9b51-edc4b08b71c5
                     "object state changed while transaction was in progress",  # Old error msg, can remove this ignore later
                     "another session modified the catalog while this DDL transaction was open",
+                    "was dropped while executing a statement",
                 ]
             )
         if exe.db.scenario == Scenario.Cancel:
@@ -1319,6 +1320,10 @@ def __init__(
             "314572800",  # 300 MiB, the production value
         ]
         self.flags_with_values["cluster"] = ["quickstart", "dont_exist"]
+        self.flags_with_values["enable_frontend_peek_sequencing"] = [
+            "true",
+            "false",
+        ]
 
         # If you are adding a new config flag in Materialize, consider using it
         # here instead of just marking it as uninteresting to silence the

src/adapter/src/client.rs

@@ -39,12 +39,14 @@ use mz_sql::catalog::{EnvironmentId, SessionCatalog};
 use mz_sql::session::hint::ApplicationNameHint;
 use mz_sql::session::metadata::SessionMetadata;
 use mz_sql::session::user::SUPPORT_USER;
-use mz_sql::session::vars::{CLUSTER, OwnedVarInput, SystemVars, Var};
+use mz_sql::session::vars::{
+    CLUSTER, ENABLE_FRONTEND_PEEK_SEQUENCING, OwnedVarInput, SystemVars, Var,
+};
 use mz_sql_parser::parser::{ParserStatementError, StatementParseResult};
 use prometheus::Histogram;
 use serde_json::json;
 use tokio::sync::{mpsc, oneshot};
-use tracing::error;
+use tracing::{debug, error};
 use uuid::Uuid;
 
 use crate::catalog::Catalog;
@@ -63,7 +65,7 @@ use crate::session::{
 use crate::statement_logging::{StatementEndedExecutionReason, StatementExecutionStrategy};
 use crate::telemetry::{self, EventDetails, SegmentClientExt, StatementFailureType};
 use crate::webhook::AppendWebhookResponse;
-use crate::{AdapterNotice, AppendWebhookError, PeekResponseUnary, StartupResponse};
+use crate::{AdapterNotice, AppendWebhookError, PeekClient, PeekResponseUnary, StartupResponse};
 
 /// A handle to a running coordinator.
 ///
@@ -252,21 +254,36 @@ impl Client {
 
         // Create the client as soon as startup succeeds (before any await points) so its `Drop` can
         // handle termination.
+        // Build the PeekClient with controller handles returned from startup.
+        let StartupResponse {
+            role_id,
+            write_notify,
+            session_defaults,
+            catalog,
+            storage_collections,
+            transient_id_gen,
+            optimizer_metrics,
+            persist_client,
+        } = response;
+
+        let peek_client = PeekClient::new(
+            self.clone(),
+            storage_collections,
+            transient_id_gen,
+            optimizer_metrics,
+            persist_client,
+        );
+
         let mut client = SessionClient {
             inner: Some(self.clone()),
             session: Some(session),
             timeouts: Timeout::new(),
             environment_id: self.environment_id.clone(),
             segment_client: self.segment_client.clone(),
+            peek_client,
+            enable_frontend_peek_sequencing: false, // initialized below, once we have a ConnCatalog
         };
 
-        let StartupResponse {
-            role_id,
-            write_notify,
-            session_defaults,
-            catalog,
-        } = response;
-
         let session = client.session();
         session.initialize_role_metadata(role_id);
         let vars_mut = session.vars_mut();
@@ -396,6 +413,10 @@ Issue a SQL query to get started. Need help?
             }
         }
 
+        client.enable_frontend_peek_sequencing = ENABLE_FRONTEND_PEEK_SEQUENCING
+            .require(catalog.system_vars())
+            .is_ok();
+
         Ok(client)
     }
 
@@ -412,7 +433,7 @@ Issue a SQL query to get started. Need help?
     pub async fn support_execute_one(
         &self,
         sql: &str,
-    ) -> Result<Pin<Box<dyn Stream<Item = PeekResponseUnary> + Send + Sync>>, anyhow::Error> {
+    ) -> Result<Pin<Box<dyn Stream<Item = PeekResponseUnary> + Send>>, anyhow::Error> {
         // Connect to the coordinator.
         let conn_id = self.new_conn_id()?;
         let session = self.new_session(SessionConfig {
@@ -503,7 +524,7 @@ Issue a SQL query to get started. Need help?
     }
 
     #[instrument(level = "debug")]
-    fn send(&self, cmd: Command) {
+    pub(crate) fn send(&self, cmd: Command) {
         self.inner_cmd_tx
             .send((OpenTelemetryContext::obtain(), cmd))
             .expect("coordinator unexpectedly gone");
@@ -524,6 +545,13 @@ pub struct SessionClient {
     timeouts: Timeout,
     segment_client: Option<mz_segment::Client>,
     environment_id: EnvironmentId,
+    /// Client for frontend peek sequencing; populated at connection startup.
+    peek_client: PeekClient,
+    /// Whether frontend peek sequencing is enabled; initialized at connection startup.
+    // TODO(peek-seq): Currently, this is initialized only at session startup. We'll be able to
+    // check the actual feature flag value at every peek (without a Coordinator call) once we'll
+    // always have a catalog snapshot at hand.
+    pub enable_frontend_peek_sequencing: bool,
 }
 
 impl SessionClient {
@@ -672,6 +700,17 @@ impl SessionClient {
         outer_ctx_extra: Option<ExecuteContextExtra>,
     ) -> Result<(ExecuteResponse, Instant), AdapterError> {
         let execute_started = Instant::now();
+
+        // Attempt peek sequencing in the session task.
+        // If unsupported, fall back to the Coordinator path.
+        // TODO(peek-seq): wire up cancel_future
+        if let Some(resp) = self.try_frontend_peek(&portal_name).await? {
+            debug!("frontend peek succeeded");
+            return Ok((resp, execute_started));
+        } else {
+            debug!("frontend peek did not happen");
+        }
+
         let response = self
             .send_with_cancel(
                 |tx, session| Command::Execute {
@@ -973,7 +1012,9 @@ impl SessionClient {
                 | Command::Terminate { .. }
                 | Command::RetireExecute { .. }
                 | Command::CheckConsistency { .. }
-                | Command::Dump { .. } => {}
+                | Command::Dump { .. }
+                | Command::GetComputeInstanceClient { .. }
+                | Command::GetOracle { .. } => {}
             };
             cmd
         });
@@ -1045,6 +1086,34 @@ impl SessionClient {
     pub async fn recv_timeout(&mut self) -> Option<TimeoutType> {
         self.timeouts.recv().await
     }
+
+    /// Returns a reference to the PeekClient used for frontend peek sequencing.
+    pub fn peek_client(&self) -> &PeekClient {
+        &self.peek_client
+    }
+
+    /// Returns a reference to the PeekClient used for frontend peek sequencing.
+    pub fn peek_client_mut(&mut self) -> &mut PeekClient {
+        &mut self.peek_client
+    }
+
+    /// Attempt to sequence a peek from the session task.
+    ///
+    /// Returns Some(response) if we handled the peek, or None to fall back to the Coordinator's
+    /// peek sequencing.
+    pub(crate) async fn try_frontend_peek(
+        &mut self,
+        portal_name: &str,
+    ) -> Result<Option<ExecuteResponse>, AdapterError> {
+        if self.enable_frontend_peek_sequencing {
+            let session = self.session.as_mut().expect("SessionClient invariant");
+            self.peek_client
+                .try_frontend_peek_inner(portal_name, session)
+                .await
+        } else {
+            Ok(None)
+        }
+    }
 }
 
 impl Drop for SessionClient {

src/adapter/src/command.rs

@@ -21,15 +21,20 @@ use mz_compute_types::ComputeInstanceId;
 use mz_ore::collections::CollectionExt;
 use mz_ore::soft_assert_no_log;
 use mz_ore::tracing::OpenTelemetryContext;
+use mz_persist_client::PersistClient;
 use mz_pgcopy::CopyFormatParams;
+use mz_repr::global_id::TransientIdGen;
 use mz_repr::role_id::RoleId;
 use mz_repr::{CatalogItemId, ColumnIndex, RowIterator};
 use mz_sql::ast::{FetchDirection, Raw, Statement};
 use mz_sql::catalog::ObjectType;
+use mz_sql::optimizer_metrics::OptimizerMetrics;
 use mz_sql::plan::{ExecuteTimeout, Plan, PlanKind};
 use mz_sql::session::user::User;
 use mz_sql::session::vars::{OwnedVarInput, SystemVars};
 use mz_sql_parser::ast::{AlterObjectRenameStatement, AlterOwnerStatement, DropObjectsStatement};
+use mz_storage_types::sources::Timeline;
+use mz_timestamp_oracle::TimestampOracle;
 use tokio::sync::{mpsc, oneshot};
 use uuid::Uuid;
 
@@ -153,6 +158,23 @@ pub enum Command {
     Dump {
         tx: oneshot::Sender<Result<serde_json::Value, anyhow::Error>>,
     },
+
+    GetComputeInstanceClient {
+        instance_id: ComputeInstanceId,
+        tx: oneshot::Sender<
+            Result<
+                mz_compute_client::controller::instance::Client<mz_repr::Timestamp>,
+                mz_compute_client::controller::error::InstanceMissing,
+            >,
+        >,
+    },
+
+    GetOracle {
+        timeline: Timeline,
+        tx: oneshot::Sender<
+            Result<Arc<dyn TimestampOracle<mz_repr::Timestamp> + Send + Sync>, AdapterError>,
+        >,
+    },
 }
 
 impl Command {
@@ -172,7 +194,9 @@ impl Command {
             | Command::SetSystemVars { .. }
             | Command::RetireExecute { .. }
             | Command::CheckConsistency { .. }
-            | Command::Dump { .. } => None,
+            | Command::Dump { .. }
+            | Command::GetComputeInstanceClient { .. }
+            | Command::GetOracle { .. } => None,
         }
     }
 
@@ -192,7 +216,9 @@ impl Command {
             | Command::SetSystemVars { .. }
             | Command::RetireExecute { .. }
             | Command::CheckConsistency { .. }
-            | Command::Dump { .. } => None,
+            | Command::Dump { .. }
+            | Command::GetComputeInstanceClient { .. }
+            | Command::GetOracle { .. } => None,
         }
     }
 }
@@ -216,6 +242,15 @@ pub struct StartupResponse {
     /// Map of (name, VarInput::Flat) tuples of session default variables that should be set.
     pub session_defaults: BTreeMap<String, OwnedVarInput>,
     pub catalog: Arc<Catalog>,
+    pub storage_collections: Arc<
+        dyn mz_storage_client::storage_collections::StorageCollections<
+                Timestamp = mz_repr::Timestamp,
+            > + Send
+            + Sync,
+    >,
+    pub transient_id_gen: Arc<TransientIdGen>,
+    pub optimizer_metrics: OptimizerMetrics,
+    pub persist_client: PersistClient,
 }
 
 /// The response to [`Client::authenticate`](crate::Client::authenticate).

src/adapter/src/coord.rs

@@ -208,26 +208,26 @@ use crate::util::{ClientTransmitter, ResultExt};
 use crate::webhook::{WebhookAppenderInvalidator, WebhookConcurrencyLimiter};
 use crate::{AdapterNotice, ReadHolds, flags};
 
+pub(crate) mod appends;
+pub(crate) mod catalog_serving;
+pub(crate) mod cluster_scheduling;
+pub(crate) mod consistency;
 pub(crate) mod id_bundle;
 pub(crate) mod in_memory_oracle;
 pub(crate) mod peek;
+pub(crate) mod read_policy;
+pub(crate) mod sequencer;
 pub(crate) mod statement_logging;
 pub(crate) mod timeline;
 pub(crate) mod timestamp_selection;
 
-pub mod appends;
-mod catalog_serving;
 mod caught_up;
-pub mod cluster_scheduling;
 mod command_handler;
-pub mod consistency;
 mod ddl;
 mod indexes;
 mod introspection;
 mod message_handler;
 mod privatelink_status;
-pub mod read_policy;
-mod sequencer;
 mod sql;
 mod validity;
 
@@ -361,6 +361,8 @@ impl Message {
                 Command::AuthenticatePassword { .. } => "command-auth_check",
                 Command::AuthenticateGetSASLChallenge { .. } => "command-auth_get_sasl_challenge",
                 Command::AuthenticateVerifySASLProof { .. } => "command-auth_verify_sasl_proof",
+                Command::GetComputeInstanceClient { .. } => "get-compute-instance-client",
+                Command::GetOracle { .. } => "get-oracle",
             },
             Message::ControllerReady {
                 controller: ControllerReadiness::Compute,

src/adapter/src/coord/appends.rs

@@ -926,6 +926,8 @@ pub struct GroupCommitPermit {
 /// So as an optimization we do not wait for these writes to complete. But if a [`Session`] tries
 /// to query any of these builtin objects, we need to block that query on the writes completing to
 /// maintain linearizability.
+///
+/// Warning: this already clears the wait flag (i.e., it calls `clear_builtin_table_updates`).
 pub(crate) fn waiting_on_startup_appends(
     catalog: &Catalog,
     session: &mut Session,