Context 值传递与链路追踪

难度：⭐⭐ 中等

考点

• context.WithValue 使用
• 自定义 key 类型避免冲突
• context 值的层级传播（子可读父，父不可读子）

提示

1. key 用自定义类型（不要用 string），避免包间冲突
2. 获取值时做类型断言，不存在返回零值
3. WithValue 创建的是子 context，不会修改父 context

---

参考答案（Go）

点击展开参考答案

//go:build ignore

package answer

import "context"

type contextKey string

const traceIDKey contextKey = "trace_id"
const userIDKey contextKey = "user_id"

func WithTraceID(ctx context.Context, traceID string) context.Context {
	return context.WithValue(ctx, traceIDKey, traceID)
}

func GetTraceID(ctx context.Context) string {
	if v, ok := ctx.Value(traceIDKey).(string); ok {
		return v
	}
	return ""
}

func WithUserID(ctx context.Context, userID string) context.Context {
	return context.WithValue(ctx, userIDKey, userID)
}

func GetUserID(ctx context.Context) string {
	if v, ok := ctx.Value(userIDKey).(string); ok {
		return v
	}
	return ""
}

func Middleware(traceID, userID string, handler func(ctx context.Context) string) string {
	ctx := context.Background()
	ctx = WithTraceID(ctx, traceID)
	ctx = WithUserID(ctx, userID)
	return handler(ctx)
}