/
session.go
90 lines (73 loc) · 1.91 KB
/
session.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
package codec
import (
"context"
"sync"
)
// define the rpc context for a request
// the session init instally in client side
// when server side receive the requset, session init in server side
// maybe call `RpcContext`
type Session interface {
RPCName() string
Request() interface{}
Response() interface{}
SetResponse(rsp interface{})
SetError(err error)
}
// BaseSession implements some basic methods defined in `Session`
type BaseSession struct {
ReqHead interface{}
RspHead interface{}
}
func (r *BaseSession) Request() interface{} {
if r != nil {
return r.ReqHead
}
return nil
}
func (r *BaseSession) Response() interface{} {
if r != nil {
return r.RspHead
}
return nil
}
func (r *BaseSession) SetResponse(rsp interface{}) {
if r != nil {
r.RspHead = rsp
}
}
var (
sessionMutex sync.RWMutex
builders = map[string]SessionBuilder{}
)
// SessionBuilder when extending protocols, SessionBuilder should be
// implemented and registered to help build the `Session`.
type SessionBuilder interface {
Build(reqHead interface{}) (Session, error)
}
// RegisterSessionBuilder register extended SessionBuilder for protocol `proto`
func RegisterSessionBuilder(proto string, builder SessionBuilder) {
sessionMutex.Lock()
defer sessionMutex.Unlock()
builders[proto] = builder
}
// GetSessionBuilder return SessionBuilder for protocol `proto`
func GetSessionBuilder(proto string) SessionBuilder {
sessionMutex.RLock()
defer sessionMutex.RUnlock()
return builders[proto]
}
const sessionKey = "qianxi"
// SessionFromContext return Session carried by `ctx`
func SessionFromContext(ctx context.Context) Session {
v := ctx.Value(sessionKey)
session, ok := v.(Session)
if !ok {
return nil
}
return session
}
// ContextWithSession return new context carrying value `session`
func ContextWithSession(ctx context.Context, session Session) context.Context {
return context.WithValue(ctx, sessionKey, session)
}