refactor(pgwire): unify query logs and slow query logs

src/common/src/util/mod.rs

@@ -28,6 +28,7 @@ pub mod memcmp_encoding;
 pub mod panic;
 pub mod pretty_bytes;
 pub mod prost;
+pub mod query_log;
 pub mod resource_util;
 pub mod row_id;
 pub mod row_serde;

src/common/src/util/query_log.rs

@@ -0,0 +1,32 @@
+// Copyright 2023 RisingWave Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::sync::LazyLock;
+use std::time::Duration;
+
+/// The target name for the root span while pgwire is handling a query.
+pub const PGWIRE_ROOT_SPAN_TARGET: &str = "pgwire_root_span";
+/// The target name for the event when pgwire finishes handling a query.
+pub const PGWIRE_QUERY_LOG: &str = "pgwire_query_log";
+/// The target name for the event when pgwire does not finish handling a query in time.
+pub const PGWIRE_SLOW_QUERY_LOG: &str = "pgwire_slow_query_log";
+
+/// The period of logging ongoing slow queries.
+pub static SLOW_QUERY_LOG_PERIOD: LazyLock<Duration> = LazyLock::new(|| {
+    std::env::var("RW_SLOW_QUERY_LOG_PERIOD_MS")
+        .ok()
+        .and_then(|s| s.parse().ok())
+        .map(Duration::from_millis)
+        .unwrap_or(Duration::from_secs(60))
+});

src/frontend/src/session.rs

@@ -825,27 +825,9 @@ impl SessionImpl {
             );
         }
         let stmt = stmts.swap_remove(0);
-        let rsp = {
-            let mut handle_fut = Box::pin(handle(self, stmt, sql.clone(), formats));
-            if cfg!(debug_assertions) {
-                // Report the SQL in the log periodically if the query is slow.
-                const SLOW_QUERY_LOG_PERIOD: Duration = Duration::from_secs(60);
-                const SLOW_QUERY_LOG: &str = "risingwave_frontend_slow_query_log";
-                loop {
-                    match tokio::time::timeout(SLOW_QUERY_LOG_PERIOD, &mut handle_fut).await {
-                        Ok(result) => break result,
-                        Err(_) => tracing::warn!(
-                            target: SLOW_QUERY_LOG,
-                            sql = sql.as_ref(),
-                            "slow query has been running for another {SLOW_QUERY_LOG_PERIOD:?}"
-                        ),
-                    }
-                }
-            } else {
-                handle_fut.await
-            }
-        }
-        .inspect_err(|e| tracing::error!(error = %e.as_report(), %sql, "failed to handle sql"))?;
+        let rsp = handle(self, stmt, sql.clone(), formats).await.inspect_err(
+            |e| tracing::error!(error = %e.as_report(), %sql, "failed to handle sql"),
+        )?;
         Ok(rsp)
     }
 
@@ -1033,25 +1015,11 @@ impl Session for SessionImpl {
         let string = stmt.to_string();
         let sql_str = string.as_str();
         let sql: Arc<str> = Arc::from(sql_str);
-        let rsp = {
-            let mut handle_fut = Box::pin(handle(self, stmt, sql.clone(), vec![format]));
-            if cfg!(debug_assertions) {
-                // Report the SQL in the log periodically if the query is slow.
-                const SLOW_QUERY_LOG_PERIOD: Duration = Duration::from_secs(60);
-                loop {
-                    match tokio::time::timeout(SLOW_QUERY_LOG_PERIOD, &mut handle_fut).await {
-                        Ok(result) => break result,
-                        Err(_) => tracing::warn!(
-                            sql_str,
-                            "slow query has been running for another {SLOW_QUERY_LOG_PERIOD:?}"
-                        ),
-                    }
-                }
-            } else {
-                handle_fut.await
-            }
-        }
-        .inspect_err(|e| tracing::error!(error = %e.as_report(), %sql, "failed to handle sql"))?;
+        let rsp = handle(self, stmt, sql.clone(), vec![format])
+            .await
+            .inspect_err(
+                |e| tracing::error!(error = %e.as_report(), %sql, "failed to handle sql"),
+            )?;
         Ok(rsp)
     }
 
@@ -1094,24 +1062,9 @@ impl Session for SessionImpl {
         self: Arc<Self>,
         portal: Portal,
     ) -> std::result::Result<PgResponse<PgResponseStream>, BoxedError> {
-        let rsp = {
-            let mut handle_fut = Box::pin(handle_execute(self, portal));
-            if cfg!(debug_assertions) {
-                // Report the SQL in the log periodically if the query is slow.
-                const SLOW_QUERY_LOG_PERIOD: Duration = Duration::from_secs(60);
-                loop {
-                    match tokio::time::timeout(SLOW_QUERY_LOG_PERIOD, &mut handle_fut).await {
-                        Ok(result) => break result,
-                        Err(_) => tracing::warn!(
-                            "slow query has been running for another {SLOW_QUERY_LOG_PERIOD:?}"
-                        ),
-                    }
-                }
-            } else {
-                handle_fut.await
-            }
-        }
-        .inspect_err(|e| tracing::error!(error=%e.as_report(), "failed to handle execute"))?;
+        let rsp = handle_execute(self, portal)
+            .await
+            .inspect_err(|e| tracing::error!(error=%e.as_report(), "failed to handle execute"))?;
         Ok(rsp)
     }