From 405969085fb6b5c8058c88ca250e930ba6092153 Mon Sep 17 00:00:00 2001 From: fatedier Date: Thu, 25 Apr 2024 20:20:39 +0800 Subject: [PATCH] client: add StatusExporter in service (#4182) --- client/service.go | 27 ++++++++++++++++++++------- pkg/ssh/server.go | 6 ++++-- 2 files changed, 24 insertions(+), 9 deletions(-) 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 {