Context Package

Theory: Understanding Context

Context carries deadlines, cancellation signals, and request-scoped values across API boundaries and between processes. It's designed to enable cancellation propagation across goroutines.

Context Fundamentals

Creating and Using Context

// Context creation and basic usage
package main

import (
    "context"
    "fmt"
    "time"
)

func main() {
    // Root contexts (never cancel these)
    ctx := context.Background() // Main context for incoming requests
    // ctx := context.TODO()    // When unsure which context to use
    
    // Create cancellable context
    ctx, cancel := context.WithCancel(ctx)
    defer cancel() // ALWAYS call cancel to prevent leaks
    
    // Create context with timeout
    ctx, cancel = context.WithTimeout(ctx, 5*time.Second)
    defer cancel()
    
    // Create context with deadline
    deadline := time.Now().Add(10 * time.Second)
    ctx, cancel = context.WithDeadline(ctx, deadline)
    defer cancel()
    
    // Check context state
    select {
    case <-ctx.Done():
        fmt.Printf("Context cancelled: %v\n", ctx.Err())
    default:
        fmt.Println("Context is active")
    }
}

// Proper context usage in functions
func doWork(ctx context.Context) error {
    // Check if context is already cancelled
    select {
    case <-ctx.Done():
        return ctx.Err()
    default:
    }
    
    // Simulate work with periodic context checks
    for i := 0; i < 10; i++ {
        select {
        case <-ctx.Done():
            return ctx.Err()
        case <-time.After(100 * time.Millisecond):
            // Do actual work
            fmt.Printf("Working... %d\n", i)
        }
    }
    
    return nil
}

Context Cancellation Patterns

// Cascading cancellation
func parentOperation(ctx context.Context) error {
    // Create child context
    childCtx, cancel := context.WithTimeout(ctx, 3*time.Second)
    defer cancel()
    
    // Start multiple operations
    errCh := make(chan error, 3)
    
    go func() {
        errCh <- operation1(childCtx)
    }()
    
    go func() {
        errCh <- operation2(childCtx)
    }()
    
    go func() {
        errCh <- operation3(childCtx)
    }()
    
    // Wait for all or cancellation
    for i := 0; i < 3; i++ {
        select {
        case err := <-errCh:
            if err != nil {
                cancel() // Cancel other operations
                return err
            }
        case <-ctx.Done():
            return ctx.Err()
        }
    }
    
    return nil
}

// Graceful shutdown pattern
func server() {
    ctx, cancel := context.WithCancel(context.Background())
    
    // Handle shutdown signals
    sigCh := make(chan os.Signal, 1)
    signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
    
    go func() {
        <-sigCh
        fmt.Println("Shutdown signal received")
        cancel()
    }()
    
    // Start services
    var wg sync.WaitGroup
    
    wg.Add(1)
    go func() {
        defer wg.Done()
        runHTTPServer(ctx)
    }()
    
    wg.Add(1)
    go func() {
        defer wg.Done()
        runBackgroundWorker(ctx)
    }()
    
    wg.Wait()
    fmt.Println("Graceful shutdown complete")
}

Context Values and Request Scoping

Context Values Best Practices

// Define typed keys to avoid collisions
type contextKey string

const (
    requestIDKey contextKey = "requestID"
    userKey      contextKey = "user"
    traceKey     contextKey = "trace"
)

// Type-safe context value setters/getters
type User struct {
    ID    string
    Name  string
    Roles []string
}

func WithUser(ctx context.Context, user *User) context.Context {
    return context.WithValue(ctx, userKey, user)
}

func UserFromContext(ctx context.Context) (*User, bool) {
    user, ok := ctx.Value(userKey).(*User)
    return user, ok
}

func WithRequestID(ctx context.Context, requestID string) context.Context {
    return context.WithValue(ctx, requestIDKey, requestID)
}

func RequestIDFromContext(ctx context.Context) string {
    if requestID, ok := ctx.Value(requestIDKey).(string); ok {
        return requestID
    }
    return ""
}

// Structured logging with context
type Logger struct {
    // logger implementation
}

func (l *Logger) WithContext(ctx context.Context) *Logger {
    logger := &Logger{}
    
    if requestID := RequestIDFromContext(ctx); requestID != "" {
        // Add request ID to all logs
        logger = logger.With("request_id", requestID)
    }
    
    if user, ok := UserFromContext(ctx); ok {
        logger = logger.With("user_id", user.ID)
    }
    
    return logger
}

// Middleware pattern
func RequestIDMiddleware(next http.HandlerFunc) http.HandlerFunc {
    return func(w http.ResponseWriter, r *http.Request) {
        requestID := r.Header.Get("X-Request-ID")
        if requestID == "" {
            requestID = generateRequestID()
        }
        
        ctx := WithRequestID(r.Context(), requestID)
        w.Header().Set("X-Request-ID", requestID)
        
        next(w, r.WithContext(ctx))
    }
}

️ Database and HTTP with Context

Database Operations

import (
    "database/sql"
    _ "github.com/lib/pq"
)

type Repository struct {
    db *sql.DB
}

// Query with timeout
func (r *Repository) GetUser(ctx context.Context, id string) (*User, error) {
    // Add query timeout
    ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
    defer cancel()
    
    query := `SELECT id, name, email FROM users WHERE id = $1`
    
    var user User
    err := r.db.QueryRowContext(ctx, query, id).Scan(
        &user.ID, &user.Name, &user.Email,
    )
    
    if err == sql.ErrNoRows {
        return nil, ErrNotFound
    }
    if err != nil {
        return nil, fmt.Errorf("query failed: %w", err)
    }
    
    return &user, nil
}

// Transaction with context
func (r *Repository) Transfer(ctx context.Context, from, to string, amount float64) error {
    tx, err := r.db.BeginTx(ctx, &sql.TxOptions{
        Isolation: sql.LevelSerializable,
    })
    if err != nil {
        return err
    }
    defer tx.Rollback()
    
    // Debit from account
    _, err = tx.ExecContext(ctx,
        `UPDATE accounts SET balance = balance - $1 WHERE id = $2`,
        amount, from,
    )
    if err != nil {
        return err
    }
    
    // Credit to account
    _, err = tx.ExecContext(ctx,
        `UPDATE accounts SET balance = balance + $1 WHERE id = $2`,
        amount, to,
    )
    if err != nil {
        return err
    }
    
    return tx.Commit()
}

// Batch operations with cancellation
func (r *Repository) BatchInsert(ctx context.Context, users []User) error {
    for i, user := range users {
        // Check context before each operation
        select {
        case <-ctx.Done():
            return fmt.Errorf("batch cancelled at item %d: %w", i, ctx.Err())
        default:
        }
        
        _, err := r.db.ExecContext(ctx,
            `INSERT INTO users (id, name, email) VALUES ($1, $2, $3)`,
            user.ID, user.Name, user.Email,
        )
        if err != nil {
            return fmt.Errorf("insert failed at item %d: %w", i, err)
        }
    }
    return nil
}

HTTP Client with Context

// HTTP client with proper context handling
func callAPI(ctx context.Context, url string) ([]byte, error) {
    // Create request with context
    req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
    if err != nil {
        return nil, err
    }
    
    // Add request ID from context
    if reqID := RequestIDFromContext(ctx); reqID != "" {
        req.Header.Set("X-Request-ID", reqID)
    }
    
    // Custom client with timeout
    client := &http.Client{
        Timeout: 10 * time.Second,
    }
    
    resp, err := client.Do(req)
    if err != nil {
        return nil, err
    }
    defer resp.Body.Close()
    
    return io.ReadAll(resp.Body)
}

// Retry with exponential backoff
func callWithRetry(ctx context.Context, url string) ([]byte, error) {
    var lastErr error
    backoff := 100 * time.Millisecond
    
    for i := 0; i < 3; i++ {
        // Create timeout for individual attempt
        attemptCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
        data, err := callAPI(attemptCtx, url)
        cancel()
        
        if err == nil {
            return data, nil
        }
        
        lastErr = err
        
        // Check if context was cancelled
        if ctx.Err() != nil {
            return nil, ctx.Err()
        }
        
        // Wait with backoff
        select {
        case <-time.After(backoff):
            backoff *= 2
        case <-ctx.Done():
            return nil, ctx.Err()
        }
    }
    
    return nil, fmt.Errorf("all retries failed: %w", lastErr)
}

Best Practices and Patterns

Practice Exercises

Exercise 1: Request Tracer

Build a distributed tracing system using context to track requests across multiple services. Include timing, logging, and error tracking.

Exercise 2: Timeout Manager

Create a service that manages different timeout policies for various operations (database, API calls, background jobs).

Exercise 3: Graceful Service

Implement a web service with proper shutdown handling, ensuring all requests complete or timeout gracefully.

Exercise 4: Context Middleware

Build HTTP middleware that enriches context with authentication, request ID, and locale information.

Challenge: Circuit Breaker

Design a circuit breaker that uses context for timeout management and cancellation propagation across failing services.