Go provides a "context" package to help keep track of data and timeouts, which is especially useful on servers where, for example, you want a piece of data to be available from when a middleware runs all throughout the request, or to limit how long a handler is allowed to run.

In the following example, if the authorize function takes too long to run (according to the handlerTimeout constant at the top), the TimeoutHandler will kick in and "cancel" the context, making the actual handler (doStuff) not actually execute any of its queries and instead output this:

doStuff(): error during query: context deadline exceeded

package main

import (
    "context"
    "fmt"
    "github.com/gorilla/mux"
    "go.mongodb.org/mongo-driver/bson"
    "go.mongodb.org/mongo-driver/mongo"
    "go.mongodb.org/mongo-driver/mongo/options"
    "log"
    "net/http"
    "time"
)

const (
    handlerTimeout = 3 * time.Second
    mongoUri       = "mongodb://localhost:27017"
    mongoDbName    = "mydb"
)

func main() {
    r := mux.NewRouter()

    mgo, _ := mongo.NewClient(options.Client().ApplyURI(mongoUri))
    err := mgo.Connect(context.Background())
    if err != nil {
        panic(err.Error())
    }

    s := service{
        mgo.Database(mongoDbName),
    }

    withMiddleware := func(f http.HandlerFunc) http.HandlerFunc {
        var f2 http.HandlerFunc = func(w http.ResponseWriter, r *http.Request) {
            key := r.URL.Query().Get("key")
            userId, err := s.authorize(r.Context(), key)
            if err != nil {
                http.Error(w, "error during authorization: "+err.Error(), 500)
                return
            }

            // context.WithValue() "wraps" the existing request context and adds a single
            // piece of data associated to a key, in our case, a user ID.
            newCtx := context.WithValue(r.Context(), "userId", userId)

            f(w, r.WithContext(newCtx))
        }

        // http.TimeoutHandler, internally, calls context.WithTimeout(), wrapping the context
        // again and therefore adding the timeout functionality we want.
        return http.TimeoutHandler(f2, handlerTimeout, "").ServeHTTP
    }

    r.HandleFunc("/", withMiddleware(s.doStuff))

    srv := &http.Server{Handler: r}

    srv.ListenAndServe()
}

type service struct {
    db *mongo.Database
}

// authorize takes a string key, runs a DB query, and returns the corresponding user ID.
func (s *service) authorize(ctx context.Context, key string) (string, error) {
    var user struct {
        Key    string
        UserId string
    }

    filter := bson.M{"Key": key}

    err := s.db.Collection("users").FindOne(ctx, filter).Decode(&user)
    if err != nil {
        log.Printf("authorize(): error during query: %v", err)
    }

    return user.UserId, err
}

// doStuff outputs every piece of data found in the DB for the current user.
func (s *service) doStuff(w http.ResponseWriter, r *http.Request) {
    ctx := r.Context()
    userId := ctx.Value("userId")

    filter := bson.M{"UserId": userId}

    cur, err := s.db.Collection("data").Find(ctx, filter)
    if err != nil {
        log.Printf("doStuff(): error during query: %v", err)
        return
    }

    for cur.Next(ctx) {
        var data struct {
            Data string
        }

        err := cur.Decode(&data)
        if err != nil {
            log.Printf("doStuff(): error during decoding: %v", err.Error())
            continue
        }

        w.Write([]byte(fmt.Sprintf("data record: %v\n", data.Data)))
    }
}
Previous on Go
Mastodon Mastodon