Exercise: Context Manager

Difficulty - Intermediate

Learning Objectives

• Master context package for cancellation and timeouts
• Implement context-aware operations
• Handle cascading cancellation
• Attach metadata to contexts
• Build context utilities

Problem Statement

Create utilities for working with contexts including timeout helpers, cancellation coordination, and value management. Your implementation should handle edge cases gracefully and follow Go's concurrency best practices.

Requirements

1. WithRetry Function

Implement a retry mechanism with exponential backoff that:

• Accepts a context, maximum retry count, and a function to execute
• Retries the function up to maxRetries times if it returns an error
• Uses exponential backoff: 1s, 2s, 4s, 8s, etc. between retries
• Respects context cancellation during retries and backoff periods
• Returns the last error if all retries fail

Example Usage:

1err := WithRetry(ctx, 3, func(ctx context.Context) error {
2    return makeAPICall(ctx)
3})
4// Will retry up to 3 times with exponential backoff

2. WithDeadline Function

Create a helper that:

• Accepts an absolute deadline time and a function to execute
• Creates a context with the deadline
• Executes the function with the deadline context
• Properly cancels the context to prevent leaks

Example Usage:

1deadline := time.Now().Add(5 * time.Second)
2err := WithDeadline(deadline, func(ctx context.Context) error {
3    return processData(ctx)
4})

3. WithTimeout Function

Similar to WithDeadline but accepts a duration:

• Accepts a timeout duration and a function to execute
• Creates a context with the timeout
• Executes the function with the timeout context
• Properly cleans up resources

Example Usage:

1err := WithTimeout(2*time.Second, func(ctx context.Context) error {
2    return fetchFromDatabase(ctx)
3})

4. Parallel Function

Execute multiple functions concurrently with fail-fast behavior:

• Accepts a context and variadic functions
• Runs all functions concurrently in separate goroutines
• If ANY function returns an error, cancels all other functions immediately
• Returns the first error encountered
• Returns nil only if ALL functions succeed

Example Usage:

1err := Parallel(ctx,
2    func(ctx context.Context) error { return task1(ctx) },
3    func(ctx context.Context) error { return task2(ctx) },
4    func(ctx context.Context) error { return task3(ctx) },
5)
6// All tasks run concurrently; first error cancels all others

5. First Function

Race multiple functions and return first success:

• Accepts a context and variadic functions
• Runs all functions concurrently
• Returns nil as soon as ANY function succeeds
• Returns an error only if ALL functions fail
• Useful for trying multiple fallback strategies

Example Usage:

1err := First(ctx,
2    func(ctx context.Context) error { return tryPrimaryAPI(ctx) },
3    func(ctx context.Context) error { return tryBackupAPI(ctx) },
4    func(ctx context.Context) error { return tryCache(ctx) },
5)
6// First successful operation wins; others are cancelled

Function Signatures

 1package contextutils
 2
 3import (
 4    "context"
 5    "time"
 6)
 7
 8// WithRetry executes fn with retries and exponential backoff
 9func WithRetry(ctx context.Context, maxRetries int, fn func(context.Context) error) error
10
11// WithDeadline runs fn with a deadline
12func WithDeadline(deadline time.Time, fn func(context.Context) error) error
13
14// WithTimeout runs fn with a timeout
15func WithTimeout(timeout time.Duration, fn func(context.Context) error) error
16
17// Parallel runs multiple functions concurrently, cancels all if one fails
18func Parallel(ctx context.Context, fns ...func(context.Context) error) error
19
20// First runs functions concurrently, returns first success
21func First(ctx context.Context, fns ...func(context.Context) error) error

Test Cases

Your implementation should pass these test scenarios:

 1// Test WithRetry succeeds after failures
 2func TestWithRetrySuccess() {
 3    attempt := 0
 4    err := WithRetry(context.Background(), 3, func(ctx context.Context) error {
 5        attempt++
 6        if attempt < 3 {
 7            return errors.New("temporary error")
 8        }
 9        return nil // Success on 3rd attempt
10    })
11    // Should succeed and attempt should be 3
12}
13
14// Test Parallel fails fast
15func TestParallelFailFast() {
16    start := time.Now()
17    err := Parallel(context.Background(),
18        func(ctx context.Context) error {
19            time.Sleep(1 * time.Second)
20            return nil
21        },
22        func(ctx context.Context) error {
23            return errors.New("immediate failure")
24        },
25    )
26    elapsed := time.Since(start)
27    // Should fail immediately, not wait for slow function
28    // elapsed should be < 100ms
29}
30
31// Test First returns first success
32func TestFirstSuccess() {
33    result := ""
34    err := First(context.Background(),
35        func(ctx context.Context) error {
36            time.Sleep(100 * time.Millisecond)
37            result = "slow"
38            return nil
39        },
40        func(ctx context.Context) error {
41            time.Sleep(10 * time.Millisecond)
42            result = "fast"
43            return nil
44        },
45    )
46    // Should succeed with result = "fast"
47}

Common Pitfalls

⚠️ Watch out for these common mistakes:

1. Goroutine leaks: Always call cancel() to clean up contexts, even if the function succeeds
2. Not checking context cancellation: In WithRetry, check ctx.Done() during backoff
3. Race conditions: Use proper synchronization when multiple goroutines write to shared state
4. Channel deadlocks: Make sure channels are properly sized or closed
5. Missing defer: Always defer cancel() immediately after creating a cancellable context

Hints

1backoff := time.Duration(1<<uint(retryNumber)) * time.Second
2// retry 0: 1 << 0 = 1 second
3// retry 1: 1 << 1 = 2 seconds
4// retry 2: 1 << 2 = 4 seconds

1errs := make(chan error, len(fns))

Solution

 1package contextutils
 2
 3import (
 4    "context"
 5    "fmt"
 6    "sync"
 7    "time"
 8)
 9
10func WithRetry(ctx context.Context, maxRetries int, fn func(context.Context) error) error {
11    var lastErr error
12
13    for i := 0; i < maxRetries; i++ {
14        if err := fn(ctx); err == nil {
15            return nil
16        } else {
17            lastErr = err
18        }
19
20        if i < maxRetries-1 {
21            backoff := time.Duration(1<<uint(i)) * time.Second
22            select {
23            case <-time.After(backoff):
24            case <-ctx.Done():
25                return ctx.Err()
26            }
27        }
28    }
29
30    return fmt.Errorf("failed after %d retries: %w", maxRetries, lastErr)
31}
32
33func WithDeadline(deadline time.Time, fn func(context.Context) error) error {
34    ctx, cancel := context.WithDeadline(context.Background(), deadline)
35    defer cancel()
36    return fn(ctx)
37}
38
39func WithTimeout(timeout time.Duration, fn func(context.Context) error) error {
40    ctx, cancel := context.WithTimeout(context.Background(), timeout)
41    defer cancel()
42    return fn(ctx)
43}
44
45func Parallel(ctx context.Context, fns ...func(context.Context) error) error {
46    ctx, cancel := context.WithCancel(ctx)
47    defer cancel()
48
49    errs := make(chan error, len(fns))
50    var wg sync.WaitGroup
51
52    for _, fn := range fns {
53        wg.Add(1)
54        go func(f func(context.Context) error) {
55            defer wg.Done()
56            if err := f(ctx); err != nil {
57                cancel() // Cancel all on first error
58                errs <- err
59            }
60        }(fn)
61    }
62
63    wg.Wait()
64    close(errs)
65
66    if err := <-errs; err != nil {
67        return err
68    }
69
70    return nil
71}
72
73func First(ctx context.Context, fns ...func(context.Context) error) error {
74    ctx, cancel := context.WithCancel(ctx)
75    defer cancel()
76
77    done := make(chan error, len(fns))
78
79    for _, fn := range fns {
80        go func(f func(context.Context) error) {
81            done <- f(ctx)
82        }(fn)
83    }
84
85    for i := 0; i < len(fns); i++ {
86        if err := <-done; err == nil {
87            cancel() // Cancel remaining
88            return nil
89        }
90    }
91
92    return fmt.Errorf("all operations failed")
93}

Key Takeaways

• Context propagates cancellation and deadlines
• Always defer cancel() to prevent leaks
• Check ctx.Done() in long-running operations
• Use context.WithValue() sparingly
• Contexts are immutable and safe for concurrent use