sdk: simplify listener retry
rabbitprincess committed
Mar 11, 2026 at 00:52 UTC
3fd12c9885ab59176e113770379b827ebebde5a1
3 files changed
+169
-97
sdk/listener.go
+109
-97
@@ -28,27 +28,29 @@ type ListenerConfig struct {
28
LeaseTTL time.Duration
29
RenewBefore time.Duration
30
ReadyTarget int
31
+ RetryCount int
32
+ RetryWait time.Duration
33
}
34
35
type Listener struct {
36
tlsCloser io.Closer
37
tlsConfig *tls.Config
38
readyTarget int
39
+ retryCount int
40
+ retryWait time.Duration
41
leaseTTL time.Duration
42
renewBefore time.Duration
43
handshakeTimeout time.Duration
44
ctx context.Context
45
cancel context.CancelFunc
46
api *relayClient
43
- signal chan struct{}
47
accepted chan net.Conn
48
leaseID string
49
hostnames []string
50
metadata types.LeaseMetadata
51
49
- activeSessions int
50
- closeOnce sync.Once
51
- mu sync.Mutex
52
+ closeOnce sync.Once
53
+ mu sync.Mutex
54
}
55
56
// NewListener creates one relay listener and its dedicated relay transport for one relay URL.
@@ -74,6 +76,10 @@ func NewListener(ctx context.Context, relayURL string, cfg ListenerConfig) (*Lis
76
if renewBefore <= 0 {
77
renewBefore = defaultRenewBefore
78
}
79
+ retryWait := cfg.RetryWait
80
+ if retryWait <= 0 {
81
+ retryWait = defaultRetryWait
82
+ }
83
84
api, err := newRelayClient(listenerCtx, relayURL, cfg)
85
if err != nil {
@@ -85,9 +91,10 @@ func NewListener(ctx context.Context, relayURL string, cfg ListenerConfig) (*Lis
91
ctx: listenerCtx,
92
cancel: cancel,
93
api: api,
88
- signal: make(chan struct{}, 1),
94
accepted: make(chan net.Conn, max(readyTarget*2, 1)),
95
readyTarget: readyTarget,
96
+ retryCount: cfg.RetryCount,
97
+ retryWait: retryWait,
98
leaseTTL: leaseTTL,
99
renewBefore: renewBefore,
100
handshakeTimeout: handshakeTimeout,
@@ -124,9 +131,10 @@ func NewListener(ctx context.Context, relayURL string, cfg ListenerConfig) (*Lis
131
l.tlsCloser = tlsCloser
132
l.mu.Unlock()
133
127
- go l.runSupervisor()
134
+ for i := 0; i < l.readyTarget; i++ {
135
+ go l.runSessionLoop()
136
+ }
137
go l.runRenewLoop()
129
- l.notify()
138
return l, nil
139
}
140
@@ -135,9 +143,6 @@ func (l *Listener) Accept() (net.Conn, error) {
143
case <-l.ctx.Done():
144
return nil, net.ErrClosed
145
case conn := <-l.accepted:
138
- if conn == nil {
139
- return nil, net.ErrClosed
140
- }
146
return conn, nil
147
}
148
}
@@ -156,7 +161,6 @@ func (l *Listener) Close() error {
161
l.leaseID = ""
162
l.tlsConfig = nil
163
l.tlsCloser = nil
159
- l.activeSessions = 0
164
l.mu.Unlock()
165
166
l.drainAccepted()
@@ -180,7 +184,7 @@ func (l *Listener) Addr() net.Addr {
184
l.mu.Lock()
185
defer l.mu.Unlock()
186
if l.leaseID == "" {
183
- return listenerAddr("portal:pending")
187
+ return listenerAddr("portal:closed")
188
}
189
return listenerAddr("portal:" + l.leaseID)
190
}
@@ -215,16 +219,25 @@ func (l *Listener) PublicURLs() []string {
219
return urls
220
}
221
218
-func (l *Listener) runSupervisor() {
222
+func (l *Listener) runSessionLoop() {
223
+ var retries int
224
+
225
for {
220
- select {
221
- case <-l.ctx.Done():
226
+ claimed, err := l.runSession()
227
+ switch {
228
+ case err == nil:
229
+ retries = 0
230
+ case errors.Is(err, context.Canceled), errors.Is(err, net.ErrClosed):
231
return
223
- case <-l.signal:
224
- }
225
-
226
- for l.reserveSessionSlot() {
227
- go l.runSession()
232
+ case claimed:
233
+ // A claimed connection already reached the data plane.
234
+ // Do not spend retry budget on browser-side TLS failures or disconnects.
235
+ retries = 0
236
+ default:
237
+ retries++
238
+ if !l.retryOrClose("reverse session connect", err, retries) {
239
+ return
240
+ }
241
}
242
}
243
}
@@ -241,65 +254,33 @@ func (l *Listener) runRenewLoop() {
254
interval = 30 * time.Second
255
}
256
244
- ticker := time.NewTicker(interval)
245
- defer ticker.Stop()
246
-
247
- var consecutiveFailures int
248
-
257
for {
250
- select {
251
- case <-l.ctx.Done():
258
+ sleepOrDone(l.context(), interval)
259
+ if l.isClosed() {
260
return
253
- case <-ticker.C:
254
- l.mu.Lock()
255
- leaseID := l.leaseID
256
- l.mu.Unlock()
257
-
258
- ctx, cancel := context.WithTimeout(l.context(), 10*time.Second)
259
- err := l.api.renewLease(ctx, leaseID, l.leaseTTL)
260
- cancel()
261
-
262
- if err != nil {
263
- if isLeaseNotFound(err) {
264
- log.Warn().
265
- Str("component", "sdk-listener").
266
- Str("lease_id", leaseID).
267
- Msg("lease not found on relay, attempting re-registration")
268
- if reregErr := l.reregister(); reregErr != nil {
269
- log.Error().Err(reregErr).
270
- Str("component", "sdk-listener").
271
- Msg("lease re-registration failed")
272
- } else {
273
- consecutiveFailures = 0
274
- log.Info().
275
- Str("component", "sdk-listener").
276
- Str("lease_id", l.LeaseID()).
277
- Strs("hostnames", l.Hostnames()).
278
- Msg("lease re-registered successfully")
279
- continue
280
- }
281
- }
261
+ }
262
283
- consecutiveFailures++
284
- event := log.Warn()
285
- if consecutiveFailures >= 3 {
286
- event = log.Error()
263
+ var retries int
264
+ for {
265
+ err := l.renewLease()
266
+ switch {
267
+ case err == nil:
268
+ goto nextRenew
269
+ case errors.Is(err, context.Canceled), errors.Is(err, net.ErrClosed):
270
+ return
271
+ default:
272
+ retries++
273
+ if !l.retryOrClose("lease renewal", err, retries) {
274
+ return
275
}
288
- event.Err(err).
289
- Str("component", "sdk-listener").
290
- Str("lease_id", l.LeaseID()).
291
- Int("consecutive_failures", consecutiveFailures).
292
- Msg("lease renewal failed")
293
- } else {
294
- consecutiveFailures = 0
276
}
277
}
278
+
279
+ nextRenew:
280
}
281
}
282
300
-func (l *Listener) runSession() {
301
- defer l.releaseSessionSlot()
302
-
283
+func (l *Listener) runSession() (bool, error) {
284
sessionCtx := l.context()
285
l.mu.Lock()
286
leaseID := l.leaseID
@@ -307,17 +288,15 @@ func (l *Listener) runSession() {
288
289
conn, err := l.api.openReverseSession(sessionCtx, leaseID)
290
if err != nil {
310
- sleepOrDone(sessionCtx, time.Second)
311
- return
291
+ return false, err
292
}
293
294
claimed, err := l.awaitActivation(conn)
295
if err != nil {
296
_ = conn.Close()
317
- if !claimed && !errors.Is(err, context.Canceled) && !errors.Is(err, net.ErrClosed) {
318
- sleepOrDone(sessionCtx, time.Second)
319
- }
297
+ return claimed, err
298
}
299
+ return claimed, nil
300
}
301
302
func (l *Listener) awaitActivation(conn net.Conn) (bool, error) {
@@ -364,6 +343,39 @@ func (l *Listener) activate(conn net.Conn) error {
343
}
344
}
345
346
+func (l *Listener) renewLease() error {
347
+ l.mu.Lock()
348
+ leaseID := l.leaseID
349
+ l.mu.Unlock()
350
+
351
+ ctx, cancel := context.WithTimeout(l.context(), 10*time.Second)
352
+ err := l.api.renewLease(ctx, leaseID, l.leaseTTL)
353
+ cancel()
354
+ if err == nil {
355
+ return nil
356
+ }
357
+ if !isLeaseNotFound(err) {
358
+ return err
359
+ }
360
+
361
+ log.Warn().
362
+ Err(err).
363
+ Str("component", "sdk-listener").
364
+ Str("lease_id", leaseID).
365
+ Msg("lease not found on relay, attempting re-registration")
366
+
367
+ if err := l.reregister(); err != nil {
368
+ return err
369
+ }
370
+
371
+ log.Info().
372
+ Str("component", "sdk-listener").
373
+ Str("lease_id", l.LeaseID()).
374
+ Strs("hostnames", l.Hostnames()).
375
+ Msg("lease re-registered successfully")
376
+ return nil
377
+}
378
+
379
func (l *Listener) reregister() error {
380
ctx, cancel := context.WithTimeout(l.context(), 10*time.Second)
381
defer cancel()
@@ -401,8 +413,6 @@ func (l *Listener) reregister() error {
413
if oldCloser != nil {
414
_ = oldCloser.Close()
415
}
404
-
405
- l.notify()
416
return nil
417
}
418
@@ -410,33 +420,35 @@ func isLeaseNotFound(err error) bool {
420
return errors.Is(err, &types.APIRequestError{Code: types.APIErrorCodeLeaseNotFound})
421
}
422
413
-func (l *Listener) reserveSessionSlot() bool {
414
- l.mu.Lock()
415
- defer l.mu.Unlock()
423
+func (l *Listener) retryOrClose(operation string, err error, retries int) bool {
424
if l.isClosed() {
425
return false
426
}
419
- if l.activeSessions >= l.readyTarget {
427
+
428
+ logger := log.With().
429
+ Str("component", "sdk-listener").
430
+ Str("operation", operation).
431
+ Str("lease_id", l.LeaseID()).
432
+ Logger()
433
+
434
+ if l.retryCount > 0 && retries > l.retryCount {
435
+ logger.Error().
436
+ Err(err).
437
+ Int("retry_count", l.retryCount).
438
+ Msg("retry budget exhausted; closing listener")
439
+ _ = l.Close()
440
return false
441
}
422
- l.activeSessions++
423
- return true
424
-}
442
426
-func (l *Listener) releaseSessionSlot() {
427
- l.mu.Lock()
428
- if l.activeSessions > 0 {
429
- l.activeSessions--
430
- }
431
- l.mu.Unlock()
432
- l.notify()
433
-}
443
+ logger.Warn().
444
+ Err(err).
445
+ Int("retry_attempt", retries).
446
+ Int("retry_count", l.retryCount).
447
+ Dur("retry_wait", l.retryWait).
448
+ Msg("operation failed; retrying")
449
435
-func (l *Listener) notify() {
436
- select {
437
- case l.signal <- struct{}{}:
438
- default:
439
- }
450
+ sleepOrDone(l.context(), l.retryWait)
451
+ return !l.isClosed()
452
}
453
454
func sleepOrDone(ctx context.Context, d time.Duration) {
sdk/relayclient.go
+1
@@ -29,6 +29,7 @@ const (
29
defaultLeaseTTL = 30 * time.Second
30
defaultRenewBefore = 30 * time.Second
31
defaultReadyTarget = 1
32
+ defaultRetryWait = 10 * time.Second
33
defaultHTTPShutdownTimeout = 5 * time.Second
34
)
35
sdk/sdk_test.go
+59
@@ -189,6 +189,65 @@ func TestNewListenerReregistersOnLeaseNotFound(t *testing.T) {
189
})
190
}
191
192
+func TestNewListenerClosesAfterReverseSessionRetryBudgetExhausted(t *testing.T) {
193
+ t.Parallel()
194
+
195
+ var connectCount atomic.Int32
196
+ var unregisterCount atomic.Int32
197
+ server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
198
+ switch r.URL.Path {
199
+ case types.PathSDKDomain:
200
+ writeSDKTestEnvelope(w, http.StatusOK, types.APIEnvelope[types.DomainResponse]{
201
+ OK: true,
202
+ Data: types.DomainResponse{
203
+ RootHost: "localhost",
204
+ Version: types.SDKProtocolVersion,
205
+ },
206
+ })
207
+ case types.PathSDKRegister:
208
+ writeSDKTestEnvelope(w, http.StatusCreated, types.APIEnvelope[types.RegisterResponse]{
209
+ OK: true,
210
+ Data: types.RegisterResponse{
211
+ LeaseID: "lease-1",
212
+ Hostnames: []string{"127.0.0.1"},
213
+ },
214
+ })
215
+ case types.PathSDKConnect:
216
+ connectCount.Add(1)
217
+ writeSDKTestEnvelope(w, http.StatusForbidden, types.APIEnvelope[any]{
218
+ OK: false,
219
+ Error: &types.APIError{Code: types.APIErrorCodeUnauthorized, Message: "reverse session denied"},
220
+ })
221
+ case types.PathSDKUnregister:
222
+ unregisterCount.Add(1)
223
+ writeSDKTestEnvelope(w, http.StatusOK, types.APIEnvelope[any]{OK: true})
224
+ default:
225
+ http.NotFound(w, r)
226
+ }
227
+ }))
228
+ defer server.Close()
229
+
230
+ listener, err := NewListener(context.Background(), server.URL, ListenerConfig{
231
+ Name: "demo",
232
+ RetryCount: 1,
233
+ RetryWait: 10 * time.Millisecond,
234
+ })
235
+ if err != nil {
236
+ t.Fatalf("NewListener() error = %v", err)
237
+ }
238
+ defer listener.Close()
239
+
240
+ waitForSDKTest(t, func() bool {
241
+ return listener.isClosed()
242
+ })
243
+ if connectCount.Load() < 2 {
244
+ t.Fatalf("connect count = %d, want at least 2", connectCount.Load())
245
+ }
246
+ if unregisterCount.Load() == 0 {
247
+ t.Fatal("expected listener to unregister lease after retry budget exhaustion")
248
+ }
249
+}
250
+
251
func TestExposeFailsFastWhenAnyRelayCannotRegister(t *testing.T) {
252
t.Parallel()
253