storage: make MySQL snapshot wait_timeout a configurable system var

src/adapter/src/flags.rs

@@ -48,6 +48,7 @@ pub fn storage_config(config: &SystemVars) -> StorageParameters {
         mysql_source_timeouts: mz_mysql_util::TimeoutConfig::build(
             config.mysql_source_snapshot_max_execution_time(),
             config.mysql_source_snapshot_lock_wait_timeout(),
+            config.mysql_source_snapshot_wait_timeout(),
             config.mysql_source_tcp_keepalive(),
             config.mysql_source_connect_timeout(),
         ),

src/mysql-util/src/lib.rs

@@ -15,8 +15,8 @@ use std::time::Duration;
 use aws_rds::RdsTokenError;
 pub use tunnel::{
     Config, DEFAULT_CONNECT_TIMEOUT, DEFAULT_SNAPSHOT_LOCK_WAIT_TIMEOUT,
-    DEFAULT_SNAPSHOT_MAX_EXECUTION_TIME, DEFAULT_TCP_KEEPALIVE, MySqlConn, TimeoutConfig,
-    TunnelConfig,
+    DEFAULT_SNAPSHOT_MAX_EXECUTION_TIME, DEFAULT_SNAPSHOT_WAIT_TIMEOUT, DEFAULT_TCP_KEEPALIVE,
+    MySqlConn, TimeoutConfig, TunnelConfig,
 };
 
 mod desc;

src/mysql-util/src/tunnel.rs

@@ -53,13 +53,16 @@ pub enum TunnelConfig {
 pub const DEFAULT_TCP_KEEPALIVE: Duration = Duration::from_secs(60);
 pub const DEFAULT_SNAPSHOT_MAX_EXECUTION_TIME: Duration = Duration::ZERO;
 pub const DEFAULT_SNAPSHOT_LOCK_WAIT_TIMEOUT: Duration = Duration::from_secs(3600);
+/// The MySQL server default `wait_timeout`.
+pub const DEFAULT_SNAPSHOT_WAIT_TIMEOUT: Duration = Duration::from_secs(28800);
 pub const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(60);
 
 #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub struct TimeoutConfig {
     // Snapshot-related configs
     pub snapshot_max_execution_time: Option<Duration>,
     pub snapshot_lock_wait_timeout: Option<Duration>,
+    pub snapshot_wait_timeout: Option<Duration>,
 
     // Socket-related configs
     pub tcp_keepalive: Option<Duration>,
@@ -78,6 +81,7 @@ impl Default for TimeoutConfig {
         Self {
             snapshot_max_execution_time: Some(DEFAULT_SNAPSHOT_MAX_EXECUTION_TIME),
             snapshot_lock_wait_timeout: Some(DEFAULT_SNAPSHOT_LOCK_WAIT_TIMEOUT),
+            snapshot_wait_timeout: Some(DEFAULT_SNAPSHOT_WAIT_TIMEOUT),
             tcp_keepalive: Some(DEFAULT_TCP_KEEPALIVE),
             connect_timeout: Some(DEFAULT_CONNECT_TIMEOUT),
         }
@@ -88,6 +92,7 @@ impl TimeoutConfig {
     pub fn build(
         snapshot_max_execution_time: Duration,
         snapshot_lock_wait_timeout: Duration,
+        snapshot_wait_timeout: Duration,
         tcp_keepalive: Duration,
         connect_timeout: Duration,
     ) -> Self {
@@ -106,6 +111,18 @@ impl TimeoutConfig {
             Some(snapshot_lock_wait_timeout)
         };
 
+        // https://dev.mysql.com/doc/refman/8.0/en/server-system-variables.html#sysvar_wait_timeout
+        let snapshot_wait_timeout =
+            if snapshot_wait_timeout.as_secs() > 31536000 || snapshot_wait_timeout.as_secs() < 1 {
+                error!(
+                    "snapshot_wait_timeout is out of range: {}. Must be within [1, 31536000].",
+                    snapshot_wait_timeout.as_secs()
+                );
+                Some(DEFAULT_SNAPSHOT_WAIT_TIMEOUT)
+            } else {
+                Some(snapshot_wait_timeout)
+            };
+
         // https://dev.mysql.com/doc/refman/8.0/en/server-system-variables.html#sysvar_max_execution_time
         let snapshot_max_execution_time = if snapshot_max_execution_time.as_millis() > 4294967295 {
             error!(
@@ -144,6 +161,7 @@ impl TimeoutConfig {
         Self {
             snapshot_max_execution_time,
             snapshot_lock_wait_timeout,
+            snapshot_wait_timeout,
             tcp_keepalive,
             connect_timeout,
         }

src/sql/src/session/vars.rs

@@ -1236,6 +1236,7 @@ impl SystemVars {
             &MYSQL_SOURCE_TCP_KEEPALIVE,
             &MYSQL_SOURCE_SNAPSHOT_MAX_EXECUTION_TIME,
             &MYSQL_SOURCE_SNAPSHOT_LOCK_WAIT_TIMEOUT,
+            &MYSQL_SOURCE_SNAPSHOT_WAIT_TIMEOUT,
             &MYSQL_SOURCE_CONNECT_TIMEOUT,
             &SSH_CHECK_INTERVAL,
             &SSH_CONNECT_TIMEOUT,
@@ -1887,6 +1888,11 @@ impl SystemVars {
         *self.expect_value(&MYSQL_SOURCE_SNAPSHOT_LOCK_WAIT_TIMEOUT)
     }
 
+    /// Returns the `mysql_source_snapshot_wait_timeout` configuration parameter.
+    pub fn mysql_source_snapshot_wait_timeout(&self) -> Duration {
+        *self.expect_value(&MYSQL_SOURCE_SNAPSHOT_WAIT_TIMEOUT)
+    }
+
     /// Returns the `mysql_source_connect_timeout` configuration parameter.
     pub fn mysql_source_connect_timeout(&self) -> Duration {
         *self.expect_value(&MYSQL_SOURCE_CONNECT_TIMEOUT)
@@ -2331,6 +2337,7 @@ impl SystemVars {
             || name == MYSQL_SOURCE_TCP_KEEPALIVE.name()
             || name == MYSQL_SOURCE_SNAPSHOT_MAX_EXECUTION_TIME.name()
             || name == MYSQL_SOURCE_SNAPSHOT_LOCK_WAIT_TIMEOUT.name()
+            || name == MYSQL_SOURCE_SNAPSHOT_WAIT_TIMEOUT.name()
             || name == MYSQL_SOURCE_CONNECT_TIMEOUT.name()
             || name == ENABLE_STORAGE_SHARD_FINALIZATION.name()
             || name == SSH_CHECK_INTERVAL.name()

src/sql/src/session/vars/definitions.rs

@@ -1047,6 +1047,15 @@ pub static MYSQL_SOURCE_SNAPSHOT_LOCK_WAIT_TIMEOUT: VarDefinition = VarDefinitio
     false,
 );
 
+/// Sets the `wait_timeout` session value on connections used during the
+/// snapshotting phase of MySQL sources.
+pub static MYSQL_SOURCE_SNAPSHOT_WAIT_TIMEOUT: VarDefinition = VarDefinition::new(
+    "mysql_source_snapshot_wait_timeout",
+    value!(Duration; mz_mysql_util::DEFAULT_SNAPSHOT_WAIT_TIMEOUT),
+    "Sets the `wait_timeout` value to use on connections during the snapshotting phase of MySQL sources (Materialize)",
+    false,
+);
+
 /// Sets the timeout for establishing an authenticated connection to MySQL
 pub static MYSQL_SOURCE_CONNECT_TIMEOUT: VarDefinition = VarDefinition::new(
     "mysql_source_connect_timeout",

src/storage/src/source/mysql/snapshot.rs

@@ -424,7 +424,14 @@ async fn lock_and_prepare_snapshot(
         )
         .await?;
 
-    restore_default_wait_timeout(&mut *lock_conn).await?;
+    if let Some(timeout) = config
+        .config
+        .parameters
+        .mysql_source_timeouts
+        .snapshot_wait_timeout
+    {
+        set_wait_timeout(&mut *lock_conn, timeout).await?;
+    }
 
     let errored_outputs = verify_output_schemas(&mut *lock_conn, tables).await?;
     let errored: BTreeSet<usize> = errored_outputs.iter().map(|(idx, _)| *idx).collect();
@@ -772,7 +779,14 @@ pub(crate) fn render<'scope>(
                     Ok(()) => (),
                 };
 
-                restore_default_wait_timeout(&mut *conn).await?;
+                if let Some(timeout) = config
+                    .config
+                    .parameters
+                    .mysql_source_timeouts
+                    .snapshot_wait_timeout
+                {
+                    set_wait_timeout(&mut *conn, timeout).await?;
+                }
 
                 let mut lock_conn = if is_snapshot_leader {
                     match lock_and_prepare_snapshot(
@@ -1100,14 +1114,19 @@ fn publish_snapshot_size(
     }
 }
 
-/// Restores the server-default wait_timeout (28800) so a lowered global value
-/// cannot reap the connection while it sits idle during snapshot setup.
-async fn restore_default_wait_timeout<Q>(conn: &mut Q) -> Result<(), mysql_async::Error>
+/// Sets the session wait_timeout so a lowered global value cannot reap the
+/// connection while it sits idle during snapshot setup.
+async fn set_wait_timeout<Q>(conn: &mut Q, timeout: Duration) -> Result<(), mysql_async::Error>
 where
     Q: Queryable,
 {
+    // Interpolating a `Duration` integer; not parameterizable in MySQL `SET`.
     #[allow(clippy::disallowed_methods)]
-    conn.query_drop("SET @@session.wait_timeout = 28800").await
+    conn.query_drop(format!(
+        "SET @@session.wait_timeout = {}",
+        timeout.as_secs()
+    ))
+    .await
 }
 async fn lock_tables_and_read_gtid_set(
     lock_conn: &mut MySqlConn,

test/launchdarkly-flag-consistency/mzcompose.py

@@ -313,6 +313,7 @@
     mysql_source_connect_timeout
     mysql_source_snapshot_lock_wait_timeout
     mysql_source_snapshot_max_execution_time
+    mysql_source_snapshot_wait_timeout
     mysql_source_tcp_keepalive
     network_policy
     oidc_audience

test/mysql-cdc/mysql-cdc.td

@@ -181,6 +181,24 @@ running
 $ postgres-execute connection=mz_system
 ALTER SYSTEM SET mysql_source_snapshot_lock_wait_timeout = 3600
 
+# Validate mysql_source_snapshot_wait_timeout
+$ postgres-execute connection=mz_system
+ALTER SYSTEM SET mysql_source_snapshot_wait_timeout = '60s'
+
+> CREATE SOURCE "test_slot_source"
+  IN CLUSTER cdc_cluster
+  FROM MYSQL CONNECTION mysql_conn;
+
+> CREATE TABLE pk_table FROM SOURCE test_slot_source (REFERENCE public.pk_table);
+
+> SELECT COUNT(*) > 0 FROM pk_table;
+true
+
+> DROP SOURCE test_slot_source CASCADE;
+
+$ postgres-execute connection=mz_system
+ALTER SYSTEM RESET mysql_source_snapshot_wait_timeout
+
 
 #
 # Error checking