Use context in your HTTP handlers

Let's consider a simple webserver that runs a task. The task can be anything like a time consuming computation, a database query.

You need to be able to do two things:

• when the client drops the connection while the task is running terminate the task so resources are released.

• when the task takes too much time return a proper HTTP status code to the client

Let's see a practical example

package main

import (
    "context"
    "errors"
    "fmt"
    "math/rand"
    "net/http"
    "time"
)

func main() {
    mux := http.NewServeMux()

    // Register the handler with the middleware, setting the timeout to 5 seconds
    mux.HandleFunc("/", longRunningTaskHandler)

    const timeout = time.Second * 5
    wrappedMux := contextMiddleware(timeout)(mux)

    // Start the server
    http.ListenAndServe(":8080", wrappedMux)
}

// contextMiddleware is a middleware that sets a timeout for the request context.
// If the request takes longer than the timeout, the middleware will cancel the context.
func contextMiddleware(timeout time.Duration) func(http.Handler) http.Handler {
    return func(next http.Handler) http.Handler {
        return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
            // Set a timeout for the request context
            ctx, cancel := context.WithTimeout(r.Context(), timeout)
            defer cancel()

            // Create a new request with the updated context
            r = r.WithContext(ctx)

            // Call the next handler with the new context
            next.ServeHTTP(w, r)
        })
    }
}

func longRunningTaskHandler(w http.ResponseWriter, r *http.Request) {
    err := longRunningTask(r.Context())
    if err != nil {
        // if the context expired, return a 504 Gateway Timeout
        if errors.Is(err, context.DeadlineExceeded) {
            w.WriteHeader(http.StatusGatewayTimeout)
            return
        }

        // if the task failed for some other reason, return a 500 Internal Server Error
        w.WriteHeader(http.StatusInternalServerError)

        return
    }

    // if the task completed successfully, return a 200 OK
    w.WriteHeader(http.StatusOK)
}

// longRunningTask is a dummy function that simulates a long running task.
// the task will take 10 seconds to complete.
// If the context is cancelled before the task completes, it will return the error.
func longRunningTask(ctx context.Context) error {
    var dur time.Duration

    // for ~50% of the time, the task will take 10 seconds to complete
    if rand.Float64() > 0.5 {
        dur = 10 * time.Second
    } else {
        dur = 1 * time.Second // and for the other ~50% of the time, the task will take 1 second to complete
    }

    // simulate the task by sleeping for the duration
    select {
    case <-time.After(dur):
        return nil
    case <-ctx.Done():
        fmt.Println("task cancelled")
        return ctx.Err()
    }
}

Run the webserver:

go run main.go

and in another terminal:

curl -i http://localhost:8080

Curl command will return at most in 5 seconds, even though the task can take up to 10 seconds.

See some runs:

➜  ~ curl -i http://localhost:8080/
HTTP/1.1 200 OK
Date: Wed, 12 Jun 2024 05:52:18 GMT
Content-Length: 0

➜  ~ curl -i http://localhost:8080/
HTTP/1.1 504 Gateway Timeout
Date: Wed, 12 Jun 2024 05:52:24 GMT
Content-Length: 0

➜  ~ curl -i http://localhost:8080/
HTTP/1.1 200 OK
Date: Wed, 12 Jun 2024 05:52:26 GMT
Content-Length: 0

As you notice in the second request the server returned 504.

Please run this in your computer and hit CTL-C when the request is running.
Notice the output of the webserver and you will see a message like:

task cancelled

See the related github repo