13 - Gin Handler Timeout Middleware

We just added the ability to store refresh tokens in Redis. Today, let's take a break from the main application features to learn how to create our first Gin middleware!

If at any point you are confused about file structure or code, go to Github repository and check out the branch for the previous lesson to be in sync with me!

I also wanted to mentioned that I made two updates in the main branch of the repository for more easily getting the application running for newcomers to this project. The steps to run the project have been updated in the README, as well.

First, I decided to add the .env.dev file to the repository. For development, I don't plan to include any API secrets directly in this file. Later on, however, we'll reference a Google Cloud Configuration file in .env.dev, so you'll need to make sure to keep that configuration file secret as always.

Second, I've added a make rule, make init, which will start up our Postgres containers and run database migrations so that you'll be ready to start creating data (users for now) with the application!

If you prefer video, check out the video version below!

Handler Timeout Middleware

Here's the ol' diagram of what we'll be working on today.

If we look at the right-hand side of our middleware, we see our application layers. The handler layer methods all receive a *gin.Context. This context is provided to us by Gin and gives us access to numerous utility methods that make working with HTTP requests and responses easier.

The problem we're trying to solve is if a call to these handler methods (including all of the downstream calls to the service/repository layers and data sources) takes a long time, we want to be able to terminate the handler and send an error to the user.

We can do this with built-in functionality from the context package of Go, by wrapping the request context using the context.WithTimeout method. When we wrap our context with this method, we get access to the context.Done() method which "returns a channel that's closed when work done on behalf of this context should be canceled" (source). After using WithTimeout, the context will be cancelled for us automatically after the duration of time we pass to the method.

We'll save discussing the timeoutWriter for when we write the middleware code.

What About http.Server Timeout Fields?

If we look in ~/main.go, we initialize our server with Go's http.Server.

  srv := &http.Server{
        Addr:    ":8080",
        Handler: router,
  }

You might be asking why we don't just make use of the net/http package's Server fields like ReadTimeout or WriteTimeout. These timeouts are concerned with the time to read and write the body to and from the client (i.e. networking I/O) and don't deal with the case of long running-handlers. There's a fine answer on StackOverflow you might want to read.

Middleware In Gin

Custom Middleware Function

Looking back at the diagram, remember that our goal is to manipulate the request context, which we extract from the gin.Context using the Request() method. I'll reiterate that the Gin and request context are two different contexts, which can be confusing.

To create a custom middleware in Gin, we create a function that returns a gin.HandlerFunc, which is defined as:

type HandlerFunc func(*Context)

In middleware HandlerFunc, we can modify this gin.Context according to our needs. In our handler timeout middleware, we will wrap our the request context with a timeout and handle concurrency issues. Another simple example is a logger, as shown in the Gin docs. Yet another case you might see is extracting a user from a session or token, and setting this user as a key on gin.Context. In fact, we'll do just that in a handful of tutorials!

Usage in Gin

I recommend checking out the Using Middleware section of the Gin docs to see how you can call a middleware function for routes, groups, and individual routes! We'll end up applying middleware to our existing route group in ~/handler/handler.go with the Use method.

  // Create a group, or base url for all routes
    g := c.R.Group(c.BaseURL)

    if gin.Mode() != gin.TestMode {
        g.Use() // middleware will be called in Use()
    }

Why Make Our Own Middleware?

I found the following open source middlewares for setting a handler timeout.

The first is the Go http.TimeoutHandler. The downside to this handler is that it's not configured to directly work with Gin routes and groups. Another thing I didn't like about the handler is that it is configured to send an HTML response in the case of a timeout, and doesn't appear to have a way to send a different Content-Type for a response.

The second middleware I found was gin-contrib/timeout. Note that this repository of middlewares is not the official one by gin-gonic, gin-gonic/contrib, even though it has a similar name. This middleware had issues trying to rewrite headers to a response writer, which we'll discuss when we write our middleware.

The third middleware I found was vearne/gin-timeout. This middleware actually addressed the above issues, and so I have no problem recommending it. However, it did not provide a way to modify the error response sent to the user.

Since getting a timeout handler middleware to work as I would like would require modifying any of the available handlers, I decided it would just be easier to create our own, even if I'm borrowing heavily from http.TimeoutHandler and vearne/gin-timeout.

Middleware Prep

Add A Pretend Delay to Signin Handler

In ~/handler/handler.go, let's fake a handler that takes a long time to complete. We'll use the Signin handler method as it currently doesn't yet have a "real" implementation. We'll set the Sleep duration to 6 seconds, as we'll end up setting our timeout limit to 5 seconds.

// Signin handler
func (h *Handler) Signin(c *gin.Context) {
    time.Sleep(6 * time.Second) // to demonstrate a timeout
    c.JSON(http.StatusOK, gin.H{
        "hello": "it's signin",
    })
}

Update App Errors

I want to add a ServiceUnavailable (HTTP 503) to our ~/model/apperrors/apperrors.go to send when our handlers run longer than 5 seconds.

// update const
const (
    // ... other types excluded for brevity
    ServiceUnavailable   Type = "SERVICE_UNAVAILABLE"    // For long running handlers
)

// update status method switch case to handle ServiceUnavailable

func (e *Error) Status() int {
    switch e.Type {
    // other cases and default excluded for brevity
    case ServiceUnavailable:
        return http.StatusServiceUnavailable
    // cases omitted
    }
}

// add an error factory
// NewServiceUnavailable to create an error for 503
func NewServiceUnavailable() *Error {
    return &Error{
        Type:    ServiceUnavailable,
        Message: fmt.Sprintf("Service unavailable or timed out"),
    }
}

Timeout Duration Environment Variable

Let's make timeout duration configurable. We'll set a handler timeout of 5 seconds for now. In reality, we could set different timeout lengths for different handlers or groups of handlers. As an example, you would probably expect the time to upload a file to a cloud service to take longer than a handler that merely updates text fields in the database. I encourage you to do this in your applications.

In .env.go add the following.

HANDLER_TIMEOUT=5

We now need to update the instantiation of our handler via handler.Config to accept this field as a duration. In ~/handler/handler.go.

// Config will hold services that will eventually be injected into this
// handler layer on handler initialization
type Config struct {
    R               *gin.Engine
    UserService     model.UserService
    TokenService    model.TokenService
    BaseURL         string
    TimeoutDuration time.Duration
}

Let's also load the environment variable and pass it to our handler config in ~/injection.go of package main.

  // read in HANDLER_TIMEOUT
    handlerTimeout := os.Getenv("HANDLER_TIMEOUT")
    ht, err := strconv.ParseInt(handlerTimeout, 0, 64)
    if err != nil {
        return nil, fmt.Errorf("could not parse HANDLER_TIMEOUT as int: %w", err)
    }

    handler.NewHandler(&handler.Config{
        R:               router,
        UserService:     userService,
        TokenService:    tokenService,
        BaseURL:         baseURL,
        TimeoutDuration: time.Duration(time.Duration(ht) * time.Second),
    })

We now have a configurable handler timeout duration! Woot woot!

Timeout Middleware Code

After quite a preface, let's add the middleware code. Our middleware will be added to a sub-package of handler. Let's add a folder and file at the path ~/handler/middleware/timeout.go. We'll go step-by-step over the code we add to this middleware function.

Middleware Function Definition

Our middleware will accept a timeout duration, which will be received from handler.Config, and an error that we'll send to the user in case of a timeout. We'll end up passing the custom ServiceUnavailable error we recently created.

package middleware

// Imports omitted

func Timeout(timeout time.Duration, errTimeout *apperrors.Error) gin.HandlerFunc {
    return func(c *gin.Context) {
    // Code will go here
  }
}

Set Gin Writer to Custom Timeout Writer

We first set the writer on Gin's context, c, to a custom writer, tw. You'll also observe that this writer is passed as a field to our custom timeoutWriter. This is a little confusing, but will hopefully become clearer after we go over the timeoutWriter definition.

    // set Gin's writer as our custom writer
        tw := &timeoutWriter{ResponseWriter: c.Writer, h: make(http.Header)}
        c.Writer = tw

Wrap Context in Timeout

In order to apply a timeout duration, we can use context.WithTimeout as follows. We defer the cancel function's execution until the end of the function. When cancel is called, all resources associated with it context, ctx, will be freed up.

    // wrap the request context with a timeout
        ctx, cancel := context.WithTimeout(c.Request.Context(), timeout)
        defer cancel()

        // update gin request context
        c.Request = c.Request.WithContext(ctx)

Channels For Finished of Panic "Signals"

We create channels that we'll send values into if our handler successfully completes or if the application panics.

    finished := make(chan struct{})        // to indicate handler finished
        panicChan := make(chan interface{}, 1) // used to handle panics if we can't recover

Call Handler Inside of a GoRoutine

We then create and invoke a goroutine which calls c.Next(). Next is a helper method which calls the next middleware or handler function (you can "cascade" handler functions and middlewares). As of yet, we only have a single middleware, so c.Next() will call a handler method. If the method completes, we pass an empty struct into the finished channel.

We also defer a call to recover from a panic and send the captured value from the panic and send it to the panicChan (it turns out we don't do anything with this value, so send and empty struct, if you prefer).

    go func() {
            defer func() {
                if p := recover(); p != nil {
                    panicChan <- p
                }
            }()

            c.Next() // calls subsequent middleware(s) and handler
            finished <- struct{}{}
        }()

Handling Either Finished, Panic, or Timeout/context.Done()

After invoking the goroutine, we then create a select statement that listens for incoming values on one of three channels. In all of these cases, we then write a header and body to a response writer embedded inside of our custom timeoutWriter in variable tw.

This may be a little confusing, writing to a Writer inside of a Writer, but this is how we prevent overwriting response bodies and headers. We basically use the timeoutWriter as a temporary/proxy writer that also stores the state of whether or not we have either timed out or already sent the response from a handler.

    select {
        case <-panicChan:
            // if we cannot recover from panic,
            // send internal server error
            e := apperrors.NewInternal()
            tw.ResponseWriter.WriteHeader(e.Status())
            eResp, _ := json.Marshal(gin.H{
                "error": e,
            })
            tw.ResponseWriter.Write(eResp)
        case <-finished:
            // if finished, set headers and write resp
            tw.mu.Lock()
            defer tw.mu.Unlock()
            // map Headers from tw.Header() (written to by gin)
            // to tw.ResponseWriter for response
            dst := tw.ResponseWriter.Header()
            for k, vv := range tw.Header() {
                dst[k] = vv
            }
            tw.ResponseWriter.WriteHeader(tw.code)
            // tw.wbuf will have been written to already when gin writes to tw.Write()
            tw.ResponseWriter.Write(tw.wbuf.Bytes())
        case <-ctx.Done():
            // timeout has occurred, send errTimeout and write headers
            tw.mu.Lock()
            defer tw.mu.Unlock()
            // ResponseWriter from gin
            tw.ResponseWriter.Header().Set("Content-Type", "application/json")
            tw.ResponseWriter.WriteHeader(errTimeout.Status())
            eResp, _ := json.Marshal(gin.H{
                "error": errTimeout,
            })
            tw.ResponseWriter.Write(eResp)
            c.Abort()
            tw.SetTimedOut()
        }

In the case of a panic, we simply write an internal server error with the appropriate status code.

If we receive on the finished channel, this means that our custom timeoutWriter will have the headers written by Gin to its h field, which will store headers. We then copy these headers to the inner, gin.ResponseWriter, referenced by dst. We then send the status code from tw, and the actual response body from the handler which is stored in tw.wbuf. We make sure to lock tw using it's mutex field to prevent writes in the case we "simultaneously" receive a on another channel.

In the case when the context times out, we receive on the channel returned by ctx.Done(). We also create a response with the *apperrors.Error passed to the middleware as a parameter. We then call c.Abort() which prevents any subsequent handlers or middlewares from being called, and set tw's timedOut field to true so that if the handler response eventually comes through, it will not be written to the client.

Let's now discuss the timeoutWriter in detail.

Custom Timeout Writer Explanation

The code for timeoutWriter is included below.

// implements http.Writer, but tracks if Writer has timed out
// or has already written its header to prevent
// header and body overwrites
// also locks access to this writer to prevent race conditions
// holds the gin.ResponseWriter which we'll manually call Write()
// on in the middleware function to send response
type timeoutWriter struct {
    gin.ResponseWriter
    h    http.Header
    wbuf bytes.Buffer // The zero value for Buffer is an empty buffer ready to use.

    mu          sync.Mutex
    timedOut    bool
    wroteHeader bool
    code        int
}

// Writes the response, but first makes sure there
// hasn't already been a timeout
// In http.ResponseWriter interface
func (tw *timeoutWriter) Write(b []byte) (int, error) {
    tw.mu.Lock()
    defer tw.mu.Unlock()
    if tw.timedOut {
        return 0, nil
    }

    return tw.wbuf.Write(b)
}

// In http.ResponseWriter interface
func (tw *timeoutWriter) WriteHeader(code int) {
    checkWriteHeaderCode(code)
    tw.mu.Lock()
    defer tw.mu.Unlock()
    // We do not write the header if we've timed out or written the header
    if tw.timedOut || tw.wroteHeader {
        return
    }
    tw.writeHeader(code)
}

// set that the header has been written
func (tw *timeoutWriter) writeHeader(code int) {
    tw.wroteHeader = true
    tw.code = code
}

// Header "relays" the header, h, set in struct
// In http.ResponseWriter interface
func (tw *timeoutWriter) Header() http.Header {
    return tw.h
}

// SetTimeOut sets timedOut field to true
func (tw *timeoutWriter) SetTimedOut() {
    tw.timedOut = true
}

func checkWriteHeaderCode(code int) {
    if code < 100 || code > 999 {
        panic(fmt.Sprintf("invalid WriteHeader code %v", code))
    }
}

This writer implements an http.ResponseWriter with the required Write, WriteHeader, and Header methods. In the middleware, we replaced the writer on *gin.Context with the custom timeoutWriter. Therefore, when the Gin handlers write a response, they will end up calling timeoutWriter.Write, timeoutWriter.WriterHeader and timeoutWriter.Header.

Let's take a look at what would happen if we use this middleware and send a response from the Signin handler.

  c.JSON(http.StatusOK, gin.H{
        "hello": "it's signin",
  })

When Gin calls the JSON method it will set the response body by calling timeoutWriter.Write with the JSON body inside of the gin.H map. If we look at timeoutWriter.Write, you can see that it checks that we haven't already set the timedOut field. This is how we prevent overwriting our response body! If we haven't timed out, we write the body to our custom field, wbuf. This wbuf temporarily stores the valid response which gets "relayed" to the client in the finished channel branch of the switch case.

The c.JSON method will also call timeoutWriter.Header() to get access to our writer's headers, timeoutWriter.h. Gin then sets headers such as Content-Type and Content-Length.

Finally, c.JSON will also set the HTTP status via the timeoutWriter.WriteHeader() method. Our custom WriteHeader() method will check to see if the header has already been written or if we've already timed out via the fields timedOut and wroteHeader. If we've already timed out or written a header, we do nothing. Otherwise, we set wroteHeader and the HTTP status code to our struct.

We also have a mu sync.Mutex in our custom writer. You can see that anytime we write to our timeoutWriter, we call tw.mu.Lock(). When we add a sync.Mutex to a struct, this allows us to lock struct to updates by concurrent routines which may want to access our timeoutWriter data (sorry if my terminology is off). Once all of the updates to our timeoutWriter data have been made, we Unlock access to the fields/data.

Apply Middleware

In ~/handler/handler.go we'll use our middleware immediately after creating our route group.

// Create a group, or base url for all routes
    g := c.R.Group(c.BaseURL)

    if gin.Mode() != gin.TestMode {
        g.Use(middleware.Timeout(c.TimeoutDuration, apperrors.NewServiceUnavailable()))
    }

You should now be able to run your app with docker-compose up and send a request to the "signin" endpoint as follows (with Curl or other client):

➜ curl --location --request POST 'http://malcorp.test/api/account/signin' \
--header 'Authorization: Bearer {{idToken}}' \
--header 'Content-Type: application/json' \
--data-raw '{}'
{"error":{"type":"SERVICE_UNAVAILABLE","message":"Service unavailable or timed out process"}}% 

If you then set the timeout in Signin to something less than 5 seconds (our application's handler timeout duration), you should receive a status 200 response with JSON.

// Signin handler
func (h *Handler) Signin(c *gin.Context) {
    time.Sleep(1 * time.Second) // to demonstrate normal use
    c.JSON(http.StatusOK, gin.H{
        "hello": "it's signin",
    })
}

Rerunning the same curl command we receive:

{"hello":"it's signin"}%   

Conclusion

That was actually a relatively complex middleware. If you're like me when I was learning from other's work, you may have to stare at it for a while until it makes sense. Yet I hope I've laid it out in a manner that saves you from a few minutes of confusion! Heck, I'm not sure I did everything right, so let me know if you think there are any blaring mistakes!

Next time, we'll write and test the Signin handler. Following that, we'll add service and repository method implementations to complete the full signin functionality.

¡Hasta pronto!