gocommon/pkg/internal/util/kv_context.go

package util

import (
	"context"
	"sync"
	"time"
)

type KVContext interface {
	context.Context
	Set(key string, value interface{})
	Get(key string) (interface{}, bool)
}

func AssertKVContext(ctx context.Context) KVContext {
	c, isOk := ctx.(KVContext)
	if isOk {
		return c
	}
	return nil
}

func NewKVContext(ctx context.Context) KVContext {
	// 防止panic
	if ctx == nil {
		ctx = context.Background()
	}
	return &kVContext{
		p:    ctx,
		_map: map[string]interface{}{},
	}
}

/**
具体实现
*/
type kVContext struct {
	/**
	不使用 并发map的原因：
	1、并不是所有的key都可以做value
	2、原生map效率高于sync.map
	*/
	sync.RWMutex
	_map map[string]interface{} // 防止并发读写问题
	p    context.Context
}

func (c *kVContext) Set(key string, value interface{}) {
	c.Lock()
	defer c.Unlock()
	c._map[key] = value
}

func (c *kVContext) Get(key string) (interface{}, bool) {
	c.RLock()
	defer c.RUnlock()
	value, isExist := c._map[key]
	return value, isExist
}

/************************************/
/***** GOLANG.ORG/X/NET/CONTEXT *****/
/************************************/

// Deadline returns the time when work done on behalf of this context
// should be canceled. Deadline returns ok==false when no deadline is
// set. Successive calls to Deadline return the same results.
func (c *kVContext) Deadline() (deadline time.Time, ok bool) {
	return
}

// Done returns a channel that's closed when work done on behalf of this
// context should be canceled. Done may return nil if this context can
// never be canceled. Successive calls to Done return the same value.
func (c *kVContext) Done() <-chan struct{} {
	return nil
}

// Err returns a non-nil error value after Done is closed,
// successive calls to Err return the same error.
// If Done is not yet closed, Err returns nil.
// If Done is closed, Err returns a non-nil error explaining why:
// Canceled if the context was canceled
// or DeadlineExceeded if the context's deadline passed.
func (c *kVContext) Err() error {
	return nil
}

// Value returns the value associated with this context for key, or nil
// if no value is associated with key. Successive calls to Value with
// the same key returns the same result.
func (c *kVContext) Value(key interface{}) interface{} {
	if k, ok := key.(string); ok {
		value, _ := c.Get(k)
		if value == nil {
			return c.p.Value(key)
		}
		return value
	}
	return c.p.Value(key)
}