@@ -398,24 +398,28 @@ func (a *agent) createTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) (_
398
398
}
399
399
}()
400
400
if err = a .trackConnGoroutine (func () {
401
+ var wg sync.WaitGroup
401
402
for {
402
403
conn , err := sshListener .Accept ()
403
404
if err != nil {
404
- return
405
+ break
405
406
}
407
+ wg .Add (1 )
406
408
closed := make (chan struct {})
407
- _ = a . trackConnGoroutine ( func () {
409
+ go func () {
408
410
select {
409
- case <- network .Closed ():
410
411
case <- closed :
412
+ case <- a .closed :
413
+ _ = conn .Close ()
411
414
}
412
- _ = conn . Close ()
413
- })
414
- _ = a . trackConnGoroutine ( func () {
415
+ wg . Done ()
416
+ }( )
417
+ go func () {
415
418
defer close (closed )
416
419
a .sshServer .HandleConn (conn )
417
- })
420
+ }( )
418
421
}
422
+ wg .Wait ()
419
423
}); err != nil {
420
424
return nil , err
421
425
}
@@ -431,35 +435,47 @@ func (a *agent) createTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) (_
431
435
}()
432
436
if err = a .trackConnGoroutine (func () {
433
437
logger := a .logger .Named ("reconnecting-pty" )
434
-
438
+ var wg sync. WaitGroup
435
439
for {
436
440
conn , err := reconnectingPTYListener .Accept ()
437
441
if err != nil {
438
442
logger .Debug (ctx , "accept pty failed" , slog .Error (err ))
439
- return
440
- }
441
- // This cannot use a JSON decoder, since that can
442
- // buffer additional data that is required for the PTY.
443
- rawLen := make ([]byte , 2 )
444
- _ , err = conn .Read (rawLen )
445
- if err != nil {
446
- continue
447
- }
448
- length := binary .LittleEndian .Uint16 (rawLen )
449
- data := make ([]byte , length )
450
- _ , err = conn .Read (data )
451
- if err != nil {
452
- continue
453
- }
454
- var msg codersdk.WorkspaceAgentReconnectingPTYInit
455
- err = json .Unmarshal (data , & msg )
456
- if err != nil {
457
- continue
443
+ break
458
444
}
445
+ wg .Add (1 )
446
+ closed := make (chan struct {})
447
+ go func () {
448
+ select {
449
+ case <- closed :
450
+ case <- a .closed :
451
+ _ = conn .Close ()
452
+ }
453
+ wg .Done ()
454
+ }()
459
455
go func () {
456
+ defer close (closed )
457
+ // This cannot use a JSON decoder, since that can
458
+ // buffer additional data that is required for the PTY.
459
+ rawLen := make ([]byte , 2 )
460
+ _ , err = conn .Read (rawLen )
461
+ if err != nil {
462
+ return
463
+ }
464
+ length := binary .LittleEndian .Uint16 (rawLen )
465
+ data := make ([]byte , length )
466
+ _ , err = conn .Read (data )
467
+ if err != nil {
468
+ return
469
+ }
470
+ var msg codersdk.WorkspaceAgentReconnectingPTYInit
471
+ err = json .Unmarshal (data , & msg )
472
+ if err != nil {
473
+ return
474
+ }
460
475
_ = a .handleReconnectingPTY (ctx , logger , msg , conn )
461
476
}()
462
477
}
478
+ wg .Wait ()
463
479
}); err != nil {
464
480
return nil , err
465
481
}
@@ -474,20 +490,29 @@ func (a *agent) createTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) (_
474
490
}
475
491
}()
476
492
if err = a .trackConnGoroutine (func () {
493
+ var wg sync.WaitGroup
477
494
for {
478
495
conn , err := speedtestListener .Accept ()
479
496
if err != nil {
480
497
a .logger .Debug (ctx , "speedtest listener failed" , slog .Error (err ))
481
- return
498
+ break
482
499
}
483
- if err = a .trackConnGoroutine (func () {
500
+ wg .Add (1 )
501
+ closed := make (chan struct {})
502
+ go func () {
503
+ select {
504
+ case <- closed :
505
+ case <- a .closed :
506
+ _ = conn .Close ()
507
+ }
508
+ wg .Done ()
509
+ }()
510
+ go func () {
511
+ defer close (closed )
484
512
_ = speedtest .ServeConn (conn )
485
- }); err != nil {
486
- a .logger .Debug (ctx , "speedtest listener failed" , slog .Error (err ))
487
- _ = conn .Close ()
488
- return
489
- }
513
+ }()
490
514
}
515
+ wg .Wait ()
491
516
}); err != nil {
492
517
return nil , err
493
518
}
@@ -511,7 +536,10 @@ func (a *agent) createTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) (_
511
536
ErrorLog : slog .Stdlib (ctx , a .logger .Named ("statistics_http_server" ), slog .LevelInfo ),
512
537
}
513
538
go func () {
514
- <- ctx .Done ()
539
+ select {
540
+ case <- ctx .Done ():
541
+ case <- a .closed :
542
+ }
515
543
_ = server .Close ()
516
544
}()
517
545
0 commit comments