From 341624dcfaaa62da4d037dbb514c0357c6bb87af Mon Sep 17 00:00:00 2001
From: jebbs <qjebbs@gmail.com>
Date: Wed, 12 Oct 2022 15:35:11 +0800
Subject: [PATCH] make sure health check is stoped

---
 balancer/healthcheck.go | 14 +++++++++++---
 outbound/balancer.go    | 15 +++++++++++++--
 outbound/leastload.go   |  7 ++++++-
 outbound/leastping.go   |  7 ++++++-
 4 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/balancer/healthcheck.go b/balancer/healthcheck.go
index 16099d48..1e9981c0 100644
--- a/balancer/healthcheck.go
+++ b/balancer/healthcheck.go
@@ -2,6 +2,7 @@ package balancer
 
 import (
 	"fmt"
+	"io"
 	"math/rand"
 	"sync"
 	"time"
@@ -9,9 +10,15 @@ import (
 	"github.com/sagernet/sing-box/adapter"
 	"github.com/sagernet/sing-box/log"
 	"github.com/sagernet/sing-box/option"
+	"github.com/sagernet/sing/common"
 	E "github.com/sagernet/sing/common/exceptions"
 )
 
+var (
+	_ common.Starter = (*HealthCheck)(nil)
+	_ io.Closer      = (*HealthCheck)(nil)
+)
+
 // HealthCheck is the health checker for balancers
 type HealthCheck struct {
 	mutex sync.Mutex
@@ -62,7 +69,7 @@ func NewHealthCheck(router adapter.Router, tags []string, logger log.Logger, con
 	}
 }
 
-// Start starts the health check service
+// Start starts the health check service, implements common.Starter
 func (h *HealthCheck) Start() error {
 	h.mutex.Lock()
 	defer h.mutex.Unlock()
@@ -86,14 +93,15 @@ func (h *HealthCheck) Start() error {
 	return nil
 }
 
-// Stop stops the health check service
-func (h *HealthCheck) Stop() {
+// Close stops the health check service, implements io.Closer
+func (h *HealthCheck) Close() error {
 	h.mutex.Lock()
 	defer h.mutex.Unlock()
 	if h.ticker != nil {
 		h.ticker.Stop()
 		h.ticker = nil
 	}
+	return nil
 }
 
 // Check does a one time health check
diff --git a/outbound/balancer.go b/outbound/balancer.go
index 60b4b48a..387e0b5f 100644
--- a/outbound/balancer.go
+++ b/outbound/balancer.go
@@ -2,6 +2,7 @@ package outbound
 
 import (
 	"context"
+	"io"
 	"math/rand"
 	"net"
 
@@ -17,6 +18,8 @@ import (
 var (
 	_ adapter.Outbound      = (*Balancer)(nil)
 	_ adapter.OutboundGroup = (*Balancer)(nil)
+	_ common.Starter        = (*Balancer)(nil)
+	_ io.Closer             = (*Balancer)(nil)
 )
 
 // Balancer is a outbound group that picks outbound with least load
@@ -112,8 +115,16 @@ func (s *Balancer) NewPacketConnection(ctx context.Context, conn N.PacketConn, m
 	return s.pick(N.NetworkUDP).NewPacketConnection(ctx, conn, metadata)
 }
 
-// initialize inits the balancer
-func (s *Balancer) initialize() error {
+// Close implements io.Closer
+func (s *Balancer) Close() error {
+	if c, ok := s.Balancer.(io.Closer); ok {
+		return c.Close()
+	}
+	return nil
+}
+
+// Start implements common.Starter
+func (s *Balancer) Start() error {
 	// the fallback is required, in case that all outbounds are not available,
 	// we can pick it instead of returning nil to avoid panic.
 	if s.fallbackTag == "" {
diff --git a/outbound/leastload.go b/outbound/leastload.go
index d942be0c..46261b24 100644
--- a/outbound/leastload.go
+++ b/outbound/leastload.go
@@ -1,16 +1,21 @@
 package outbound
 
 import (
+	"io"
+
 	"github.com/sagernet/sing-box/adapter"
 	"github.com/sagernet/sing-box/balancer"
 	C "github.com/sagernet/sing-box/constant"
 	"github.com/sagernet/sing-box/log"
 	"github.com/sagernet/sing-box/option"
+	"github.com/sagernet/sing/common"
 )
 
 var (
 	_ adapter.Outbound      = (*LeastLoad)(nil)
 	_ adapter.OutboundGroup = (*LeastLoad)(nil)
+	_ common.Starter        = (*LeastLoad)(nil)
+	_ io.Closer             = (*LeastLoad)(nil)
 )
 
 // LeastLoad is a outbound group that picks outbound with least load
@@ -33,7 +38,7 @@ func NewLeastLoad(router adapter.Router, logger log.ContextLogger, tag string, o
 
 // Start implements common.Starter
 func (s *LeastLoad) Start() error {
-	err := s.Balancer.initialize()
+	err := s.Balancer.Start()
 	if err != nil {
 		return err
 	}
diff --git a/outbound/leastping.go b/outbound/leastping.go
index c26affb1..2a222a7a 100644
--- a/outbound/leastping.go
+++ b/outbound/leastping.go
@@ -1,16 +1,21 @@
 package outbound
 
 import (
+	"io"
+
 	"github.com/sagernet/sing-box/adapter"
 	"github.com/sagernet/sing-box/balancer"
 	C "github.com/sagernet/sing-box/constant"
 	"github.com/sagernet/sing-box/log"
 	"github.com/sagernet/sing-box/option"
+	"github.com/sagernet/sing/common"
 )
 
 var (
 	_ adapter.Outbound      = (*LeastPing)(nil)
 	_ adapter.OutboundGroup = (*LeastPing)(nil)
+	_ common.Starter        = (*LeastPing)(nil)
+	_ io.Closer             = (*LeastPing)(nil)
 )
 
 // LeastPing is a outbound group that picks outbound with least load
@@ -33,7 +38,7 @@ func NewLeastPing(router adapter.Router, logger log.ContextLogger, tag string, o
 
 // Start implements common.Starter
 func (s *LeastPing) Start() error {
-	err := s.Balancer.initialize()
+	err := s.Balancer.Start()
 	if err != nil {
 		return err
 	}