Implement a Bounded Per-Key Ordered Task Executor

Problem Statement

Design and implement a thread-safe executor, KeyedTaskExecutor, that executes submitted tasks asynchronously subject to strict concurrency and ordering invariants:

1. Per-Key Mutual Exclusion & Ordering: Tasks associated with the same key must never run concurrently. Tasks for a given key must execute strictly in the order they were admitted.
2. Global Concurrency Bound: At most max_workers (e.g., 4) tasks across all keys may execute concurrently at any point in time.
3. Bounded Capacity: The executor accepts at most max_capacity total pending tasks across all keys. Submissions exceeding this capacity must be rejected immediately.
4. Task Cancellation: Pending (non-running) tasks can be cancelled via a task handle or identifier.
5. Graceful Shutdown: The executor must support a shutdown(wait=True, timeout=None) method. Once shut down, new submissions are rejected. If wait=True, the executor waits up to timeout seconds for in-flight tasks to complete.

Interface

from typing import Callable, Any, Optional

class TaskRejectedException(Exception):
    pass

class TaskHandle:
    def result(self, timeout: Optional[float] = None) -> Any:
        """Blocks until completion and returns value or raises task exception."""
        pass

    def cancel(self) -> bool:
        """Attempts to cancel task. Returns True if cancelled before execution started."""
        pass

class KeyedTaskExecutor:
    def __init__(self, max_workers: int = 4, max_capacity: int = 100):
        pass

    def submit(self, key: str, fn: Callable[[], Any]) -> TaskHandle:
        """Atomically admits a task for `key`. Raises TaskRejectedException if capacity is reached or executor is shut down."""
        pass

    def shutdown(self, wait: bool = True, timeout: Optional[float] = None) -> bool:
        """Stops accepting tasks and waits for active/queued tasks up to `timeout`. Returns True if all finished."""
        pass

Example

executor = KeyedTaskExecutor(max_workers=2, max_capacity=10)

# Submissions
h1 = executor.submit("user_1", task_a) # Starts immediately (worker 1)
h2 = executor.submit("user_2", task_b) # Starts immediately (worker 2)
h3 = executor.submit("user_1", task_c) # Queued: user_1 already running
h4 = executor.submit("user_3", task_d) # Queued: max_workers reached

# When task_a finishes, user_1 becomes eligible again.
# task_c runs next on user_1, maintaining FIFO order for user_1.

Constraints & Edge Cases

• 1 <= max_workers <= 64
• 1 <= max_capacity <= 100,000
• Task execution may raise arbitrary exceptions; failure of one task must not stall subsequent tasks for that key or leak worker threads.
• Empty keys or keys whose queues drain completely must clean up their internal state to prevent memory leaks.
• Submissions and cancellations must be thread-safe.