Priority-Based Async Task Scheduler

Problem Statement

Build an async task scheduler that manages and executes async tasks with different priorities, deadlines, and retry policies. The system should accept tasks, queue them by priority, execute them with worker pools, handle timeouts and failures, support task cancellation, and shutdown gracefully. Tasks represent units of work like API calls, database queries, email sends, or report generation.

Use Cases

• Background job processing - Process user uploads, generate thumbnails, send emails
• API rate-limited requests - Queue API calls, respect rate limits, retry on failure
• Distributed task queues - Similar to Celery (Python), Sidekiq (Ruby), Bull (Node.js)
• Scheduled jobs - Cron-like scheduling (run at specific times)
• ETL pipelines - Extract, transform, load operations with dependencies
• Webhook delivery - Retry failed webhook calls with exponential backoff
• Report generation - Queue long-running report jobs, notify on completion
• Batch processing - Process large datasets in chunks

Why It Matters

Priority Scheduling: Without priorities, low-priority bulk tasks block critical operations. Example: 1000 thumbnail generations (low priority, 1s each) block a password reset email (high priority, 100ms) for 16 minutes. Priority queues ensure critical tasks run first.

Worker Pools: Single-threaded task execution means 10 concurrent tasks at 1s each = 10s total. Worker pool with 10 workers = 1s total (10x faster). Worker pools maximize throughput while limiting concurrency to prevent resource exhaustion.

Timeout & Retry: Async operations can hang (unresponsive servers). Without timeouts, hung tasks block workers forever, degrading throughput. Retries handle transient failures—95% success rate becomes 99.9% with 3 retries.

Graceful Shutdown: Killing workers mid-task loses work. Graceful shutdown stops accepting new tasks, waits for in-progress tasks, then exits cleanly. Critical for deployments and restarts.

Example performance:

No priorities:     Critical task waits behind 1000 low-priority tasks = 16min latency
With priorities:   Critical task executes immediately = 100ms latency (960x faster)

1 worker:          100 tasks × 1s = 100s
10 workers:        100 tasks ÷ 10 × 1s = 10s (10x throughput)
100 workers:       100 tasks ÷ 100 × 1s = 1s (limited by CPU/network)

Key Concepts Explained

This project requires understanding several advanced Rust async programming concepts. These concepts enable building high-performance concurrent task schedulers that would be difficult to implement safely in other languages.

Async/Await and Futures

The Core Concept: Async functions don’t execute immediately—they return Future trait objects that represent pending computation.

#![allow(unused)]
fn main() {
// Synchronous (blocks the thread)
fn download_file(url: &str) -> String {
    // Blocks for seconds/minutes
    http_client.get(url).text() // Thread waits here
}

// Asynchronous (suspends, allows other work)
async fn download_file(url: &str) -> String {
    // Returns immediately with a Future
    http_client.get(url).await // Suspends here, thread does other work
}
}

How Futures Work:

#![allow(unused)]
fn main() {
// What async fn actually returns:
fn download_file(url: &str) -> impl Future<Output = String> {
    // Returns a state machine that implements Future trait
}

// The Future trait:
trait Future {
    type Output;
    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output>;
}

enum Poll<T> {
    Ready(T),     // Computation complete
    Pending,      // Not ready yet, will notify when ready
}
}

Execution Model:

#![allow(unused)]
fn main() {
// Without await (nothing happens):
let future = download_file("http://example.com");
// File not downloaded! Future is lazy

// With await (executes):
let content = download_file("http://example.com").await;
// Now it executes when polled by the runtime
}

Why This Matters for Schedulers:

• Tasks are futures that can be stored, moved, and executed later
• Async operations don’t block threads—one thread can handle thousands of tasks
• Futures compose: You can race them (select!), join them, timeout them

Pin and Pinned Futures

The Problem: Futures can be self-referential (contain pointers to their own data). Moving them invalidates these pointers.

#![allow(unused)]
fn main() {
// Simplified self-referential future:
struct MyFuture {
    data: String,
    pointer: *const String,  // Points to self.data
}

// If we move this:
let mut f1 = MyFuture { data: "hello".into(), pointer: &f1.data };
let f2 = f1;  // Moved!
// f2.pointer now points to f1's old location (dangling pointer!)
}

The Solution: Pin

Pin<P> is a wrapper that prevents moving the pointed-to value.

#![allow(unused)]
fn main() {
use std::pin::Pin;

// Pin prevents moving the future
let pinned: Pin<Box<MyFuture>> = Box::pin(MyFuture { ... });
// Can't move out of pinned now!
}

Why Futures Need Pin:

#![allow(unused)]
fn main() {
async fn example() {
    let x = 42;
    let y = &x;  // y borrows x
    other_async_fn().await;  // Suspends here
    println!("{}", y);  // Resumes here
}

// Desugars to state machine:
enum ExampleFuture {
    State1 { x: i32, y: *const i32 },  // y points to x!
    State2 { ... },
}

// If this future moves, y becomes invalid!
}

Pin in This Project:

#![allow(unused)]
fn main() {
pub struct Task {
    // Future must be pinned to prevent invalidating internal pointers
    pub future: Pin<Box<dyn Future<Output = TaskResult> + Send>>,
}

impl Task {
    pub fn new<F>(name: String, future: F) -> Self
    where
        F: Future<Output = TaskResult> + Send + 'static,
    {
        Task {
            future: Box::pin(future),  // Pin immediately
            // ...
        }
    }
}
}

Key Rules:

• Once pinned, value can’t be moved
• Pin<Box<T>> is most common pattern (heap-allocated, pinned)
• Most code doesn’t need to think about Pin—just use Box::pin()

Trait Objects: Box

The Problem: Different async functions return different concrete Future types.

#![allow(unused)]
fn main() {
async fn task_a() -> i32 { 42 }
async fn task_b() -> i32 { 100 }

// These return DIFFERENT types!
type FutureA = impl Future<Output = i32>;  // Compiler-generated type A
type FutureB = impl Future<Output = i32>;  // Compiler-generated type B

// Can't store them in the same Vec:
let tasks = vec![task_a(), task_b()];  // ❌ Error: type mismatch
}

The Solution: Trait Objects

Use dyn Future to erase the specific type:

#![allow(unused)]
fn main() {
// Store any future that returns i32:
let tasks: Vec<Pin<Box<dyn Future<Output = i32>>>> = vec![
    Box::pin(task_a()),  // Type A erased to dyn Future
    Box::pin(task_b()),  // Type B erased to dyn Future
];

// Now they're compatible!
}

Type Erasure:

Before (concrete types):
task_a() → CompilerGeneratedFutureA { ... }  // 16 bytes
task_b() → CompilerGeneratedFutureB { ... }  // 24 bytes

After (trait objects):
Box::pin(task_a()) → Box<dyn Future> { ptr: *, vtable: * }  // 16 bytes (fat pointer)
Box::pin(task_b()) → Box<dyn Future> { ptr: *, vtable: * }  // 16 bytes (fat pointer)

All trait objects have the same size (2 pointers)!

The + Send Bound:

#![allow(unused)]
fn main() {
// For multithreaded execution:
Pin<Box<dyn Future<Output = T> + Send>>
//                                 ^^^^ Required to send across threads

// Send means "safe to transfer ownership between threads"
// Without Send, can't spawn on tokio runtime:
tokio::spawn(future);  // Requires future: Send
}

Performance Cost:

• Dynamic dispatch: Virtual function call through vtable (~2-5ns overhead)
• Heap allocation: Box requires allocation
• Trade-off: Flexibility vs. slight performance hit

Priority Queues and BinaryHeap

The Concept: A data structure where the “largest” or “smallest” element is always accessible in O(1) time.

How BinaryHeap Works:

#![allow(unused)]
fn main() {
use std::collections::BinaryHeap;

let mut heap = BinaryHeap::new();
heap.push(5);
heap.push(1);
heap.push(10);
heap.push(3);

// Internal representation (max-heap):
//       10
//      /  \
//     5    3
//    /
//   1

// Pop always returns the maximum:
assert_eq!(heap.pop(), Some(10));
assert_eq!(heap.pop(), Some(5));
assert_eq!(heap.pop(), Some(3));
assert_eq!(heap.pop(), Some(1));
}

Heap Properties:

• Complete binary tree: Stored in a Vec, very cache-friendly
• Heap property: Parent ≥ children (for max-heap)
• O(1) peek: Top element always at index 0
• O(log n) insert/remove: Bubble up/down to maintain heap property

Using Ord for Custom Priority:

#![allow(unused)]
fn main() {
#[derive(Eq, PartialEq)]
struct PriorityTask {
    priority: u8,  // Lower value = higher priority
    sequence: u64, // FIFO within same priority
}

impl Ord for PriorityTask {
    fn cmp(&self, other: &Self) -> Ordering {
        // BinaryHeap is max-heap, so reverse to get min-heap behavior
        other.priority.cmp(&self.priority)  // Reversed!
            .then_with(|| other.sequence.cmp(&self.sequence))  // Reversed!
    }
}

// Now BinaryHeap pops highest priority (lowest priority value) first
let mut queue = BinaryHeap::new();
queue.push(PriorityTask { priority: 2, sequence: 1 });  // Normal
queue.push(PriorityTask { priority: 0, sequence: 2 });  // Critical
queue.push(PriorityTask { priority: 1, sequence: 3 });  // High

assert_eq!(queue.pop().unwrap().priority, 0);  // Critical first
assert_eq!(queue.pop().unwrap().priority, 1);  // High second
assert_eq!(queue.pop().unwrap().priority, 2);  // Normal third
}

Why This Matters:

• Task priority scheduling: Critical tasks execute before low-priority ones
• Efficient: O(log n) is fast even for millions of tasks
• Fair within priority: Sequence number ensures FIFO for same-priority tasks

Channels: Communication Between Async Tasks

Channels enable safe message passing between tasks. Tokio provides three main types:

1. mpsc (Multi-Producer Single-Consumer)

Use Case: Many workers sending results to one coordinator.

#![allow(unused)]
fn main() {
use tokio::sync::mpsc;

let (tx, mut rx) = mpsc::channel(100);  // Buffer 100 messages

// Producer
let tx1 = tx.clone();
tokio::spawn(async move {
    tx1.send(42).await.unwrap();
});

// Another producer
let tx2 = tx.clone();
tokio::spawn(async move {
    tx2.send(100).await.unwrap();
});

// Single consumer
while let Some(value) = rx.recv().await {
    println!("Received: {}", value);
}
}

Properties:

• Multiple senders: Clone tx for each producer
• Single receiver: Only one rx
• Buffered: Senders can send up to capacity before blocking
• Backpressure: send().await blocks when buffer full

Why For Task Scheduler:

• Workers are producers (send results)
• Scheduler is consumer (collects results)

2. oneshot (Single-Use Channel)

Use Case: One-time signals, like cancellation.

#![allow(unused)]
fn main() {
use tokio::sync::oneshot;

let (tx, rx) = oneshot::channel();

tokio::spawn(async move {
    // Long-running task
    tokio::select! {
        result = do_work() => {
            println!("Work done: {}", result);
        }
        _ = rx => {
            println!("Cancelled!");
            return;
        }
    }
});

// Cancel after 1 second
tokio::time::sleep(Duration::from_secs(1)).await;
tx.send(()).unwrap();  // Cancels the task
}

Properties:

• Single message: Can only send once
• Zero overhead: Optimized for one-shot use
• Cancellation-friendly: Closing sender/receiver signals cancel

3. watch (Broadcast State Changes)

Use Case: Broadcasting status updates to multiple observers.

#![allow(unused)]
fn main() {
use tokio::sync::watch;

let (tx, mut rx1) = watch::channel("initial");
let mut rx2 = rx1.clone();  // Multiple receivers

tokio::spawn(async move {
    while rx1.changed().await.is_ok() {
        println!("rx1 sees: {}", *rx1.borrow());
    }
});

tokio::spawn(async move {
    while rx2.changed().await.is_ok() {
        println!("rx2 sees: {}", *rx2.borrow());
    }
});

tx.send("updated").unwrap();  // Both receivers notified
tx.send("final").unwrap();    // Both receivers notified again
}

Properties:

• Broadcast: All receivers get updates
• Latest value only: Old updates are overwritten
• Cheap cloning: Receivers can be cloned

Why For Task Scheduler:

• Broadcast task status (Queued → Running → Completed)
• Multiple observers can monitor same task

Timeout and Racing Futures

tokio::time::timeout: Race a future against a timer.

#![allow(unused)]
fn main() {
use tokio::time::{timeout, Duration};

async fn slow_operation() -> i32 {
    tokio::time::sleep(Duration::from_secs(10)).await;
    42
}

// Timeout after 1 second:
match timeout(Duration::from_secs(1), slow_operation()).await {
    Ok(result) => println!("Completed: {}", result),
    Err(_) => println!("Timed out!"),  // After 1 second
}
}

How It Works:

#![allow(unused)]
fn main() {
// Simplified implementation:
async fn timeout<F>(duration: Duration, future: F) -> Result<F::Output, Elapsed>
where
    F: Future,
{
    tokio::select! {
        result = future => Ok(result),        // Future completed first
        _ = tokio::time::sleep(duration) => Err(Elapsed),  // Timer won
    }
}
}

Why This Matters:

• Prevents hung tasks: Unresponsive services don’t block workers forever
• Guarantees bounded latency: Task fails fast rather than hang indefinitely
• Resource protection: Limits time spent on any single task

tokio::select!: Race multiple futures, return first to complete.

#![allow(unused)]
fn main() {
tokio::select! {
    result = task1 => println!("Task1 won: {}", result),
    result = task2 => println!("Task2 won: {}", result),
    _ = shutdown_signal => println!("Shutting down"),
}
// Only ONE branch executes (first to complete)
}

Atomic Types: Lock-Free Synchronization

The Problem: Sharing counters/flags between threads requires synchronization.

#![allow(unused)]
fn main() {
// Bad: Data race
static mut COUNTER: u64 = 0;
COUNTER += 1;  // ❌ UNDEFINED BEHAVIOR if multiple threads access

// Good: Atomic operations
use std::sync::atomic::{AtomicU64, Ordering};

static COUNTER: AtomicU64 = AtomicU64::new(0);
COUNTER.fetch_add(1, Ordering::Relaxed);  // ✓ Thread-safe
}

Common Atomic Types:

#![allow(unused)]
fn main() {
use std::sync::atomic::*;

let count = AtomicU64::new(0);
let flag = AtomicBool::new(false);
let size = AtomicUsize::new(100);

// Operations:
count.fetch_add(1, Ordering::Relaxed);   // count++
count.fetch_sub(1, Ordering::Relaxed);   // count--
count.load(Ordering::Relaxed);           // Read value
count.store(42, Ordering::Relaxed);      // Write value

flag.store(true, Ordering::Relaxed);     // Set flag
flag.load(Ordering::Relaxed);            // Check flag

let old = count.swap(10, Ordering::Relaxed);  // Atomic exchange
}

Memory Ordering (simplified):

For most metrics/counters: Use Relaxed.

Why For Task Scheduler:

• Metrics: Track tasks completed, failed, queue depth
• Shutdown flag: Signal workers to stop
• No locks needed: Atomic operations are lock-free (faster)

Performance:

Atomic increment (Relaxed):  ~1-2ns
Mutex lock/unlock:           ~20-50ns
Speedup: 10-25x faster for simple counters

Graceful Shutdown Pattern

The Problem: Abruptly killing workers loses in-progress work and leaves system in inconsistent state.

The Pattern:

#![allow(unused)]
fn main() {
// 1. Set shutdown flag (stop accepting new work)
shutdown_flag.store(true, Ordering::Relaxed);

// 2. Close input channel (signals workers no more tasks coming)
drop(task_sender);

// 3. Wait for workers to finish current tasks
for worker in workers {
    worker.await.unwrap();
}

// 4. Drain any remaining results
while let Ok(result) = result_receiver.try_recv() {
    process(result);
}
}

Worker Loop with Graceful Shutdown:

#![allow(unused)]
fn main() {
async fn worker(mut task_rx: mpsc::Receiver<Task>) {
    // Loop until channel closed
    while let Some(task) = task_rx.recv().await {
        // Process task completely
        let result = execute_task(task).await;
        send_result(result).await;
    }
    // Channel closed → exit gracefully
}
}

Benefits:

• No lost work: In-flight tasks complete
• Clean state: All resources properly released
• Safe restarts: Can restart without corrupting data

Timeout-Based Shutdown:

#![allow(unused)]
fn main() {
async fn shutdown_with_timeout(workers: Vec<JoinHandle<()>>, timeout: Duration) {
    let shutdown_future = async {
        for worker in workers {
            worker.await.unwrap();
        }
    };

    match tokio::time::timeout(timeout, shutdown_future).await {
        Ok(_) => println!("Clean shutdown"),
        Err(_) => println!("Forced shutdown after timeout"),
    }
}
}

Send and Sync Traits

Send: Type can be transferred across thread boundaries.

#![allow(unused)]
fn main() {
// i32 is Send:
let x = 42;
std::thread::spawn(move || {
    println!("{}", x);  // ✓ i32 moved to new thread
});

// Rc<T> is NOT Send:
use std::rc::Rc;
let rc = Rc::new(42);
std::thread::spawn(move || {
    println!("{}", rc);  // ❌ Error: Rc is not Send
});

// Arc<T> IS Send:
use std::sync::Arc;
let arc = Arc::new(42);
std::thread::spawn(move || {
    println!("{}", arc);  // ✓ Arc is Send
});
}

Sync: Type can be shared across threads (via &T).

#![allow(unused)]
fn main() {
// Equivalent: &T is Send
trait Sync {}

// RefCell is NOT Sync:
use std::cell::RefCell;
let cell = RefCell::new(42);
let cell_ref = &cell;
std::thread::spawn(move || {
    cell_ref.borrow_mut();  // ❌ Error: RefCell is not Sync
});

// Mutex IS Sync:
use std::sync::Mutex;
let mutex = Mutex::new(42);
let mutex_ref = &mutex;
std::thread::spawn(move || {
    mutex_ref.lock().unwrap();  // ✓ Mutex is Sync
});
}

Why This Matters:

#![allow(unused)]
fn main() {
// Futures must be Send to spawn on tokio:
tokio::spawn(async {
    // This entire future must be Send
    // All variables must be Send
});

// Task definition requires Send:
pub struct Task {
    pub future: Pin<Box<dyn Future<Output = TaskResult> + Send>>,
    //                                                      ^^^^ Required
}
}

Auto Traits:

• Types are automatically Send/Sync if all fields are Send/Sync
• Compiler prevents unsafe cross-thread access at compile time

Connection to This Project

Now that you understand the core concepts, here’s how they map to the milestones:

Milestone 1: Basic Task Definition

• Concepts Used: Futures, Pin<Box>, trait objects, Send bound
• Why: Tasks are heterogeneous futures that need to be stored and executed later
• Key Insight: Box::pin() enables storing different async function types in the same struct

Milestone 2: Priority Queue

• Concepts Used: BinaryHeap, Ord trait, sequence numbers for FIFO
• Why: Schedule tasks by importance, not just arrival order
• Key Insight: Custom Ord implementation controls heap ordering, reversed comparisons give min-heap behavior

Milestone 3: Worker Pool

• Concepts Used: mpsc channels, tokio::spawn, concurrent execution
• Why: Execute multiple tasks simultaneously, maximize throughput
• Key Insight: Channels decouple task submission from execution, workers pull from shared queue

Milestone 4: Timeout and Retry

• Concepts Used: tokio::time::timeout, tokio::select!, exponential backoff
• Why: Prevent hung tasks, handle transient failures automatically
• Key Insight: Racing futures against timers enables bounded execution time

Milestone 5: Cancellation and Tracking

• Concepts Used: oneshot channels, watch channels, task status broadcasting
• Why: Stop obsolete work, monitor task lifecycle
• Key Insight: Cancellation is cooperative—tasks must check cancel signal via select!

Milestone 6: Graceful Shutdown and Metrics

• Concepts Used: AtomicU64/AtomicBool, graceful shutdown pattern, channel closure
• Why: Clean termination, observability into system performance
• Key Insight: Atomic types enable lock-free metrics tracking, channel closure signals workers to exit

Putting It All Together:

The complete scheduler combines all concepts:

1. Tasks as futures enable lazy execution and composition
2. Priority queues ensure critical work runs first
3. Channels coordinate producer-consumer communication
4. Timeouts bound execution time
5. Cancellation stops obsolete work
6. Atomics track metrics without locks
7. Graceful shutdown ensures clean termination

Each milestone builds on the previous one, progressively adding more sophisticated async patterns until you have a production-ready task scheduler.

Milestone 1: Basic Task Definition and Execution

Introduction

Before building a scheduler, you need to define what a “task” is. This milestone teaches you to represent async work as trait objects (boxed futures) and execute them.

Why Start Here: Rust’s async traits are tricky—you can’t store async fn directly because the future type is unnameable. Using Box<dyn Future> (trait objects) solves this, allowing heterogeneous task storage.

Architecture

Structs:

• Task - Represents a unit of async work

  • Field id: Uuid - Unique task identifier
  • Field name: String - Human-readable task name
  • Field future: Pin<Box<dyn Future<Output = TaskResult> + Send>> - The actual async work
  • Field created_at: Instant - When task was created

• TaskResult - Result of task execution

  • Variant Success(String) - Task completed successfully
  • Variant Failure(String) - Task failed with error
  • Variant Timeout - Task exceeded deadline

Key Functions:

• impl Task::new<F>(name: String, future: F) -> Self where F: Future<Output = TaskResult> + Send + 'static - Creates task
• async fn execute_task(task: Task) -> (Uuid, TaskResult) - Runs task and returns result
• async fn example_task(name: &str, duration_ms: u64) -> TaskResult - Sample task

Role Each Plays:

• Pin<Box>: Heap-allocated future that can be moved safely (required for async trait objects)
• Send bound: Ensures future can cross thread boundaries (required for tokio::spawn)
• Uuid: Unique identifier for tracking tasks

Checkpoint Tests

#![allow(unused)]
fn main() {
#[tokio::test]
async fn test_task_creation() {
    let task = Task::new(
        "test_task".to_string(),
        async { TaskResult::Success("done".to_string()) },
    );

    assert_eq!(task.name, "test_task");
    assert!(task.id != Uuid::nil());
}

#[tokio::test]
async fn test_task_execution() {
    let task = Task::new(
        "success_task".to_string(),
        async { TaskResult::Success("completed".to_string()) },
    );

    let id = task.id;
    let (result_id, result) = execute_task(task).await;

    assert_eq!(result_id, id);
    assert!(matches!(result, TaskResult::Success(_)));
}

#[tokio::test]
async fn test_task_failure() {
    let task = Task::new(
        "fail_task".to_string(),
        async { TaskResult::Failure("error occurred".to_string()) },
    );

    let (_, result) = execute_task(task).await;

    assert!(matches!(result, TaskResult::Failure(_)));
}

#[tokio::test]
async fn test_example_task() {
    let result = example_task("test", 10).await;

    match result {
        TaskResult::Success(msg) => assert!(msg.contains("test")),
        _ => panic!("Expected success"),
    }
}

#[tokio::test]
async fn test_async_task_with_work() {
    let task = Task::new(
        "fetch_task".to_string(),
        async {
            // Simulate async work
            tokio::time::sleep(Duration::from_millis(50)).await;
            TaskResult::Success("fetched data".to_string())
        },
    );

    let start = Instant::now();
    let (_, result) = execute_task(task).await;
    let elapsed = start.elapsed();

    assert!(elapsed >= Duration::from_millis(50));
    assert!(matches!(result, TaskResult::Success(_)));
}
}

Starter Code

use tokio::time::{sleep, Duration, Instant};
use std::future::Future;
use std::pin::Pin;
use uuid::Uuid;

#[derive(Debug, Clone)]
pub enum TaskResult {
    Success(String),
    Failure(String),
    Timeout,
}

pub struct Task {
    pub id: Uuid,
    pub name: String,
    pub future: Pin<Box<dyn Future<Output = TaskResult> + Send>>,
    pub created_at: Instant,
}

impl Task {
    pub fn new<F>(name: String, future: F) -> Self
    where
        F: Future<Output = TaskResult> + Send + 'static,
    {
        // TODO: Generate unique ID
        // TODO: Box and pin the future
        // TODO: Record creation time

        todo!("Implement Task::new")
    }
}

pub async fn execute_task(task: Task) -> (Uuid, TaskResult) {
    // TODO: Await the task's future
    // TODO: Return (id, result)
    // Hint: You may need to use .await on task.future

    todo!("Implement task execution")
}

pub async fn example_task(name: &str, duration_ms: u64) -> TaskResult {
    // TODO: Simulate work with sleep
    // TODO: Return success with message including name

    todo!("Implement example task")
}

#[tokio::main]
async fn main() {
    let task1 = Task::new(
        "greet".to_string(),
        example_task("Alice", 100),
    );

    let task2 = Task::new(
        "calculate".to_string(),
        async {
            sleep(Duration::from_millis(50)).await;
            TaskResult::Success("42".to_string())
        },
    );

    println!("Executing tasks...");

    let (id1, result1) = execute_task(task1).await;
    println!("Task {}: {:?}", id1, result1);

    let (id2, result2) = execute_task(task2).await;
    println!("Task {}: {:?}", id2, result2);
}

Implementation Hints:

1. Use uuid::Uuid::new_v4() to generate IDs
2. Box and pin: Box::pin(future) returns Pin<Box<dyn Future>>
3. To await a pinned future: task.future.await
4. Use Instant::now() for timestamps
5. For trait objects: ensure + Send + 'static bounds

Milestone 2: Priority Queue for Task Scheduling

Introduction

Why Milestone 1 Isn’t Enough: FIFO (first-in-first-out) execution is unfair. Critical tasks wait behind low-priority bulk operations. We need priority-based scheduling.

The Improvement: Implement a priority queue using BinaryHeap where high-priority tasks execute first. Tasks with same priority use FIFO.

Optimization: Priority scheduling eliminates head-of-line blocking. Without it, one slow low-priority task delays all high-priority tasks queued behind it. With priorities, urgent tasks jump the queue.

Architecture

Structs:

• PriorityTask - Task wrapper with priority

  • Field task: Task - The actual task
  • Field priority: Priority - Execution priority
  • Field sequence: u64 - FIFO order for same priority

• Priority - Priority levels

  • Variant Critical = 0, High = 1, Normal = 2, Low = 3
  • (Lower number = higher priority)

• TaskQueue - Priority-based queue

  • Field queue: BinaryHeap<PriorityTask> - Max-heap (highest priority first)
  • Field sequence: AtomicU64 - Counter for FIFO within priority

Key Functions:

• impl TaskQueue::new() -> Self - Creates queue
• fn push(&mut self, task: Task, priority: Priority) - Enqueues task
• fn pop(&mut self) -> Option<Task> - Dequeues highest priority task
• fn len(&self) -> usize - Queue size

Role Each Plays:

• BinaryHeap: Max-heap data structure (O(log n) insert/remove)
• Ord trait: Defines ordering (by priority, then sequence)
• Sequence number: Breaks ties within same priority (maintains FIFO)

Checkpoint Tests

#![allow(unused)]
fn main() {
#[test]
fn test_priority_ordering() {
    assert!(Priority::Critical < Priority::High);
    assert!(Priority::High < Priority::Normal);
    assert!(Priority::Normal < Priority::Low);
}

#[test]
fn test_task_queue_push_pop() {
    let mut queue = TaskQueue::new();

    let task1 = Task::new("t1".into(), async { TaskResult::Success("1".into()) });
    let task2 = Task::new("t2".into(), async { TaskResult::Success("2".into()) });

    queue.push(task1, Priority::Normal);
    queue.push(task2, Priority::High);

    assert_eq!(queue.len(), 2);

    // High priority should come out first
    let first = queue.pop().unwrap();
    assert_eq!(first.name, "t2");

    let second = queue.pop().unwrap();
    assert_eq!(second.name, "t1");

    assert!(queue.pop().is_none());
}

#[test]
fn test_priority_queue_ordering() {
    let mut queue = TaskQueue::new();

    queue.push(
        Task::new("low".into(), async { TaskResult::Success("".into()) }),
        Priority::Low,
    );
    queue.push(
        Task::new("critical".into(), async { TaskResult::Success("".into()) }),
        Priority::Critical,
    );
    queue.push(
        Task::new("normal".into(), async { TaskResult::Success("".into()) }),
        Priority::Normal,
    );
    queue.push(
        Task::new("high".into(), async { TaskResult::Success("".into()) }),
        Priority::High,
    );

    // Should pop in order: critical, high, normal, low
    assert_eq!(queue.pop().unwrap().name, "critical");
    assert_eq!(queue.pop().unwrap().name, "high");
    assert_eq!(queue.pop().unwrap().name, "normal");
    assert_eq!(queue.pop().unwrap().name, "low");
}

#[test]
fn test_fifo_within_priority() {
    let mut queue = TaskQueue::new();

    // Add 3 tasks with same priority
    queue.push(
        Task::new("first".into(), async { TaskResult::Success("".into()) }),
        Priority::Normal,
    );
    queue.push(
        Task::new("second".into(), async { TaskResult::Success("".into()) }),
        Priority::Normal,
    );
    queue.push(
        Task::new("third".into(), async { TaskResult::Success("".into()) }),
        Priority::Normal,
    );

    // Should come out in FIFO order
    assert_eq!(queue.pop().unwrap().name, "first");
    assert_eq!(queue.pop().unwrap().name, "second");
    assert_eq!(queue.pop().unwrap().name, "third");
}
}

Starter Code

use std::collections::BinaryHeap;
use std::cmp::Ordering;
use std::sync::atomic::{AtomicU64, Ordering as AtomicOrdering};

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Priority {
    Critical = 0,
    High = 1,
    Normal = 2,
    Low = 3,
}

pub struct PriorityTask {
    pub task: Task,
    pub priority: Priority,
    pub sequence: u64,
}

impl Ord for PriorityTask {
    fn cmp(&self, other: &Self) -> Ordering {
        // TODO: Compare by priority first (lower value = higher priority)
        // TODO: If same priority, compare by sequence (lower = earlier)
        // Hint: Use .reverse() to flip ordering for max-heap

        todo!("Implement ordering")
    }
}

impl PartialOrd for PriorityTask {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Eq for PriorityTask {}

impl PartialEq for PriorityTask {
    fn eq(&self, other: &Self) -> bool {
        self.priority == other.priority && self.sequence == other.sequence
    }
}

pub struct TaskQueue {
    queue: BinaryHeap<PriorityTask>,
    sequence: AtomicU64,
}

impl TaskQueue {
    pub fn new() -> Self {
        // TODO: Initialize empty queue

        todo!("Implement TaskQueue::new")
    }

    pub fn push(&mut self, task: Task, priority: Priority) {
        // TODO: Get next sequence number
        // TODO: Wrap task in PriorityTask
        // TODO: Push to heap

        todo!("Implement push")
    }

    pub fn pop(&mut self) -> Option<Task> {
        // TODO: Pop from heap
        // TODO: Extract task from PriorityTask

        todo!("Implement pop")
    }

    pub fn len(&self) -> usize {
        self.queue.len()
    }

    pub fn is_empty(&self) -> bool {
        self.queue.is_empty()
    }
}

#[tokio::main]
async fn main() {
    let mut queue = TaskQueue::new();

    // Enqueue tasks with different priorities
    queue.push(
        Task::new("Low priority task".into(), example_task("low", 100)),
        Priority::Low,
    );
    queue.push(
        Task::new("Critical task".into(), example_task("critical", 50)),
        Priority::Critical,
    );
    queue.push(
        Task::new("Normal task".into(), example_task("normal", 75)),
        Priority::Normal,
    );

    // Execute in priority order
    while let Some(task) = queue.pop() {
        println!("Executing: {}", task.name);
        let (id, result) = execute_task(task).await;
        println!("  Result: {:?}", result);
    }
}

Implementation Hints:

1. For Ord: self.priority.cmp(&other.priority).reverse().then_with(|| self.sequence.cmp(&other.sequence).reverse())
2. Sequence: self.sequence.fetch_add(1, AtomicOrdering::Relaxed)
3. BinaryHeap is max-heap, so reverse comparisons to get desired order
4. Use queue.pop().map(|pt| pt.task) to extract task

Milestone 3: Worker Pool for Concurrent Execution

Introduction

Why Milestone 2 Isn’t Enough: Sequential task execution is slow. With 100 tasks at 1s each, completion takes 100s. A worker pool executes multiple tasks concurrently.

The Improvement: Create a worker pool with N workers that pull tasks from the queue and execute them in parallel. Use channels for communication.

Optimization (Parallelism): 10 workers × 1s per task = 10 tasks/second throughput vs 1 task/second sequential. Worker pool saturates available concurrency (CPU cores, network connections).

Architecture

Structs:

• WorkerPool - Manages concurrent task execution

  • Field workers: Vec<JoinHandle<()>> - Worker task handles
  • Field task_tx: mpsc::Sender<PriorityTask> - Channel to send tasks to workers
  • Field result_rx: mpsc::Receiver<(Uuid, TaskResult)> - Channel for results

• WorkerConfig - Worker pool configuration

  • Field worker_count: usize - Number of concurrent workers
  • Field queue_capacity: usize - Task queue buffer size

Key Functions:

• async fn WorkerPool::new(config: WorkerConfig) -> Self - Creates pool
• async fn submit_task(&self, task: Task, priority: Priority) -> Result<(), String> - Enqueues task
• async fn collect_results(&mut self) -> Vec<(Uuid, TaskResult)> - Gets completed results
• async fn shutdown(self) - Stops workers gracefully

Role Each Plays:

• mpsc channel: Task queue (multiple producers = submitters, single consumer = distributor)
• JoinHandle: Reference to spawned worker tasks
• Worker loop: Continuously pulls tasks and executes them

Checkpoint Tests

#![allow(unused)]
fn main() {
#[tokio::test]
async fn test_worker_pool_creation() {
    let config = WorkerConfig {
        worker_count: 4,
        queue_capacity: 100,
    };

    let pool = WorkerPool::new(config).await;

    assert_eq!(pool.workers.len(), 4);
}

#[tokio::test]
async fn test_submit_and_execute() {
    let config = WorkerConfig {
        worker_count: 2,
        queue_capacity: 10,
    };

    let mut pool = WorkerPool::new(config).await;

    let task = Task::new(
        "test".into(),
        async { TaskResult::Success("done".into()) },
    );

    pool.submit_task(task, Priority::Normal).await.unwrap();

    // Give workers time to execute
    tokio::time::sleep(Duration::from_millis(100)).await;

    let results = pool.collect_results().await;

    assert_eq!(results.len(), 1);
    assert!(matches!(results[0].1, TaskResult::Success(_)));
}

#[tokio::test]
async fn test_concurrent_execution() {
    let config = WorkerConfig {
        worker_count: 5,
        queue_capacity: 100,
    };

    let mut pool = WorkerPool::new(config).await;

    // Submit 10 tasks
    for i in 0..10 {
        let task = Task::new(
            format!("task-{}", i),
            async move {
                tokio::time::sleep(Duration::from_millis(100)).await;
                TaskResult::Success(format!("result-{}", i))
            },
        );
        pool.submit_task(task, Priority::Normal).await.unwrap();
    }

    // With 5 workers, 10 tasks @ 100ms should take ~200ms (2 batches)
    let start = Instant::now();

    tokio::time::sleep(Duration::from_millis(250)).await;

    let results = pool.collect_results().await;
    let elapsed = start.elapsed();

    assert_eq!(results.len(), 10);
    assert!(elapsed < Duration::from_millis(300)); // Faster than sequential (1000ms)
}

#[tokio::test]
async fn test_priority_execution() {
    let config = WorkerConfig {
        worker_count: 1, // Single worker to see ordering
        queue_capacity: 100,
    };

    let mut pool = WorkerPool::new(config).await;

    // Submit in reverse priority order
    pool.submit_task(
        Task::new("low".into(), example_task("low", 10)),
        Priority::Low,
    )
    .await
    .unwrap();

    pool.submit_task(
        Task::new("critical".into(), example_task("critical", 10)),
        Priority::Critical,
    )
    .await
    .unwrap();

    tokio::time::sleep(Duration::from_millis(100)).await;

    let results = pool.collect_results().await;

    // Critical should execute first (would need task name tracking to verify fully)
    assert_eq!(results.len(), 2);
}
}

Starter Code

use tokio::sync::mpsc;
use tokio::task::JoinHandle;

#[derive(Clone)]
pub struct WorkerConfig {
    pub worker_count: usize,
    pub queue_capacity: usize,
}

pub struct WorkerPool {
    workers: Vec<JoinHandle<()>>,
    task_tx: mpsc::Sender<PriorityTask>,
    result_rx: mpsc::Receiver<(Uuid, TaskResult)>,
}

impl WorkerPool {
    pub async fn new(config: WorkerConfig) -> Self {
        // TODO: Create task channel
        // TODO: Create result channel
        // TODO: Spawn N worker tasks
        // Each worker:
        //   - Receives PriorityTask from task_rx
        //   - Executes task
        //   - Sends (id, result) to result_tx

        todo!("Implement WorkerPool::new")
    }

    pub async fn submit_task(&self, task: Task, priority: Priority) -> Result<(), String> {
        // TODO: Wrap task in PriorityTask with sequence number
        // TODO: Send to task channel
        // Hint: Use try_send or send with error handling

        todo!("Implement submit_task")
    }

    pub async fn collect_results(&mut self) -> Vec<(Uuid, TaskResult)> {
        // TODO: Drain all available results from result_rx
        // TODO: Use try_recv in loop until empty

        todo!("Implement collect_results")
    }

    pub async fn shutdown(self) {
        // TODO: Drop task_tx to close channel (signals workers to exit)
        // TODO: Await all worker JoinHandles

        todo!("Implement shutdown")
    }
}

async fn worker_loop(
    mut task_rx: mpsc::Receiver<PriorityTask>,
    result_tx: mpsc::Sender<(Uuid, TaskResult)>,
) {
    // TODO: Loop while receiving tasks
    // TODO: Execute each task
    // TODO: Send result back

    todo!("Implement worker loop")
}

#[tokio::main]
async fn main() {
    let config = WorkerConfig {
        worker_count: 4,
        queue_capacity: 100,
    };

    let mut pool = WorkerPool::new(config).await;

    // Submit tasks
    for i in 0..20 {
        let priority = match i % 3 {
            0 => Priority::Critical,
            1 => Priority::Normal,
            _ => Priority::Low,
        };

        let task = Task::new(
            format!("task-{}", i),
            example_task(&format!("work-{}", i), 50 + (i * 10)),
        );

        pool.submit_task(task, priority).await.unwrap();
    }

    // Collect results periodically
    tokio::time::sleep(Duration::from_secs(2)).await;

    let results = pool.collect_results().await;
    println!("Completed {} tasks", results.len());

    pool.shutdown().await;
}

Implementation Hints:

1. Spawn workers: tokio::spawn(worker_loop(task_rx.clone(), result_tx.clone()))
2. Worker loop: while let Some(pt) = task_rx.recv().await { ... }
3. Execute: let result = pt.task.future.await;
4. Send result: result_tx.send((pt.task.id, result)).await;
5. Collect: Use while let Ok(result) = self.result_rx.try_recv() { ... }

Milestone 4: Timeout and Retry Mechanisms

Introduction

Why Milestone 3 Isn’t Enough: Tasks can hang indefinitely (unresponsive servers, infinite loops). Transient failures (network hiccups) should retry automatically.

The Improvement: Wrap task execution with tokio::time::timeout, implement retry logic with exponential backoff.

Optimization: Timeouts prevent worker starvation. One hung task without timeout blocks that worker forever, reducing pool capacity from N to N-1. Retries improve success rates from ~95% to 99.9%.

Architecture

Structs:

• TaskConfig - Per-task execution configuration

  • Field timeout: Duration - Maximum execution time
  • Field max_retries: u32 - Retry attempts on failure
  • Field retry_delay: Duration - Base delay between retries

• TaskWithConfig - Task + configuration bundle

  • Field task: Task - The task
  • Field priority: Priority - Execution priority
  • Field config: TaskConfig - Execution config

Key Functions:

• async fn execute_with_timeout(task: Task, timeout: Duration) -> TaskResult - Executes with deadline
• async fn execute_with_retry(task: Task, config: TaskConfig) -> TaskResult - Retries on failure
• fn should_retry(result: &TaskResult) -> bool - Determines if retry needed

Role Each Plays:

• tokio::time::timeout: Races future against timer, cancels if too slow
• Retry loop: Attempts execution multiple times with delays
• Exponential backoff: Increases delay between retries (1s, 2s, 4s…)

Checkpoint Tests

#![allow(unused)]
fn main() {
#[tokio::test]
async fn test_task_timeout() {
    let task = Task::new(
        "slow".into(),
        async {
            tokio::time::sleep(Duration::from_secs(10)).await;
            TaskResult::Success("done".into())
        },
    );

    let result = execute_with_timeout(task, Duration::from_millis(100)).await;

    assert!(matches!(result, TaskResult::Timeout));
}

#[tokio::test]
async fn test_task_completes_within_timeout() {
    let task = Task::new(
        "fast".into(),
        async {
            tokio::time::sleep(Duration::from_millis(50)).await;
            TaskResult::Success("done".into())
        },
    );

    let result = execute_with_timeout(task, Duration::from_millis(200)).await;

    assert!(matches!(result, TaskResult::Success(_)));
}

#[tokio::test]
async fn test_retry_on_failure() {
    use std::sync::atomic::{AtomicU32, Ordering};
    use std::sync::Arc;

    let attempt = Arc::new(AtomicU32::new(0));
    let attempt_clone = Arc::clone(&attempt);

    let task = Task::new(
        "retry_test".into(),
        async move {
            let current = attempt_clone.fetch_add(1, Ordering::Relaxed);
            if current < 2 {
                // Fail first 2 attempts
                TaskResult::Failure("not yet".into())
            } else {
                TaskResult::Success("finally".into())
            }
        },
    );

    let config = TaskConfig {
        timeout: Duration::from_secs(1),
        max_retries: 3,
        retry_delay: Duration::from_millis(10),
    };

    let result = execute_with_retry(task, config).await;

    assert!(matches!(result, TaskResult::Success(_)));
    assert_eq!(attempt.load(Ordering::Relaxed), 3); // Took 3 attempts
}

#[tokio::test]
async fn test_retry_exhaustion() {
    let task = Task::new(
        "always_fail".into(),
        async { TaskResult::Failure("nope".into()) },
    );

    let config = TaskConfig {
        timeout: Duration::from_secs(1),
        max_retries: 2,
        retry_delay: Duration::from_millis(10),
    };

    let result = execute_with_retry(task, config).await;

    assert!(matches!(result, TaskResult::Failure(_)));
}

#[test]
fn test_should_retry() {
    assert!(should_retry(&TaskResult::Failure("error".into())));
    assert!(should_retry(&TaskResult::Timeout));
    assert!(!should_retry(&TaskResult::Success("ok".into())));
}
}

Starter Code

use tokio::time::{timeout, Duration};

#[derive(Clone)]
pub struct TaskConfig {
    pub timeout: Duration,
    pub max_retries: u32,
    pub retry_delay: Duration,
}

impl Default for TaskConfig {
    fn default() -> Self {
        Self {
            timeout: Duration::from_secs(30),
            max_retries: 3,
            retry_delay: Duration::from_secs(1),
        }
    }
}

pub struct TaskWithConfig {
    pub task: Task,
    pub priority: Priority,
    pub config: TaskConfig,
}

pub async fn execute_with_timeout(task: Task, timeout_duration: Duration) -> TaskResult {
    // TODO: Wrap task.future with tokio::time::timeout
    // TODO: If timeout, return TaskResult::Timeout
    // TODO: Otherwise return result

    todo!("Implement timeout execution")
}

pub fn should_retry(result: &TaskResult) -> bool {
    // TODO: Return true for Failure and Timeout
    // TODO: Return false for Success

    todo!("Implement retry decision")
}

pub async fn execute_with_retry(mut task: Task, config: TaskConfig) -> TaskResult {
    // TODO: Loop up to max_retries times
    // TODO: Execute with timeout
    // TODO: If success, return immediately
    // TODO: If should_retry, wait and try again
    // TODO: Implement exponential backoff (delay * 2^attempt)

    todo!("Implement retry logic")
}

#[tokio::main]
async fn main() {
    let task = Task::new(
        "flaky_api_call".into(),
        async {
            // Simulate flaky API (50% failure rate)
            if rand::random::<bool>() {
                TaskResult::Failure("API error".into())
            } else {
                TaskResult::Success("API response".into())
            }
        },
    );

    let config = TaskConfig {
        timeout: Duration::from_secs(5),
        max_retries: 5,
        retry_delay: Duration::from_millis(100),
    };

    println!("Executing with retry...");
    let result = execute_with_retry(task, config).await;
    println!("Result: {:?}", result);
}

Implementation Hints:

1. Timeout: match timeout(duration, task.future).await { Ok(result) => result, Err(_) => TaskResult::Timeout }
2. Retry loop: for attempt in 0..=max_retries { ... }
3. Exponential backoff: let delay = retry_delay * 2u32.pow(attempt);
4. Sleep between retries: tokio::time::sleep(delay).await;
5. Note: Can’t retry a consumed future—need to redesign Task to be callable multiple times (use Arc<dyn Fn() -> Future> instead)

Milestone 5: Task Cancellation and Tracking

Introduction

Why Milestone 4 Isn’t Enough: Running tasks may become obsolete (user cancels request, data invalidated). We need to cancel in-flight tasks and track their status.

The Improvement: Use tokio::sync::oneshot channels for cancellation signals. Track task states (Queued, Running, Completed, Cancelled).

Optimization (Resource Efficiency): Cancelling obsolete tasks frees workers for useful work. Without cancellation, workers waste time on tasks whose results will be discarded.

Architecture

Structs:

• TaskStatus - Task lifecycle state

  • Variant Queued - Waiting in queue
  • Variant Running - Being executed
  • Variant Completed(TaskResult) - Finished
  • Variant Cancelled - Aborted before completion

• TaskHandle - Reference to submitted task

  • Field id: Uuid - Task identifier
  • Field cancel_tx: oneshot::Sender<()> - Send to cancel
  • Field status_rx: watch::Receiver<TaskStatus> - Monitor status

• CancellableTask - Task with cancellation support

  • Field task: Task - The task
  • Field cancel_rx: oneshot::Receiver<()> - Receives cancel signal

Key Functions:

• async fn submit_cancellable_task(...) -> TaskHandle - Submit with cancellation ability
• async fn execute_cancellable(task: CancellableTask, status_tx: watch::Sender<TaskStatus>) - Execute with cancel check
• fn cancel(&self) - Cancels task

Role Each Plays:

• oneshot channel: One-time signal for cancellation
• watch channel: Broadcast current task status to observers
• tokio::select!: Race task execution against cancel signal

Checkpoint Tests

#![allow(unused)]
fn main() {
#[tokio::test]
async fn test_task_cancellation() {
    use tokio::sync::oneshot;

    let (cancel_tx, cancel_rx) = oneshot::channel();

    let task = Task::new(
        "long_task".into(),
        async {
            tokio::time::sleep(Duration::from_secs(10)).await;
            TaskResult::Success("done".into())
        },
    );

    let cancellable = CancellableTask {
        task,
        cancel_rx,
    };

    let (status_tx, mut status_rx) = tokio::sync::watch::channel(TaskStatus::Queued);

    let exec_handle = tokio::spawn(execute_cancellable(cancellable, status_tx));

    // Wait a bit then cancel
    tokio::time::sleep(Duration::from_millis(50)).await;
    cancel_tx.send(()).unwrap();

    exec_handle.await.unwrap();

    // Status should be Cancelled
    assert!(matches!(*status_rx.borrow(), TaskStatus::Cancelled));
}

#[tokio::test]
async fn test_task_completes_before_cancel() {
    use tokio::sync::oneshot;

    let (cancel_tx, cancel_rx) = oneshot::channel();

    let task = Task::new(
        "fast_task".into(),
        async {
            tokio::time::sleep(Duration::from_millis(10)).await;
            TaskResult::Success("done".into())
        },
    );

    let cancellable = CancellableTask {
        task,
        cancel_rx,
    };

    let (status_tx, mut status_rx) = tokio::sync::watch::channel(TaskStatus::Queued);

    execute_cancellable(cancellable, status_tx).await;

    // Should complete normally
    assert!(matches!(
        *status_rx.borrow(),
        TaskStatus::Completed(TaskResult::Success(_))
    ));
}

#[tokio::test]
async fn test_task_handle() {
    // This would require full implementation with WorkerPool integration
    // For now, verify handle structure

    let (_cancel_tx, cancel_rx) = tokio::sync::oneshot::channel();
    let (_status_tx, status_rx) = tokio::sync::watch::channel(TaskStatus::Queued);

    let handle = TaskHandle {
        id: Uuid::new_v4(),
        cancel_tx: _cancel_tx,
        status_rx,
    };

    assert!(!handle.id.is_nil());
}
}

Starter Code

use tokio::sync::{oneshot, watch};

#[derive(Debug, Clone)]
pub enum TaskStatus {
    Queued,
    Running,
    Completed(TaskResult),
    Cancelled,
}

pub struct TaskHandle {
    pub id: Uuid,
    cancel_tx: oneshot::Sender<()>,
    pub status_rx: watch::Receiver<TaskStatus>,
}

impl TaskHandle {
    pub fn cancel(self) {
        // TODO: Send cancel signal
        // Hint: self.cancel_tx.send(())

        todo!("Implement cancel")
    }

    pub async fn wait(mut self) -> TaskStatus {
        // TODO: Wait for status to be Completed or Cancelled
        // Hint: Use status_rx.changed().await

        todo!("Implement wait")
    }

    pub fn status(&self) -> TaskStatus {
        // TODO: Get current status
        // Hint: self.status_rx.borrow().clone()

        todo!("Implement status check")
    }
}

pub struct CancellableTask {
    pub task: Task,
    pub cancel_rx: oneshot::Receiver<()>,
}

pub async fn execute_cancellable(
    mut task: CancellableTask,
    status_tx: watch::Sender<TaskStatus>,
) {
    // TODO: Update status to Running
    // TODO: Use tokio::select! to race between:
    //   - task.task.future.await => update to Completed
    //   - task.cancel_rx => update to Cancelled
    // TODO: Send final status

    todo!("Implement cancellable execution")
}

pub async fn submit_cancellable_task(
    pool: &WorkerPool,
    task: Task,
    priority: Priority,
) -> TaskHandle {
    // TODO: Create cancel and status channels
    // TODO: Wrap task in CancellableTask
    // TODO: Submit to pool
    // TODO: Return TaskHandle

    todo!("Implement cancellable submit")
}

#[tokio::main]
async fn main() {
    use tokio::sync::oneshot;

    let (cancel_tx, cancel_rx) = oneshot::channel();

    let task = Task::new(
        "cancellable_task".into(),
        async {
            for i in 0..10 {
                println!("Working... {}/10", i + 1);
                tokio::time::sleep(Duration::from_millis(500)).await;
            }
            TaskResult::Success("completed".into())
        },
    );

    let cancellable = CancellableTask { task, cancel_rx };

    let (status_tx, mut status_rx) = watch::channel(TaskStatus::Queued);

    let exec_handle = tokio::spawn(execute_cancellable(cancellable, status_tx));

    // Cancel after 1.5 seconds
    tokio::spawn(async move {
        tokio::time::sleep(Duration::from_millis(1500)).await;
        println!("Cancelling task...");
        cancel_tx.send(()).unwrap();
    });

    exec_handle.await.unwrap();

    println!("Final status: {:?}", *status_rx.borrow());
}

Implementation Hints:

1. Update status: status_tx.send(TaskStatus::Running).unwrap();
2. select!: tokio::select! { result = task.future => {...}, _ = cancel_rx => {...} }
3. Wait for changes: while status_rx.changed().await.is_ok() { if matches!(...) { break; } }
4. Note: Cancellation is cooperative—task must check cancel signal periodically

Milestone 6: Graceful Shutdown and Metrics

Introduction

Why Milestone 5 Isn’t Enough: Production systems need clean shutdown (deployments, restarts) and observability (throughput, queue depth, latency).

The Improvement: Implement graceful shutdown (stop accepting tasks, drain queue, wait for workers) and collect metrics (tasks completed, avg latency, queue depth).

Optimization (Observability): Without metrics, performance problems are invisible. Metrics reveal bottlenecks—high queue depth means workers are overloaded, high latency means tasks are too slow.

Architecture

Structs:

• SchedulerMetrics - Performance metrics

  • Field tasks_submitted: AtomicU64 - Total submitted
  • Field tasks_completed: AtomicU64 - Successfully completed
  • Field tasks_failed: AtomicU64 - Failed tasks
  • Field tasks_cancelled: AtomicU64 - Cancelled tasks
  • Field total_latency_ms: AtomicU64 - Sum of all latencies
  • Field current_queue_depth: AtomicUsize - Tasks waiting

• Scheduler - Complete task scheduler

  • Field pool: WorkerPool - Worker pool
  • Field metrics: Arc<SchedulerMetrics> - Metrics tracker
  • Field shutdown: Arc<AtomicBool> - Shutdown flag

Key Functions:

• async fn Scheduler::new(config: WorkerConfig) -> Self - Creates scheduler
• async fn submit(&self, task: Task, priority: Priority) -> Result<TaskHandle, String> - Submit with metrics
• async fn shutdown_gracefully(&mut self) - Clean shutdown
• fn get_metrics_report(&self) -> String - Formats metrics

Role Each Plays:

• AtomicBool: Thread-safe shutdown flag
• Metrics: Track system performance
• Graceful shutdown: Ensures no work is lost

Checkpoint Tests

#![allow(unused)]
fn main() {
#[tokio::test]
async fn test_metrics_tracking() {
    let metrics = Arc::new(SchedulerMetrics::new());

    metrics.increment_submitted();
    metrics.increment_submitted();
    metrics.increment_completed();
    metrics.increment_failed();
    metrics.record_latency(Duration::from_millis(100));
    metrics.set_queue_depth(5);

    let report = metrics.get_report();

    assert!(report.contains("Submitted: 2"));
    assert!(report.contains("Completed: 1"));
    assert!(report.contains("Failed: 1"));
    assert!(report.contains("Queue depth: 5"));
}

#[tokio::test]
async fn test_graceful_shutdown() {
    let config = WorkerConfig {
        worker_count: 2,
        queue_capacity: 10,
    };

    let mut scheduler = Scheduler::new(config).await;

    // Submit tasks
    for i in 0..5 {
        let task = Task::new(
            format!("task-{}", i),
            async move {
                tokio::time::sleep(Duration::from_millis(100)).await;
                TaskResult::Success(format!("{}", i))
            },
        );
        scheduler.submit(task, Priority::Normal).await.unwrap();
    }

    // Initiate shutdown
    scheduler.shutdown_gracefully().await;

    // All tasks should complete
    let metrics = scheduler.get_metrics_report();
    println!("{}", metrics);
}

#[tokio::test]
async fn test_reject_after_shutdown() {
    let config = WorkerConfig {
        worker_count: 1,
        queue_capacity: 10,
    };

    let mut scheduler = Scheduler::new(config).await;

    scheduler.shutdown.store(true, Ordering::Relaxed);

    let task = Task::new("test".into(), example_task("test", 10));
    let result = scheduler.submit(task, Priority::Normal).await;

    assert!(result.is_err());
}
}

Starter Code

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};

pub struct SchedulerMetrics {
    tasks_submitted: AtomicU64,
    tasks_completed: AtomicU64,
    tasks_failed: AtomicU64,
    tasks_cancelled: AtomicU64,
    total_latency_ms: AtomicU64,
    current_queue_depth: AtomicUsize,
}

impl SchedulerMetrics {
    pub fn new() -> Self {
        // TODO: Initialize atomic counters

        todo!("Implement SchedulerMetrics::new")
    }

    pub fn increment_submitted(&self) {
        self.tasks_submitted.fetch_add(1, Ordering::Relaxed);
    }

    pub fn increment_completed(&self) {
        self.tasks_completed.fetch_add(1, Ordering::Relaxed);
    }

    pub fn increment_failed(&self) {
        self.tasks_failed.fetch_add(1, Ordering::Relaxed);
    }

    pub fn increment_cancelled(&self) {
        self.tasks_cancelled.fetch_add(1, Ordering::Relaxed);
    }

    pub fn record_latency(&self, latency: Duration) {
        self.total_latency_ms
            .fetch_add(latency.as_millis() as u64, Ordering::Relaxed);
    }

    pub fn set_queue_depth(&self, depth: usize) {
        self.current_queue_depth.store(depth, Ordering::Relaxed);
    }

    pub fn get_report(&self) -> String {
        // TODO: Format all metrics into readable report
        // Include: submitted, completed, failed, cancelled, avg latency, queue depth

        todo!("Implement metrics report")
    }
}

pub struct Scheduler {
    pool: WorkerPool,
    metrics: Arc<SchedulerMetrics>,
    shutdown: Arc<AtomicBool>,
}

impl Scheduler {
    pub async fn new(config: WorkerConfig) -> Self {
        // TODO: Create worker pool
        // TODO: Initialize metrics
        // TODO: Set shutdown flag to false

        todo!("Implement Scheduler::new")
    }

    pub async fn submit(
        &self,
        task: Task,
        priority: Priority,
    ) -> Result<TaskHandle, String> {
        // TODO: Check shutdown flag
        // TODO: If shutdown, reject with error
        // TODO: Increment submitted metric
        // TODO: Submit to pool
        // TODO: Return handle

        todo!("Implement submit with metrics")
    }

    pub async fn shutdown_gracefully(&mut self) {
        // TODO: Set shutdown flag
        // TODO: Stop accepting new tasks
        // TODO: Wait for pool to drain
        // TODO: Shutdown pool

        todo!("Implement graceful shutdown")
    }

    pub fn get_metrics_report(&self) -> String {
        self.metrics.get_report()
    }
}

pub async fn monitor_metrics(metrics: Arc<SchedulerMetrics>) {
    // TODO: Periodically print metrics
    // TODO: Loop with sleep, print every N seconds

    todo!("Implement metrics monitoring")
}

#[tokio::main]
async fn main() {
    let config = WorkerConfig {
        worker_count: 8,
        queue_capacity: 1000,
    };

    let mut scheduler = Scheduler::new(config).await;

    // Spawn metrics monitor
    let metrics_clone = Arc::clone(&scheduler.metrics);
    tokio::spawn(monitor_metrics(metrics_clone));

    // Submit many tasks
    for i in 0..100 {
        let priority = match i % 4 {
            0 => Priority::Critical,
            1 => Priority::High,
            2 => Priority::Normal,
            _ => Priority::Low,
        };

        let task = Task::new(
            format!("task-{}", i),
            async move {
                let duration = Duration::from_millis(50 + (i % 10) * 10);
                tokio::time::sleep(duration).await;

                if i % 20 == 0 {
                    TaskResult::Failure("simulated failure".into())
                } else {
                    TaskResult::Success(format!("result-{}", i))
                }
            },
        );

        if let Ok(handle) = scheduler.submit(task, priority).await {
            // Could track handles for cancellation
        }
    }

    // Let tasks run
    tokio::time::sleep(Duration::from_secs(3)).await;

    println!("\nInitiating graceful shutdown...");
    scheduler.shutdown_gracefully().await;

    println!("\n=== Final Metrics ===");
    println!("{}", scheduler.get_metrics_report());
}

Implementation Hints:

1. Check shutdown: if self.shutdown.load(Ordering::Relaxed) { return Err(...); }
2. Graceful shutdown: self.shutdown.store(true, ...); self.pool.shutdown().await;
3. Metrics report: format submitted, completed, failed, success rate, avg latency
4. Avg latency: total_latency / completed (handle divide by zero)
5. Monitor: loop { sleep(Duration::from_secs(5)).await; println!(metrics); }

Complete Working Example

// Cargo.toml:
// [dependencies]
// tokio = { version = "1.35", features = ["full"] }
// uuid = { version = "1.6", features = ["v4"] }
// rand = "0.8"

use tokio::time::{sleep, timeout, Duration, Instant};
use tokio::sync::{mpsc, oneshot, watch};
use tokio::task::JoinHandle;
use std::future::Future;
use std::pin::Pin;
use std::collections::BinaryHeap;
use std::cmp::Ordering;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering as AtomicOrdering};
use uuid::Uuid;

// Basic task types
#[derive(Debug, Clone)]
pub enum TaskResult {
    Success(String),
    Failure(String),
    Timeout,
}

pub struct Task {
    pub id: Uuid,
    pub name: String,
    pub future: Pin<Box<dyn Future<Output = TaskResult> + Send>>,
    pub created_at: Instant,
}

impl Task {
    pub fn new<F>(name: String, future: F) -> Self
    where
        F: Future<Output = TaskResult> + Send + 'static,
    {
        Self {
            id: Uuid::new_v4(),
            name,
            future: Box::pin(future),
            created_at: Instant::now(),
        }
    }
}

// Priority types
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Priority {
    Critical = 0,
    High = 1,
    Normal = 2,
    Low = 3,
}

pub struct PriorityTask {
    pub task: Task,
    pub priority: Priority,
    pub sequence: u64,
    pub config: TaskConfig,
}

impl Ord for PriorityTask {
    fn cmp(&self, other: &Self) -> Ordering {
        other
            .priority
            .cmp(&self.priority)
            .then_with(|| other.sequence.cmp(&self.sequence))
    }
}

impl PartialOrd for PriorityTask {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Eq for PriorityTask {}

impl PartialEq for PriorityTask {
    fn eq(&self, other: &Self) -> bool {
        self.priority == other.priority && self.sequence == other.sequence
    }
}

// Task configuration
#[derive(Clone)]
pub struct TaskConfig {
    pub timeout: Duration,
    pub max_retries: u32,
    pub retry_delay: Duration,
}

impl Default for TaskConfig {
    fn default() -> Self {
        Self {
            timeout: Duration::from_secs(30),
            max_retries: 3,
            retry_delay: Duration::from_secs(1),
        }
    }
}

// Metrics
pub struct SchedulerMetrics {
    tasks_submitted: AtomicU64,
    tasks_completed: AtomicU64,
    tasks_failed: AtomicU64,
    current_queue_depth: AtomicUsize,
}

impl SchedulerMetrics {
    pub fn new() -> Self {
        Self {
            tasks_submitted: AtomicU64::new(0),
            tasks_completed: AtomicU64::new(0),
            tasks_failed: AtomicU64::new(0),
            current_queue_depth: AtomicUsize::new(0),
        }
    }

    pub fn increment_submitted(&self) {
        self.tasks_submitted
            .fetch_add(1, AtomicOrdering::Relaxed);
    }

    pub fn increment_completed(&self) {
        self.tasks_completed
            .fetch_add(1, AtomicOrdering::Relaxed);
    }

    pub fn increment_failed(&self) {
        self.tasks_failed.fetch_add(1, AtomicOrdering::Relaxed);
    }

    pub fn set_queue_depth(&self, depth: usize) {
        self.current_queue_depth
            .store(depth, AtomicOrdering::Relaxed);
    }

    pub fn get_report(&self) -> String {
        let submitted = self.tasks_submitted.load(AtomicOrdering::Relaxed);
        let completed = self.tasks_completed.load(AtomicOrdering::Relaxed);
        let failed = self.tasks_failed.load(AtomicOrdering::Relaxed);
        let queue = self.current_queue_depth.load(AtomicOrdering::Relaxed);

        let success_rate = if completed + failed > 0 {
            (completed as f64 / (completed + failed) as f64) * 100.0
        } else {
            0.0
        };

        format!(
            "Submitted: {}, Completed: {}, Failed: {}, Queue: {}, Success rate: {:.1}%",
            submitted, completed, failed, queue, success_rate
        )
    }
}

// Worker pool
pub struct WorkerPool {
    task_tx: mpsc::Sender<PriorityTask>,
    metrics: Arc<SchedulerMetrics>,
    workers: Vec<JoinHandle<()>>,
}

impl WorkerPool {
    pub async fn new(worker_count: usize, metrics: Arc<SchedulerMetrics>) -> Self {
        let (task_tx, task_rx) = mpsc::channel(1000);
        let task_rx = Arc::new(tokio::sync::Mutex::new(task_rx));

        let mut workers = Vec::new();

        for worker_id in 0..worker_count {
            let task_rx = Arc::clone(&task_rx);
            let metrics = Arc::clone(&metrics);

            let worker = tokio::spawn(async move {
                loop {
                    let priority_task = {
                        let mut rx = task_rx.lock().await;
                        rx.recv().await
                    };

                    match priority_task {
                        Some(pt) => {
                            let task_id = pt.task.id;
                            let task_name = pt.task.name.clone();

                            let result = timeout(pt.config.timeout, pt.task.future).await;

                            match result {
                                Ok(TaskResult::Success(_)) => {
                                    metrics.increment_completed();
                                }
                                _ => {
                                    metrics.increment_failed();
                                }
                            }
                        }
                        None => break,
                    }
                }
            });

            workers.push(worker);
        }

        Self {
            task_tx,
            metrics,
            workers,
        }
    }

    pub async fn submit(&self, task: Task, priority: Priority, config: TaskConfig) {
        static SEQUENCE: AtomicU64 = AtomicU64::new(0);
        let sequence = SEQUENCE.fetch_add(1, AtomicOrdering::Relaxed);

        let pt = PriorityTask {
            task,
            priority,
            sequence,
            config,
        };

        self.metrics.increment_submitted();
        let _ = self.task_tx.send(pt).await;
    }

    pub async fn shutdown(self) {
        drop(self.task_tx);

        for worker in self.workers {
            let _ = worker.await;
        }
    }
}

// Complete scheduler
pub struct Scheduler {
    pool: WorkerPool,
    metrics: Arc<SchedulerMetrics>,
    shutdown: Arc<AtomicBool>,
}

impl Scheduler {
    pub async fn new(worker_count: usize) -> Self {
        let metrics = Arc::new(SchedulerMetrics::new());
        let pool = WorkerPool::new(worker_count, Arc::clone(&metrics)).await;

        Self {
            pool,
            metrics,
            shutdown: Arc::new(AtomicBool::new(false)),
        }
    }

    pub async fn submit(
        &self,
        task: Task,
        priority: Priority,
        config: TaskConfig,
    ) -> Result<Uuid, String> {
        if self.shutdown.load(AtomicOrdering::Relaxed) {
            return Err("Scheduler is shutting down".to_string());
        }

        let id = task.id;
        self.pool.submit(task, priority, config).await;
        Ok(id)
    }

    pub async fn shutdown_gracefully(self) {
        self.shutdown.store(true, AtomicOrdering::Relaxed);
        println!("Shutdown initiated. Draining queue...");

        self.pool.shutdown().await;

        println!("All workers stopped. Shutdown complete.");
    }

    pub fn get_metrics(&self) -> String {
        self.metrics.get_report()
    }
}

// Example task
async fn example_work(name: &str, duration_ms: u64, fail: bool) -> TaskResult {
    sleep(Duration::from_millis(duration_ms)).await;

    if fail {
        TaskResult::Failure(format!("{} failed", name))
    } else {
        TaskResult::Success(format!("{} completed", name))
    }
}

// Main
#[tokio::main]
async fn main() {
    println!("=== Priority-Based Async Task Scheduler ===\n");

    let scheduler = Scheduler::new(4).await;

    // Spawn metrics monitor
    let metrics_clone = Arc::clone(&scheduler.metrics);
    tokio::spawn(async move {
        loop {
            sleep(Duration::from_secs(2)).await;
            println!("[METRICS] {}", metrics_clone.get_report());
        }
    });

    // Submit various tasks
    for i in 0..50 {
        let priority = match i % 4 {
            0 => Priority::Critical,
            1 => Priority::High,
            2 => Priority::Normal,
            _ => Priority::Low,
        };

        let name = format!("task-{}", i);
        let duration = 100 + (i * 20);
        let fail = i % 10 == 0;

        let task = Task::new(name.clone(), example_work(&name, duration, fail));

        let config = TaskConfig {
            timeout: Duration::from_secs(5),
            max_retries: 2,
            retry_delay: Duration::from_millis(100),
        };

        if let Ok(id) = scheduler.submit(task, priority, config).await {
            if i < 5 {
                println!("Submitted {} with priority {:?} (ID: {})", name, priority, id);
            }
        }
    }

    // Run for a bit
    sleep(Duration::from_secs(5)).await;

    println!("\n{}", scheduler.get_metrics());
    scheduler.shutdown_gracefully().await;
}