Context Propagation for Application-wide Concerns

Your Goal for Today

Learning Outcome: Use context propagation for request-scoped values, timeouts, and cancellation throughout the application.

As we build more complex applications, especially distributed ones like our target mini-Kafka, we face a new set of challenges. An incoming request might trigger a chain of operations: an API handler, a database query, maybe even a call to another service. What happens if the original client hangs up? Or if one step takes too long? How do we trace a single request's path through all these operations?

The context package is Go's standard solution for managing this entire lifecycle. Today, you will learn how to use it to control timeouts, propagate cancellation signals, and pass request-scoped data like trace IDs through your application's call stack.

1. What is context.Context?

At its core, context.Context is an interface that carries deadlines, cancellation signals, and other request-scoped values across API boundaries and between processes. It is an indispensable tool for any concurrent or networked Go program.

A context.Context provides four main capabilities:

1. Cancellation: It can signal to functions in its call chain that the operation should be abandoned.
2. Timeouts and Deadlines: It can automatically cancel an operation after a specific duration or at a set time.
3. Request-Scoped Values: It provides a way to pass data like tracing IDs or user authentication information down a request's call stack.
4. Error Information: It reports why it was canceled (e.g., context.DeadlineExceeded).

For a concise introduction to these concepts, the following video is an excellent starting point.

2. The Context Tree: Derivation and Propagation

A fundamental concept to grasp is that contexts are immutable and form a tree-like structure. You never modify a context; instead, you create a new child context derived from a parent.

• The Root: The chain always starts with a root context, typically context.Background(), which is an empty context with no values, deadlines, or cancellation capabilities.
• Derivation: You create new contexts using functions like context.WithCancel, context.WithTimeout, and context.WithValue. Each of these functions takes a parent context and returns a new child context and, in some cases, a cancel function.
• Propagation: Cancellation signals flow down the tree from parent to child. If a parent context is canceled, all its descendants are also canceled. However, values are looked up by traversing up the tree from child to parent.

This diagram illustrates how new contexts are derived to form a hierarchy.

The video you watched earlier has a great segment explaining this immutability and the tree structure.

3. Applying Context: Timeouts and Cancellation

The most common use case for context is managing the lifetime of an operation. In our REST API, we want to ensure that no request can hang indefinitely. We can achieve this by creating a context with a timeout at the beginning of our HTTP handler and passing it down to our storage layer.

How it Works:

1. In the Handler: An incoming http.Request already contains a context (r.Context()). We derive a new context from it with a timeout.
2. defer cancel(): When you create a cancellable context (with WithCancel, WithTimeout, or WithDeadline), you also get a cancel function. It is critical to call this function when the operation is complete to release the resources associated with the context. The defer statement is the perfect tool for this.
3. Propagation: We pass this new ctx as the first argument to any downstream functions (e.g., storage.Get(ctx, ...)).
4. Listening for Cancellation: Downstream functions can then listen for the context's cancellation signal in two ways:
   • Implicitly: By passing the context to library functions that are context-aware. For example, the pgx driver's query methods (QueryRowContext, ExecContext, etc.) will automatically abort a database query if the context is canceled.
   • Explicitly: In long-running goroutines, you can use a select statement to listen on both your work channels and the context's Done() channel: case <-ctx.Done():.

The following article provides a superb, production-grade example of this entire flow within an HTTP middleware.

4. Applying Context: Request-Scoped Values

Context is also the idiomatic way to pass request-scoped metadata, such as tracing IDs, through your application. This avoids cluttering function signatures with extra parameters and ensures the data is available wherever it's needed for logging or instrumentation.

Best Practices for WithValue:

• Use it sparingly: context.WithValue is not for passing optional function parameters or core business data. It's for cross-cutting concerns like tracing, authentication info, etc.
• Use custom key types: To avoid key collisions between different packages, never use built-in types like string for context keys. Define your own unexported key type.

// Define an unexported key type.
type key string

// Define an unexported key for our request ID.
const requestIDKey key = "requestID"

// In a middleware or handler:
requestID := "some-unique-id"
ctx := context.WithValue(r.Context(), requestIDKey, requestID)

// In a downstream function:
if reqID, ok := ctx.Value(requestIDKey).(string); ok {
    // Use the request ID for logging...
}

This pattern is very common in HTTP middleware, where you can "enrich" the context for all subsequent handlers.

This video provides a great walkthrough of using WithValue correctly, including the custom key pattern.

5. Project Integration: Making Our API Context-Aware

Now, let's apply these patterns to our REST API project.

Step 1: Update the Storage Interface

First, we must update our core abstraction. By convention, functions that should be cancellable or accept request-scoped data should accept context.Context as their first parameter.

Modify your internal/storage/storage.go file:

package storage

import "context"

type Storage interface {
    // Add context.Context to all methods
    Get(ctx context.Context, id string) (*Item, error)
    Create(ctx context.Context, item *Item) error
    // ... update other methods as well
}

Step 2: Update the Implementations (PostgresStorage and InMemoryStorage)

Now, update both PostgresStorage and InMemoryStorage to conform to the new interface signature.

For PostgresStorage, you can now use the context-aware methods from pgx. This gives you timeout/cancellation handling for your database queries for free.

In internal/storage/postgres.go:

// In your PostgresStorage Get method
func (p *PostgresStorage) Get(ctx context.Context, id string) (*Item, error) {
    // ... your query and item setup
    
    // Use QueryRow, which is an alias for QueryRowContext. It respects the context.
    err := p.db.QueryRow(ctx, query, id).Scan(...)
    if err != nil {
        if errors.Is(err, pgx.ErrNoRows) {
            return nil, ErrNotFound
        }
        // Check if the error was caused by the context being canceled.
        if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
             return nil, err
        }
        return nil, fmt.Errorf("failed to get item: %w", err)
    }
    return &item, nil
}

When the context passed to p.db.QueryRow is canceled, the pgx driver will abort the query and return an error. By checking errors.Is(err, context.DeadlineExceeded), your repository can now correctly report that a timeout occurred.

Step 3: Update the API Handlers

Finally, wire everything together in your API handlers. Here, you'll create the context with a timeout and pass it down to the storage layer.

// In your API handler for getting an item
func(s *Server) handleGetItem(w http.ResponseWriter, r *http.Request) {
    // Create a new context with a 3-second timeout.
    ctx, cancel := context.WithTimeout(r.Context(), 3*time.Second)
    defer cancel() // IMPORTANT: ensure cancellation function is called.

    // Get item ID from request...
    id := chi.URLParam(r, "id")

    // Pass the derived context to the storage layer.
    item, err := s.storage.Get(ctx, id)
    if err != nil {
        if errors.Is(err, storage.ErrNotFound) {
            // handle 404
            return
        }
        // Check if the error was a timeout from our context.
        if errors.Is(err, context.DeadlineExceeded) {
            // handle 504 Gateway Timeout
            log.Println("Request timed out")
            return
        }
        // handle other 500 errors
        return
    }

    // handle success
}

With this structure, if s.storage.Get takes longer than 3 seconds, context.WithTimeout will cancel the ctx, pgx will abort the query, and your handler will correctly identify and handle the timeout error.

Conclusion

You have now integrated one of the most important packages in the Go standard library into our application. By propagating context, you've made our service more robust, resilient, and observable.

Key Takeaways:

• context.Context is the standard mechanism for propagating cancellation, deadlines, and request-scoped values.
• Contexts are immutable and form a tree, with cancellation flowing downwards.
• Always accept context.Context as the first argument to functions that perform I/O or other long-running work.
• When creating a context with WithCancel, WithTimeout, or WithDeadline, always defer the call to the returned cancel function.
• Use context.WithValue sparingly, primarily for request-scoped metadata like trace IDs, and always use a custom key type.

Next Up

Now that our application understands how to handle cancellation signals via context, we are perfectly positioned for the next lesson. We will implement graceful shutdown, allowing our server to finish processing in-flight requests before shutting down when it receives an OS signal (like SIGINT from Ctrl+C).