Implement a Time-Aware Task Management System

Problem Statement

Many services at Airbnb require time-sensitive asynchronous actions, such as sending checkout reminder messages to guests or releasing expired reservation holds.

Implement an in-memory class TaskScheduler that manages tasks with scheduled execution times and priorities.

Class Interface

class TaskScheduler:
    def __init__(self):
        pass

    def schedule(self, task_id: str, priority: int, execute_at: int) -> bool:
        """Schedules a new task."""
        pass

    def reschedule(self, task_id: str, new_execute_at: int) -> bool:
        """Updates the execution timestamp of an existing pending task."""
        pass

    def cancel(self, task_id: str) -> bool:
        """Cancels an existing pending task."""
        pass

    def poll(self, current_time: int, max_tasks: int) -> list[str]:
        """Fetches and marks as executed up to max_tasks that are due."""
        pass

Behavior & Invariants

1. schedule(task_id, priority, execute_at):

   • Schedules a task. Returns True if successfully scheduled.
   • If a pending task with task_id already exists, returns False without modifying it.

2. reschedule(task_id, new_execute_at):

   • Updates execute_at of an existing pending task, keeping its original priority.
   • Returns True on success, or False if task_id does not exist or has already been executed.

3. cancel(task_id):

   • Removes the task from pending tasks. Returns True if successfully cancelled, or False if it does not exist or was already executed.

4. poll(current_time, max_tasks):

   • Returns up to max_tasks task IDs whose execute_at <= current_time.
   • Ordering: Higher priority first. For tasks with equal priority, earlier execute_at first. If both tie, lexicographically smaller task_id first.
   • Polled tasks are transitioned to executed and can no longer be cancelled, rescheduled, or re-polled.

Constraints

• Number of operations: up to 10^5 total calls across all methods.
• 1 = 1

Example

scheduler = TaskScheduler()

scheduler.schedule("taskA", priority=1, execute_at=100) # True
scheduler.schedule("taskB", priority=2, execute_at=120) # True
scheduler.schedule("taskC", priority=2, execute_at=90)  # True
scheduler.schedule("taskA", priority=5, execute_at=50)  # False (taskA already exists)

scheduler.reschedule("taskA", new_execute_at=80)        # True
scheduler.cancel("taskC")                               # True

# At time 100, pending tasks are:
# - taskA: priority=1, execute_at=80 (due)
# - taskB: priority=2, execute_at=120 (not due yet)
scheduler.poll(current_time=100, max_tasks=5)          # ["taskA"]

# At time 150:
# - taskB: priority=2, execute_at=120 (due)
scheduler.poll(current_time=150, max_tasks=1)          # ["taskB"]
scheduler.cancel("taskB")                               # False (already executed)

Edge Cases

• Polling when no tasks are due or when queue is empty (return empty list).
• max_tasks is larger than the number of available eligible tasks.
• Tasks rescheduled to a later or earlier time that alters their readiness at current_time.