Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 19 additions & 3 deletions crawl4ai/async_dispatcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -534,6 +534,7 @@ async def run_urls_stream(
config: Union[CrawlerRunConfig, List[CrawlerRunConfig]],
) -> AsyncGenerator[CrawlerTaskResult, None]:
self.crawler = crawler
active_tasks = []

# Start the memory monitor task
memory_monitor = asyncio.create_task(self._memory_monitor_task())
Expand All @@ -550,7 +551,6 @@ async def run_urls_stream(
# Add to queue with initial priority 0, retry count 0, and current time
await self.task_queue.put((0, (url, task_id, 0, time.time())))

active_tasks = []
completed_count = 0
total_urls = len(urls)

Expand Down Expand Up @@ -614,8 +614,24 @@ async def run_urls_stream(
await self._update_queue_priorities()

finally:
# Clean up
# Cancel and await every task owned by this stream before returning
# control to the caller. Otherwise a closed stream can leave crawls
# using browser pages and contexts in the background.
for task in active_tasks:
if not task.done():
task.cancel()
if active_tasks:
await asyncio.gather(*active_tasks, return_exceptions=True)

# Discard URLs that were queued by this stream but never started.
while True:
try:
self.task_queue.get_nowait()
except asyncio.QueueEmpty:
break

memory_monitor.cancel()
await asyncio.gather(memory_monitor, return_exceptions=True)
if self.monitor:
self.monitor.stop()

Expand Down Expand Up @@ -769,4 +785,4 @@ async def run_urls(
return await asyncio.gather(*tasks, return_exceptions=True)
finally:
if self.monitor:
self.monitor.stop()
self.monitor.stop()
44 changes: 43 additions & 1 deletion tests/async/test_dispatchers.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,8 @@
import pytest
import asyncio
import time
from types import SimpleNamespace

import pytest
from crawl4ai import (
AsyncWebCrawler,
BrowserConfig,
Expand Down Expand Up @@ -63,6 +66,45 @@ async def test_memory_adaptive_with_rate_limit(
assert len(results) == len(test_urls)
assert all(r.success for r in results)

async def test_memory_adaptive_stream_closure_cleans_up_tasks(self, run_config):
class BlockingCrawler:
async def arun(self, url, config=None, session_id=None):
if url == "fast":
return SimpleNamespace(
success=True,
status_code=200,
error_message="",
)
await asyncio.Event().wait()

class TrackingDispatcher(MemoryAdaptiveDispatcher):
def __init__(self):
super().__init__(max_session_permit=3)
self.tasks = {}

async def crawl_url(self, url, config, task_id, retry_count=0):
self.tasks[url] = asyncio.current_task()
return await super().crawl_url(url, config, task_id, retry_count)

dispatcher = TrackingDispatcher()
stream = dispatcher.run_urls_stream(
["fast", "blocked-1", "blocked-2", "queued"],
BlockingCrawler(),
run_config,
)

first = await asyncio.wait_for(stream.__anext__(), timeout=1)
assert first.url == "fast"

await asyncio.wait_for(stream.aclose(), timeout=1)

assert set(dispatcher.tasks) == {"fast", "blocked-1", "blocked-2"}
assert all(task.done() for task in dispatcher.tasks.values())
assert dispatcher.tasks["blocked-1"].cancelled()
assert dispatcher.tasks["blocked-2"].cancelled()
assert dispatcher.concurrent_sessions == 0
assert dispatcher.task_queue.empty()

async def test_semaphore_basic(self, browser_config, run_config, test_urls):
async with AsyncWebCrawler(config=browser_config) as crawler:
dispatcher = SemaphoreDispatcher(semaphore_count=2)
Expand Down