go context value保存map

go context虽然只能保存一对key和value,但是支持any;虽然获取value时支持不断从父context中找,但是无法获取所有kv。

package main

import (
	"maps"
	"context"
	"fmt"
)

// context推荐使用自定义结构体类型作为key,与go内置类型有所区分
type ContextKey struct{}

type contextValue struct {
	// 向前回溯
	pre *contextValue
	kvs map[any]any
}

func AddKvs(ctx context.Context, kvs ...any) context.Context {
	if len(kvs) == 0 || len(kvs)%2 == 1 {
		return ctx
	}

	kvsInContext := make(map[any]any, len(kvs))
	for i := 0; i < len(kvs); i++ {
		kvsInContext[kvs[i]] = kvs[i+1]
		i++
	}

	return context.WithValue(ctx, ContextKey{}, &contextValue{
		pre: getContextValue(ctx),
		kvs: kvsInContext,
	})
}

func getContextValue(ctx context.Context) *contextValue {
	val, ok := ctx.Value(ContextKey{}).(*contextValue)
	if !ok {
		return nil
	}
	return val
}

func GetAllKvs(ctx context.Context) []any {
	m := make(map[any]any, 0)
	ctxVal := getContextValue(ctx)
	if ctxVal == nil {
		return []any{}
	}
	maps.Copy(m, ctxVal.kvs)

	for ctxVal.pre != nil {
		ctxVal = ctxVal.pre
		maps.Copy(m, ctxVal.kvs)
	}

	res := make([]any, 0, len(m))
	for k, v := range m {
		res = append(res, k)
		res = append(res, v)
	}
	return res
}

func main() {
	ctx := context.Background()
	fmt.Println(GetAllKvs(ctx))

	ctx1 := AddKvs(ctx, "ctx1", "1")
	fmt.Println(GetAllKvs(ctx1))

	ctx2 := AddKvs(ctx1)
	fmt.Println(GetAllKvs(ctx2))

	ctx3 := AddKvs(ctx2, "ctx3", "3")
	fmt.Println(GetAllKvs(ctx3))
}

posted on 2025-06-21 11:00  王景迁  阅读(15)  评论(0)    收藏  举报

导航