From 95434a3ebdd208ef3f27a34087d6f5fff16e1fc0 Mon Sep 17 00:00:00 2001
From: baixl <baixl>
Date: Wed, 5 Jul 2023 17:10:38 +0800
Subject: [PATCH] add stop pxy.StopCheck

---
 main.go       |  1 +
 rule/group.go | 70 ++++++++++++++++++++++++++++++---------------------
 rule/proxy.go |  9 +++++++
 3 files changed, 51 insertions(+), 29 deletions(-)

diff --git a/main.go b/main.go
index 7403a0a..ab241bb 100644
--- a/main.go
+++ b/main.go
@@ -68,6 +68,7 @@ func main() {
 
 	// enable checkers
 	pxy.Check()
+	defer pxy.StopCheck()
 
 	// run proxy servers
 	for _, listen := range config.Listens {
diff --git a/rule/group.go b/rule/group.go
index fc9607b..e1a8fc0 100644
--- a/rule/group.go
+++ b/rule/group.go
@@ -33,6 +33,7 @@ type FwdrGroup struct {
 	index    uint32
 	priority uint32
 	next     func(addr string) *Forwarder
+	stopChan chan bool
 }
 
 // NewFwdrGroup returns a new forward group.
@@ -235,41 +236,47 @@ func (p *FwdrGroup) check(fwdr *Forwarder, checker Checker) {
 	intval := time.Duration(p.config.CheckInterval) * time.Second
 
 	for {
-		time.Sleep(intval * time.Duration(wait))
+		select {
+		case <-p.stopChan:
+			log.F("[check] %s: stop checking", p.name)
+			return
+		default:
+			time.Sleep(intval * time.Duration(wait))
 
-		// check all forwarders at least one time
-		if wait > 0 && (fwdr.Priority() < p.Priority()) {
-			continue
-		}
-
-		if fwdr.Enabled() && p.config.CheckDisabledOnly {
-			continue
-		}
-
-		elapsed, err := checker.Check(fwdr)
-		if err != nil {
-			if errors.Is(err, proxy.ErrNotSupported) {
-				fwdr.SetMaxFailures(0)
-				log.F("[check] %s: %s(%d), %s, stop checking", p.name, fwdr.Addr(), fwdr.Priority(), err)
-				fwdr.Enable()
-				break
+			// check all forwarders at least one time
+			if wait > 0 && (fwdr.Priority() < p.Priority()) {
+				continue
 			}
 
-			wait++
-			if wait > 16 {
-				wait = 16
+			if fwdr.Enabled() && p.config.CheckDisabledOnly {
+				continue
 			}
 
-			log.F("[check] %s: %s(%d), FAILED. error: %s", p.name, fwdr.Addr(), fwdr.Priority(), err)
-			fwdr.Disable()
-			continue
-		}
+			elapsed, err := checker.Check(fwdr)
+			if err != nil {
+				if errors.Is(err, proxy.ErrNotSupported) {
+					fwdr.SetMaxFailures(0)
+					log.F("[check] %s: %s(%d), %s, stop checking", p.name, fwdr.Addr(), fwdr.Priority(), err)
+					fwdr.Enable()
+					break
+				}
 
-		wait = 1
-		p.setLatency(fwdr, elapsed)
-		log.F("[check] %s: %s(%d), SUCCESS. Elapsed: %dms, Latency: %dms.",
-			p.name, fwdr.Addr(), fwdr.Priority(), elapsed.Milliseconds(), time.Duration(fwdr.Latency()).Milliseconds())
-		fwdr.Enable()
+				wait++
+				if wait > 16 {
+					wait = 16
+				}
+
+				log.F("[check] %s: %s(%d), FAILED. error: %s", p.name, fwdr.Addr(), fwdr.Priority(), err)
+				fwdr.Disable()
+				continue
+			}
+
+			wait = 1
+			p.setLatency(fwdr, elapsed)
+			log.F("[check] %s: %s(%d), SUCCESS. Elapsed: %dms, Latency: %dms.",
+				p.name, fwdr.Addr(), fwdr.Priority(), elapsed.Milliseconds(), time.Duration(fwdr.Latency()).Milliseconds())
+			fwdr.Enable()
+		}
 	}
 }
 
@@ -316,3 +323,8 @@ func (p *FwdrGroup) scheduleDH(dstAddr string) *Forwarder {
 	fnv1a.Write([]byte(dstAddr))
 	return p.avail[fnv1a.Sum32()%uint32(len(p.avail))]
 }
+
+// Stop Check
+func (p *FwdrGroup) StopCheck() {
+	p.stopChan <- true
+}
diff --git a/rule/proxy.go b/rule/proxy.go
index 495b83e..b3a3a0a 100644
--- a/rule/proxy.go
+++ b/rule/proxy.go
@@ -154,3 +154,12 @@ func (p *Proxy) Check() {
 		fwdrGroup.Check()
 	}
 }
+
+// Stop All checkers.
+func (p *Proxy) StopCheck() {
+	p.main.StopCheck()
+
+	for _, fwdrGroup := range p.all {
+		fwdrGroup.StopCheck()
+	}
+}