【Reprint】Go Standard Library Context

title: 【转载】Go Standard Library Context
date: 2021-08-09 16:38:33
comment: false
toc: true
category:

Golang
tags:
转载
Go
Standard Library
Context

This article is reprinted from: Go Standard Library Context | Li Wenzhou's Blog

In the Go http package's Server, each request has a corresponding goroutine to handle it. The request handling function typically starts additional goroutines to access backend services, such as databases and RPC services. The goroutine handling a request usually needs to access some request-specific data, such as the terminal user's authentication information, verification-related tokens, and the request's deadline. When a request is canceled or times out, all goroutines handling that request should exit quickly, allowing the system to release the resources occupied by those goroutines.

Why Context is Needed#

Basic Example#

package main  

import (  
    "fmt"  
    "sync"  

    "time"  
)  

var wg sync.WaitGroup  

// Initial example  

func worker() {  
    for {  
        fmt.Println("worker")  
        time.Sleep(time.Second)  
    }  
    // How to receive external commands to exit  
    wg.Done()  
}  

func main() {  
    wg.Add(1)  
    go worker()  
    // How to gracefully end the child goroutine  
    wg.Wait()  
    fmt.Println("over")  
}

Global Variable Method#

package main  

import (  
    "fmt"  
    "sync"  

    "time"  
)  

var wg sync.WaitGroup  
var exit bool  

// Problems with the global variable method:  
// 1. Using global variables is not easy to unify when called across packages  
// 2. If a goroutine is started in worker, it becomes difficult to control.  

func worker() {  
    for {  
        fmt.Println("worker")  
        time.Sleep(time.Second)  
        if exit {  
            break  
        }  
    }  
    wg.Done()  
}  

func main() {  
    wg.Add(1)  
    go worker()  
    time.Sleep(time.Second * 3) // sleep for 3 seconds to prevent the program from exiting too quickly  
    exit = true                 // Modify the global variable to exit the child goroutine  
    wg.Wait()  
    fmt.Println("over")  
}

Channel Method#

package main  

import (  
    "fmt"  
    "sync"  

    "time"  
)  

var wg sync.WaitGroup  

// Problems with the channel method:  
// 1. Using global variables is not easy to implement standards and unification when called across packages, requiring maintenance of a shared channel  

func worker(exitChan chan struct{}) {  
LOOP:  
    for {  
        fmt.Println("worker")  
        time.Sleep(time.Second)  
        select {  
        case <-exitChan: // Wait for notification from the parent  
            break LOOP  
        default:  
        }  
    }  
    wg.Done()  
}  

func main() {  
    var exitChan = make(chan struct{})  
    wg.Add(1)  
    go worker(exitChan)  
    time.Sleep(time.Second * 3) // sleep for 3 seconds to prevent the program from exiting too quickly  
    exitChan <- struct{}{}      // Send exit signal to the child goroutine  
    close(exitChan)  
    wg.Wait()  
    fmt.Println("over")  
}

Official Version Solution#

package main  

import (  
    "context"  
    "fmt"  
    "sync"  

    "time"  
)  

var wg sync.WaitGroup  

func worker(ctx context.Context) {  
LOOP:  
    for {  
        fmt.Println("worker")  
        time.Sleep(time.Second)  
        select {  
        case <-ctx.Done(): // Wait for notification from the parent  
            break LOOP  
        default:  
        }  
    }  
    wg.Done()  
}  

func main() {  
    ctx, cancel := context.WithCancel(context.Background())  
    wg.Add(1)  
    go worker(ctx)  
    time.Sleep(time.Second * 3)  
    cancel() // Notify the child goroutine to end  
    wg.Wait()  
    fmt.Println("over")  
}

When a child goroutine starts another goroutine, it only needs to pass the ctx:

package main  

import (  
    "context"  
    "fmt"  
    "sync"  

    "time"  
)  

var wg sync.WaitGroup  

func worker(ctx context.Context) {  
    go worker2(ctx)  
LOOP:  
    for {  
        fmt.Println("worker")  
        time.Sleep(time.Second)  
        select {  
        case <-ctx.Done(): // Wait for notification from the parent  
            break LOOP  
        default:  
        }  
    }  
    wg.Done()  
}  

func worker2(ctx context.Context) {  
LOOP:  
    for {  
        fmt.Println("worker2")  
        time.Sleep(time.Second)  
        select {  
        case <-ctx.Done(): // Wait for notification from the parent  
            break LOOP  
        default:  
        }  
    }  
}  
func main() {  
    ctx, cancel := context.WithCancel(context.Background())  
    wg.Add(1)  
    go worker(ctx)  
    time.Sleep(time.Second * 3)  
    cancel() // Notify the child goroutine to end  
    wg.Wait()  
    fmt.Println("over")  
}

Introduction to Context#

Go 1.7 introduced a new standard library context, which defines the Context type, specifically designed to simplify operations related to handling multiple goroutines for a single request, including request domain data, cancellation signals, deadlines, etc., which may involve multiple API calls.

Contexts should be created for incoming requests to the server, and outgoing calls to the server should accept contexts. The function call chain between them must pass the context, or derived contexts created with WithCancel, WithDeadline, WithTimeout, or WithValue can be used. When a context is canceled, all derived contexts are also canceled.

Context Interface#

context.Context is an interface that defines four methods that need to be implemented. The specific signatures are as follows:

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

Among them:

The Deadline method needs to return the time when the current Context is canceled, which is the deadline for completing the work;
The Done method needs to return a Channel, which will close after the current work is completed or the context is canceled. Multiple calls to the Done method will return the same channel;
The Err method will return the reason for the current Context ending, and it will only return a non-empty value when the channel returned by Done is closed;
- If the current Context is canceled, it will return the Canceled error;
- If the current Context times out, it will return the DeadlineExceeded error;
The Value method will return the value corresponding to the key from the Context. For the same context, multiple calls to Value with the same Key will return the same result. This method is only used to pass request domain data across APIs and processes;

Background() and TODO()#

Go has two built-in functions: Background() and TODO(), which return a background and todo that implement the Context interface, respectively. In our code, these two built-in context objects are initially used as the top-level parent context, from which more child context objects are derived.

Background() is mainly used in the main function, initialization, and test code as the top-level Context of the Context tree, which is the root Context.

TODO(), it currently does not know the specific usage scenario. If we do not know what context to use, we can use this.

Both background and todo are essentially emptyCtx struct types, which are non-cancelable, have no deadline set, and carry no values.

With Series Functions#

In addition, the context package also defines four With series functions.

WithCancel#

The function signature of WithCancel is as follows:

func WithCancel(parent Context) (ctx Context, cancel CancelFunc)

WithCancel returns a copy of the parent with a new Done channel. When the returned cancel function is called or when the Done channel of the parent context is closed, the Done channel of the returned context will be closed, regardless of which happens first.

Cancelling this context will release the resources associated with it, so the code should call cancel immediately after the operations running in this context are completed.

func gen(ctx context.Context) <-chan int {  
    dst := make(chan int)  
    n := 1  
    go func() {  
        for {  
            select {  
            case <-ctx.Done():  
                return // return to end this goroutine and prevent leaks  
            case dst <- n:  
                n++  
            }  
        }  
    }()  
    return dst  
}  
func main() {  
    ctx, cancel := context.WithCancel(context.Background())  
    defer cancel() // Call cancel after we have taken the needed integers  

    for n := range gen(ctx) {  
        fmt.Println(n)  
        if n == 5 {  
            break  
        }  
    }  
}

In the above example code, the gen function generates integers in a separate goroutine and sends them to the returned channel. The caller of gen needs to cancel the context after using the generated integers to prevent leaks in the internal goroutine started by gen.

WithDeadline#

The function signature of WithDeadline is as follows:

func WithDeadline(parent Context, deadline time.Time) (Context, CancelFunc)

Returns a copy of the parent context and adjusts the deadline to not be later than d. If the deadline of the parent context is already earlier than d, WithDeadline(parent, d) is semantically equivalent to the parent context. When the deadline expires, the Done channel of the returned context will be closed when the returned cancel function is called or when the Done channel of the parent context is closed, whichever happens first.

Cancelling this context will release the resources associated with it, so the code should call cancel immediately after the operations running in this context are completed.

func main() {  
    d := time.Now().Add(50 * time.Millisecond)  
    ctx, cancel := context.WithDeadline(context.Background(), d)  

    // Although ctx will expire, it is good practice to call its cancel function in any case.  
    // If not, it may cause the context and its parent to live longer than necessary.  
    defer cancel()  

    select {  
    case <-time.After(1 * time.Second):  
        fmt.Println("overslept")  
    case <-ctx.Done():  
        fmt.Println(ctx.Err())  
    }  
}

In the above code, a deadline is defined to expire after 50 milliseconds, and then we call context.WithDeadline(context.Background(), d) to get a context (ctx) and a cancel function (cancel), and then use a select statement to make the main program wait: either wait for 1 second and print overslept to exit or wait for ctx to expire and exit.

In the above example code, since ctx will expire after 50 milliseconds, ctx.Done() will receive the context expiration notification first and print the content of ctx.Err().

WithTimeout#

The function signature of WithTimeout is as follows:

func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc)

WithTimeout returns WithDeadline(parent, time.Now().Add(timeout)).

Cancelling this context will release the resources associated with it, so the code should call cancel immediately after the operations running in this context are completed, and it is usually used for timeout control of database or network connections. The specific example is as follows:

package main  

import (  
    "context"  
    "fmt"  
    "sync"  

    "time"  
)  

// context.WithTimeout  

var wg sync.WaitGroup  

func worker(ctx context.Context) {  
LOOP:  
    for {  
        fmt.Println("db connecting ...")  
        time.Sleep(time.Millisecond * 10) // Assume normal database connection takes 10 milliseconds  
        select {  
        case <-ctx.Done(): // Automatically called after 50 milliseconds  
            break LOOP  
        default:  
        }  
    }  
    fmt.Println("worker done!")  
    wg.Done()  
}  

func main() {  
    // Set a timeout of 50 milliseconds  
    ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*50)  
    wg.Add(1)  
    go worker(ctx)  
    time.Sleep(time.Second * 5)  
    cancel() // Notify the child goroutine to end  
    wg.Wait()  
    fmt.Println("over")  
}

WithValue#

The WithValue function can establish a relationship between request-scoped data and the Context object. It is declared as follows:

func WithValue(parent Context, key, val interface{}) Context

WithValue returns a copy of the parent node, where the value associated with the key is val.

Context values should only be used to pass necessary request domain data across APIs and processes, not to pass optional parameters to functions.

The provided key must be comparable and should not be of type string or any other built-in type to avoid conflicts when using context between packages. Users of WithValue should define their own types for keys. To avoid allocation when assigning to interface{}, context keys typically have a concrete type of struct{}. Alternatively, the static type of exported context key variables should be a pointer or interface.

package main  

import (  
    "context"  
    "fmt"  
    "sync"  

    "time"  
)  

// context.WithValue  

type TraceCode string  

var wg sync.WaitGroup  

func worker(ctx context.Context) {  
    key := TraceCode("TRACE_CODE")  
    traceCode, ok := ctx.Value(key).(string) // Get trace code in the child goroutine  
    if !ok {  
        fmt.Println("invalid trace code")  
    }  
LOOP:  
    for {  
        fmt.Printf("worker, trace code:%s\n", traceCode)  
        time.Sleep(time.Millisecond * 10) // Assume normal database connection takes 10 milliseconds  
        select {  
        case <-ctx.Done(): // Automatically called after 50 milliseconds  
            break LOOP  
        default:  
        }  
    }  
    fmt.Println("worker done!")  
    wg.Done()  
}  

func main() {  
    // Set a timeout of 50 milliseconds  
    ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*50)  
    // Set trace code in the system entry to pass to subsequent started goroutines for log data aggregation  
    ctx = context.WithValue(ctx, TraceCode("TRACE_CODE"), "12512312234")  
    wg.Add(1)  
    go worker(ctx)  
    time.Sleep(time.Second * 5)  
    cancel() // Notify the child goroutine to end  
    wg.Wait()  
    fmt.Println("over")  
}

Notes on Using Context#

It is recommended to pass Context explicitly as a parameter
Functions that take Context as a parameter should have Context as the first parameter.
When passing Context to a function, do not pass nil; if you do not know what to pass, use context.TODO()
The Value-related methods of Context should pass necessary data from the request domain and should not be used to pass optional parameters
Context is thread-safe and can be safely passed among multiple goroutines

Client Timeout Cancellation Example#

How to implement timeout control on the client side when calling server APIs?

Server Side#

// context_timeout/server/main.go  
package main  

import (  
    "fmt"  
    "math/rand"  
    "net/http"  

    "time"  
)  

// Server side, randomly slow response  

func indexHandler(w http.ResponseWriter, r *http.Request) {  
    number := rand.Intn(2)  
    if number == 0 {  
        time.Sleep(time.Second * 10) // Slow response taking 10 seconds  
        fmt.Fprintf(w, "slow response")  
        return  
    }  
    fmt.Fprint(w, "quick response")  
}  

func main() {  
    http.HandleFunc("/", indexHandler)  
    err := http.ListenAndServe(":8000", nil)  
    if err != nil {  
        panic(err)  
    }  
}

Client Side#

// context_timeout/client/main.go  
package main  

import (  
    "context"  
    "fmt"  
    "io/ioutil"  
    "net/http"  
    "sync"  
    "time"  
)  

// Client  

type respData struct {  
    resp *http.Response  
    err  error  
}  

func doCall(ctx context.Context) {  
    transport := http.Transport{  
       // Define a global client object and enable long connections for frequent requests  
       // Use short connections for infrequent requests  
       DisableKeepAlives: true,  
    }  
    client := http.Client{  
        Transport: &transport,  
    }  

    respChan := make(chan *respData, 1)  
    req, err := http.NewRequest("GET", "http://127.0.0.1:8000/", nil)  
    if err != nil {  
        fmt.Printf("new request failed, err:%v\n", err)  
        return  
    }  
    req = req.WithContext(ctx) // Create a new client request using the ctx with timeout  
    var wg sync.WaitGroup  
    wg.Add(1)  
    defer wg.Wait()  
    go func() {  
        resp, err := client.Do(req)  
        fmt.Printf("client.do resp:%v, err:%v\n", resp, err)  
        rd := &respData{  
            resp: resp,  
            err:  err,  
        }  
        respChan <- rd  
        wg.Done()  
    }()  

    select {  
    case <-ctx.Done():  
        //transport.CancelRequest(req)  
        fmt.Println("call api timeout")  
    case result := <-respChan:  
        fmt.Println("call server api success")  
        if result.err != nil {  
            fmt.Printf("call server api failed, err:%v\n", result.err)  
            return  
        }  
        defer result.resp.Body.Close()  
        data, _ := ioutil.ReadAll(result.resp.Body)  
        fmt.Printf("resp:%v\n", string(data))  
    }  
}  

func main() {  
    // Define a timeout of 100 milliseconds  
    ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*100)  
    defer cancel() // Call cancel to release child goroutine resources  
    doCall(ctx)  
}