Golang——Context.WithValue的基本使用

在Golang中，可以通过Context对协程做同步，或者传递上下文变量给其他协程。这样可以避免在协程之间传递大量的变量，代码更整洁可维护。下面的例子通过WithValue传递给协程一个变量，并且通过channel在协程之间通信。

package main

import (
	"context"
	"fmt"
	"sync"
)

func worker(cancelCtx context.Context, ch chan int, wg *sync.WaitGroup) {
	defer wg.Done()

	fmt.Println(fmt.Sprintf("context value: %v", cancelCtx.Value("key1")))
	for {
		select {
		case val := <-ch:
			fmt.Println(fmt.Sprintf("read from ch value: %d", val))
		case <-cancelCtx.Done():
			fmt.Println("worker is cancelled")
			return
		}
	}
}

func main() {
	rootCtx := context.Background()
	childCtx := context.WithValue(rootCtx, "key1", "value1")
	cancelCtx, cancelFunc := context.WithCancel(childCtx)

	ch := make(chan int)
	wg := &sync.WaitGroup{}
	wg.Add(1)
	go worker(cancelCtx, ch, wg)

	for i := 0; i < 10; i++ {
		ch <- i
	}

	cancelFunc()
	wg.Wait()
	close(ch)

}

输出：

context value: value1
read from ch value: 0
read from ch value: 1
read from ch value: 2
read from ch value: 3
read from ch value: 4
read from ch value: 5
read from ch value: 6
read from ch value: 7
read from ch value: 8
read from ch value: 9
worker is cancelled

在实际的生产环境中，比如web服务器http请求处理器中，可以通过WithValue传递通用的字段给请求处理协程，比如用于多个请求之间的链路追踪：

package main

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

	"github.com/google/uuid"
)

func welcome(w http.ResponseWriter, r *http.Request) {
	traceid := ""
	if m := r.Context().Value("traceid"); m != nil {
		if value, ok := m.(string); ok {
			traceid = value
		}
	}
	w.Header().Add("traceid", traceid)
	w.Write([]byte("Welcome to China"))
}

func traceID(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		traceid := uuid.New().String()
		ctx := context.WithValue(r.Context(), "traceid", traceid)
		req := r.WithContext(ctx)
		next.ServeHTTP(w, req)
	})
}

func main() {
	welcomeHandler := http.HandlerFunc(welcome)
	http.Handle("/welcome", traceID(welcomeHandler))
	http.ListenAndServe(":9090", nil)
}