# 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

```go
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:

```bash
go run main.go
```

and in another terminal:

```bash
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:

```bash
➜  ~ 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](https://github.com/gosom/gkomdev/blob/main/examples/http-handler-context/main.go) repo