/
context.go
108 lines (89 loc) ยท 2.1 KB
/
context.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
package share
import (
"context"
"fmt"
"reflect"
"sync"
)
// var _ context.Context = &Context{}
// Context is a rpcx customized Context that can contains multiple values.
type Context struct {
tagsLock *sync.Mutex
tags map[interface{}]interface{}
context.Context
}
func NewContext(ctx context.Context) *Context {
tagsLock := &sync.Mutex{}
ctx = context.WithValue(ctx, ContextTagsLock, tagsLock)
return &Context{
tagsLock: tagsLock,
Context: ctx,
tags: map[interface{}]interface{}{isShareContext: true},
}
}
func (c *Context) Lock() {
c.tagsLock.Lock()
}
func (c *Context) Unlock() {
c.tagsLock.Unlock()
}
func (c *Context) Value(key interface{}) interface{} {
c.tagsLock.Lock()
defer c.tagsLock.Unlock()
if c.tags == nil {
c.tags = make(map[interface{}]interface{})
}
if v, ok := c.tags[key]; ok {
return v
}
return c.Context.Value(key)
}
func (c *Context) SetValue(key, val interface{}) {
c.tagsLock.Lock()
defer c.tagsLock.Unlock()
if c.tags == nil {
c.tags = make(map[interface{}]interface{})
}
c.tags[key] = val
}
// DeleteKey delete the kv pair by key.
func (c *Context) DeleteKey(key interface{}) {
c.tagsLock.Lock()
defer c.tagsLock.Unlock()
if c.tags == nil || key == nil {
return
}
delete(c.tags, key)
}
func (c *Context) String() string {
return fmt.Sprintf("%v.WithValue(%v)", c.Context, c.tags)
}
func WithValue(parent context.Context, key, val interface{}) *Context {
if key == nil {
panic("nil key")
}
if !reflect.TypeOf(key).Comparable() {
panic("key is not comparable")
}
tags := make(map[interface{}]interface{})
tags[key] = val
return &Context{Context: parent, tags: tags, tagsLock: &sync.Mutex{}}
}
func WithLocalValue(ctx *Context, key, val interface{}) *Context {
if key == nil {
panic("nil key")
}
if !reflect.TypeOf(key).Comparable() {
panic("key is not comparable")
}
if ctx.tags == nil {
ctx.tags = make(map[interface{}]interface{})
}
ctx.tags[key] = val
return ctx
}
// IsShareContext checks whether a context is share.Context.
func IsShareContext(ctx context.Context) bool {
ok := ctx.Value(isShareContext)
return ok != nil
}