Limit coroutines using a pool instead of chunks

ci/environment-friends.yml

@@ -21,6 +21,7 @@ dependencies:
   - flake8
   - black
   - google-cloud-core
+  - google-cloud-storage
   - google-api-core
   - google-api-python-client
   - httpretty

fsspec/asyn.py

@@ -239,20 +239,38 @@ async def _run_coros_in_chunks(
         batch_size = len(coros)
 
     assert batch_size > 0
+
+    async def _run_coro(coro):
+        try:
+            return await asyncio.wait_for(coro, timeout=timeout)
+        except Exception as e:
+            if not return_exceptions:
+                raise
+            return e
+        finally:
+            callback.relative_update(1)
+
+    coros = iter(coros)
+    stop = False
+    pending = set()
     results = []
-    for start in range(0, len(coros), batch_size):
-        chunk = [
-            asyncio.Task(asyncio.wait_for(c, timeout=timeout))
-            for c in coros[start : start + batch_size]
-        ]
-        if callback is not DEFAULT_CALLBACK:
-            [
-                t.add_done_callback(lambda *_, **__: callback.relative_update(1))
-                for t in chunk
-            ]
-        results.extend(
-            await asyncio.gather(*chunk, return_exceptions=return_exceptions),
-        )
+
+    while pending or not stop:
+        while len(pending) < batch_size and not stop:
+            try:
+                coro = next(coros)
+            except StopIteration:
+                stop = True
+            else:
+                pending.add(asyncio.ensure_future(_run_coro(coro)))
+
+        if not pending:
+            break
+
+        done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
+        while done:
+            results.append(await done.pop())
+
     return results