Skip to content

Commit b823fac

Browse files
calmkartcalmkart
authored
fix(gpud): back off session reconnects (#1256)
## Summary - Add structured reconnect signals for gpud session reader, writer, and health-check failures. - Parse `Retry-After` from 429/503 responses and use it as the lower bound for reconnect delay. - Replace fixed reconnect timing with exponential backoff plus jitter, including startup jitter. - Add jittered reconnect after `updateToken` so token rotation does not reconnect all agents at once. - Preserve auth failure persistence for 401/403 and token-validation 500 responses. ## Testing - `go test -gcflags="all=-N -l" ./pkg/session` - E2E single-machine replacement on worker node `fargate-ip-10-50-115-131.ec2.internal` - Verified real gpud service health, states, components, metrics, and control/worker machine status - Verified fake manager scenarios for health 503, reader 503, writer 429, auth 500, and token update
1 parent a72085b commit b823fac

9 files changed

Lines changed: 940 additions & 123 deletions

‎pkg/session/session.go‎

Lines changed: 70 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -246,13 +246,21 @@ type Session struct {
246246
// In production: time.Sleep, in tests: can be mocked to skip sleep
247247
timeSleepFunc func(d time.Duration)
248248

249+
// nowFunc returns the current time.
250+
// In production: time.Now, in tests: can be mocked for deterministic backoff.
251+
nowFunc func() time.Time
252+
253+
// jitterFunc returns a delay in [0, max) for reconnect spreading.
254+
// In production: random jitter, in tests: can be mocked for deterministic backoff.
255+
jitterFunc func(max time.Duration) time.Duration
256+
249257
// startReaderFunc starts a reader connection
250258
// In production: s.startReader, in tests: can be mocked
251-
startReaderFunc func(ctx context.Context, readerExit chan any, jar *cookiejar.Jar)
259+
startReaderFunc func(ctx context.Context, readerExit chan reconnectSignal, jar *cookiejar.Jar)
252260

253261
// startWriterFunc starts a writer connection
254262
// In production: s.startWriter, in tests: can be mocked
255-
startWriterFunc func(ctx context.Context, writerExit chan any, jar *cookiejar.Jar)
263+
startWriterFunc func(ctx context.Context, writerExit chan reconnectSignal, jar *cookiejar.Jar)
256264

257265
// checkServerHealthFunc checks server health
258266
// In production: s.checkServerHealth, in tests: can be mocked
@@ -344,6 +352,8 @@ func NewSession(ctx context.Context, epLocalGPUdServer string, epControlPlane st
344352

345353
s.timeAfterFunc = time.After
346354
s.timeSleepFunc = time.Sleep
355+
s.nowFunc = time.Now
356+
s.jitterFunc = defaultJitter
347357

348358
s.startReaderFunc = s.startReader
349359
s.startWriterFunc = s.startWriter
@@ -443,31 +453,53 @@ func controlPlaneOrigin(epControlPlane string) (string, error) {
443453
return host, nil
444454
}
445455

446-
func (s *Session) startWriter(ctx context.Context, writerExit chan any, jar *cookiejar.Jar) {
456+
func (s *Session) startWriter(ctx context.Context, writerExit chan reconnectSignal, jar *cookiejar.Jar) {
457+
sig := reconnectSignal{side: reconnectSideWriter}
447458
pipeFinishCh := make(chan any)
448459
goroutineCloseCh := make(chan any)
449460
defer func() {
450461
close(goroutineCloseCh)
451462
s.closer.Close()
452463
<-pipeFinishCh
453-
close(writerExit)
464+
sendReconnectSignal(writerExit, sig)
454465
}()
455466
reader, writer := io.Pipe()
456467
go s.handleWriterPipe(writer, goroutineCloseCh, pipeFinishCh)
457468

458469
req, err := createSessionRequest(ctx, s.epControlPlane, s.machineID, "write", s.token, reader)
459470
if err != nil {
471+
sig.err = err
460472
log.Logger.Warnw("session writer: error creating request", "error", err)
461473
return
462474
}
463475

464476
client := createHTTPClient(jar)
465477
resp, err := client.Do(req)
466478
if err != nil {
479+
if ctx.Err() != nil {
480+
sig.err = ctx.Err()
481+
} else {
482+
sig.err = err
483+
}
467484
log.Logger.Warnw("session writer: error making request", "error", err)
468485
return
469486
}
470487

488+
if resp.StatusCode != http.StatusOK {
489+
sig = classifySessionHTTPResponse(reconnectSideWriter, resp, s.now())
490+
log.Logger.Warnw(
491+
"session writer: request resp not ok -- retrying",
492+
"status", resp.Status,
493+
"statusCode", resp.StatusCode,
494+
"retryAfter", sig.retryAfter.String(),
495+
"reason", sig.reason,
496+
)
497+
return
498+
}
499+
defer func() {
500+
_ = resp.Body.Close()
501+
}()
502+
471503
serverID := resp.Header.Get("X-GPUD-Server-ID")
472504
log.Logger.Infow("session writer: http closed", "status", resp.Status, "statusCode", resp.StatusCode, "serverID", serverID)
473505
}
@@ -515,18 +547,20 @@ func (s *Session) writeBodyToPipe(writer *io.PipeWriter, body Body) error {
515547
return nil
516548
}
517549

518-
func (s *Session) startReader(ctx context.Context, readerExit chan any, jar *cookiejar.Jar) {
550+
func (s *Session) startReader(ctx context.Context, readerExit chan reconnectSignal, jar *cookiejar.Jar) {
551+
sig := reconnectSignal{side: reconnectSideReader}
519552
goroutineCloseCh := make(chan any)
520553
pipeFinishCh := make(chan any)
521554
defer func() {
522555
close(goroutineCloseCh)
523556
s.closer.Close()
524557
<-pipeFinishCh
525-
close(readerExit)
558+
sendReconnectSignal(readerExit, sig)
526559
}()
527560

528561
req, err := createSessionRequest(ctx, s.epControlPlane, s.machineID, "read", s.token, nil)
529562
if err != nil {
563+
sig.err = err
530564
log.Logger.Debugw("session reader: error creating request", "error", err)
531565
close(pipeFinishCh)
532566
return
@@ -535,19 +569,29 @@ func (s *Session) startReader(ctx context.Context, readerExit chan any, jar *coo
535569
client := createHTTPClient(jar)
536570
resp, err := client.Do(req)
537571
if err != nil {
572+
if ctx.Err() != nil {
573+
sig.err = ctx.Err()
574+
} else {
575+
sig.err = err
576+
}
538577
log.Logger.Warnw("session reader: error making request -- retrying", "error", err)
539578
close(pipeFinishCh)
540579
return
541580
}
542581

543582
if resp.StatusCode != http.StatusOK {
544-
log.Logger.Warnw("session reader: request resp not ok -- retrying", "status", resp.Status, "statusCode", resp.StatusCode)
545-
546-
if resp.StatusCode == http.StatusForbidden || resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusInternalServerError {
547-
bodyBytes, _ := io.ReadAll(resp.Body)
548-
if message, ok := loginFailureStatusMessage(resp.StatusCode, string(bodyBytes)); ok {
549-
s.persistLoginStatus(ctx, false, message)
550-
}
583+
status := resp.Status
584+
statusCode := resp.StatusCode
585+
sig = classifySessionHTTPResponse(reconnectSideReader, resp, s.now())
586+
log.Logger.Warnw(
587+
"session reader: request resp not ok -- retrying",
588+
"status", status,
589+
"statusCode", statusCode,
590+
"retryAfter", sig.retryAfter.String(),
591+
"reason", sig.reason,
592+
)
593+
if sig.authFailure {
594+
s.persistLoginStatus(ctx, false, sig.reason)
551595
}
552596

553597
close(pipeFinishCh)
@@ -558,6 +602,9 @@ func (s *Session) startReader(ctx context.Context, readerExit chan any, jar *coo
558602
log.Logger.Infow("session reader got X-GPUD-Server-ID", "serverID", serverID)
559603

560604
s.processReaderResponse(resp, goroutineCloseCh, pipeFinishCh)
605+
if ctx.Err() != nil {
606+
sig.err = ctx.Err()
607+
}
561608
}
562609

563610
// healthCheckHTTPError is returned by checkServerHealth when the control plane
@@ -566,6 +613,7 @@ func (s *Session) startReader(ctx context.Context, readerExit chan any, jar *coo
566613
type healthCheckHTTPError struct {
567614
statusCode int
568615
body string
616+
retryAfter time.Duration
569617
}
570618

571619
func (e *healthCheckHTTPError) Error() string {
@@ -639,18 +687,21 @@ func (s *Session) checkServerHealth(ctx context.Context, jar *cookiejar.Jar, tok
639687
if err != nil {
640688
return err
641689
}
642-
defer func() {
643-
_ = resp.Body.Close()
644-
}()
645690

646691
if resp.StatusCode != http.StatusOK {
647-
bodyBytes, _ := io.ReadAll(resp.Body)
648-
body := string(bodyBytes)
692+
body := readLimitedResponseBody(resp, sessionErrorBodyReadLimit)
649693
if len(body) > 500 {
650694
body = body[:500]
651695
}
652-
return &healthCheckHTTPError{statusCode: resp.StatusCode, body: body}
696+
return &healthCheckHTTPError{
697+
statusCode: resp.StatusCode,
698+
body: body,
699+
retryAfter: parseRetryAfter(resp.Header.Get("Retry-After"), s.now()),
700+
}
653701
}
702+
defer func() {
703+
_ = resp.Body.Close()
704+
}()
654705
return nil
655706
}
656707

‎pkg/session/session_keepalive.go‎

Lines changed: 42 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -4,44 +4,22 @@ import (
44
"context"
55
"errors"
66
"net/http/cookiejar"
7-
"time"
87

98
"github.com/leptonai/gpud/pkg/log"
109
)
1110

1211
func (s *Session) keepAlive() {
13-
// Start the initial connection immediately
14-
firstConnection := true
12+
backoff := reconnectBackoff{}
13+
if !s.waitReconnectDelay(s.ctx, s.jitter(startupJitterMax)) {
14+
return
15+
}
1516

1617
for {
1718
select {
1819
case <-s.ctx.Done():
1920
log.Logger.Debug("session keep alive: closing keep alive")
2021
return
2122
default:
22-
// CRITICAL: Reconnection delay to prevent race conditions
23-
//
24-
// Without this delay, rapid connection failures would create multiple
25-
// overlapping reader/writer goroutines that all write to the same channels,
26-
// causing "reader channel full" errors and making GPUd unresponsive.
27-
//
28-
// The 3-second delay ensures:
29-
// - Previous connections have time to fully clean up
30-
// - We don't overwhelm the control plane with rapid reconnection attempts
31-
// - Only one set of reader/writer goroutines exists at a time
32-
if !firstConnection {
33-
select {
34-
case <-s.ctx.Done():
35-
return
36-
case <-s.timeAfterFunc(3 * time.Second):
37-
log.Logger.Debug("session keep alive: attempting reconnection after delay")
38-
}
39-
}
40-
firstConnection = false
41-
42-
readerExit := make(chan any)
43-
writerExit := make(chan any)
44-
4523
// CLEANUP: Ensure previous connections are fully terminated
4624
//
4725
// This cleanup is essential to prevent the "reader channel full" bug:
@@ -55,7 +33,7 @@ func (s *Session) keepAlive() {
5533
s.closer.Close()
5634

5735
// Give old goroutines time to detect closer signal and exit
58-
s.timeSleepFunc(100 * time.Millisecond)
36+
s.sleep(cleanupDrainDelay)
5937

6038
// Drain any stale messages left in the reader channel
6139
s.drainReaderChannel()
@@ -78,15 +56,29 @@ func (s *Session) keepAlive() {
7856
// "failed to validate token" when the token is invalid.
7957
var httpErr *healthCheckHTTPError
8058
if errors.As(err, &httpErr) {
81-
if message, ok := loginFailureStatusMessage(httpErr.statusCode, httpErr.body); ok {
82-
s.persistLoginStatus(ctx, false, message)
59+
sig := classifyHealthCheckError(httpErr)
60+
if sig.authFailure {
61+
s.persistLoginStatus(ctx, false, sig.reason)
8362
}
8463
}
8564

8665
cancel()
66+
if s.ctx.Err() != nil {
67+
return
68+
}
69+
sig := classifyHealthCheckError(err)
70+
delay := backoff.nextDelay(s, sig)
71+
log.Logger.Debugw("session keep alive: attempting reconnection after delay", "delay", delay.String(), "reason", sig.reason, "retryAfter", sig.retryAfter.String())
72+
if !s.waitReconnectDelay(s.ctx, delay) {
73+
return
74+
}
8775
continue
8876
}
8977

78+
readerExit := make(chan reconnectSignal, 1)
79+
writerExit := make(chan reconnectSignal, 1)
80+
connectionStartedAt := s.now()
81+
9082
go s.startReaderFunc(ctx, readerExit, jar)
9183
go s.startWriterFunc(ctx, writerExit, jar)
9284

@@ -107,18 +99,33 @@ func (s *Session) keepAlive() {
10799
// - No deadlock: We handle both possible exit orders
108100
// - Clean shutdown: Both goroutines exit before we create new ones
109101
// - No goroutine leaks: Context cancellation ensures termination
102+
var firstExit reconnectSignal
103+
var secondExit reconnectSignal
110104
select {
111-
case <-readerExit:
105+
case firstExit = <-readerExit:
112106
log.Logger.Debug("session reader: reader exited first")
113-
cancel() // Signal writer to exit
114-
<-writerExit // Wait for writer cleanup
107+
cancel() // Signal writer to exit
108+
secondExit = <-writerExit // Wait for writer cleanup
115109
log.Logger.Debug("session writer: writer exited after cancellation")
116-
case <-writerExit:
110+
case firstExit = <-writerExit:
117111
log.Logger.Debug("session writer: writer exited first")
118-
cancel() // Signal reader to exit
119-
<-readerExit // Wait for reader cleanup
112+
cancel() // Signal reader to exit
113+
secondExit = <-readerExit // Wait for reader cleanup
120114
log.Logger.Debug("session reader: reader exited after cancellation")
121115
}
116+
if s.ctx.Err() != nil {
117+
return
118+
}
119+
120+
if s.now().Sub(connectionStartedAt) >= reconnectStableWindow {
121+
backoff.reset()
122+
}
123+
sig := chooseReconnectSignal(firstExit, secondExit)
124+
delay := backoff.nextDelay(s, sig)
125+
log.Logger.Debugw("session keep alive: attempting reconnection after delay", "delay", delay.String(), "reason", sig.reason, "retryAfter", sig.retryAfter.String())
126+
if !s.waitReconnectDelay(s.ctx, delay) {
127+
return
128+
}
122129
}
123130
}
124131
}

‎pkg/session/session_keepalive_health_check_test.go‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -69,10 +69,10 @@ func TestKeepAliveHealthCheck403PersistsFailure(t *testing.T) {
6969
s.timeSleepFunc = func(d time.Duration) {}
7070

7171
// These should never be called since health check fails before reader/writer start
72-
s.startReaderFunc = func(ctx context.Context, readerExit chan any, jar *cookiejar.Jar) {
72+
s.startReaderFunc = func(ctx context.Context, readerExit chan reconnectSignal, jar *cookiejar.Jar) {
7373
t.Fatal("startReader should not be called when health check fails")
7474
}
75-
s.startWriterFunc = func(ctx context.Context, writerExit chan any, jar *cookiejar.Jar) {
75+
s.startWriterFunc = func(ctx context.Context, writerExit chan reconnectSignal, jar *cookiejar.Jar) {
7676
t.Fatal("startWriter should not be called when health check fails")
7777
}
7878

@@ -146,10 +146,10 @@ func TestKeepAliveHealthCheck500TokenValidationPersistsFailure(t *testing.T) {
146146
return ch
147147
}
148148
s.timeSleepFunc = func(d time.Duration) {}
149-
s.startReaderFunc = func(ctx context.Context, readerExit chan any, jar *cookiejar.Jar) {
149+
s.startReaderFunc = func(ctx context.Context, readerExit chan reconnectSignal, jar *cookiejar.Jar) {
150150
t.Fatal("startReader should not be called when health check fails")
151151
}
152-
s.startWriterFunc = func(ctx context.Context, writerExit chan any, jar *cookiejar.Jar) {
152+
s.startWriterFunc = func(ctx context.Context, writerExit chan reconnectSignal, jar *cookiejar.Jar) {
153153
t.Fatal("startWriter should not be called when health check fails")
154154
}
155155

@@ -210,7 +210,7 @@ func TestStartReader500TokenValidationPersistsAuthenticationFailure(t *testing.T
210210
jar, err := cookiejar.New(nil)
211211
require.NoError(t, err)
212212

213-
readerExit := make(chan any)
213+
readerExit := make(chan reconnectSignal, 1)
214214
s.startReader(ctx, readerExit, jar)
215215

216216
dbRO, err := sqlite.Open(stateFile, sqlite.WithReadOnly(true))
@@ -323,10 +323,10 @@ func TestKeepAliveHealthCheckNon403DoesNotPersist(t *testing.T) {
323323
return ch
324324
}
325325
s.timeSleepFunc = func(d time.Duration) {}
326-
s.startReaderFunc = func(ctx context.Context, readerExit chan any, jar *cookiejar.Jar) {
326+
s.startReaderFunc = func(ctx context.Context, readerExit chan reconnectSignal, jar *cookiejar.Jar) {
327327
t.Fatal("startReader should not be called when health check fails")
328328
}
329-
s.startWriterFunc = func(ctx context.Context, writerExit chan any, jar *cookiejar.Jar) {
329+
s.startWriterFunc = func(ctx context.Context, writerExit chan reconnectSignal, jar *cookiejar.Jar) {
330330
t.Fatal("startWriter should not be called when health check fails")
331331
}
332332

@@ -390,10 +390,10 @@ func TestKeepAliveHealthCheckNetworkErrorDoesNotPersist(t *testing.T) {
390390
return ch
391391
}
392392
s.timeSleepFunc = func(d time.Duration) {}
393-
s.startReaderFunc = func(ctx context.Context, readerExit chan any, jar *cookiejar.Jar) {
393+
s.startReaderFunc = func(ctx context.Context, readerExit chan reconnectSignal, jar *cookiejar.Jar) {
394394
t.Fatal("startReader should not be called when health check fails")
395395
}
396-
s.startWriterFunc = func(ctx context.Context, writerExit chan any, jar *cookiejar.Jar) {
396+
s.startWriterFunc = func(ctx context.Context, writerExit chan reconnectSignal, jar *cookiejar.Jar) {
397397
t.Fatal("startWriter should not be called when health check fails")
398398
}
399399

0 commit comments

Comments
 (0)