@@ -27,11 +27,54 @@ import (
27
27
"sigs.k8s.io/apiserver-network-proxy/konnectivity-client/proto/client"
28
28
)
29
29
30
+ // connContext tracks a connection from agent to node network.
31
+ type connContext struct {
32
+ conn net.Conn
33
+ cleanFunc func ()
34
+ dataCh chan []byte
35
+ cleanOnce sync.Once
36
+ }
37
+
38
+ func (c * connContext ) cleanup () {
39
+ c .cleanOnce .Do (c .cleanFunc )
40
+ }
41
+
42
+ type connectionManager struct {
43
+ mu sync.RWMutex
44
+ connections map [int64 ]* connContext
45
+ }
46
+
47
+ func (cm * connectionManager ) Add (connID int64 , ctx * connContext ) {
48
+ cm .mu .Lock ()
49
+ defer cm .mu .Unlock ()
50
+ cm .connections [connID ] = ctx
51
+ }
52
+
53
+ func (cm * connectionManager ) Get (connID int64 ) (* connContext , bool ) {
54
+ cm .mu .RLock ()
55
+ defer cm .mu .RUnlock ()
56
+ ctx , ok := cm .connections [connID ]
57
+ return ctx , ok
58
+ }
59
+
60
+ func (cm * connectionManager ) Delete (connID int64 ) {
61
+ cm .mu .Lock ()
62
+ defer cm .mu .Unlock ()
63
+ delete (cm .connections , connID )
64
+ }
65
+
66
+ func newConnectionManager () * connectionManager {
67
+ return & connectionManager {
68
+ connections : make (map [int64 ]* connContext ),
69
+ }
70
+ }
71
+
30
72
// AgentClient runs on the node network side. It connects to proxy server and establishes
31
73
// a stream connection from which it sends and receives network traffic.
32
74
type AgentClient struct {
33
- nextConnID int64
34
- connContext map [int64 ]* connContext
75
+ nextConnID int64
76
+
77
+ connManager * connectionManager
35
78
36
79
stream * RedialableAgentClient
37
80
stopCh <- chan struct {}
@@ -47,7 +90,7 @@ func newAgentClient(address, agentID string, cs *ClientSet, opts ...grpc.DialOpt
47
90
48
91
func newAgentClientWithRedialableAgentClient (rac * RedialableAgentClient ) * AgentClient {
49
92
return & AgentClient {
50
- connContext : make ( map [ int64 ] * connContext ),
93
+ connManager : newConnectionManager ( ),
51
94
stream : rac ,
52
95
stopCh : rac .stopCh ,
53
96
}
@@ -62,18 +105,6 @@ func (c *AgentClient) Close() {
62
105
c .stream .Close ()
63
106
}
64
107
65
- // connContext tracks a connection from agent to node network.
66
- type connContext struct {
67
- conn net.Conn
68
- cleanFunc func ()
69
- dataCh chan []byte
70
- cleanOnce sync.Once
71
- }
72
-
73
- func (c * connContext ) cleanup () {
74
- c .cleanOnce .Do (c .cleanFunc )
75
- }
76
-
77
108
// Connect connects to proxy server to establish a gRPC stream,
78
109
// on which the proxied traffic is multiplexed through the stream
79
110
// and piped to the local connection. It register itself as a
@@ -146,7 +177,7 @@ func (a *AgentClient) Serve() {
146
177
147
178
connID := atomic .AddInt64 (& a .nextConnID , 1 )
148
179
dataCh := make (chan []byte , 5 )
149
- a . connContext [ connID ] = & connContext {
180
+ ctx : = & connContext {
150
181
conn : conn ,
151
182
dataCh : dataCh ,
152
183
cleanFunc : func () {
@@ -167,24 +198,26 @@ func (a *AgentClient) Serve() {
167
198
}
168
199
169
200
close (dataCh )
170
- delete ( a . connContext , connID )
201
+ a . connManager . Delete ( connID )
171
202
},
172
203
}
204
+ a .connManager .Add (connID , ctx )
173
205
174
206
resp .GetDialResponse ().ConnectID = connID
175
207
if err := a .stream .RetrySend (resp ); err != nil {
176
208
klog .Warningf ("stream send error: %v" , err )
177
209
continue
178
210
}
179
211
180
- go a .remoteToProxy (conn , connID )
181
- go a .proxyToRemote (conn , connID )
212
+ go a .remoteToProxy (connID , ctx )
213
+ go a .proxyToRemote (connID , ctx )
182
214
183
215
case client .PacketType_DATA :
184
216
data := pkt .GetData ()
185
217
klog .Infof ("received DATA(id=%d)" , data .ConnectID )
186
218
187
- if ctx , ok := a .connContext [data .ConnectID ]; ok {
219
+ ctx , ok := a .connManager .Get (data .ConnectID )
220
+ if ok {
188
221
ctx .dataCh <- data .Data
189
222
}
190
223
@@ -194,7 +227,8 @@ func (a *AgentClient) Serve() {
194
227
195
228
klog .Infof ("received CLOSE_REQ(id=%d)" , connID )
196
229
197
- if ctx , ok := a .connContext [connID ]; ok {
230
+ ctx , ok := a .connManager .Get (connID )
231
+ if ok {
198
232
ctx .cleanup ()
199
233
} else {
200
234
resp := & client.Packet {
@@ -215,12 +249,7 @@ func (a *AgentClient) Serve() {
215
249
}
216
250
}
217
251
218
- func (a * AgentClient ) remoteToProxy (conn net.Conn , connID int64 ) {
219
- ctx := a .connContext [connID ]
220
- if ctx == nil {
221
- return
222
- }
223
-
252
+ func (a * AgentClient ) remoteToProxy (connID int64 , ctx * connContext ) {
224
253
defer ctx .cleanup ()
225
254
226
255
var buf [1 << 12 ]byte
@@ -229,7 +258,7 @@ func (a *AgentClient) remoteToProxy(conn net.Conn, connID int64) {
229
258
}
230
259
231
260
for {
232
- n , err := conn .Read (buf [:])
261
+ n , err := ctx . conn .Read (buf [:])
233
262
klog .Infof ("received %d bytes from remote for connID[%d]" , n , connID )
234
263
235
264
if err == io .EOF {
@@ -251,18 +280,13 @@ func (a *AgentClient) remoteToProxy(conn net.Conn, connID int64) {
251
280
}
252
281
}
253
282
254
- func (a * AgentClient ) proxyToRemote (conn net.Conn , connID int64 ) {
255
- ctx := a .connContext [connID ]
256
- if ctx == nil {
257
- return
258
- }
259
-
283
+ func (a * AgentClient ) proxyToRemote (connID int64 , ctx * connContext ) {
260
284
defer ctx .cleanup ()
261
285
262
286
for d := range ctx .dataCh {
263
287
pos := 0
264
288
for {
265
- n , err := conn .Write (d [pos :])
289
+ n , err := ctx . conn .Write (d [pos :])
266
290
if err == nil {
267
291
klog .Infof ("[connID: %d] write last %d data to remote" , connID , n )
268
292
break
0 commit comments