From 405969085fb6b5c8058c88ca250e930ba6092153 Mon Sep 17 00:00:00 2001
From: fatedier <fatedier@gmail.com>
Date: Thu, 25 Apr 2024 20:20:39 +0800
Subject: [PATCH] client: add StatusExporter in service (#4182)

---
 client/service.go | 29 +++++++++++++++++++++--------
 pkg/ssh/server.go |  6 ++++--
 2 files changed, 25 insertions(+), 10 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 ws, nil
+	return ctl.pm.GetProxyStatus(name)
+}
+
+func (svr *Service) StatusExporter() StatusExporter {
+	return &statusExporterImpl{
+		getProxyStatusFunc: svr.getProxyStatus,
+	}
+}
+
+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 {