Go: Context for Lifecycle Management

Table of Contents

Following the Learn Go with Tests guide, exploring Go’s context package, it’s managing request lifecycles, cancellation, and timeouts in concurrent applications.

Request Lifecycle Management
#

In server applications, you need to:

Cancel long-running operations when clients disconnect
Set timeouts for operations
Pass request-scoped values (user ID, trace ID)
Coordinate cancellation across goroutines
Handle graceful shutdowns

The context package provides a standardized way to handle these.

Context Fundamentals
#

Go Blog About Context

type Context interface {
    Deadline() (deadline time.Time, ok bool)
    Done() <-chan struct{}
    Err() error
    Value(key interface{}) interface{}
}

A Context carries a deadline, cancellation signal, and request-scoped values across API boundaries. Its methods are safe for simultaneous use by multiple goroutines.

Deadline(): Returns time When the context should be cancelled (if any)
Done(): Channel that closes when context is cancelled
Err(): Why the context was cancelled, after Done channel is closed
Value(key): Request-scoped data (nil if none)

GoDoc Context

HTTP Server with Timeout
#

HTTP server that respects context cancellation:

type Store interface {
    Fetch(ctx context.Context) (string, error)
}

// Returns a handler for calling store
func Server(store Store) http.HandlerFunc {
    return func(w http.ResponseWriter, r *http.Request) {
        // Passes Context to Fetch
        // Let's it handle it rather than handling it in Server
        data, err := store.Fetch(r.Context())
        if err != nil {
            return // TODO: log error reason
        }
        fmt.Fprint(w, data)
    }
}

Store implementation with context support:

func (s *Store) Fetch(ctx context.Context) (string, error) {
    data := make(chan string, 1)

    go func() {
        var result string
        for _, c := range s.response {
            select {
            case <-ctx.Done():
                log.Println("spy store got cancelled")
                return
            default:
                time.Sleep(10 * time.Millisecond) 
                // Simulate work
                result += string(c)
            }
        }
        data <- result 
        // Send completed result
    }()

    select {
    case <-ctx.Done():
        return "", ctx.Err() 
        // Return cancellation error
    case res := <-data:
        return res, nil 
        // Return completed work
    }
}

Testing Context Cancellation
#

Test 1: Normal Operation
#

func TestServer(t *testing.T) {
    data := "hello, world"

    t.Run("returns data from store", func(t *testing.T) {
        // Create a spy store that returns test data
        store := &SpyStore{response: data}
        svr := Server(store)

        // Create HTTP test request and response recorder
        request := httptest.NewRequest(http.MethodGet, "/", nil)
        response := httptest.NewRecorder()

        // Execute the handler
        svr.ServeHTTP(response, request)

        if response.Body.String() != data {
            t.Errorf(`got "%s", want "%s"`, response.Body.String(), data)
        }
    })
}

Uses httptest.NewRequest() to create a test HTTP request
Uses httptest.NewRecorder() to capture the response

Test 2: Cancellation Behavior
#

t.Run(
    "tells store to cancel work if request is cancelled", 
    func(t *testing.T
    ) {
    // Setup spy store and server
    store := &SpyStore{response: data}
    svr := Server(store)

    // Create base HTTP request
    request := httptest.NewRequest(http.MethodGet, "/", nil)

    // Create cancelling context from request's context
    cancellingCtx, cancel := context.WithCancel(request.Context())
    
    // Schedule cancellation after 5ms
    time.AfterFunc(5*time.Millisecond, cancel)
    
    // Replace request's context with the cancelling one
    request = request.WithContext(cancellingCtx)

    // Use spy response writer to track if response was written
    response := &SpyResponseWriter{}

    // Execute handler - should be cancelled before completion
    svr.ServeHTTP(response, request)

    if response.written {
        t.Error("a response should not have been written")
    }
})

Creates a cancelling context using context.WithCancel()
cancellingCtx (Context) - The new context that can be cancelled
- Gets passed to functions that need to respect cancellation
- When cancelled, its Done() channel closes
cancel (Function) - The cancellation trigger function
- When called, it cancels the cancellingCtx
- You call this function when you want to trigger cancellation
Uses time.AfterFunc() to trigger cancellation after 5ms
Replaces the request’s context with the cancelling context using WithContext()

Request Lifecycle Management #

Context Fundamentals #

HTTP Server with Timeout #

Testing Context Cancellation #

Test 1: Normal Operation #

Test 2: Cancellation Behavior #