From a6bcf69ea0e19e4f8c22c526a0ba96d3fbd2ed58 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?H=C3=A9ctor=20Ram=C3=B3n=20Jim=C3=A9nez?= <hector@hecrj.dev>
Date: Tue, 8 Apr 2025 22:12:47 +0200
Subject: [PATCH] Fix `Task` unit calculation in `Task::batch`

---
 runtime/src/task.rs | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/runtime/src/task.rs b/runtime/src/task.rs
index 81973e23..756f592b 100644
--- a/runtime/src/task.rs
+++ b/runtime/src/task.rs
@@ -89,11 +89,16 @@ impl<T> Task<T> {
     where
         T: 'static,
     {
-        let select_all = stream::select_all(
-            tasks.into_iter().filter_map(|task| task.stream),
-        );
+        let mut select_all = stream::SelectAll::new();
+        let mut units = 0;
 
-        let units = select_all.len();
+        for task in tasks.into_iter() {
+            if let Some(stream) = task.stream {
+                select_all.push(stream);
+            }
+
+            units += task.units;
+        }
 
         Self {
             stream: Some(boxed_stream(select_all)),