diff --git a/client/service.go b/client/service.go index 7a7f6dc..0cbd875 100644 --- a/client/service.go +++ b/client/service.go @@ -380,18 +380,31 @@ func (svr *Service) stop() { } } -// TODO(fatedier): Use StatusExporter to provide query interfaces instead of directly using methods from the Service. -func (svr *Service) GetProxyStatus(name string) (*proxy.WorkingStatus, error) { +func (svr *Service) getProxyStatus(name string) (*proxy.WorkingStatus, bool) { svr.ctlMu.RLock() ctl := svr.ctl svr.ctlMu.RUnlock() if ctl == nil { - return nil, fmt.Errorf("control is not running") + return nil, false } - ws, ok := ctl.pm.GetProxyStatus(name) - if !ok { - return nil, fmt.Errorf("proxy [%s] is not found", name) + return ctl.pm.GetProxyStatus(name) +} + +func (svr *Service) StatusExporter() StatusExporter { + return &statusExporterImpl{ + getProxyStatusFunc: svr.getProxyStatus, } - return ws, nil +} + +type StatusExporter interface { + GetProxyStatus(name string) (*proxy.WorkingStatus, bool) +} + +type statusExporterImpl struct { + getProxyStatusFunc func(name string) (*proxy.WorkingStatus, bool) +} + +func (s *statusExporterImpl) GetProxyStatus(name string) (*proxy.WorkingStatus, bool) { + return s.getProxyStatusFunc(name) } diff --git a/pkg/ssh/server.go b/pkg/ssh/server.go index 2adab17..84b744f 100644 --- a/pkg/ssh/server.go +++ b/pkg/ssh/server.go @@ -363,11 +363,13 @@ func (s *TunnelServer) waitProxyStatusReady(name string, timeout time.Duration) timer := time.NewTimer(timeout) defer timer.Stop() + statusExporter := s.vc.Service().StatusExporter() + for { select { case <-ticker.C: - ps, err := s.vc.Service().GetProxyStatus(name) - if err != nil { + ps, ok := statusExporter.GetProxyStatus(name) + if !ok { continue } switch ps.Phase {