227. Efficient Task Cancellation in Asyncio

🔹 1. Basic Task Cancellation

This example demonstrates how to cancel an asyncio task gracefully.

Copy

import asyncio

async def long_task():
    print("Task started")
    try:
        await asyncio.sleep(5)
    except asyncio.CancelledError:
        print("Task was cancelled")
        raise

async def main():
    task = asyncio.create_task(long_task())
    await asyncio.sleep(2)  # Wait for a while
    task.cancel()  # Cancel the task
    try:
        await task
    except asyncio.CancelledError:
        print("Handled cancellation")

asyncio.run(main())

✅ Fix: The CancelledError is caught, and the task is properly cleaned up after cancellation.

---

🔹 2. Task Cancellation with Timeout

You can cancel a task after a specific timeout using asyncio.wait_for.

Copy

import asyncio

async def long_running_task():
    await asyncio.sleep(10)  # Simulating a long-running task
    print("Task completed")

async def main():
    try:
        await asyncio.wait_for(long_running_task(), timeout=3)
    except asyncio.TimeoutError:
        print("Task timed out and was cancelled")

asyncio.run(main())

✅ Fix: If the task takes longer than the specified timeout, it’s cancelled automatically.

---

🔹 3. Canceling Multiple Tasks in a Group

You can cancel multiple tasks within a task group.

Copy

import asyncio

async def task_1():
    await asyncio.sleep(3)
    print("Task 1 finished")

async def task_2():
    await asyncio.sleep(4)
    print("Task 2 finished")

async def main():
    tasks = [asyncio.create_task(task_1()), asyncio.create_task(task_2())]
    await asyncio.sleep(2)
    for task in tasks:
        task.cancel()  # Cancel all tasks after 2 seconds

    await asyncio.gather(*tasks, return_exceptions=True)  # Gather tasks to handle exceptions

asyncio.run(main())

✅ Fix: All tasks in the group are cancelled after a certain period, and exceptions are handled with return_exceptions=True.

---

🔹 4. Cancelling Task with Cleanup

It’s essential to perform cleanup operations when a task is cancelled.

Copy

import asyncio

async def cleanup_task():
    try:
        print("Cleaning up...")
        await asyncio.sleep(2)  # Simulate cleanup operation
        print("Cleanup complete")
    except asyncio.CancelledError:
        print("Task cancelled during cleanup")
        raise

async def main():
    task = asyncio.create_task(cleanup_task())
    await asyncio.sleep(1)
    task.cancel()
    await task

asyncio.run(main())

✅ Fix: The task handles cancellation during cleanup and ensures resources are released properly.

---

🔹 5. Handling Cancellation in Long-Running Loops

If your task runs in a loop, cancel the loop gracefully by catching CancelledError.

Copy

import asyncio

async def long_running_loop():
    i = 0
    while True:
        await asyncio.sleep(1)
        i += 1
        print(f"Task running... iteration {i}")
        if i == 5:
            break

async def main():
    task = asyncio.create_task(long_running_loop())
    await asyncio.sleep(3)  # Let it run for a while
    task.cancel()  # Cancel the task
    try:
        await task
    except asyncio.CancelledError:
        print("Task was cancelled during loop execution")

asyncio.run(main())

✅ Fix: The task is cancelled during its execution, with proper handling of CancelledError.

---

🔹 6. Using asyncio.shield() to Prevent Task Cancellation

You can use asyncio.shield() to protect a task from cancellation.

Copy

import asyncio

async def long_running_task():
    await asyncio.sleep(5)
    print("Task finished")

async def main():
    task = asyncio.create_task(long_running_task())
    task = asyncio.shield(task)  # Prevent cancellation
    await asyncio.sleep(2)
    task.cancel()  # This will not cancel the task
    await task  # Wait for the task to complete

asyncio.run(main())

✅ Fix: asyncio.shield() ensures that the task cannot be cancelled, even when task.cancel() is called.

---

🔹 7. Canceling Tasks in a Context Manager

Using an async context manager allows you to manage task cancellation in a clean way.

Copy

import asyncio

async def long_task():
    print("Long task started")
    await asyncio.sleep(5)
    print("Long task finished")

async def main():
    async with asyncio.TaskGroup() as group:
        group.create_task(long_task())
        await asyncio.sleep(2)
        group.cancel_all()  # Cancel all tasks in the group

asyncio.run(main())

✅ Fix: The context manager ensures that all tasks in the group are cancelled as needed.

---

🔹 8. Task Cancellation with CancelledError Handling

You can handle task cancellations explicitly using try-except blocks.

Copy

import asyncio

async def risky_task():
    print("Starting risky task")
    await asyncio.sleep(5)
    print("Task completed")

async def main():
    task = asyncio.create_task(risky_task())
    await asyncio.sleep(2)
    task.cancel()  # Cancel the task
    try:
        await task
    except asyncio.CancelledError:
        print("Risky task was cancelled gracefully")

asyncio.run(main())

✅ Fix: The CancelledError is explicitly handled to cleanly cancel the task.

---

🔹 9. Handling Task Cancellation in Multiple Coroutine Functions

You can manage cancellations in tasks that are waiting for multiple async functions.

Copy

import asyncio

async def task_1():
    await asyncio.sleep(3)
    print("Task 1 completed")

async def task_2():
    await asyncio.sleep(4)
    print("Task 2 completed")

async def main():
    task1 = asyncio.create_task(task_1())
    task2 = asyncio.create_task(task_2())
    
    await asyncio.sleep(2)
    task1.cancel()  # Cancel task 1
    task2.cancel()  # Cancel task 2

    await asyncio.gather(task1, task2, return_exceptions=True)  # Handle cancellations

asyncio.run(main())

✅ Fix: Both tasks are cancelled, and exceptions are managed with return_exceptions=True.

---

🔹 10. Task Cancellation with Custom Timeout

Implementing custom timeout handling with task cancellation.

Copy

import asyncio

async def fetch_data():
    print("Fetching data...")
    await asyncio.sleep(10)

async def main():
    task = asyncio.create_task(fetch_data())
    try:
        await asyncio.wait_for(task, timeout=3)
    except asyncio.TimeoutError:
        print("Fetching data took too long, task was cancelled")

asyncio.run(main())

✅ Fix: The task is cancelled automatically after the specified timeout.

---

💡 Best Practices for Efficient Task Cancellation:

1. Catch CancelledError: Always handle the CancelledError to cleanly stop the task and release any resources.

2. Use asyncio.wait_for: Set a timeout for tasks that may take longer than expected.

3. Graceful Cleanup: Ensure tasks clean up resources when they are cancelled.

4. Task Groups: Use asyncio.TaskGroup for managing and cancelling multiple tasks together.

5. Shield Important Tasks: Protect tasks from cancellation with asyncio.shield() when necessary.

By properly handling task cancellations, you can avoid resource leaks, improve the reliability of your async code, and ensure that tasks don’t leave your system in an inconsistent state.