Improved healthcheck, idlewatcher support for loadbalanced routes, bug fixes

internal/route/http.go

@@ -18,6 +18,7 @@ import (
 	url "github.com/yusing/go-proxy/internal/net/types"
 	P "github.com/yusing/go-proxy/internal/proxy"
 	PT "github.com/yusing/go-proxy/internal/proxy/fields"
+	"github.com/yusing/go-proxy/internal/types"
 	F "github.com/yusing/go-proxy/internal/utils/functional"
 	"github.com/yusing/go-proxy/internal/watcher/health"
 )
@@ -26,12 +27,12 @@ type (
 	HTTPRoute struct {
 		*P.ReverseProxyEntry `json:"entry"`
 
-		LoadBalancer *loadbalancer.LoadBalancer `json:"load_balancer,omitempty"`
-		HealthMon    health.HealthMonitor       `json:"health"`
+		HealthMon health.HealthMonitor `json:"health,omitempty"`
 
-		server  *loadbalancer.Server
-		handler http.Handler
-		rp      *gphttp.ReverseProxy
+		loadBalancer *loadbalancer.LoadBalancer
+		server       *loadbalancer.Server
+		handler      http.Handler
+		rp           *gphttp.ReverseProxy
 	}
 
 	SubdomainKey = PT.Alias
@@ -102,10 +103,6 @@ func (r *HTTPRoute) URL() url.URL {
 }
 
 func (r *HTTPRoute) Start() E.NestedError {
-	if r.handler != nil {
-		return nil
-	}
-
 	if r.ShouldNotServe() {
 		return nil
 	}
@@ -113,6 +110,10 @@ func (r *HTTPRoute) Start() E.NestedError {
 	httpRoutesMu.Lock()
 	defer httpRoutesMu.Unlock()
 
+	if r.handler != nil {
+		return nil
+	}
+
 	if r.HealthCheck.Disabled && (r.UseIdleWatcher() || r.UseLoadBalance()) {
 		logrus.Warnf("%s.healthCheck.disabled cannot be false when loadbalancer or idlewatcher is enabled", r.Alias)
 		r.HealthCheck.Disabled = true
@@ -129,15 +130,23 @@ func (r *HTTPRoute) Start() E.NestedError {
 		r.HealthMon = waker
 	case !r.HealthCheck.Disabled:
 		r.HealthMon = health.NewHTTPHealthMonitor(common.GlobalTask("Reverse proxy "+r.String()), r.URL(), r.HealthCheck)
-		fallthrough
-	case len(r.PathPatterns) == 1 && r.PathPatterns[0] == "/":
-		r.handler = ReverseProxyHandler{r.rp}
-	default:
-		mux := http.NewServeMux()
-		for _, p := range r.PathPatterns {
-			mux.HandleFunc(string(p), r.rp.ServeHTTP)
+	}
+
+	if r.handler == nil {
+		switch {
+		case len(r.PathPatterns) == 1 && r.PathPatterns[0] == "/":
+			r.handler = ReverseProxyHandler{r.rp}
+		default:
+			mux := http.NewServeMux()
+			for _, p := range r.PathPatterns {
+				mux.HandleFunc(string(p), r.rp.ServeHTTP)
+			}
+			r.handler = mux
 		}
-		r.handler = mux
+	}
+
+	if r.HealthMon != nil {
+		r.HealthMon.Start()
 	}
 
 	if r.UseLoadBalance() {
@@ -146,9 +155,6 @@ func (r *HTTPRoute) Start() E.NestedError {
 		httpRoutes.Store(string(r.Alias), r)
 	}
 
-	if r.HealthMon != nil {
-		r.HealthMon.Start()
-	}
 	return nil
 }
 
@@ -160,7 +166,7 @@ func (r *HTTPRoute) Stop() (_ E.NestedError) {
 	httpRoutesMu.Lock()
 	defer httpRoutesMu.Unlock()
 
-	if r.LoadBalancer != nil {
+	if r.loadBalancer != nil {
 		r.removeFromLoadBalancer()
 	} else {
 		httpRoutes.Delete(string(r.Alias))
@@ -184,29 +190,40 @@ func (r *HTTPRoute) addToLoadBalancer() {
 	var lb *loadbalancer.LoadBalancer
 	linked, ok := httpRoutes.Load(r.LoadBalance.Link)
 	if ok {
-		lb = linked.LoadBalancer
+		lb = linked.loadBalancer
+		lb.UpdateConfigIfNeeded(r.LoadBalance)
+		if linked.Raw.Homepage == nil && r.Raw.Homepage != nil {
+			linked.Raw.Homepage = r.Raw.Homepage
+		}
 	} else {
 		lb = loadbalancer.New(r.LoadBalance)
 		lb.Start()
 		linked = &HTTPRoute{
-			LoadBalancer: lb,
+			ReverseProxyEntry: &P.ReverseProxyEntry{
+				Raw: &types.RawEntry{
+					Homepage: r.Raw.Homepage,
+				},
+				Alias: PT.Alias(lb.Link),
+			},
+			HealthMon:    lb,
+			loadBalancer: lb,
 			handler:      lb,
 		}
 		httpRoutes.Store(r.LoadBalance.Link, linked)
 	}
-	r.LoadBalancer = lb
+	r.loadBalancer = lb
 	r.server = loadbalancer.NewServer(string(r.Alias), r.rp.TargetURL, r.LoadBalance.Weight, r.handler, r.HealthMon)
 	lb.AddServer(r.server)
 }
 
 func (r *HTTPRoute) removeFromLoadBalancer() {
-	r.LoadBalancer.RemoveServer(r.server)
-	if r.LoadBalancer.IsEmpty() {
+	r.loadBalancer.RemoveServer(r.server)
+	if r.loadBalancer.IsEmpty() {
 		httpRoutes.Delete(r.LoadBalance.Link)
 		logrus.Debugf("loadbalancer %q removed from route table", r.LoadBalance.Link)
 	}
 	r.server = nil
-	r.LoadBalancer = nil
+	r.loadBalancer = nil
 }
 
 func ProxyHandler(w http.ResponseWriter, r *http.Request) {

internal/route/route.go

@@ -1,6 +1,7 @@
 package route
 
 import (
+	"github.com/yusing/go-proxy/internal/docker"
 	E "github.com/yusing/go-proxy/internal/error"
 	url "github.com/yusing/go-proxy/internal/net/types"
 	P "github.com/yusing/go-proxy/internal/proxy"
@@ -36,6 +37,13 @@ const (
 // function alias.
 var NewRoutes = F.NewMap[Routes]
 
+func (rt *Route) Container() *docker.Container {
+	if rt.Entry.Container == nil {
+		return docker.DummyContainer
+	}
+	return rt.Entry.Container
+}
+
 func NewRoute(en *types.RawEntry) (*Route, E.NestedError) {
 	entry, err := P.ValidateEntry(en)
 	if err != nil {

internal/route/stream.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"net"
 	"sync"
 
 	"github.com/sirupsen/logrus"
@@ -24,12 +25,11 @@ type StreamRoute struct {
 
 	url url.URL
 
-	wg     sync.WaitGroup
 	task   common.Task
 	cancel context.CancelFunc
+	done   chan struct{}
 
-	connCh chan any
-	l      logrus.FieldLogger
+	l logrus.FieldLogger
 
 	mu sync.Mutex
 }
@@ -61,7 +61,6 @@ func NewStreamRoute(entry *P.StreamEntry) (*StreamRoute, E.NestedError) {
 	base := &StreamRoute{
 		StreamEntry: entry,
 		url:         url,
-		connCh:      make(chan any, 100),
 	}
 	if entry.Scheme.ListeningScheme.IsTCP() {
 		base.StreamImpl = NewTCPRoute(base)
@@ -73,7 +72,7 @@ func NewStreamRoute(entry *P.StreamEntry) (*StreamRoute, E.NestedError) {
 }
 
 func (r *StreamRoute) String() string {
-	return fmt.Sprintf("%s stream: %s", r.Scheme, r.Alias)
+	return fmt.Sprintf("stream %s", r.Alias)
 }
 
 func (r *StreamRoute) URL() url.URL {
@@ -88,14 +87,12 @@ func (r *StreamRoute) Start() E.NestedError {
 		return nil
 	}
 	r.task, r.cancel = common.NewTaskWithCancel(r.String())
-	r.wg.Wait()
 	if err := r.Setup(); err != nil {
 		return E.FailWith("setup", err)
 	}
+	r.done = make(chan struct{})
 	r.l.Infof("listening on port %d", r.Port.ListeningPort)
-	r.wg.Add(2)
 	go r.acceptConnections()
-	go r.handleConnections()
 	if !r.Healthcheck.Disabled {
 		r.HealthMon = health.NewRawHealthMonitor(r.task, r.URL(), r.Healthcheck)
 		r.HealthMon.Start()
@@ -122,11 +119,7 @@ func (r *StreamRoute) Stop() E.NestedError {
 	r.cancel()
 	r.CloseListeners()
 
-	r.wg.Wait()
-	r.task.Finished()
-
-	r.task, r.cancel = nil, nil
-
+	<-r.done
 	return nil
 }
 
@@ -135,41 +128,45 @@ func (r *StreamRoute) Started() bool {
 }
 
 func (r *StreamRoute) acceptConnections() {
-	defer r.wg.Done()
+	var connWg sync.WaitGroup
+
+	task := r.task.Subtask("%s accept connections", r.String())
+
+	defer func() {
+		connWg.Wait()
+		task.Finished()
+		r.task.Finished()
+		r.task, r.cancel = nil, nil
+		close(r.done)
+		r.done = nil
+	}()
 
 	for {
 		select {
-		case <-r.task.Context().Done():
+		case <-task.Context().Done():
 			return
 		default:
 			conn, err := r.Accept()
 			if err != nil {
 				select {
-				case <-r.task.Context().Done():
+				case <-task.Context().Done():
 					return
 				default:
-					r.l.Error(err)
+					var nErr *net.OpError
+					ok := errors.As(err, &nErr)
+					if !(ok && nErr.Timeout()) {
+						r.l.Error(err)
+					}
 					continue
 				}
 			}
-			r.connCh <- conn
-		}
-	}
-}
-
-func (r *StreamRoute) handleConnections() {
-	defer r.wg.Done()
-
-	for {
-		select {
-		case <-r.task.Context().Done():
-			return
-		case conn := <-r.connCh:
+			connWg.Add(1)
 			go func() {
 				err := r.Handle(conn)
 				if err != nil && !errors.Is(err, context.Canceled) {
 					r.l.Error(err)
 				}
+				connWg.Done()
 			}()
 		}
 	}

internal/route/tcp.go

@@ -4,31 +4,25 @@ import (
 	"context"
 	"fmt"
 	"net"
-	"sync"
 	"time"
 
 	T "github.com/yusing/go-proxy/internal/proxy/fields"
 	U "github.com/yusing/go-proxy/internal/utils"
+	F "github.com/yusing/go-proxy/internal/utils/functional"
 )
 
 const tcpDialTimeout = 5 * time.Second
 
 type (
-	Pipes []U.BidirectionalPipe
-
-	TCPRoute struct {
+	TCPConnMap = F.Map[net.Conn, struct{}]
+	TCPRoute   struct {
 		*StreamRoute
-		listener net.Listener
-		pipe     Pipes
-		mu       sync.Mutex
+		listener *net.TCPListener
 	}
 )
 
 func NewTCPRoute(base *StreamRoute) StreamImpl {
-	return &TCPRoute{
-		StreamRoute: base,
-		pipe:        make(Pipes, 0),
-	}
+	return &TCPRoute{StreamRoute: base}
 }
 
 func (route *TCPRoute) Setup() error {
@@ -38,11 +32,12 @@ func (route *TCPRoute) Setup() error {
 	}
 	//! this read the allocated port from original ':0'
 	route.Port.ListeningPort = T.Port(in.Addr().(*net.TCPAddr).Port)
-	route.listener = in
+	route.listener = in.(*net.TCPListener)
 	return nil
 }
 
 func (route *TCPRoute) Accept() (any, error) {
+	route.listener.SetDeadline(time.Now().Add(time.Second))
 	return route.listener.Accept()
 }
 
@@ -50,24 +45,23 @@ func (route *TCPRoute) Handle(c any) error {
 	clientConn := c.(net.Conn)
 
 	defer clientConn.Close()
+	go func() {
+		<-route.task.Context().Done()
+		clientConn.Close()
+	}()
 
 	ctx, cancel := context.WithTimeout(route.task.Context(), tcpDialTimeout)
-	defer cancel()
 
 	serverAddr := fmt.Sprintf("%s:%v", route.Host, route.Port.ProxyPort)
 	dialer := &net.Dialer{}
 
 	serverConn, err := dialer.DialContext(ctx, string(route.Scheme.ProxyScheme), serverAddr)
+	cancel()
 	if err != nil {
 		return err
 	}
 
-	route.mu.Lock()
-
 	pipe := U.NewBidirectionalPipe(route.task.Context(), clientConn, serverConn)
-	route.pipe = append(route.pipe, pipe)
-
-	route.mu.Unlock()
 	return pipe.Start()
 }
 

internal/route/udp.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"time"
 
 	T "github.com/yusing/go-proxy/internal/proxy/fields"
 	U "github.com/yusing/go-proxy/internal/utils"
@@ -67,6 +68,7 @@ func (route *UDPRoute) Accept() (any, error) {
 	in := route.listeningConn
 
 	buffer := make([]byte, udpBufferSize)
+	route.listeningConn.SetReadDeadline(time.Now().Add(time.Second))
 	nRead, srcAddr, err := in.ReadFromUDP(buffer)
 	if err != nil {
 		return nil, err