如何用aiohttp实现异步批量文件下载？求队列/信号量等更优方案

Hey there! Let's work through your async file download problem together. Your current setup is hitting issues because you're sending too many requests at once, and your manual batch approach has a small indexing bug that's causing empty coroutine errors. Let's fix that with two robust solutions—using semaphores (the simpler option) and queues (for more control)—plus proper error handling for server rejections.

First, let's address the small bugs in your existing code:

• In your do_downloads function, you're calling download(link) but the function expects both link and name—that's a typo that would cause errors right away!
• When doing manual batches, if the last batch has fewer items than BATCH_LENGTH, you might end up with empty lists, which triggers that "Set of coroutines is empty" error.

Solution 1: Using Semaphores (Simpler Concurrency Control) #

Semaphores act like a gatekeeper—they only let a set number of coroutines run at the same time. This prevents flooding the server with too many requests, which is likely why your full batch was crashing.

Step 1: Updated Download Function with Error Handling #

First, let's beef up the download function to handle server rejections and file IO errors:

import asyncio
import aiohttp
from aiohttp import ClientError
import os

async def download(link, name, proxy=None):
    try:
        async with aiohttp.ClientSession() as session:
            async with session.get(link, proxy=proxy) as resp:
                resp.raise_for_status()  # Catch HTTP errors like 403/404/500
                content = await resp.read()
                file_name = f"files/{name}.pdf"
                os.makedirs("files", exist_ok=True)  # Ensure output folder exists
                with open(file_name, 'wb') as f:
                    f.write(content)
                print(f"✅ Downloaded {file_name}")
    except ClientError as e:
        print(f"❌ Failed to fetch {link}: {str(e)}")
    except IOError as e:
        print(f"❌ Failed to save {name}.pdf: {str(e)}")

Step 2: Bounded Downloads with Semaphore #

Now, use a semaphore to limit concurrent requests in do_downloads:

async def do_downloads(links, names, max_concurrent=5):
    # Limit to 5 simultaneous requests (adjust based on server tolerance)
    sem = asyncio.Semaphore(max_concurrent)
    
    async def bounded_download(link, name):
        # Only run this coroutine if the semaphore has an available slot
        async with sem:
            await download(link, name)
    
    tasks = [bounded_download(link, name) for link, name in zip(links, names)]
    # Wait for all tasks to finish (even if some fail)
    await asyncio.gather(*tasks)

Step 3: Run the Code #

if __name__ == "__main__":
    links = ["https://example.com/file1.pdf", "https://example.com/file2.pdf"]  # Your links
    names = ["report1", "report2"]  # Corresponding file names
    
    # Ensure links and names lists match in length
    assert len(links) == len(names), "Links and names must be the same length!"
    
    asyncio.run(do_downloads(links, names))

Solution 2: Using Queues (Flexible for Complex Workflows) #

If you need more control (like dynamic task addition or retries), a queue-based approach is better. We'll use a producer to add tasks to a queue, and multiple consumers to process them in parallel.

Full Queue-Based Code #

import asyncio
import aiohttp
from aiohttp import ClientError
import os

async def download(link, name, proxy=None):
    # Same error-handled download function as above
    try:
        async with aiohttp.ClientSession() as session:
            async with session.get(link, proxy=proxy) as resp:
                resp.raise_for_status()
                content = await resp.read()
                file_name = f"files/{name}.pdf"
                os.makedirs("files", exist_ok=True)
                with open(file_name, 'wb') as f:
                    f.write(content)
                print(f"✅ Downloaded {file_name}")
    except ClientError as e:
        print(f"❌ Failed to fetch {link}: {str(e)}")
    except IOError as e:
        print(f"❌ Failed to save {name}.pdf: {str(e)}")

async def consumer(queue):
    # Pull tasks from the queue until told to stop
    while True:
        task = await queue.get()
        if task is None:  # Signal to exit
            break
        link, name = task
        await download(link, name)
        queue.task_done()  # Mark task as completed

async def producer(queue, links, names):
    # Add all download tasks to the queue
    for link, name in zip(links, names):
        await queue.put((link, name))

async def main(links, names, num_consumers=5):
    queue = asyncio.Queue()
    
    # Start consumer coroutines
    consumers = [asyncio.create_task(consumer(queue)) for _ in range(num_consumers)]
    
    # Add all tasks to the queue
    await producer(queue, links, names)
    
    # Wait for all tasks to be processed
    await queue.join()
    
    # Tell consumers to exit
    for _ in range(num_consumers):
        await queue.put(None)
    
    # Wait for all consumers to finish
    await asyncio.gather(*consumers)

if __name__ == "__main__":
    links = ["https://example.com/file1.pdf", "https://example.com/file2.pdf"]
    names = ["report1", "report2"]
    assert len(links) == len(names), "Links and names must match!"
    
    asyncio.run(main(links, names))

Fixing Your Manual Batch Error #

If you still want to use manual batches, add a check to avoid empty batches:

if __name__ == "__main__":
    links = ["https://example.com/file1.pdf", ...]
    names = ["report1", ...]
    BATCH_LENGTH = 10
    current_start = 0
    total = len(links)
    
    loop = asyncio.get_event_loop()
    while current_start < total:
        end = min(current_start + BATCH_LENGTH, total)
        batch_links = links[current_start:end]
        batch_names = names[current_start:end]
        
        # Only run if the batch isn't empty
        if batch_links:
            loop.run_until_complete(do_downloads(batch_links, batch_names))
        
        current_start = end
    loop.close()

内容的提问来源于stack exchange，提问作者liamod