fix: Hacky fix for running the experimental scheduler when spawning t…

…ask (#2284)

Until I figure out the best solution (which hopefully is close).

Pushing just to benchmark and test stuff successfully.

To run the tokio console debugger:

1. Download the CLI:
   ```
   cargo --locked tokio-console
   ```
2. Build/Run local with new just commands (along with `--debug-tokio`
flag):
   ```
   just run-debug-tokio local --debug-tokio
   ```

**NOTE:** This builds in a new `target` directory:
`unstable-tokio-target` so as to not mess up with existing target
directory (which can cause issues with LSP etc.).

---------

Signed-off-by: Vaibhav <vrongmeal@gmail.com>

.gitignore

@@ -1,8 +1,8 @@
 # Adding something new? It's probably good to add to .dockerignore too.
 
 .pre-commit-config.yaml
-# nix build results
-/target
+# Cargo target directories.
+/*target
 /result
 glaredb_image
 pgsrv_image
@@ -31,4 +31,4 @@ __pycache__
 bench_data/
 
 # Logs
-*.log
+*.log

crates/glaredb/Cargo.toml

@@ -36,6 +36,7 @@ reedline = "0.27.1"
 nu-ansi-term = "0.49.0"
 url.workspace = true
 atty = "0.2.14"
+console-subscriber = "0.2.0"
 
 [dev-dependencies]
 predicates = "3.0.4"

crates/glaredb/src/args/local.rs

@@ -21,6 +21,11 @@ pub struct LocalArgs {
     /// File for logs to be written to
     #[arg(long, value_parser)]
     pub log_file: Option<PathBuf>,
+
+    /// Start the tokio console subscriber to debug runtime.
+    #[cfg(not(release))]
+    #[arg(long, value_parser, hide = true)]
+    pub debug_tokio: bool,
 }
 
 #[derive(Debug, Clone, Args)]

crates/glaredb/src/commands.rs

@@ -52,8 +52,21 @@ trait RunCommand {
     fn run(self) -> Result<()>;
 }
 
+impl LocalArgs {
+    #[cfg(not(release))]
+    fn start_tokio_debugger(&self) {
+        if self.debug_tokio {
+            eprintln!("> Starting tokio debugger");
+            console_subscriber::init();
+        }
+    }
+}
+
 impl RunCommand for LocalArgs {
     fn run(self) -> Result<()> {
+        #[cfg(not(release))]
+        self.start_tokio_debugger();
+
         let runtime = build_runtime("local")?;
         runtime.block_on(async move {
             let query = match (self.file, self.query) {

crates/sqlexec/src/distexec/scheduler.rs

@@ -25,8 +25,28 @@ impl Scheduler {
     /// Returns an executor builder along-with to create different kinds of
     /// task executors.
     pub fn new() -> (Self, ExecutorBuilder) {
-        let (task_sender, task_receiver) = async_channel::unbounded();
-        (Self { task_sender }, ExecutorBuilder { task_receiver })
+        let (actual_task_sender, actual_task_receiver) = async_channel::unbounded();
+
+        // Due to our implementation, the runtime is blocked somehow and unable
+        // to spawn new tasks.
+        //
+        // HACK: Introducing a new task here doesn't block the runtime from
+        // spawning new tasks.
+        let (dummy_task_sender, dummy_task_receiver) = async_channel::unbounded();
+        tokio::spawn(async move {
+            while let Ok(task) = dummy_task_receiver.recv().await {
+                actual_task_sender.send(task).await.unwrap();
+            }
+        });
+
+        (
+            Self {
+                task_sender: dummy_task_sender,
+            },
+            ExecutorBuilder {
+                task_receiver: actual_task_receiver,
+            },
+        )
     }
 
     /// Schedule a plan for execution.

justfile

@@ -29,6 +29,14 @@ run *args: protoc
 build *args: protoc
   cargo build --bin glaredb {{args}}
 
+# Build glaredb with unstable_tokio flag.
+#
+# Need to run the local command with `--debug-tokio` argument.
+build-debug-tokio *args:
+  RUSTFLAGS="--cfg tokio_unstable" CARGO_TARGET_DIR=tokio-unstable-target just build {{args}}
+
+run-debug-tokio *args:
+  RUSTFLAGS="--cfg tokio_unstable" CARGO_TARGET_DIR=tokio-unstable-target just run {{args}}
 
 # A zip archive will be placed in `target/dist` containing the release binary.