diff --git a/supervisor/supervisor.go b/supervisor/supervisor.go index b2c19e79f28..1843212896e 100644 --- a/supervisor/supervisor.go +++ b/supervisor/supervisor.go @@ -124,11 +124,17 @@ func (s *Supervisor) Run( s.log.Logger().Error().Err(err).Msg("initial tunnel connection failed") return err } + + return s.runLoop(ctx) +} + +func (s *Supervisor) runLoop(ctx context.Context) error { var tunnelsWaiting []int tunnelsActive := s.config.HAConnections backoff := retry.NewBackoff(s.config.Retries, tunnelRetryDuration, true) var backoffTimer <-chan time.Time + gracefulShutdownC := s.gracefulShutdownC shuttingDown := false for { @@ -175,8 +181,9 @@ func (s *Supervisor) Run( // No more tunnels outstanding, clear backoff timer backoff.SetGracePeriod() } - case <-s.gracefulShutdownC: + case <-gracefulShutdownC: shuttingDown = true + gracefulShutdownC = nil } } } diff --git a/supervisor/tunnel_test.go b/supervisor/tunnel_test.go index 66229c543a7..96b4c20d92b 100644 --- a/supervisor/tunnel_test.go +++ b/supervisor/tunnel_test.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "sync/atomic" "testing" "time" @@ -298,6 +299,68 @@ func newTestSupervisor(t *testing.T, server TunnelServer) *Supervisor { } } +type countingDoneContext struct { + context.Context + doneCalls atomic.Int32 + doneCallC chan struct{} +} + +func (c *countingDoneContext) Done() <-chan struct{} { + c.doneCalls.Add(1) + select { + case c.doneCallC <- struct{}{}: + default: + } + return nil +} + +func TestRunLoopHandlesGracefulShutdownOnce(t *testing.T) { + gracefulShutdownC := make(chan struct{}) + s := newTestSupervisor(t, nil) + s.config.HAConnections = 1 + s.gracefulShutdownC = gracefulShutdownC + + ctx := &countingDoneContext{ + Context: context.Background(), + doneCallC: make(chan struct{}, 3), + } + runDone := make(chan error, 1) + go func() { + runDone <- s.runLoop(ctx) + }() + + select { + case <-ctx.doneCallC: + case <-time.After(time.Second): + t.Fatal("run loop did not start") + } + + close(gracefulShutdownC) + select { + case <-ctx.doneCallC: + case <-time.After(time.Second): + t.Fatal("run loop did not handle graceful shutdown") + } + + busyLoop := false + select { + case <-ctx.doneCallC: + busyLoop = true + case <-time.After(100 * time.Millisecond): + } + + s.tunnelErrors <- tunnelError{index: 0} + select { + case err := <-runDone: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("run loop did not exit after the active tunnel stopped") + } + + assert.False(t, busyLoop, "run loop repeatedly selected the closed graceful shutdown channel") + assert.Equal(t, int32(2), ctx.doneCalls.Load()) +} + func TestStartFirstTunnelRetryLoop(t *testing.T) { t.Parallel()