Managing SQL query timeouts

So far in our code we’ve been using Go’s Exec() and QueryRow() methods to run our SQL queries. But Go also provides context-aware variants of these two methods: ExecContext() and QueryRowContext(). These variants accept a context.Context instance as the first parameter which you can leverage to terminate running database queries.

This feature can be useful when you have a SQL query that is taking longer to run than expected. When this happens, it suggests a problem — either with that particular query or your database or application more generally — and you probably want to cancel the query (in order to free up resources), log an error for further investigation, and return a 500 Internal Server Error response to the client.

In this chapter, we’ll update our application to do exactly that.

Mimicking a long-running query

To help demonstrate how this all works, let’s start by adapting our database model’s Get() method so that it mimics a long-running query. Specifically, we’ll update our SQL query to return a pg_sleep(8) value, which will make PostgreSQL sleep for 8 seconds before returning its result.

package data

...

func (m MovieModel) Get(id int64) (*Movie, error) {
    if id < 1 {
        return nil, ErrRecordNotFound
    }

    // Update the query to return pg_sleep(8) as the first value.
    query := `
        SELECT pg_sleep(8), id, created_at, title, year, runtime, genres, version
        FROM movies
        WHERE id = $1`

    var movie Movie

    // Importantly, update the Scan() arguments so that the pg_sleep(8) return value 
    // is scanned into a []byte slice.
    err := m.DB.QueryRow(query, id).Scan(
        &[]byte{}, // Add this line.
        &movie.ID,
        &movie.CreatedAt,
        &movie.Title,
        &movie.Year,
        &movie.Runtime,
        pq.Array(&movie.Genres),
        &movie.Version,
    )

    if err != nil {
        switch {
        case errors.Is(err, sql.ErrNoRows):
            return nil, ErrRecordNotFound
        default:
            return nil, err
        }
    }

    return &movie, nil
}

...

If you restart the application and make a request to the GET /v1/movies/:id endpoint, you should find that the request hangs for 8 seconds before you finally get a successful response displaying the movie information. Similar to this:

$ curl -w '\nTime: %{time_total}s \n' localhost:4000/v1/movies/1
{
    "movie": {
        "id": 1,
        "title": "Moana",
        "year": 2015,
        "runtime": "107 mins",
        "genres": [
            "animation",
            "adventure"
        ],
        "version": 1
    }
}

Time: 8.013534s

Adding a query timeout

Now that we’ve got some code that mimics a long-running query, let’s enforce a timeout so that the SQL query is automatically canceled if it doesn’t complete within 3 seconds.

To do this we need to:

1. Use the context.WithTimeout() function to create a context.Context instance with a 3-second timeout deadline.
2. Execute the SQL query using the QueryRowContext() method, passing the context.Context instance as a parameter.

I’ll demonstrate:

package data

import (
    "context" // New import
    "database/sql"
    "errors"
    "time"

    "greenlight.alexedwards.net/internal/validator"

    "github.com/lib/pq"
)

...

func (m MovieModel) Get(id int64) (*Movie, error) {
    if id < 1 {
        return nil, ErrRecordNotFound
    }

    query := `
        SELECT pg_sleep(8), id, created_at, title, year, runtime, genres, version
        FROM movies
        WHERE id = $1`

    var movie Movie

    // Use the context.WithTimeout() function to create a context.Context which carries a
    // 3-second timeout deadline. Note that we're using the empty context.Background() 
    // as the 'parent' context.
    ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)

    // Importantly, use defer to make sure that we cancel the context before the Get()
    // method returns.
    defer cancel()

    // Use the QueryRowContext() method to execute the query, passing in the context 
    // (including the deadline) as the first argument. 
    err := m.DB.QueryRowContext(ctx, query, id).Scan(
        &[]byte{},
        &movie.ID,
        &movie.CreatedAt,
        &movie.Title,
        &movie.Year,
        &movie.Runtime,
        pq.Array(&movie.Genres),
        &movie.Version,
    )

    if err != nil {
        switch {
        case errors.Is(err, sql.ErrNoRows):
            return nil, ErrRecordNotFound
        default:
            return nil, err
        }
    }

    return &movie, nil
}

...

There are a couple of things in the code above that I’d like to emphasize and explain:

• The defer cancel() line is necessary because it ensures that the resources associated with our context will always be released before the Get() method returns, thereby preventing a memory leak. Without it, the resources won’t be released until either the 3-second timeout is hit or the parent context (which in this specific example is context.Background()) is canceled.

• The timeout countdown begins from the moment that the context is created using context.WithTimeout(). Any time spent executing code between creating the context and calling QueryRowContext() will count towards the timeout.

OK, let’s try this out.

If you restart the application and make another request to the GET /v1/movies/:id endpoint, you should now get an error response similar to this after a 3-second delay:

$ curl -w '\nTime: %{time_total}s \n' localhost:4000/v1/movies/1
{
    "error": "the server encountered a problem and could not process your request"
}

Time: 3.025179s

If you go back to the terminal window which is running the application, you should also see a log line with the error message "pq: canceling statement due to user request". Like so:

$ go run ./cmd/api
time=2023-09-10T10:59:13.722+02:00 level=INFO msg="database connection pool established"
time=2023-09-10T10:59:13.722+02:00 level=INFO msg="starting server" addr=:4000 env=development
time=2023-09-10T10:59:16.722+02:00 level=ERROR msg="pq: canceling statement due to user request"

At first the wording of this error message might seem odd… until you learn that the message “canceling statement due to user request” is coming from PostgreSQL. In that light it makes sense: our application is the user of PostgreSQL, and we’re purposefully canceling the query after 3 seconds.

So, this is actually really good, and things are working as we would hope.

After 3 seconds, the context timeout is reached and our pq database driver sends a cancellation signal to PostgreSQL^†. PostgreSQL then terminates the running query, the corresponding resources are freed-up, and it returns the error message that we see above. The client is then sent a 500 Internal Server Error response, and the error is logged so that we know something has gone wrong.

Timeouts outside of PostgreSQL

There’s another important thing to point out here: it’s possible that the timeout deadline will be hit before the PostgreSQL query even starts.

You might remember that earlier in the book we configured our sql.DB connection pool to allow a maximum of 25 open connections. If all those connections are in-use, then any additional queries will be ‘queued’ by sql.DB until a connection becomes available. In this scenario — or any other which causes a delay — it’s possible that the timeout deadline will be hit before a free database connection even becomes available. If this happens then QueryRowContext() will return a context.DeadlineExceeded error.

In fact, we can demonstrate this in our application by setting the maximum open connections to 1 and making two concurrent requests to our endpoint. Go ahead and restart the API using a -db-max-open-conns=1 flag, like so:

$ go run ./cmd/api -db-max-open-conns=1
time=2023-09-10T10:59:13.722+02:00 level=INFO msg="database connection pool established"
time=2023-09-10T10:59:13.722+02:00 level=INFO msg="starting server" addr=:4000 env=development

Then in another terminal window make two requests to the GET /v1/movies/:id endpoint at the same time. At the moment that the 3-second timeout is reached, we should have one running SQL query and the other still ‘queued’ in the sql.DB connection pool. You should get two error responses which look like this:

$ curl localhost:4000/v1/movies/1 & curl localhost:4000/v1/movies/1 &
[1] 33221
[2] 33222
$ {
    "error": "the server encountered a problem and could not process your request"
}
{
    "error": "the server encountered a problem and could not process your request"
}

When you now head back to your original terminal, you should see two different error messages:

$ go run ./cmd/api -db-max-open-conns=1
time=2023-09-10T10:59:13.722+02:00 level=INFO msg="database connection pool established"
time=2023-09-10T10:59:13.722+02:00 level=INFO msg="starting server" addr=:4000 env=development
time=2023-09-10T10:59:27.722+02:00 level=ERROR msg="context deadline exceeded"
time=2023-09-10T10:59:27.722+02:00 level=ERROR msg="pq: canceling statement due to user request" addr=:4000 env=development

Here the pq: canceling statement due to user request error message relates to the running SQL query being terminated, whereas the context deadline exceeded message relates to the queued SQL query being canceled before a free database connection even became available.

In a similar vein, it’s also possible that the timeout deadline will be hit later on when the data returned from the query is being processed with Scan(). If this happens, then Scan() will also return a context.DeadlineExceeded error.

Updating our database model

Let’s quickly update our database model to use a 3-second timeout deadline for all our operations. While we’re at it, we’ll remove the pg_sleep(8) clause from our Get() method too.

package data

...

func (m MovieModel) Insert(movie *Movie) error {
    query := `
        INSERT INTO movies (title, year, runtime, genres) 
        VALUES ($1, $2, $3, $4)
        RETURNING id, created_at, version`

    args := []any{movie.Title, movie.Year, movie.Runtime, pq.Array(movie.Genres)}

    // Create a context with a 3-second timeout.
    ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
    defer cancel()

    // Use QueryRowContext() and pass the context as the first argument.
    return m.DB.QueryRowContext(ctx, query, args...).Scan(&movie.ID, &movie.CreatedAt, &movie.Version)
}

func (m MovieModel) Get(id int64) (*Movie, error) {
    if id < 1 {
        return nil, ErrRecordNotFound
    }

    // Remove the pg_sleep(8) clause.
    query := `
        SELECT id, created_at, title, year, runtime, genres, version
        FROM movies
        WHERE id = $1`

    var movie Movie

    ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
    defer cancel()

    // Remove &[]byte{} from the first Scan() destination.
    err := m.DB.QueryRowContext(ctx, query, id).Scan(
        &movie.ID,
        &movie.CreatedAt,
        &movie.Title,
        &movie.Year,
        &movie.Runtime,
        pq.Array(&movie.Genres),
        &movie.Version,
    )

    if err != nil {
        switch {
        case errors.Is(err, sql.ErrNoRows):
            return nil, ErrRecordNotFound
        default:
            return nil, err
        }
    }

    return &movie, nil
}

func (m MovieModel) Update(movie *Movie) error {
    query := `
        UPDATE movies 
        SET title = $1, year = $2, runtime = $3, genres = $4, version = version + 1
        WHERE id = $5 AND version = $6
        RETURNING version`

    args := []any{
        movie.Title,
        movie.Year,
        movie.Runtime,
        pq.Array(movie.Genres),
        movie.ID,
        movie.Version,
    }

    // Create a context with a 3-second timeout.
    ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
    defer cancel()

    // Use QueryRowContext() and pass the context as the first argument.
    err := m.DB.QueryRowContext(ctx, query, args...).Scan(&movie.Version)
    if err != nil {
        switch {
        case errors.Is(err, sql.ErrNoRows):
            return ErrEditConflict
        default:
            return err
        }
    }

    return nil
}

func (m MovieModel) Delete(id int64) error {
    if id < 1 {
        return ErrRecordNotFound
    }

    query := `
        DELETE FROM movies
        WHERE id = $1`

    // Create a context with a 3-second timeout.
    ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
    defer cancel()

    // Use ExecContext() and pass the context as the first argument.
    result, err := m.DB.ExecContext(ctx, query, id)
    if err != nil {
        return err
    }

    rowsAffected, err := result.RowsAffected()
    if err != nil {
        return err
    }

    if rowsAffected == 0 {
        return ErrRecordNotFound
    }

    return nil
}

Additional information

Using the request context

As an alternative to the pattern that we’ve used in the code above, we could create a context with a timeout in our handlers using the request context as the parent — and then pass that on to our database model.

But — and it’s a big but — doing this introduces a lot of behavioral complexity, and for most applications the benefits aren’t large enough to make the trade-off worthwhile.

The details behind this are very interesting, but also quite intricate and heavy-going. For that reason I’ve discussed it further in this appendix.