feat(tracing): make minitrace a cli option (#586)

* feat(tracing): make minitrace a cli option Signed-off-by: PhilipsPot <lichengamoy@gmail.com> * delete enable_tracing feature in cargo.toml Signed-off-by: PhilipsPot <lichengamoy@gmail.com> * add enable_tracing to tests and bench mods Signed-off-by: PhilipsPot <lichengamoy@gmail.com>
risinglightdb · Mar 27, 2022 · 7c5e0d5 · 7c5e0d5
1 parent 3a3910e
commit 7c5e0d5
Show file tree

Hide file tree

Showing 6 changed files with 76 additions and 29 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -8,7 +8,7 @@ version = "0.1.1"
 default = ["jemalloc"]
 simd = []
 jemalloc = ["tikv-jemallocator"]
-enable_tracing = []
+
 
 [dependencies]
 anyhow = "1"

diff --git a/benches/e2e.rs b/benches/e2e.rs
@@ -9,7 +9,8 @@ fn create_table(c: &mut Criterion) {
         b.to_async(&runtime).iter_batched(
             Database::new_in_memory,
             |db| async move {
-                db.run("create table t(v1 int, v2 int, v3 int)")
+                let enable_tracing = false;
+                db.run("create table t(v1 int, v2 int, v3 int)", enable_tracing)
                     .await
                     .unwrap()
             },
@@ -29,16 +30,17 @@ fn insert(c: &mut Criterion) {
                 .chain(std::iter::repeat("(1,10,100),").take(size - 1))
                 .chain(std::iter::once("(1,10,100)"))
                 .collect::<String>();
+            let enable_tracing = false;
             b.to_async(&runtime).iter_batched(
                 || async {
                     let db = Database::new_in_memory();
-                    db.run("create table t(v1 int, v2 int, v3 int)")
+                    db.run("create table t(v1 int, v2 int, v3 int)", enable_tracing)
                         .await
                         .unwrap();
                     db
                 },
                 |db| async {
-                    db.await.run(&sql).await.unwrap();
+                    db.await.run(&sql, enable_tracing).await.unwrap();
                 },
                 BatchSize::LargeInput,
             );
@@ -58,15 +60,21 @@ fn select_add(c: &mut Criterion) {
                 .chain(std::iter::repeat("(1,10),").take(size - 1))
                 .chain(std::iter::once("(1,10)"))
                 .collect::<String>();
+            let enable_tracing = false;
             b.to_async(&runtime).iter_batched(
                 || async {
                     let db = Database::new_in_memory();
-                    db.run("create table t(v1 int, v2 int)").await.unwrap();
-                    db.run(&insert_sql).await.unwrap();
+                    db.run("create table t(v1 int, v2 int)", enable_tracing)
+                        .await
+                        .unwrap();
+                    db.run(&insert_sql, enable_tracing).await.unwrap();
                     db
                 },
                 |db| async {
-                    db.await.run("select v1 + v2 from t").await.unwrap();
+                    db.await
+                        .run("select v1 + v2 from t", enable_tracing)
+                        .await
+                        .unwrap();
                 },
                 BatchSize::LargeInput,
             );

diff --git a/src/db.rs b/src/db.rs
@@ -158,14 +158,16 @@ impl Database {
 
     /// Run SQL queries and return the outputs.
 
-    pub async fn run(&self, sql: &str) -> Result<Vec<Chunk>, Error> {
-        self.run_with_context(Default::default(), sql).await
+    pub async fn run(&self, sql: &str, enable_tracing: bool) -> Result<Vec<Chunk>, Error> {
+        self.run_with_context(Default::default(), sql, enable_tracing)
+            .await
     }
 
     pub async fn run_with_context(
         &self,
         context: Arc<Context>,
         sql: &str,
+        enable_tracing: bool,
     ) -> Result<Vec<Chunk>, Error> {
         if let Some(cmdline) = sql.trim().strip_prefix('\\') {
             return self.run_internal(cmdline).await;
@@ -198,7 +200,7 @@ impl Database {
             let executor = executor_builder.build(optimized_plan);
 
             let (root, _collector) = Span::root("root");
-            let output: Vec<DataChunk> = if cfg!(feature = "enable_tracing") {
+            let output: Vec<DataChunk> = if enable_tracing {
                 executor.try_collect().in_span(root).await.map_err(|e| {
                     debug!("error: {}", e);
                     e
@@ -217,10 +219,11 @@ impl Database {
                 chunk.set_header(column_names);
             }
             outputs.push(chunk);
-            #[cfg(feature = "enable_tracing")]
-            let records: Vec<SpanRecord> = _collector.collect().await;
-            #[cfg(feature = "enable_tracing")]
-            println!("{records:#?}");
+
+            if enable_tracing {
+                let records: Vec<SpanRecord> = _collector.collect().await;
+                println!("{records:#?}");
+            }
         }
         Ok(outputs)
     }

diff --git a/src/main.rs b/src/main.rs
@@ -40,6 +40,10 @@ struct Args {
     /// - `human`: human readable format
     #[clap(long)]
     output_format: Option<String>,
+
+    /// Whether to use minitrace
+    #[clap(long)]
+    enable_tracing: bool,
 }
 
 // human-readable message
@@ -89,12 +93,17 @@ fn print_execution_time(start_time: Instant) {
     }
 }
 
-async fn run_query_in_background(db: Arc<Database>, sql: String, output_format: Option<String>) {
+async fn run_query_in_background(
+    db: Arc<Database>,
+    sql: String,
+    output_format: Option<String>,
+    enable_tracing: bool,
+) {
     let context: Arc<Context> = Default::default();
     let start_time = Instant::now();
     let handle = tokio::spawn({
         let context = context.clone();
-        async move { db.run_with_context(context, &sql).await }
+        async move { db.run_with_context(context, &sql, enable_tracing).await }
     });
 
     select! {
@@ -123,7 +132,11 @@ async fn run_query_in_background(db: Arc<Database>, sql: String, output_format:
 }
 
 /// Run RisingLight interactive mode
-async fn interactive(db: Database, output_format: Option<String>) -> Result<()> {
+async fn interactive(
+    db: Database,
+    output_format: Option<String>,
+    enable_tracing: bool,
+) -> Result<()> {
     let mut rl = Editor::<()>::new();
     let history_path = dirs::cache_dir().map(|p| {
         let cache_dir = p.join("risinglight");
@@ -149,7 +162,13 @@ async fn interactive(db: Database, output_format: Option<String>) -> Result<()>
             Ok(line) => {
                 if !line.trim().is_empty() {
                     rl.add_history_entry(line.as_str());
-                    run_query_in_background(db.clone(), line, output_format.clone()).await;
+                    run_query_in_background(
+                        db.clone(),
+                        line,
+                        output_format.clone(),
+                        enable_tracing,
+                    )
+                    .await;
                 }
             }
             Err(ReadlineError::Interrupted) => {
@@ -176,11 +195,16 @@ async fn interactive(db: Database, output_format: Option<String>) -> Result<()>
 }
 
 /// Run a SQL file in RisingLight
-async fn run_sql(db: Database, path: &str, output_format: Option<String>) -> Result<()> {
+async fn run_sql(
+    db: Database,
+    path: &str,
+    output_format: Option<String>,
+    enable_tracing: bool,
+) -> Result<()> {
     let lines = std::fs::read_to_string(path)?;
 
     info!("{}", lines);
-    let chunks = db.run(&lines).await?;
+    let chunks = db.run(&lines, enable_tracing).await?;
     for chunk in chunks {
         print_chunk(&chunk, &output_format);
     }
@@ -192,14 +216,15 @@ async fn run_sql(db: Database, path: &str, output_format: Option<String>) -> Res
 struct DatabaseWrapper {
     db: Database,
     output_format: Option<String>,
+    enable_tracing: bool,
 }
 
 #[async_trait]
 impl sqllogictest::AsyncDB for DatabaseWrapper {
     type Error = risinglight::Error;
     async fn run(&mut self, sql: &str) -> Result<String, Self::Error> {
         info!("{}", sql);
-        let chunks = self.db.run(sql).await?;
+        let chunks = self.db.run(sql, self.enable_tracing).await?;
         for chunk in &chunks {
             print_chunk(chunk, &self.output_format);
         }
@@ -211,8 +236,17 @@ impl sqllogictest::AsyncDB for DatabaseWrapper {
 }
 
 /// Run a sqllogictest file in RisingLight
-async fn run_sqllogictest(db: Database, path: &str, output_format: Option<String>) -> Result<()> {
-    let mut tester = sqllogictest::Runner::new(DatabaseWrapper { db, output_format });
+async fn run_sqllogictest(
+    db: Database,
+    path: &str,
+    output_format: Option<String>,
+    enable_tracing: bool,
+) -> Result<()> {
+    let mut tester = sqllogictest::Runner::new(DatabaseWrapper {
+        db,
+        output_format,
+        enable_tracing,
+    });
     let path = path.to_string();
     tester
         .run_file_async(path)
@@ -245,15 +279,15 @@ async fn main() -> Result<()> {
 
     if let Some(file) = args.file {
         if file.ends_with(".sql") {
-            run_sql(db, &file, args.output_format).await?;
+            run_sql(db, &file, args.output_format, args.enable_tracing).await?;
         } else if file.ends_with(".slt") {
-            run_sqllogictest(db, &file, args.output_format).await?;
+            run_sqllogictest(db, &file, args.output_format, args.enable_tracing).await?;
         } else {
             warn!("No suffix detected, assume sql file");
-            run_sql(db, &file, args.output_format).await?;
+            run_sql(db, &file, args.output_format, args.enable_tracing).await?;
         }
     } else {
-        interactive(db, args.output_format).await?;
+        interactive(db, args.output_format, args.enable_tracing).await?;
     }
 
     Ok(())

diff --git a/src/optimizer/logical_plan_rewriter/arith_expr_simplification.rs b/src/optimizer/logical_plan_rewriter/arith_expr_simplification.rs
@@ -122,7 +122,8 @@ mod tests {
         let create_stmt = "create table t(a int)";
         let sql0 = "select a + 0 from t";
         // a + 0 should be converted to a
-        let _ = db.run(create_stmt).await;
+        let enable_tracing = false;
+        let _ = db.run(create_stmt, enable_tracing).await;
 
         let plans = db.generate_execution_plan(sql0).unwrap();
         assert_eq!(plans.len(), 1);

diff --git a/tests/sqllogictest/src/lib.rs b/tests/sqllogictest/src/lib.rs
@@ -69,7 +69,8 @@ struct DatabaseWrapper {
 impl sqllogictest::AsyncDB for DatabaseWrapper {
     type Error = Error;
     async fn run(&mut self, sql: &str) -> Result<String, Self::Error> {
-        let chunks = self.db.run(sql).await?;
+        let enable_tracing = false;
+        let chunks = self.db.run(sql, enable_tracing).await?;
         let output = chunks
             .iter()
             .map(datachunk_to_sqllogictest_string)