diff --git a/internal/server/server.go b/internal/server/server.go index 455ee27..3ef25c2 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -282,7 +282,7 @@ func (s *Server) serve(ctx context.Context) { default: } logger.Infof("AcceptStream returned %v — waiting for new session", err) - time.Sleep(100 * time.Millisecond) + s.waitForNewSession(ctx, sess) continue } @@ -294,6 +294,22 @@ func (s *Server) serve(ctx context.Context) { } } +func (s *Server) waitForNewSession(ctx context.Context, dead *smux.Session) { + for { + select { + case <-ctx.Done(): + return + case <-time.After(50 * time.Millisecond): + } + s.sessMu.RLock() + current := s.session + s.sessMu.RUnlock() + if current != dead { + return + } + } +} + func (s *Server) shutdown() { s.sessMu.Lock() if s.session != nil {