Mastering Contexts in Go

In Go, Contexts provide a standard way to pass metadata and control signals between goroutines. They are mainly used to manage task execution time, data passing, and operation cancellation. This article covers different types of contexts in Go and examples of how to use them.

Introduction to Contexts

Contexts in Go are represented by the context.Context interface, which includes methods for getting deadlines, cancellation, values, and done channels. The primary package for working with contexts is context.

package context

type Context interface {
    Deadline() (deadline time.Time, ok bool)
    Done() <-chan struct{}
    Err() error
    Value(key interface{}) interface{}
}

Context types

There are four main functions to create contexts:

1. context.Background(): Returns an empty context. It is usually used as the root context for the entire application.
2. context.TODO(): Returns a context that can be used when a context is required but not yet defined. It signals that the context needs further work.
3. context.WithCancel(parent Context): Returns a derived context that can be canceled by calling the cancel function.
4. context.WithDeadline(parent Context, d time.Time): Returns a derived context that automatically cancels at a specified time (deadline).
5. context.WithTimeout(parent Context, timeout time.Duration): Similar to the WithDeadline, but the deadline is set by a duration.
6. context.WithValue(parent Context, key, val interface{}): Returns a derived context that contains a key-value pair.

Contexts in Action

Context with Cancellation

A context with cancellation is useful when you need to stop a goroutine based on an event.

package main

import (
    "context"
    "fmt"
    "time"
)

func main() {
    ctx, cancel := context.WithCancel(context.Background())
    go func() {
        select {
        case <-time.After(2 * time.Second):
            fmt.Println("Operation completed")
        case <-ctx.Done():
            fmt.Println("Operation canceled")
        }
    }()
    
    // try to change this value to 3 and execute again
    time.Sleep(1 * time.Second)
    cancel()
    time.Sleep(2 * time.Second)
}

Context with Timeout

This context automatically cancels after a specified duration.

package main

import (
    "context"
    "fmt"
    "time"
)

func main() {
    ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
    defer cancel()
    
    go func() {
        select {
        case <-time.After(3 * time.Second):
            fmt.Println("Operation completed")
        case <-ctx.Done():
            fmt.Println("Operation timed out")
        }
    }()
    
    // try to change this value to 2 and execute again
    time.Sleep(4 * time.Second)
}

Context with Deadline

A context with a deadline is similar to a context with a timeout, but the time is set as a specific value.

package main

import (
    "context"
    "fmt"
    "time"
)

func main() {
    // try to change this value to 3 and execute again
    deadline := time.Now().Add(2 * time.Second)
    ctx, cancel := context.WithDeadline(context.Background(), deadline)
    defer cancel()
    
    go func() {
        select {
        case <-time.After(3 * time.Second):
            fmt.Println("Operation completed")
        case <-ctx.Done():
            fmt.Println("Operation reached deadline")
        }
    }()
    
    time.Sleep(4 * time.Second)
}

Context with Values

Contexts can store arbitrary data as key-value pairs. This is useful for passing parameters and settings to handlers.

package main

import (
	"context"
	"fmt"
	"time"
)

func main() {
	ctx := context.WithValue(context.Background(), "key", "value")

	go func(ctx context.Context) {
		if v := ctx.Value("key"); v != nil {
			fmt.Println("Value found:", v)
		} else {
			fmt.Println("No value found")
		}
	}(ctx)

	time.Sleep(1 * time.Second)
}

Applying Contexts

Contexts are widely used in various parts of Go applications, including network servers, databases, and client requests. They help properly manage task execution time, cancel unnecessary operations, and pass data between goroutines.

Using in HTTP Servers

package main

import (
    "context"
    "fmt"
    "net/http"
    "time"
)

func handler(w http.ResponseWriter, r *http.Request) {
    ctx := r.Context()
    select {
    case <-time.After(5 * time.Second):
        fmt.Fprintf(w, "Request processed")
    case <-ctx.Done():
        fmt.Fprintf(w, "Request canceled")
    }
}

func main() {
    http.HandleFunc("/", handler)
    http.ListenAndServe(":8080", nil)
}

This code sets up an HTTP server that handles requests with a context-aware handler. It either completes after 5 seconds or responds if the request is canceled.

Using in Databases

package main

import (
    "context"
    "database/sql"
    "fmt"
    "time"

    _ "github.com/go-sql-driver/mysql"
)

func queryDatabase(ctx context.Context, db *sql.DB) {
    query := "SELECT sleep(5)"
    rows, err := db.QueryContext(ctx, query)
    if err != nil {
        fmt.Println("Query error:", err)
        return
    }
    defer rows.Close()

    for rows.Next() {
        var result string
        if err := rows.Scan(&result); err != nil {
            fmt.Println("Scan error:", err)
            return
        }
        fmt.Println("Result:", result)
    }
}

func main() {
    db, err := sql.Open("mysql", "user:password@tcp(localhost:3306)/dbname")
    if err != nil {
        fmt.Println("Database connection error:", err)
        return
    }
    defer db.Close()

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

Here we connect to a MySQL database and execute a query with a context timeout of 3 seconds. If the query takes longer, it is canceled, and an error message is printed.

Using in Goroutines

package main

import (
    "context"
    "fmt"
    "time"
)

func worker(ctx context.Context, id int) {
    for {
        select {
        case <-ctx.Done():
            fmt.Printf("Worker %d stopped\n", id)
            return
        case <-time.After(1 * time.Second):
            fmt.Printf("Worker %d working\n", id)
        }
    }
}

func main() {
    ctx, cancel := context.WithCancel(context.Background())

    for i := 1; i <= 3; i++ {
        go worker(ctx, i)
    }

    time.Sleep(3 * time.Second)
    cancel()
    time.Sleep(1 * time.Second)
}

In this example, the code spawns three worker goroutines that print status messages every second. The workers stop when the main function cancels the context after 3 seconds.

Using in an API request with a deadline

package main

import (
    "context"
    "fmt"
    "net/http"
    "time"
)

func fetchAPI(ctx context.Context, url string) {
    req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
    if err != nil {
        fmt.Println("Request creation error:", err)
        return
    }

    client := &http.Client{}
    resp, err := client.Do(req)
    if err != nil {
        fmt.Println("Request error:", err)
        return
    }
    defer resp.Body.Close()

    if resp.StatusCode == http.StatusOK {
        fmt.Println("API request succeeded")
    } else {
        fmt.Println("API request failed with status:", resp.StatusCode)
    }
}

func main() {
    ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(2*time.Second))
    defer cancel()

    fetchAPI(ctx, "http://example.com/api")
}

This example demonstrates making an API request with a 2-second deadline. If the request is not completed within this timeframe, it is canceled, ensuring that the program does not wait indefinitely.

Conclusion

Contexts in Go are a powerful tool for managing execution time, cancellation, and data passing between goroutines. Using contexts correctly helps avoid resource leaks, ensures timely task completion, and improves code structure and readability. Various types of contexts, such as those with cancellation, timeout, deadline, and values, provide flexible task management in Go applications.