-
Notifications
You must be signed in to change notification settings - Fork 607
/
request_context.go
185 lines (160 loc) · 4.66 KB
/
request_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
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package thrift
import (
"context"
"net"
"sync"
"time"
)
// RPCOptions is a mirror of C++ apache::thrift::RPCOptions
// Not all options are guaranteed to be implemented by a client
type RPCOptions struct {
Timeout time.Duration
ChunkTimeout time.Duration
QueueTimeout time.Duration
OverallTimeout time.Duration
ProcessingTimeout time.Duration
Priority Priority
ClientOnlyTimeouts bool
InteractionID int64
// For sending and receiving headers.
contextHeaders
}
// Identity represents a secure peer identity
type Identity struct {
Type string
Data string
}
// RequestContext is a mirror of C++ apache::thrift::RPCOptions
// Not all options are guaranteed to be implemented by a client
type RequestContext struct {
RequestTimeout time.Duration
Method string
SequenceID int32
InteractionID int64
Priority Priority
LocalAddress net.Addr
PeerAddress net.Addr
SecurityProtocol string
PeerIdentities []Identity
PeerCommonName string
// For sending and receiving headers.
contextHeaders
}
// WithRPCOptions sets the RPCOptions in a client request go context
func WithRPCOptions(ctx context.Context, opts *RPCOptions) context.Context {
return context.WithValue(ctx, rpcOptionsKey, opts)
}
// GetRPCOptions returns the RPCOptions in a go context, or nil if there is nothing
func GetRPCOptions(ctx context.Context) *RPCOptions {
v := ctx.Value(rpcOptionsKey)
if v == nil {
return nil
}
return v.(*RPCOptions)
}
// GetRequestContext returns the RequestContext in a go context, or nil if there is nothing
func GetRequestContext(ctx context.Context) *RequestContext {
v := ctx.Value(reqContextKey)
if v == nil {
return nil
}
return v.(*RequestContext)
}
// WithRequestContext sets the RequestContext in a given go context
func WithRequestContext(ctx context.Context, reqCtx *RequestContext) context.Context {
return context.WithValue(ctx, reqContextKey, reqCtx)
}
type requestContextKey int
const (
rpcOptionsKey requestContextKey = 1
reqContextKey requestContextKey = 2
)
// Priority maps to C++ apache::thrift::concurrency::PRIORITY
type Priority uint8
// Priority maps to C++ apache::thrift::concurrency::PRIORITY
const (
PriorityNone Priority = 0
PriorityHighImportant Priority = 1
PriorityHigh Priority = 2
PriorityImportant Priority = 3
PriorityNormal Priority = 4
PriorityBestEffort Priority = 5
PriorityEnd Priority = 6
)
type contextHeaders struct {
writeHeaders map[string]string
readHeaders map[string]string
writeHeaderLock sync.Mutex
readHeaderLock sync.Mutex
}
func (c *contextHeaders) ReadHeaders() map[string]string {
c.readHeaderLock.Lock()
defer c.readHeaderLock.Unlock()
res := map[string]string{}
for k, v := range c.readHeaders {
res[k] = v
}
return res
}
func (c *contextHeaders) WriteHeaders() map[string]string {
c.writeHeaderLock.Lock()
defer c.writeHeaderLock.Unlock()
res := map[string]string{}
for k, v := range c.writeHeaders {
res[k] = v
}
return res
}
func (c *contextHeaders) SetWriteHeader(k, v string) {
c.writeHeaderLock.Lock()
defer c.writeHeaderLock.Unlock()
if c.writeHeaders == nil {
c.writeHeaders = map[string]string{}
}
c.writeHeaders[k] = v
}
func (c *contextHeaders) SetReadHeader(k, v string) {
c.readHeaderLock.Lock()
defer c.readHeaderLock.Unlock()
if c.readHeaders == nil {
c.readHeaders = map[string]string{}
}
c.readHeaders[k] = v
}
func (c *contextHeaders) SetWriteHeaders(headers map[string]string) {
c.writeHeaderLock.Lock()
defer c.writeHeaderLock.Unlock()
c.writeHeaders = headers
}
func (c *contextHeaders) SetReadHeaders(headers map[string]string) {
c.readHeaderLock.Lock()
defer c.readHeaderLock.Unlock()
c.readHeaders = headers
}
func (c *contextHeaders) GetWriteHeader(k string) (string, bool) {
c.writeHeaderLock.Lock()
defer c.writeHeaderLock.Unlock()
v, ok := c.writeHeaders[k]
return v, ok
}
func (c *contextHeaders) GetReadHeader(k string) (string, bool) {
c.readHeaderLock.Lock()
defer c.readHeaderLock.Unlock()
v, ok := c.readHeaders[k]
return v, ok
}