WIP

Author: SangJunBak

src/adapter/src/coord/statement_logging.rs

@@ -180,14 +180,13 @@ impl StatementLogging {
     /// Check if we need to drop a statement
     /// due to throttling, and update internal data structures appropriately.
     ///
-    /// Returns `None` if we must throttle this statement, and `Some(n)` otherwise, where `n`
-    /// is the number of statements that were dropped due to throttling before this one.
+    /// Returns `false` if we must throttle this statement, and `true` otherwise.
     fn throttling_check(
         &mut self,
         cost: u64,
         target_data_rate: u64,
         max_data_credit: Option<u64>,
-    ) -> Option<usize> {
+    ) -> bool {
         let ts = (self.now)() / 1000;
         // We use saturating_sub here because system time isn't monotonic, causing cases
         // when last_logged_ts_seconds is greater than ts.
@@ -202,14 +201,14 @@ impl StatementLogging {
         if let Some(remaining) = self.tokens.checked_sub(cost) {
             debug!("throttling check passed. tokens remaining: {remaining}; cost: {cost}");
             self.tokens = remaining;
-            Some(std::mem::take(&mut self.throttled_count))
+            true
         } else {
             debug!(
                 "throttling check failed. tokens available: {}; cost: {cost}",
                 self.tokens
             );
             self.throttled_count += 1;
-            None
+            false
         }
     }
 }
@@ -276,15 +275,14 @@ impl Coordinator {
     /// Check whether we need to do throttling (i.e., whether `STATEMENT_LOGGING_TARGET_DATA_RATE` is set).
     /// If so, actually do the check.
     ///
-    /// Returns `None` if we must throttle this statement, and `Some(n)` otherwise, where `n`
-    /// is the number of statements that were dropped due to throttling before this one.
-    fn statement_logging_throttling_check(&mut self, cost: usize) -> Option<usize> {
+    /// Returns `false` if we must throttle this statement, and `true` otherwise.
+    fn statement_logging_throttling_check(&mut self, cost: usize) -> bool {
         let Some(target_data_rate) = self
             .catalog
             .system_config()
             .statement_logging_target_data_rate()
         else {
-            return Some(std::mem::take(&mut self.statement_logging.throttled_count));
+            return true;
         };
         let max_data_credit = self
             .catalog
@@ -297,20 +295,32 @@ impl Coordinator {
         )
     }
 
+    fn log_prepared_statement(
+        &mut self,
+        uuid: Uuid,
+        session: &mut Session,
+        logging: &Arc<QCell<PreparedStatementLoggingInfo>>,
+    ) {
+        let logging = session.qcell_rw(&*logging);
+        if let PreparedStatementLoggingInfo::StillToLog { .. } = logging {
+            *logging = PreparedStatementLoggingInfo::AlreadyLogged { uuid };
+            self.statement_logging.throttled_count = 0;
+        }
+    }
+
     /// Returns any statement logging events needed for a particular
     /// prepared statement. Possibly mutates the `PreparedStatementLoggingInfo` metadata.
     ///
     /// This function does not do a sampling check, and assumes we did so in a higher layer.
     ///
-    /// It _does_ do a throttling check, and returns `None` if we must not log due to throttling.
-    pub(crate) fn log_prepared_statement(
+    pub(crate) fn get_prepared_statement_info(
         &mut self,
         session: &mut Session,
         logging: &Arc<QCell<PreparedStatementLoggingInfo>>,
-    ) -> Option<(
+    ) -> (
         Option<(StatementPreparedRecord, PreparedStatementEvent)>,
         Uuid,
-    )> {
+    ) {
         let logging = session.qcell_rw(&*logging);
         let mut out = None;
 
@@ -357,22 +367,20 @@ impl Coordinator {
                     Datum::String(redacted_sql.as_str()),
                 ]);
 
-                let cost = mpsh_packer.byte_len() + sql_row.byte_len();
-                let throttled_count = self.statement_logging_throttling_check(cost)?;
+                let throttled_count = self.statement_logging.throttled_count;
                 mpsh_packer.push(Datum::UInt64(throttled_count.try_into().expect("must fit")));
+
                 out = Some((
                     record,
                     PreparedStatementEvent {
                         prepared_statement: mpsh_row,
                         sql_text: sql_row,
                     },
                 ));
-
-                *logging = PreparedStatementLoggingInfo::AlreadyLogged { uuid };
                 uuid
             }
         };
-        Some((out, uuid))
+        (out, uuid)
     }
     /// The rate at which statement execution should be sampled.
     /// This is the value of the session var `statement_logging_sample_rate`,
@@ -700,6 +708,8 @@ impl Coordinator {
             distribution.sample(&mut thread_rng())
         };
 
+        // Figure out the cost of everything before we log.
+
         // Track how many statements we're recording.
         let sampled_label = sample.then_some("true").unwrap_or("false");
         self.metrics
@@ -728,7 +738,8 @@ impl Coordinator {
         if !sample {
             return None;
         }
-        let (ps_record, ps_uuid) = self.log_prepared_statement(session, logging)?;
+
+        let (maybe_ps, ps_uuid) = self.get_prepared_statement_info(session, logging);
 
         let began_at = if let Some(lifecycle_timestamps) = lifecycle_timestamps {
             lifecycle_timestamps.received
@@ -737,11 +748,6 @@ impl Coordinator {
         };
         let now = self.now();
         let execution_uuid = epoch_to_uuid_v7(&now);
-        self.record_statement_lifecycle_event(
-            &StatementLoggingId(execution_uuid),
-            &StatementLifecycleEvent::ExecutionBegan,
-            began_at,
-        );
 
         let params = std::iter::zip(params.execute_types.iter(), params.datums.iter())
             .map(|(r#type, datum)| {
@@ -785,28 +791,72 @@ impl Coordinator {
                 .map(|s| s.as_str().to_string())
                 .collect(),
         };
+
         let mseh_update = Self::pack_statement_began_execution_update(&record);
+
+        let (maybe_ps_event, maybe_sh_update) = if let Some((ps_record, ps_event)) = maybe_ps {
+            if let Some(sh) = self
+                .statement_logging
+                .unlogged_sessions
+                .get(&ps_record.session_id)
+            {
+                (Some(ps_event), Some(Self::pack_session_history_update(sh)))
+            } else {
+                (None, None)
+            }
+        } else {
+            (None, None)
+        };
+
+        let statement_lifecycle_event_update = Self::pack_statement_lifecycle_event(
+            &StatementLoggingId(execution_uuid),
+            &StatementLifecycleEvent::ExecutionBegan,
+            began_at,
+        );
+
+        let cost = mseh_update.byte_len().saturating_add(
+            maybe_ps_event
+                .as_ref()
+                .map(|e| {
+                    e.prepared_statement
+                        .byte_len()
+                        .saturating_add(e.sql_text.byte_len())
+                })
+                .unwrap_or(0)
+                .saturating_add(maybe_sh_update.as_ref().map(|u| u.byte_len()).unwrap_or(0))
+                .saturating_add(statement_lifecycle_event_update.byte_len()),
+        );
+
+        if !self.statement_logging_throttling_check(cost) {
+            return None;
+        }
+
+        self.log_prepared_statement(ps_uuid, session, logging);
+
+        self.record_statement_lifecycle_event(
+            &StatementLoggingId(execution_uuid),
+            &StatementLifecycleEvent::ExecutionBegan,
+            began_at,
+        );
+
         self.statement_logging
             .pending_statement_execution_events
             .push((mseh_update, Diff::ONE));
         self.statement_logging
             .executions_begun
             .insert(execution_uuid, record);
-        if let Some((ps_record, ps_update)) = ps_record {
+
+        if let Some(sh_update) = maybe_sh_update {
+            self.statement_logging
+                .pending_session_events
+                .push(sh_update);
+        }
+        if let Some(ps_event) = maybe_ps_event {
             self.statement_logging
                 .pending_prepared_statement_events
-                .push(ps_update);
-            if let Some(sh) = self
-                .statement_logging
-                .unlogged_sessions
-                .remove(&ps_record.session_id)
-            {
-                let sh_update = Self::pack_session_history_update(&sh);
-                self.statement_logging
-                    .pending_session_events
-                    .push(sh_update);
-            }
+                .push(ps_event);
         }
+
         Some(StatementLoggingId(execution_uuid))
     }
 

src/environmentd/tests/server.rs

@@ -518,7 +518,7 @@ fn test_statement_logging_throttling() {
         1.0,
         test_util::TestHarness::default().with_system_parameter_default(
             "statement_logging_target_data_rate".to_string(),
-            "100".to_string(),
+            "1000".to_string(),
         ),
     );
     thread::sleep(Duration::from_secs(2));
@@ -536,7 +536,72 @@ fn test_statement_logging_throttling() {
                     "SELECT
     sql,
     throttled_count
-FROM mz_internal.mz_prepared_statement_history mpsh
+FROM
+mz_internal.mz_statement_execution_history mseh
+JOIN mz_internal.mz_prepared_statement_history mpsh
+ON mseh.prepared_statement_id = mpsh.id
+JOIN (SELECT DISTINCT sql, sql_hash, redacted_sql FROM mz_internal.mz_sql_text) mst
+ON mpsh.sql_hash = mst.sql_hash
+WHERE sql IN ('SELECT 1', 'SELECT 2')",
+                    &[],
+                )
+                .unwrap();
+
+            if sl_results.iter().any(|stmt| {
+                let sql: String = stmt.get(0);
+                sql == "SELECT 2"
+            }) {
+                Ok(sl_results)
+            } else {
+                Err(())
+            }
+        })
+        .expect("Never saw last statement (`SELECT 2`)");
+    let throttled_count = logs
+        .iter()
+        .map(|log| {
+            let UInt8(throttled_count) = log.get(1);
+            throttled_count
+        })
+        .sum::<u64>();
+    assert!(
+        throttled_count > 0,
+        "at least some statements should have been throttled"
+    );
+
+    assert_eq!(logs.len() + usize::cast_from(throttled_count), 101);
+}
+
+#[mz_ore::test]
+fn test_statement_logging_prepared_statement_throttling() {
+    let (server, mut client) = setup_statement_logging_core(
+        1.0,
+        1.0,
+        test_util::TestHarness::default().with_system_parameter_default(
+            "statement_logging_target_data_rate".to_string(),
+            "1000".to_string(),
+        ),
+    );
+    thread::sleep(Duration::from_secs(2));
+    let statement = client.prepare("SELECT 1").unwrap();
+
+    for _ in 0..100 {
+        client.execute(&statement, &[]).unwrap();
+    }
+    thread::sleep(Duration::from_secs(4));
+    client.execute("SELECT 2", &[]).unwrap();
+    let mut client = server.connect_internal(postgres::NoTls).unwrap();
+    let logs = Retry::default()
+        .max_duration(Duration::from_secs(60))
+        .retry(|_| {
+            let sl_results = client
+                .query(
+                    "SELECT
+    sql,
+    throttled_count
+FROM mz_internal.mz_statement_execution_history mseh
+JOIN mz_internal.mz_prepared_statement_history mpsh
+ON mseh.prepared_statement_id = mpsh.id
 JOIN (SELECT DISTINCT sql, sql_hash, redacted_sql FROM mz_internal.mz_sql_text) mst
 ON mpsh.sql_hash = mst.sql_hash
 WHERE sql IN ('SELECT 1', 'SELECT 2')",
@@ -557,7 +622,9 @@ WHERE sql IN ('SELECT 1', 'SELECT 2')",
     let throttled_count = logs
         .iter()
         .map(|log| {
+            let txt = log.get::<_, String>(0);
             let UInt8(throttled_count) = log.get(1);
+            println!("throttled_count for row {}: {}", txt, throttled_count);
             throttled_count
         })
         .sum::<u64>();