wip

src/coordinator/distributed.rs

@@ -196,7 +196,7 @@ impl ExecutionPlan for DistributedExec {
         let tx = builder.tx();
 
         builder.spawn(async move {
-            let _guard = query_coordinator.end_query_guard();
+            let guard = query_coordinator.end_query_guard();
 
             let result = prepare_static_plan(&query_coordinator, &base_plan)?;
 
@@ -214,6 +214,8 @@ impl ExecutionPlan for DistributedExec {
                     break; // channel closed
                 }
             }
+            drop(stream);
+            drop(guard);
             drop(tx);
             query_coordinator.drain_pending_tasks().await?;
             Ok(())

src/execution_plans/mod.rs

@@ -7,6 +7,9 @@ mod network_broadcast;
 mod network_coalesce;
 mod network_shuffle;
 
+#[cfg(test)]
+mod orphaned_task_cleanup_tests;
+
 #[cfg(any(test, feature = "integration"))]
 pub mod benchmarks;