fix(route): add mutex lock for load balancer updates to prevent race conditions

internal/idlewatcher/watcher.go

@@ -69,8 +69,6 @@ type (
 
 const ContextKey = "idlewatcher.watcher"
 
-// TODO: replace -1 with neverTick
-
 var (
 	watcherMap   = make(map[string]*Watcher)
 	watcherMapMu sync.RWMutex
@@ -246,11 +244,14 @@ func NewWatcher(parent task.Parent, r routes.Route, cfg *idlewatcher.Config) (*W
 		kind = "proxmox"
 	}
 	w.l = log.With().
-		Stringer("idle_timeout", cfg.IdleTimeout).
 		Str("kind", kind).
 		Str("container", cfg.ContainerName()).
 		Logger()
 
+	if cfg.IdleTimeout != neverTick {
+		w.l = w.l.With().Stringer("idle_timeout", cfg.IdleTimeout).Logger()
+	}
+
 	if err != nil {
 		return nil, err
 	}
@@ -315,6 +316,8 @@ func NewWatcher(parent task.Parent, r routes.Route, cfg *idlewatcher.Config) (*W
 
 	w.dedupDependencies()
 
+	r.SetHealthMonitor(w)
+
 	w.l = w.l.With().Strs("deps", cfg.DependsOn).Logger()
 	if exists {
 		w.l.Debug().Msg("idlewatcher reloaded")

internal/route/fileserver.go

@@ -11,7 +11,6 @@ import (
 	"github.com/yusing/go-proxy/internal/net/gphttp/middleware"
 	"github.com/yusing/go-proxy/internal/route/routes"
 	"github.com/yusing/go-proxy/internal/task"
-	"github.com/yusing/go-proxy/internal/watcher/health"
 	"github.com/yusing/go-proxy/internal/watcher/health/monitor"
 )
 
@@ -19,9 +18,6 @@ type (
 	FileServer struct {
 		*Route
 
-		Health *monitor.FileServerHealthMonitor `json:"health"`
-
-		task         *task.Task
 		middleware   *middleware.Middleware
 		handler      http.Handler
 		accessLogger *accesslog.AccessLogger
@@ -90,8 +86,8 @@ func (s *FileServer) Start(parent task.Parent) gperr.Error {
 	}
 
 	if s.UseHealthCheck() {
-		s.Health = monitor.NewFileServerHealthMonitor(s.HealthCheck, s.Root)
-		if err := s.Health.Start(s.task); err != nil {
+		s.HealthMon = monitor.NewFileServerHealthMonitor(s.HealthCheck, s.Root)
+		if err := s.HealthMon.Start(s.task); err != nil {
 			return err
 		}
 	}
@@ -111,15 +107,6 @@ func (s *FileServer) Start(parent task.Parent) gperr.Error {
 	return nil
 }
 
-func (s *FileServer) Task() *task.Task {
-	return s.task
-}
-
-// Finish implements task.TaskFinisher.
-func (s *FileServer) Finish(reason any) {
-	s.task.Finish(reason)
-}
-
 // ServeHTTP implements http.Handler.
 func (s *FileServer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 	s.handler.ServeHTTP(w, req)
@@ -127,7 +114,3 @@ func (s *FileServer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 		s.accessLogger.Log(req, req.Response)
 	}
 }
-
-func (s *FileServer) HealthMonitor() health.HealthMonitor {
-	return s.Health
-}

internal/route/reverse_proxy.go

@@ -3,6 +3,7 @@ package route
 import (
 	"crypto/tls"
 	"net/http"
+	"sync"
 
 	"github.com/yusing/go-proxy/agent/pkg/agent"
 	"github.com/yusing/go-proxy/agent/pkg/agentproxy"
@@ -18,20 +19,15 @@ import (
 	"github.com/yusing/go-proxy/internal/net/types"
 	"github.com/yusing/go-proxy/internal/route/routes"
 	"github.com/yusing/go-proxy/internal/task"
-	"github.com/yusing/go-proxy/internal/watcher/health"
 	"github.com/yusing/go-proxy/internal/watcher/health/monitor"
 )
 
 type ReveseProxyRoute struct {
 	*Route
 
-	HealthMon health.HealthMonitor `json:"health,omitempty"`
-
 	loadBalancer *loadbalancer.LoadBalancer
 	handler      http.Handler
 	rp           *reverseproxy.ReverseProxy
-
-	task *task.Task
 }
 
 var _ routes.ReverseProxyRoute = (*ReveseProxyRoute)(nil)
@@ -157,30 +153,21 @@ func (r *ReveseProxyRoute) Start(parent task.Parent) gperr.Error {
 	return nil
 }
 
-// Task implements task.TaskStarter.
-func (r *ReveseProxyRoute) Task() *task.Task {
-	return r.task
-}
-
-// Finish implements task.TaskFinisher.
-func (r *ReveseProxyRoute) Finish(reason any) {
-	r.task.Finish(reason)
-}
-
 func (r *ReveseProxyRoute) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 	r.handler.ServeHTTP(w, req)
 }
 
-func (r *ReveseProxyRoute) HealthMonitor() health.HealthMonitor {
-	return r.HealthMon
-}
+var lbLock sync.Mutex
 
 func (r *ReveseProxyRoute) addToLoadBalancer(parent task.Parent) {
 	var lb *loadbalancer.LoadBalancer
 	cfg := r.LoadBalance
+	lbLock.Lock()
+
 	l, ok := routes.HTTP.Get(cfg.Link)
 	var linked *ReveseProxyRoute
 	if ok {
+		lbLock.Unlock()
 		linked = l.(*ReveseProxyRoute)
 		lb = linked.loadBalancer
 		lb.UpdateConfigIfNeeded(cfg)
@@ -192,10 +179,10 @@ func (r *ReveseProxyRoute) addToLoadBalancer(parent task.Parent) {
 		_ = lb.Start(parent) // always return nil
 		linked = &ReveseProxyRoute{
 			Route: &Route{
-				Alias:    cfg.Link + "-loadbalancer",
-				Homepage: r.Homepage,
+				Alias:     cfg.Link + "-loadbalancer",
+				Homepage:  r.Homepage,
+				HealthMon: lb,
 			},
-			HealthMon:    lb,
 			loadBalancer: lb,
 			handler:      lb,
 		}
@@ -205,6 +192,7 @@ func (r *ReveseProxyRoute) addToLoadBalancer(parent task.Parent) {
 			routes.HTTP.DelKey(cfg.Link)
 			routes.All.DelKey(cfg.Link)
 		})
+		lbLock.Unlock()
 	}
 	r.loadBalancer = lb
 

internal/route/route.go

@@ -51,7 +51,8 @@ type (
 		Homepage     *homepage.ItemConfig           `json:"homepage,omitempty"`
 		AccessLog    *accesslog.RequestLoggerConfig `json:"access_log,omitempty"`
 
-		Idlewatcher *idlewatcher.Config `json:"idlewatcher,omitempty"`
+		Idlewatcher *idlewatcher.Config  `json:"idlewatcher,omitempty"`
+		HealthMon   health.HealthMonitor `json:"health,omitempty"`
 
 		Metadata `deserialize:"-"`
 	}
@@ -68,7 +69,9 @@ type (
 
 		Excluded *bool `json:"excluded"`
 
-		impl        routes.Route
+		impl routes.Route
+		task *task.Task
+
 		isValidated bool
 		lastError   gperr.Error
 		provider    routes.Provider
@@ -243,7 +246,7 @@ func (r *Route) Impl() routes.Route {
 }
 
 func (r *Route) Task() *task.Task {
-	return r.impl.Task()
+	return r.task
 }
 
 func (r *Route) Start(parent task.Parent) (err gperr.Error) {
@@ -265,12 +268,12 @@ func (r *Route) start(parent task.Parent) gperr.Error {
 
 	if conflict, added := routes.All.AddIfNotExists(r.impl); !added {
 		err := gperr.Errorf("route %s already exists: from %s and %s", r.Alias, r.ProviderName(), conflict.ProviderName())
-		r.impl.Task().FinishAndWait(err)
+		r.task.FinishAndWait(err)
 		return err
 	} else {
 		// reference here because r.impl will be nil after Finish() is called.
 		impl := r.impl
-		impl.Task().OnCancel("remove_routes_from_all", func() {
+		r.task.OnCancel("remove_routes_from_all", func() {
 			routes.All.Del(impl)
 		})
 	}
@@ -285,7 +288,7 @@ func (r *Route) FinishAndWait(reason any) {
 	if r.impl == nil {
 		return
 	}
-	r.impl.Task().FinishAndWait(reason)
+	r.task.FinishAndWait(reason)
 	r.impl = nil
 }
 
@@ -360,7 +363,14 @@ func (r *Route) IsAgent() bool {
 }
 
 func (r *Route) HealthMonitor() health.HealthMonitor {
-	return r.impl.HealthMonitor()
+	return r.HealthMon
+}
+
+func (r *Route) SetHealthMonitor(m health.HealthMonitor) {
+	if r.HealthMon != nil && r.HealthMon != m {
+		r.HealthMon.Finish("health monitor replaced")
+	}
+	r.HealthMon = m
 }
 
 func (r *Route) IdlewatcherConfig() *idlewatcher.Config {

internal/route/routes/route.go

@@ -26,6 +26,7 @@ type (
 		GetProvider() Provider
 		TargetURL() *net.URL
 		HealthMonitor() health.HealthMonitor
+		SetHealthMonitor(m health.HealthMonitor)
 		References() []string
 
 		Started() <-chan struct{}

internal/route/stream.go

@@ -11,20 +11,14 @@ import (
 	net "github.com/yusing/go-proxy/internal/net/types"
 	"github.com/yusing/go-proxy/internal/route/routes"
 	"github.com/yusing/go-proxy/internal/task"
-	"github.com/yusing/go-proxy/internal/watcher/health"
 	"github.com/yusing/go-proxy/internal/watcher/health/monitor"
 )
 
 // TODO: support stream load balance.
 type StreamRoute struct {
 	*Route
-
 	net.Stream `json:"-"`
 
-	HealthMon health.HealthMonitor `json:"health"`
-
-	task *task.Task
-
 	l zerolog.Logger
 }
 
@@ -88,20 +82,6 @@ func (r *StreamRoute) Start(parent task.Parent) gperr.Error {
 	return nil
 }
 
-// Task implements task.TaskStarter.
-func (r *StreamRoute) Task() *task.Task {
-	return r.task
-}
-
-// Finish implements task.TaskFinisher.
-func (r *StreamRoute) Finish(reason any) {
-	r.task.Finish(reason)
-}
-
-func (r *StreamRoute) HealthMonitor() health.HealthMonitor {
-	return r.HealthMon
-}
-
 func (r *StreamRoute) acceptConnections() {
 	defer r.task.Finish("listener closed")