refactor(lease, relay): implement event loops for LeaseManager and RelayServer to centralize state and remove explicit locks
cognitive-glitch committed
Dec 9, 2025 at 10:46 UTC
50c5584555292efd9372aef761d3ca7cd44f4aa2
5 files changed
+548
-235
golangci.yml
+4
@@ -13,6 +13,10 @@ linters:
13
- govet # Standard vet
14
- errcheck # Unchecked errors are fatal
15
- staticcheck # Dominant static analysis
16
+ - gochecknoglobals # Disallows global variables (improves concurrency safety)
17
+ - gochecknoinits # Disallows init functions (can have concurrency implications)
18
+ - gocritic # General linter with many checks, including some concurrency-related ones like deferInLoop
19
+ - containedctx # Ensures context.Context is the first argument (improves context propagation)
20
21
# --- The Dark Arts (Performance & Concurrency) ---
22
- prealloc # Encourages slice pre-allocation (critical for low latency)
portal/handlers.go
+39
-41
@@ -66,11 +66,11 @@ func (g *RelayServer) handleLeaseUpdateRequest(ctx *StreamContext, packet *rdver
66
if g.leaseManager.UpdateLease(req.Lease, ctx.ConnectionID) {
67
resp.Code = rdverb.ResponseCode_RESPONSE_CODE_ACCEPTED
68
69
- // Register lease connection
69
+ // Register lease connection via event loop
70
leaseID := string(req.Lease.Identity.Id)
71
- g.leaseConnectionsLock.Lock()
72
- g.leaseConnections[leaseID] = ctx.Connection
73
- g.leaseConnectionsLock.Unlock()
71
+ done := make(chan struct{}, 1)
72
+ g.cmdCh <- &cmdRegisterLeaseConn{leaseID: leaseID, conn: ctx.Connection, done: done}
73
+ <-done
74
75
// Log lease update completion
76
log.Debug().
@@ -123,11 +123,11 @@ func (g *RelayServer) handleLeaseDeleteRequest(ctx *StreamContext, packet *rdver
123
if g.leaseManager.DeleteLease(req.Identity) {
124
resp.Code = rdverb.ResponseCode_RESPONSE_CODE_ACCEPTED
125
126
- // Remove lease connection
126
+ // Remove lease connection via event loop
127
leaseID := string(req.Identity.Id)
128
- g.leaseConnectionsLock.Lock()
129
- delete(g.leaseConnections, leaseID)
130
- g.leaseConnectionsLock.Unlock()
128
+ done := make(chan struct{}, 1)
129
+ g.cmdCh <- &cmdUnregisterLeaseConn{leaseID: leaseID, done: done}
130
+ <-done
131
132
// Log lease deletion completion
133
log.Debug().
@@ -168,12 +168,12 @@ func (g *RelayServer) handleConnectionRequest(ctx *StreamContext, packet *rdverb
168
return g.sendConnectionResponse(ctx.Stream, rdverb.ResponseCode_RESPONSE_CODE_INVALID_IDENTITY)
169
}
170
171
- // Get the lease connection
172
- g.connectionsLock.RLock()
173
- leaseConn, leaseExists := g.connections[leaseEntry.ConnectionID]
174
- g.connectionsLock.RUnlock()
171
+ // Get the lease connection via event loop
172
+ reply := make(chan *Connection, 1)
173
+ g.cmdCh <- &cmdGetConnByLeaseEntry{connID: leaseEntry.ConnectionID, reply: reply}
174
+ leaseConn := <-reply
175
176
- if !leaseExists {
176
+ if leaseConn == nil {
177
return g.sendConnectionResponse(ctx.Stream, rdverb.ResponseCode_RESPONSE_CODE_INVALID_IDENTITY)
178
}
179
@@ -190,11 +190,12 @@ func (g *RelayServer) handleConnectionRequest(ctx *StreamContext, packet *rdverb
190
// Enforce relayed connection limits
191
if respCode == rdverb.ResponseCode_RESPONSE_CODE_ACCEPTED {
192
leaseID := string(leaseEntry.Lease.Identity.Id)
193
- g.limitsLock.Lock()
194
- overPerLease := g.maxRelayedPerLease > 0 && g.relayedPerLeaseCount[leaseID] >= g.maxRelayedPerLease
195
- g.limitsLock.Unlock()
193
+ // Atomically check and increment limit via event loop
194
+ reply := make(chan bool, 1)
195
+ g.cmdCh <- &cmdCheckAndIncLimit{leaseID: leaseID, reply: reply}
196
+ underLimit := <-reply
197
197
- if overPerLease {
198
+ if !underLimit {
199
log.Warn().Str("lease_id", leaseID).Msg("[RelayServer] Relayed connection per-lease limit reached")
200
respCode = rdverb.ResponseCode_RESPONSE_CODE_REJECTED
201
leaseStream.Close()
@@ -203,6 +204,13 @@ func (g *RelayServer) handleConnectionRequest(ctx *StreamContext, packet *rdverb
204
205
// Send response to client
206
if err := g.sendConnectionResponse(ctx.Stream, respCode); err != nil {
207
+ // If we already incremented the limit, decrement it
208
+ if respCode == rdverb.ResponseCode_RESPONSE_CODE_ACCEPTED {
209
+ leaseID := string(leaseEntry.Lease.Identity.Id)
210
+ done := make(chan struct{}, 1)
211
+ g.cmdCh <- &cmdDecLimit{leaseID: leaseID, done: done}
212
+ <-done
213
+ }
214
leaseStream.Close()
215
return err
216
}
@@ -273,33 +281,23 @@ func (g *RelayServer) sendConnectionResponse(stream *yamux.Stream, code rdverb.R
281
}
282
283
func (g *RelayServer) establishRelayedConnection(clientStream, leaseStream *yamux.Stream, leaseID string) {
276
- // Register connection for tracking
277
- g.limitsLock.Lock()
278
- g.relayedPerLeaseCount[leaseID]++
279
- g.limitsLock.Unlock()
280
-
281
- g.relayedConnectionsLock.Lock()
282
- g.relayedConnections[leaseID] = append(g.relayedConnections[leaseID], clientStream)
283
- g.relayedConnectionsLock.Unlock()
284
+ // Register stream for tracking via event loop
285
+ // Note: limit was already incremented in handleConnectionRequest
286
+ done := make(chan struct{}, 1)
287
+ g.cmdCh <- &cmdAddRelayed{leaseID: leaseID, stream: clientStream, done: done}
288
+ <-done
289
290
// Cleanup function
291
defer func() {
287
- g.limitsLock.Lock()
288
- if g.relayedPerLeaseCount[leaseID] > 0 {
289
- g.relayedPerLeaseCount[leaseID]--
290
- }
291
- g.limitsLock.Unlock()
292
-
293
- g.relayedConnectionsLock.Lock()
294
- if streams, ok := g.relayedConnections[leaseID]; ok {
295
- for i, s := range streams {
296
- if s == clientStream {
297
- g.relayedConnections[leaseID] = append(streams[:i], streams[i+1:]...)
298
- break
299
- }
300
- }
301
- }
302
- g.relayedConnectionsLock.Unlock()
292
+ // Decrement limit via event loop
293
+ done := make(chan struct{}, 1)
294
+ g.cmdCh <- &cmdDecLimit{leaseID: leaseID, done: done}
295
+ <-done
296
+
297
+ // Remove stream from tracking via event loop
298
+ done2 := make(chan struct{}, 1)
299
+ g.cmdCh <- &cmdRemoveRelayed{leaseID: leaseID, stream: clientStream, done: done2}
300
+ <-done2
301
}()
302
303
// Use callback for actual relay (handles BPS limiting in relay-server)
portal/lease.go
+296
-117
@@ -9,6 +9,7 @@ import (
9
"gosuda.org/portal/portal/core/proto/rdverb"
10
)
11
12
+// LeaseEntry represents a registered lease with expiration tracking.
13
type LeaseEntry struct {
14
Lease *rdverb.Lease
15
Expires time.Time
@@ -16,12 +17,110 @@ type LeaseEntry struct {
17
ConnectionID int64 // Store the connection ID
18
}
19
20
+// leaseCmd is the command interface for LeaseManager event loop.
21
+type leaseCmd interface{ leaseCmd() }
22
+
23
+type cmdUpdateLease struct {
24
+ lease *rdverb.Lease
25
+ connID int64
26
+ reply chan<- bool
27
+}
28
+
29
+type cmdDeleteLease struct {
30
+ identity *rdsec.Identity
31
+ reply chan<- bool
32
+}
33
+
34
+type cmdGetLease struct {
35
+ identity *rdsec.Identity
36
+ reply chan<- leaseResult
37
+}
38
+
39
+type cmdGetLeaseByID struct {
40
+ leaseID string
41
+ reply chan<- leaseResult
42
+}
43
+
44
+type cmdGetAllLeases struct {
45
+ reply chan<- []*rdverb.Lease
46
+}
47
+
48
+type cmdCleanupByConnID struct {
49
+ connID int64
50
+ reply chan<- []string
51
+}
52
+
53
+type cmdBan struct {
54
+ leaseID string
55
+ done chan<- struct{}
56
+}
57
+
58
+type cmdUnban struct {
59
+ leaseID string
60
+ done chan<- struct{}
61
+}
62
+
63
+type cmdGetBanned struct {
64
+ reply chan<- [][]byte
65
+}
66
+
67
+type cmdSetPattern struct {
68
+ pattern string
69
+ reply chan<- error
70
+}
71
+
72
+type cmdSetTTL struct {
73
+ min time.Duration
74
+ max time.Duration
75
+ done chan<- struct{}
76
+}
77
+
78
+type cmdCleanExpired struct {
79
+ done chan<- struct{}
80
+}
81
+
82
+type cmdGetAllEntries struct {
83
+ reply chan<- []*LeaseEntry
84
+}
85
+
86
+type cmdGetLeaseALPNs struct {
87
+ leaseID string
88
+ reply chan<- []string
89
+}
90
+
91
+// leaseResult carries a LeaseEntry lookup result.
92
+type leaseResult struct {
93
+ entry *LeaseEntry
94
+ exists bool
95
+}
96
+
97
+// Command interface markers
98
+func (cmdUpdateLease) leaseCmd() {}
99
+func (cmdDeleteLease) leaseCmd() {}
100
+func (cmdGetLease) leaseCmd() {}
101
+func (cmdGetLeaseByID) leaseCmd() {}
102
+func (cmdGetAllLeases) leaseCmd() {}
103
+func (cmdCleanupByConnID) leaseCmd() {}
104
+func (cmdBan) leaseCmd() {}
105
+func (cmdUnban) leaseCmd() {}
106
+func (cmdGetBanned) leaseCmd() {}
107
+func (cmdSetPattern) leaseCmd() {}
108
+func (cmdSetTTL) leaseCmd() {}
109
+func (cmdCleanExpired) leaseCmd() {}
110
+func (cmdGetAllEntries) leaseCmd() {}
111
+func (cmdGetLeaseALPNs) leaseCmd() {}
112
+
113
+// LeaseManager manages lease registrations using a single-threaded event loop.
114
+// All state mutations are processed sequentially via the command channel.
115
type LeaseManager struct {
116
leases map[string]*LeaseEntry // Key: identity ID
21
- leasesLock sync.RWMutex
117
stopCh chan struct{}
118
ttlInterval time.Duration
119
120
+ // Event loop
121
+ cmdCh chan leaseCmd
122
+ runWg sync.WaitGroup
123
+
124
// policy controls
125
bannedLeases map[string]struct{}
126
namePattern *regexp.Regexp
@@ -34,48 +133,156 @@ func NewLeaseManager(ttlInterval time.Duration) *LeaseManager {
133
leases: make(map[string]*LeaseEntry),
134
stopCh: make(chan struct{}),
135
ttlInterval: ttlInterval,
136
+ cmdCh: make(chan leaseCmd, 256),
137
bannedLeases: make(map[string]struct{}),
138
}
139
}
140
141
func (lm *LeaseManager) Start() {
142
+ lm.runWg.Add(1)
143
+ go lm.run()
144
go lm.ttlWorker()
145
}
146
147
func (lm *LeaseManager) Stop() {
148
close(lm.stopCh)
149
+ lm.runWg.Wait()
150
}
151
49
-func (lm *LeaseManager) ttlWorker() {
50
- ticker := time.NewTicker(lm.ttlInterval)
51
- defer ticker.Stop()
52
-
152
+// run is the main event loop that processes all commands sequentially.
153
+func (lm *LeaseManager) run() {
154
+ defer lm.runWg.Done()
155
for {
156
select {
55
- case <-ticker.C:
56
- lm.cleanupExpiredLeases()
157
+ case cmd := <-lm.cmdCh:
158
+ lm.handleCmd(cmd)
159
case <-lm.stopCh:
160
return
161
}
162
}
163
}
164
63
-func (lm *LeaseManager) cleanupExpiredLeases() {
64
- lm.leasesLock.Lock()
65
- defer lm.leasesLock.Unlock()
165
+// handleCmd dispatches commands to their handlers.
166
+func (lm *LeaseManager) handleCmd(cmd leaseCmd) {
167
+ switch c := cmd.(type) {
168
+ case *cmdUpdateLease:
169
+ c.reply <- lm.updateLeaseInternal(c.lease, c.connID)
170
+
171
+ case *cmdDeleteLease:
172
+ identityID := string(c.identity.Id)
173
+ if _, exists := lm.leases[identityID]; exists {
174
+ delete(lm.leases, identityID)
175
+ c.reply <- true
176
+ } else {
177
+ c.reply <- false
178
+ }
179
+
180
+ case *cmdGetLease:
181
+ identityID := string(c.identity.Id)
182
+ entry, exists := lm.leases[identityID]
183
+ if !exists || time.Now().After(entry.Expires) {
184
+ c.reply <- leaseResult{nil, false}
185
+ } else {
186
+ c.reply <- leaseResult{entry, true}
187
+ }
188
+
189
+ case *cmdGetLeaseByID:
190
+ if _, banned := lm.bannedLeases[c.leaseID]; banned {
191
+ c.reply <- leaseResult{nil, false}
192
+ return
193
+ }
194
+ entry, exists := lm.leases[c.leaseID]
195
+ if !exists || time.Now().After(entry.Expires) {
196
+ c.reply <- leaseResult{nil, false}
197
+ } else {
198
+ c.reply <- leaseResult{entry, true}
199
+ }
200
+
201
+ case *cmdGetAllLeases:
202
+ now := time.Now()
203
+ var valid []*rdverb.Lease
204
+ for _, entry := range lm.leases {
205
+ if now.Before(entry.Expires) {
206
+ valid = append(valid, entry.Lease)
207
+ }
208
+ }
209
+ c.reply <- valid
210
+
211
+ case *cmdCleanupByConnID:
212
+ var cleaned []string
213
+ for leaseID, entry := range lm.leases {
214
+ if entry.ConnectionID == c.connID {
215
+ delete(lm.leases, leaseID)
216
+ cleaned = append(cleaned, leaseID)
217
+ }
218
+ }
219
+ c.reply <- cleaned
220
+
221
+ case *cmdBan:
222
+ lm.bannedLeases[c.leaseID] = struct{}{}
223
+ c.done <- struct{}{}
224
+
225
+ case *cmdUnban:
226
+ delete(lm.bannedLeases, c.leaseID)
227
+ c.done <- struct{}{}
228
+
229
+ case *cmdGetBanned:
230
+ banned := make([][]byte, 0, len(lm.bannedLeases))
231
+ for id := range lm.bannedLeases {
232
+ banned = append(banned, []byte(id))
233
+ }
234
+ c.reply <- banned
235
67
- now := time.Now()
68
- for id, lease := range lm.leases {
69
- if now.After(lease.Expires) {
70
- delete(lm.leases, id)
236
+ case *cmdSetPattern:
237
+ if c.pattern == "" {
238
+ lm.namePattern = nil
239
+ c.reply <- nil
240
+ return
241
+ }
242
+ re, err := regexp.Compile(c.pattern)
243
+ if err != nil {
244
+ c.reply <- err
245
+ return
246
+ }
247
+ lm.namePattern = re
248
+ c.reply <- nil
249
+
250
+ case *cmdSetTTL:
251
+ lm.minTTL = c.min
252
+ lm.maxTTL = c.max
253
+ c.done <- struct{}{}
254
+
255
+ case *cmdCleanExpired:
256
+ now := time.Now()
257
+ for id, entry := range lm.leases {
258
+ if now.After(entry.Expires) {
259
+ delete(lm.leases, id)
260
+ }
261
+ }
262
+ c.done <- struct{}{}
263
+
264
+ case *cmdGetAllEntries:
265
+ now := time.Now()
266
+ var entries []*LeaseEntry
267
+ for _, entry := range lm.leases {
268
+ if now.Before(entry.Expires) {
269
+ entries = append(entries, entry)
270
+ }
271
+ }
272
+ c.reply <- entries
273
+
274
+ case *cmdGetLeaseALPNs:
275
+ entry, exists := lm.leases[c.leaseID]
276
+ if !exists || time.Now().After(entry.Expires) {
277
+ c.reply <- nil
278
+ } else {
279
+ c.reply <- entry.Lease.Alpn
280
}
281
}
282
}
283
75
-func (lm *LeaseManager) UpdateLease(lease *rdverb.Lease, connectionID int64) bool {
76
- lm.leasesLock.Lock()
77
- defer lm.leasesLock.Unlock()
78
-
284
+// updateLeaseInternal contains the lease update logic, called within the event loop.
285
+func (lm *LeaseManager) updateLeaseInternal(lease *rdverb.Lease, connectionID int64) bool {
286
identityID := string(lease.Identity.Id)
287
expires := time.Unix(lease.Expires, 0)
288
@@ -91,7 +298,6 @@ func (lm *LeaseManager) UpdateLease(lease *rdverb.Lease, connectionID int64) boo
298
if lm.namePattern != nil && lease.Name != "" && !lm.namePattern.MatchString(lease.Name) {
299
return false
300
}
94
- // reserved prefix check removed
301
if lm.minTTL > 0 || lm.maxTTL > 0 {
302
ttl := time.Until(expires)
303
if lm.minTTL > 0 && ttl < lm.minTTL {
@@ -105,13 +311,10 @@ func (lm *LeaseManager) UpdateLease(lease *rdverb.Lease, connectionID int64) boo
311
// Check for name conflicts (only if name is not empty)
312
if lease.Name != "" && lease.Name != "(unnamed)" {
313
for existingID, existingEntry := range lm.leases {
108
- // Skip if it's the same identity (updating own lease)
314
if existingID == identityID {
315
continue
316
}
112
- // Check if another identity is using the same name
317
if existingEntry.Lease.Name == lease.Name {
114
- // Name conflict with a different identity
318
return false
319
}
320
}
@@ -127,133 +330,109 @@ func (lm *LeaseManager) UpdateLease(lease *rdverb.Lease, connectionID int64) boo
330
return true
331
}
332
130
-func (lm *LeaseManager) DeleteLease(identity *rdsec.Identity) bool {
131
- lm.leasesLock.Lock()
132
- defer lm.leasesLock.Unlock()
333
+func (lm *LeaseManager) ttlWorker() {
334
+ ticker := time.NewTicker(lm.ttlInterval)
335
+ defer ticker.Stop()
336
134
- identityID := string(identity.Id)
135
- if _, exists := lm.leases[identityID]; exists {
136
- delete(lm.leases, identityID)
137
- return true
337
+ for {
338
+ select {
339
+ case <-ticker.C:
340
+ done := make(chan struct{}, 1)
341
+ select {
342
+ case lm.cmdCh <- &cmdCleanExpired{done: done}:
343
+ <-done
344
+ case <-lm.stopCh:
345
+ return
346
+ }
347
+ case <-lm.stopCh:
348
+ return
349
+ }
350
}
139
- return false
351
}
352
142
-func (lm *LeaseManager) GetLease(identity *rdsec.Identity) (*LeaseEntry, bool) {
143
- lm.leasesLock.RLock()
144
- defer lm.leasesLock.RUnlock()
145
-
146
- identityID := string(identity.Id)
147
-
148
- lease, exists := lm.leases[identityID]
149
- if !exists {
150
- return nil, false
151
- }
353
+func (lm *LeaseManager) UpdateLease(lease *rdverb.Lease, connectionID int64) bool {
354
+ reply := make(chan bool, 1)
355
+ lm.cmdCh <- &cmdUpdateLease{lease: lease, connID: connectionID, reply: reply}
356
+ return <-reply
357
+}
358
153
- // Check if lease is expired
154
- if time.Now().After(lease.Expires) {
155
- return nil, false
156
- }
359
+func (lm *LeaseManager) DeleteLease(identity *rdsec.Identity) bool {
360
+ reply := make(chan bool, 1)
361
+ lm.cmdCh <- &cmdDeleteLease{identity: identity, reply: reply}
362
+ return <-reply
363
+}
364
158
- return lease, true
365
+func (lm *LeaseManager) GetLease(identity *rdsec.Identity) (*LeaseEntry, bool) {
366
+ reply := make(chan leaseResult, 1)
367
+ lm.cmdCh <- &cmdGetLease{identity: identity, reply: reply}
368
+ result := <-reply
369
+ return result.entry, result.exists
370
}
371
372
func (lm *LeaseManager) GetLeaseByID(leaseID string) (*LeaseEntry, bool) {
162
- lm.leasesLock.RLock()
163
- defer lm.leasesLock.RUnlock()
164
-
165
- // Check if banned
166
- if _, banned := lm.bannedLeases[leaseID]; banned {
167
- return nil, false
168
- }
169
-
170
- lease, exists := lm.leases[leaseID]
171
- if !exists {
172
- return nil, false
173
- }
174
-
175
- // Check if lease is expired
176
- if time.Now().After(lease.Expires) {
177
- return nil, false
178
- }
179
-
180
- return lease, true
373
+ reply := make(chan leaseResult, 1)
374
+ lm.cmdCh <- &cmdGetLeaseByID{leaseID: leaseID, reply: reply}
375
+ result := <-reply
376
+ return result.entry, result.exists
377
}
378
379
func (lm *LeaseManager) GetAllLeases() []*rdverb.Lease {
184
- lm.leasesLock.RLock()
185
- defer lm.leasesLock.RUnlock()
186
-
187
- now := time.Now()
188
- var validLeases []*rdverb.Lease
189
-
190
- for _, lease := range lm.leases {
191
- if now.Before(lease.Expires) {
192
- validLeases = append(validLeases, lease.Lease)
193
- }
194
- }
195
-
196
- return validLeases
380
+ reply := make(chan []*rdverb.Lease, 1)
381
+ lm.cmdCh <- &cmdGetAllLeases{reply: reply}
382
+ return <-reply
383
}
384
385
// Lease policy configuration helpers
386
func (lm *LeaseManager) BanLease(leaseID string) {
201
- lm.leasesLock.Lock()
202
- lm.bannedLeases[leaseID] = struct{}{}
203
- lm.leasesLock.Unlock()
387
+ done := make(chan struct{}, 1)
388
+ lm.cmdCh <- &cmdBan{leaseID: leaseID, done: done}
389
+ <-done
390
}
391
392
func (lm *LeaseManager) UnbanLease(leaseID string) {
207
- lm.leasesLock.Lock()
208
- delete(lm.bannedLeases, leaseID)
209
- lm.leasesLock.Unlock()
393
+ done := make(chan struct{}, 1)
394
+ lm.cmdCh <- &cmdUnban{leaseID: leaseID, done: done}
395
+ <-done
396
}
397
398
func (lm *LeaseManager) GetBannedLeases() [][]byte {
213
- lm.leasesLock.RLock()
214
- defer lm.leasesLock.RUnlock()
215
- banned := make([][]byte, 0, len(lm.bannedLeases))
216
- for id := range lm.bannedLeases {
217
- banned = append(banned, []byte(id))
218
- }
219
- return banned
399
+ reply := make(chan [][]byte, 1)
400
+ lm.cmdCh <- &cmdGetBanned{reply: reply}
401
+ return <-reply
402
}
403
404
func (lm *LeaseManager) SetNamePattern(pattern string) error {
223
- lm.leasesLock.Lock()
224
- defer lm.leasesLock.Unlock()
225
- if pattern == "" {
226
- lm.namePattern = nil
227
- return nil
228
- }
229
- re, err := regexp.Compile(pattern)
230
- if err != nil {
231
- return err
232
- }
233
- lm.namePattern = re
234
- return nil
405
+ reply := make(chan error, 1)
406
+ lm.cmdCh <- &cmdSetPattern{pattern: pattern, reply: reply}
407
+ return <-reply
408
}
409
410
// SetReservedPrefixes removed: reserved prefix policy no longer supported
411
412
func (lm *LeaseManager) SetTTLBounds(min, max time.Duration) {
240
- lm.leasesLock.Lock()
241
- lm.minTTL = min
242
- lm.maxTTL = max
243
- lm.leasesLock.Unlock()
413
+ done := make(chan struct{}, 1)
414
+ lm.cmdCh <- &cmdSetTTL{min: min, max: max, done: done}
415
+ <-done
416
}
417
418
func (lm *LeaseManager) CleanupLeasesByConnectionID(connectionID int64) []string {
247
- lm.leasesLock.Lock()
248
- defer lm.leasesLock.Unlock()
249
-
250
- var cleanedLeaseIDs []string
251
- for leaseID, lease := range lm.leases {
252
- if lease.ConnectionID == connectionID {
253
- delete(lm.leases, leaseID)
254
- cleanedLeaseIDs = append(cleanedLeaseIDs, leaseID)
255
- }
256
- }
419
+ reply := make(chan []string, 1)
420
+ lm.cmdCh <- &cmdCleanupByConnID{connID: connectionID, reply: reply}
421
+ return <-reply
422
+}
423
+
424
+// GetAllEntries returns all valid (non-expired) lease entries.
425
+// This method is used by RelayServer.GetAllLeaseEntries.
426
+func (lm *LeaseManager) GetAllEntries() []*LeaseEntry {
427
+ reply := make(chan []*LeaseEntry, 1)
428
+ lm.cmdCh <- &cmdGetAllEntries{reply: reply}
429
+ return <-reply
430
+}
431
258
- return cleanedLeaseIDs
432
+// GetLeaseALPNs returns the ALPN identifiers for a given lease ID.
433
+// This method is used by RelayServer.GetLeaseALPNs.
434
+func (lm *LeaseManager) GetLeaseALPNs(leaseID string) []string {
435
+ reply := make(chan []string, 1)
436
+ lm.cmdCh <- &cmdGetLeaseALPNs{leaseID: leaseID, reply: reply}
437
+ return <-reply
438
}
portal/lease_test.go
+5
@@ -10,6 +10,7 @@ import (
10
11
func TestLeaseManager_NameConflict(t *testing.T) {
12
lm := NewLeaseManager(30 * time.Second)
13
+ lm.Start()
14
defer lm.Stop()
15
16
// Create two different identities
@@ -67,6 +68,7 @@ func TestLeaseManager_NameConflict(t *testing.T) {
68
69
func TestLeaseManager_SameIdentityUpdate(t *testing.T) {
70
lm := NewLeaseManager(30 * time.Second)
71
+ lm.Start()
72
defer lm.Stop()
73
74
identity := &rdsec.Identity{
@@ -112,6 +114,7 @@ func TestLeaseManager_SameIdentityUpdate(t *testing.T) {
114
115
func TestLeaseManager_EmptyNameAllowed(t *testing.T) {
116
lm := NewLeaseManager(30 * time.Second)
117
+ lm.Start()
118
defer lm.Stop()
119
120
identity1 := &rdsec.Identity{
@@ -150,6 +153,7 @@ func TestLeaseManager_EmptyNameAllowed(t *testing.T) {
153
154
func TestLeaseManager_UnnamedAllowed(t *testing.T) {
155
lm := NewLeaseManager(30 * time.Second)
156
+ lm.Start()
157
defer lm.Stop()
158
159
identity1 := &rdsec.Identity{
@@ -188,6 +192,7 @@ func TestLeaseManager_UnnamedAllowed(t *testing.T) {
192
193
func TestLeaseManager_UnicodeNameConflict(t *testing.T) {
194
lm := NewLeaseManager(30 * time.Second)
195
+ lm.Start()
196
defer lm.Stop()
197
198
identity1 := &rdsec.Identity{
portal/relay.go
+204
-77
@@ -3,6 +3,7 @@ package portal
3
import (
4
"io"
5
"sync"
6
+ "sync/atomic"
7
"time"
8
9
"github.com/hashicorp/yamux"
@@ -20,30 +21,109 @@ type Connection struct {
21
streamsLock sync.Mutex
22
}
23
24
+// relayCmd is the command interface for RelayServer event loop.
25
+type relayCmd interface{ relayCmd() }
26
+
27
+// Connection management commands
28
+type cmdRegisterConn struct {
29
+ id int64
30
+ conn *Connection
31
+ done chan<- struct{}
32
+}
33
+
34
+type cmdRemoveConn struct {
35
+ id int64
36
+ done chan<- []string // returns cleaned lease IDs
37
+}
38
+
39
+type cmdIsConnActive struct {
40
+ id int64
41
+ reply chan<- bool
42
+}
43
+
44
+type cmdGetConnByLeaseEntry struct {
45
+ connID int64
46
+ reply chan<- *Connection
47
+}
48
+
49
+// Lease connection management commands
50
+type cmdRegisterLeaseConn struct {
51
+ leaseID string
52
+ conn *Connection
53
+ done chan<- struct{}
54
+}
55
+
56
+type cmdUnregisterLeaseConn struct {
57
+ leaseID string
58
+ done chan<- struct{}
59
+}
60
+
61
+// Relayed connection tracking commands
62
+type cmdAddRelayed struct {
63
+ leaseID string
64
+ stream *yamux.Stream
65
+ done chan<- struct{}
66
+}
67
+
68
+type cmdRemoveRelayed struct {
69
+ leaseID string
70
+ stream *yamux.Stream
71
+ done chan<- struct{}
72
+}
73
+
74
+// Limit management commands
75
+type cmdCheckAndIncLimit struct {
76
+ leaseID string
77
+ reply chan<- bool // true if under limit and incremented
78
+}
79
+
80
+type cmdDecLimit struct {
81
+ leaseID string
82
+ done chan<- struct{}
83
+}
84
+
85
+type cmdSetMaxRelayed struct {
86
+ max int
87
+ done chan<- struct{}
88
+}
89
+
90
+// Command interface markers
91
+func (cmdRegisterConn) relayCmd() {}
92
+func (cmdRemoveConn) relayCmd() {}
93
+func (cmdIsConnActive) relayCmd() {}
94
+func (cmdGetConnByLeaseEntry) relayCmd() {}
95
+func (cmdRegisterLeaseConn) relayCmd() {}
96
+func (cmdUnregisterLeaseConn) relayCmd() {}
97
+func (cmdAddRelayed) relayCmd() {}
98
+func (cmdRemoveRelayed) relayCmd() {}
99
+func (cmdCheckAndIncLimit) relayCmd() {}
100
+func (cmdDecLimit) relayCmd() {}
101
+func (cmdSetMaxRelayed) relayCmd() {}
102
+
103
+// RelayServer handles relay connections using a single-threaded event loop.
104
+// All state mutations are processed sequentially via the command channel.
105
type RelayServer struct {
106
credential *cryptoops.Credential
107
identity *rdsec.Identity
108
address []string
109
28
- connidCounter int64
29
- connections map[int64]*Connection
30
- connectionsLock sync.RWMutex
31
-
32
- leaseConnections map[string]*Connection // Key: lease ID, Value: Connection
33
- leaseConnectionsLock sync.RWMutex
34
-
35
- relayedConnections map[string][]*yamux.Stream // Key: lease ID, Value: slice of relayed streams
36
- relayedConnectionsLock sync.RWMutex
110
+ connidCounter int64
111
+ connections map[int64]*Connection
112
+ leaseConnections map[string]*Connection // Key: lease ID, Value: Connection
113
+ relayedConnections map[string][]*yamux.Stream // Key: lease ID, Value: slice of relayed streams
114
115
leaseManager *LeaseManager
116
117
+ // Event loop
118
+ cmdCh chan relayCmd
119
+ runWg sync.WaitGroup
120
+
121
stopch chan struct{}
122
waitgroup sync.WaitGroup
123
124
// Traffic control limits and counters
125
maxRelayedPerLease int
126
relayedPerLeaseCount map[string]int
46
- limitsLock sync.Mutex
127
128
// Callback for relay connection establishment (set by relay-server for BPS handling)
129
onEstablishRelay func(clientStream, leaseStream *yamux.Stream, leaseID string)
@@ -62,6 +142,7 @@ func NewRelayServer(credential *cryptoops.Credential, address []string) *RelaySe
142
leaseConnections: make(map[string]*Connection),
143
relayedConnections: make(map[string][]*yamux.Stream),
144
leaseManager: NewLeaseManager(30 * time.Second), // TTL check every 30 seconds
145
+ cmdCh: make(chan relayCmd, 256),
146
stopch: make(chan struct{}),
147
relayedPerLeaseCount: make(map[string]int),
148
}
@@ -75,8 +156,10 @@ func (g *RelayServer) handleConn(id int64, connection *Connection) {
156
defer func() {
157
log.Debug().Int64("conn_id", id).Msg("[RelayServer] Connection closing, cleaning up")
158
78
- // Clean up leases associated with this connection when it closes
79
- cleanedLeaseIDs := g.leaseManager.CleanupLeasesByConnectionID(id)
159
+ // Remove connection and clean up all associated state via event loop
160
+ done := make(chan []string, 1)
161
+ g.cmdCh <- &cmdRemoveConn{id: id, done: done}
162
+ cleanedLeaseIDs := <-done
163
164
if len(cleanedLeaseIDs) > 0 {
165
log.Debug().
@@ -85,31 +168,6 @@ func (g *RelayServer) handleConn(id int64, connection *Connection) {
168
Msg("[RelayServer] Cleaned up leases for connection")
169
}
170
88
- // Also clean up lease connections mapping
89
- g.leaseConnectionsLock.Lock()
90
- for _, leaseID := range cleanedLeaseIDs {
91
- delete(g.leaseConnections, leaseID)
92
- }
93
- g.leaseConnectionsLock.Unlock()
94
-
95
- // Clean up relayed connections for these leases
96
- g.relayedConnectionsLock.Lock()
97
- for _, leaseID := range cleanedLeaseIDs {
98
- if streams, exists := g.relayedConnections[leaseID]; exists {
99
- // Close all relayed streams
100
- for _, stream := range streams {
101
- stream.Close()
102
- }
103
- delete(g.relayedConnections, leaseID)
104
- }
105
- }
106
- g.relayedConnectionsLock.Unlock()
107
-
108
- // Remove the connection itself
109
- g.connectionsLock.Lock()
110
- delete(g.connections, id)
111
- g.connectionsLock.Unlock()
112
-
171
// Close the underlying connection
172
connection.conn.Close()
173
@@ -233,16 +291,17 @@ func (g *RelayServer) HandleConnection(conn io.ReadWriteCloser) error {
291
return err
292
}
293
236
- g.connectionsLock.Lock()
237
- g.connidCounter++
238
- connID := g.connidCounter
294
+ connID := atomic.AddInt64(&g.connidCounter, 1)
295
connection := &Connection{
296
conn: conn,
297
sess: sess,
298
streams: make(map[uint32]*yamux.Stream),
299
}
244
- g.connections[connID] = connection
245
- g.connectionsLock.Unlock()
300
+
301
+ // Register connection via event loop
302
+ done := make(chan struct{}, 1)
303
+ g.cmdCh <- &cmdRegisterConn{id: connID, conn: connection, done: done}
304
+ <-done
305
306
log.Debug().Int64("conn_id", connID).Msg("[RelayServer] Connection registered, starting handler")
307
go g.handleConn(connID, connection)
@@ -265,63 +324,131 @@ func (g *RelayServer) GetLeaseManager() *LeaseManager {
324
325
// IsConnectionActive checks if a connection with the given ID is still active
326
func (g *RelayServer) IsConnectionActive(connectionID int64) bool {
268
- g.connectionsLock.RLock()
269
- defer g.connectionsLock.RUnlock()
270
-
271
- _, exists := g.connections[connectionID]
272
- return exists
327
+ reply := make(chan bool, 1)
328
+ g.cmdCh <- &cmdIsConnActive{id: connectionID, reply: reply}
329
+ return <-reply
330
}
331
332
// GetAllLeaseEntries returns all lease entries from the lease manager
333
func (g *RelayServer) GetAllLeaseEntries() []*LeaseEntry {
277
- g.leaseManager.leasesLock.RLock()
278
- defer g.leaseManager.leasesLock.RUnlock()
279
-
280
- var entries []*LeaseEntry
281
- now := time.Now()
282
-
283
- for _, entry := range g.leaseManager.leases {
284
- if now.Before(entry.Expires) {
285
- entries = append(entries, entry)
286
- }
287
- }
288
-
289
- return entries
334
+ return g.leaseManager.GetAllEntries()
335
}
336
337
// GetLeaseALPNs returns the ALPN identifiers for a given lease ID
338
func (g *RelayServer) GetLeaseALPNs(leaseID string) []string {
294
- g.leaseManager.leasesLock.RLock()
295
- defer g.leaseManager.leasesLock.RUnlock()
296
-
297
- entry, exists := g.leaseManager.leases[leaseID]
298
- if !exists {
299
- return nil
300
- }
301
-
302
- now := time.Now()
303
- if now.After(entry.Expires) {
304
- return nil
305
- }
306
-
307
- return entry.Lease.Alpn
339
+ return g.leaseManager.GetLeaseALPNs(leaseID)
340
}
341
342
func (g *RelayServer) Start() {
343
+ g.runWg.Add(1)
344
+ go g.run()
345
g.leaseManager.Start()
346
}
347
348
func (g *RelayServer) Stop() {
349
close(g.stopch)
350
+ g.runWg.Wait()
351
g.leaseManager.Stop()
352
g.waitgroup.Wait()
353
}
354
355
+// run is the main event loop that processes all commands sequentially.
356
+func (g *RelayServer) run() {
357
+ defer g.runWg.Done()
358
+ for {
359
+ select {
360
+ case cmd := <-g.cmdCh:
361
+ g.handleCmd(cmd)
362
+ case <-g.stopch:
363
+ return
364
+ }
365
+ }
366
+}
367
+
368
+// handleCmd dispatches commands to their handlers.
369
+func (g *RelayServer) handleCmd(cmd relayCmd) {
370
+ switch c := cmd.(type) {
371
+ case *cmdRegisterConn:
372
+ g.connections[c.id] = c.conn
373
+ c.done <- struct{}{}
374
+
375
+ case *cmdRemoveConn:
376
+ delete(g.connections, c.id)
377
+ // Clean up leases associated with this connection
378
+ cleanedLeaseIDs := g.leaseManager.CleanupLeasesByConnectionID(c.id)
379
+ // Clean up lease connections and relayed connections
380
+ for _, leaseID := range cleanedLeaseIDs {
381
+ delete(g.leaseConnections, leaseID)
382
+ if streams, exists := g.relayedConnections[leaseID]; exists {
383
+ for _, stream := range streams {
384
+ stream.Close()
385
+ }
386
+ delete(g.relayedConnections, leaseID)
387
+ }
388
+ delete(g.relayedPerLeaseCount, leaseID)
389
+ }
390
+ c.done <- cleanedLeaseIDs
391
+
392
+ case *cmdIsConnActive:
393
+ _, exists := g.connections[c.id]
394
+ c.reply <- exists
395
+
396
+ case *cmdGetConnByLeaseEntry:
397
+ conn, exists := g.connections[c.connID]
398
+ if exists {
399
+ c.reply <- conn
400
+ } else {
401
+ c.reply <- nil
402
+ }
403
+
404
+ case *cmdRegisterLeaseConn:
405
+ g.leaseConnections[c.leaseID] = c.conn
406
+ c.done <- struct{}{}
407
+
408
+ case *cmdUnregisterLeaseConn:
409
+ delete(g.leaseConnections, c.leaseID)
410
+ c.done <- struct{}{}
411
+
412
+ case *cmdAddRelayed:
413
+ g.relayedConnections[c.leaseID] = append(g.relayedConnections[c.leaseID], c.stream)
414
+ c.done <- struct{}{}
415
+
416
+ case *cmdRemoveRelayed:
417
+ if streams, ok := g.relayedConnections[c.leaseID]; ok {
418
+ for i, s := range streams {
419
+ if s == c.stream {
420
+ g.relayedConnections[c.leaseID] = append(streams[:i], streams[i+1:]...)
421
+ break
422
+ }
423
+ }
424
+ }
425
+ c.done <- struct{}{}
426
+
427
+ case *cmdCheckAndIncLimit:
428
+ if g.maxRelayedPerLease > 0 && g.relayedPerLeaseCount[c.leaseID] >= g.maxRelayedPerLease {
429
+ c.reply <- false
430
+ } else {
431
+ g.relayedPerLeaseCount[c.leaseID]++
432
+ c.reply <- true
433
+ }
434
+
435
+ case *cmdDecLimit:
436
+ if g.relayedPerLeaseCount[c.leaseID] > 0 {
437
+ g.relayedPerLeaseCount[c.leaseID]--
438
+ }
439
+ c.done <- struct{}{}
440
+
441
+ case *cmdSetMaxRelayed:
442
+ g.maxRelayedPerLease = c.max
443
+ c.done <- struct{}{}
444
+ }
445
+}
446
+
447
// Traffic control setters
448
func (g *RelayServer) SetMaxRelayedPerLease(n int) {
322
- g.limitsLock.Lock()
323
- g.maxRelayedPerLease = n
324
- g.limitsLock.Unlock()
449
+ done := make(chan struct{}, 1)
450
+ g.cmdCh <- &cmdSetMaxRelayed{max: n, done: done}
451
+ <-done
452
}
453
454
// SetEstablishRelayCallback sets the callback for relay connection establishment