fixed loadbalancer with idlewatcher, fixed reload issue

internal/net/http/loadbalancer/dummy_response_writer.go

@@ -0,0 +1,15 @@
+package loadbalancer
+
+import "net/http"
+
+type DummyResponseWriter struct{}
+
+func (w *DummyResponseWriter) Header() (_ http.Header) {
+	return
+}
+
+func (w *DummyResponseWriter) Write([]byte) (_ int, _ error) {
+	return
+}
+
+func (w *DummyResponseWriter) WriteHeader(int) {}

internal/net/http/loadbalancer/ip_hash.go

@@ -21,7 +21,7 @@ func (lb *LoadBalancer) newIPHash() impl {
 	if len(lb.Options) == 0 {
 		return impl
 	}
-	var err E.NestedError
+	var err E.Error
 	impl.realIP, err = middleware.NewRealIP(lb.Options)
 	if err != nil {
 		logger.Errorf("loadbalancer %s invalid real_ip options: %s, ignoring", lb.Link, err)

internal/net/http/loadbalancer/loadbalancer.go

@@ -1,10 +1,13 @@
 package loadbalancer
 
 import (
+	"context"
 	"net/http"
 	"sync"
 	"time"
 
+	"github.com/yusing/go-proxy/internal/common"
+	idlewatcher "github.com/yusing/go-proxy/internal/docker/idlewatcher/types"
 	E "github.com/yusing/go-proxy/internal/error"
 	"github.com/yusing/go-proxy/internal/net/http/middleware"
 	"github.com/yusing/go-proxy/internal/task"
@@ -54,10 +57,10 @@ func New(cfg *Config) *LoadBalancer {
 }
 
 // Start implements task.TaskStarter.
-func (lb *LoadBalancer) Start(routeSubtask task.Task) E.NestedError {
+func (lb *LoadBalancer) Start(routeSubtask task.Task) E.Error {
 	lb.startTime = time.Now()
 	lb.task = routeSubtask
-	lb.task.OnComplete("loadbalancer cleanup", func() {
+	lb.task.OnFinished("loadbalancer cleanup", func() {
 		if lb.impl != nil {
 			lb.pool.RangeAll(func(k string, v *Server) {
 				lb.impl.OnRemoveServer(v)
@@ -69,7 +72,7 @@ func (lb *LoadBalancer) Start(routeSubtask task.Task) E.NestedError {
 }
 
 // Finish implements task.TaskFinisher.
-func (lb *LoadBalancer) Finish(reason string) {
+func (lb *LoadBalancer) Finish(reason any) {
 	lb.task.Finish(reason)
 }
 
@@ -128,7 +131,7 @@ func (lb *LoadBalancer) AddServer(srv *Server) {
 
 	lb.rebalance()
 	lb.impl.OnAddServer(srv)
-	logger.Infof("[add] %s to loadbalancer %s: %d servers available", srv.Name, lb.Link, lb.pool.Size())
+	logger.Debugf("[add] %s to loadbalancer %s: %d servers available", srv.Name, lb.Link, lb.pool.Size())
 }
 
 func (lb *LoadBalancer) RemoveServer(srv *Server) {
@@ -147,11 +150,11 @@ func (lb *LoadBalancer) RemoveServer(srv *Server) {
 
 	if lb.pool.Size() == 0 {
 		lb.task.Finish("no server left")
-		logger.Infof("[remove] loadbalancer %s stopped", lb.Link)
+		logger.Infof("loadbalancer %s stopped", lb.Link)
 		return
 	}
 
-	logger.Infof("[remove] %s from loadbalancer %s: %d servers left", srv.Name, lb.Link, lb.pool.Size())
+	logger.Debugf("[remove] %s from loadbalancer %s: %d servers left", srv.Name, lb.Link, lb.pool.Size())
 }
 
 func (lb *LoadBalancer) rebalance() {
@@ -211,6 +214,21 @@ func (lb *LoadBalancer) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
 		http.Error(rw, "Service unavailable", http.StatusServiceUnavailable)
 		return
 	}
+	if r.Header.Get(common.HeaderCheckRedirect) != "" {
+		ctx, cancel := context.WithTimeout(r.Context(), 1*time.Second)
+		defer cancel()
+		// send dummy request to wake all servers
+		var dummyRW *DummyResponseWriter
+		for _, srv := range srvs {
+			// wake only if server implements Waker
+			_, ok := srv.handler.(idlewatcher.Waker)
+			if !ok {
+				continue
+			}
+			wakeReq := r.Clone(ctx)
+			srv.ServeHTTP(dummyRW, wakeReq)
+		}
+	}
 	lb.impl.ServeHTTP(srvs, rw, r)
 }
 
@@ -261,10 +279,9 @@ func (lb *LoadBalancer) String() string {
 func (lb *LoadBalancer) availServers() []*Server {
 	avail := make([]*Server, 0, lb.pool.Size())
 	lb.pool.RangeAll(func(_ string, srv *Server) {
-		if srv.Status().Bad() {
-			return
+		if srv.Status().Good() {
+			avail = append(avail, srv)
 		}
-		avail = append(avail, srv)
 	})
 	return avail
 }

internal/net/http/loadbalancer/round_robin.go

@@ -14,8 +14,8 @@ func (lb *roundRobin) OnAddServer(srv *Server)    {}
 func (lb *roundRobin) OnRemoveServer(srv *Server) {}
 
 func (lb *roundRobin) ServeHTTP(srvs servers, rw http.ResponseWriter, r *http.Request) {
-	index := lb.index.Add(1)
-	srvs[index%uint32(len(srvs))].ServeHTTP(rw, r)
+	index := lb.index.Add(1) % uint32(len(srvs))
+	srvs[index].ServeHTTP(rw, r)
 	if lb.index.Load() >= 2*uint32(len(srvs)) {
 		lb.index.Store(0)
 	}

internal/net/http/middleware/cidr_whitelist.go

@@ -35,7 +35,7 @@ var cidrWhitelistDefaults = func() *cidrWhitelistOpts {
 	}
 }
 
-func NewCIDRWhitelist(opts OptionsRaw) (*Middleware, E.NestedError) {
+func NewCIDRWhitelist(opts OptionsRaw) (*Middleware, E.Error) {
 	wl := new(cidrWhitelist)
 	wl.m = &Middleware{
 		impl:   wl,

internal/net/http/middleware/cloudflare_real_ip.go

@@ -33,7 +33,7 @@ var CloudflareRealIP = &realIP{
 	m: &Middleware{withOptions: NewCloudflareRealIP},
 }
 
-func NewCloudflareRealIP(_ OptionsRaw) (*Middleware, E.NestedError) {
+func NewCloudflareRealIP(_ OptionsRaw) (*Middleware, E.Error) {
 	cri := new(realIP)
 	cri.m = &Middleware{
 		impl: cri,

internal/net/http/middleware/forward_auth.go

@@ -36,7 +36,7 @@ var ForwardAuth = &forwardAuth{
 	m: &Middleware{withOptions: NewForwardAuthfunc},
 }
 
-func NewForwardAuthfunc(optsRaw OptionsRaw) (*Middleware, E.NestedError) {
+func NewForwardAuthfunc(optsRaw OptionsRaw) (*Middleware, E.Error) {
 	fa := new(forwardAuth)
 	fa.forwardAuthOpts = new(forwardAuthOpts)
 	err := Deserialize(optsRaw, fa.forwardAuthOpts)

internal/net/http/middleware/middleware.go

@@ -11,7 +11,7 @@ import (
 )
 
 type (
-	Error = E.NestedError
+	Error = E.Error
 
 	ReverseProxy   = gphttp.ReverseProxy
 	ProxyRequest   = gphttp.ProxyRequest
@@ -24,7 +24,7 @@ type (
 	BeforeFunc         func(next http.HandlerFunc, w ResponseWriter, r *Request)
 	RewriteFunc        func(req *Request)
 	ModifyResponseFunc func(resp *Response) error
-	CloneWithOptFunc   func(opts OptionsRaw) (*Middleware, E.NestedError)
+	CloneWithOptFunc   func(opts OptionsRaw) (*Middleware, E.Error)
 
 	OptionsRaw = map[string]any
 	Options    any
@@ -77,7 +77,7 @@ func (m *Middleware) MarshalJSON() ([]byte, error) {
 	}, "", "  ")
 }
 
-func (m *Middleware) WithOptionsClone(optsRaw OptionsRaw) (*Middleware, E.NestedError) {
+func (m *Middleware) WithOptionsClone(optsRaw OptionsRaw) (*Middleware, E.Error) {
 	if len(optsRaw) != 0 && m.withOptions != nil {
 		return m.withOptions(optsRaw)
 	}
@@ -108,7 +108,7 @@ func (m *Middleware) ModifyResponse(resp *Response) error {
 }
 
 // TODO: check conflict or duplicates.
-func createMiddlewares(middlewaresMap map[string]OptionsRaw) (middlewares []*Middleware, res E.NestedError) {
+func createMiddlewares(middlewaresMap map[string]OptionsRaw) (middlewares []*Middleware, res E.Error) {
 	middlewares = make([]*Middleware, 0, len(middlewaresMap))
 
 	invalidM := E.NewBuilder("invalid middlewares")
@@ -136,7 +136,7 @@ func createMiddlewares(middlewaresMap map[string]OptionsRaw) (middlewares []*Mid
 	return
 }
 
-func PatchReverseProxy(rpName string, rp *ReverseProxy, middlewaresMap map[string]OptionsRaw) (err E.NestedError) {
+func PatchReverseProxy(rpName string, rp *ReverseProxy, middlewaresMap map[string]OptionsRaw) (err E.Error) {
 	var middlewares []*Middleware
 	middlewares, err = createMiddlewares(middlewaresMap)
 	if err != nil {

internal/net/http/middleware/middleware_builder.go

@@ -10,7 +10,7 @@ import (
 	"gopkg.in/yaml.v3"
 )
 
-func BuildMiddlewaresFromComposeFile(filePath string) (map[string]*Middleware, E.NestedError) {
+func BuildMiddlewaresFromComposeFile(filePath string) (map[string]*Middleware, E.Error) {
 	fileContent, err := os.ReadFile(filePath)
 	if err != nil {
 		return nil, E.FailWith("read middleware compose file", err)
@@ -18,7 +18,7 @@ func BuildMiddlewaresFromComposeFile(filePath string) (map[string]*Middleware, E
 	return BuildMiddlewaresFromYAML(fileContent)
 }
 
-func BuildMiddlewaresFromYAML(data []byte) (middlewares map[string]*Middleware, outErr E.NestedError) {
+func BuildMiddlewaresFromYAML(data []byte) (middlewares map[string]*Middleware, outErr E.Error) {
 	b := E.NewBuilder("middlewares compile errors")
 	defer b.To(&outErr)
 

internal/net/http/middleware/modify_request.go

@@ -22,7 +22,7 @@ var ModifyRequest = &modifyRequest{
 	m: &Middleware{withOptions: NewModifyRequest},
 }
 
-func NewModifyRequest(optsRaw OptionsRaw) (*Middleware, E.NestedError) {
+func NewModifyRequest(optsRaw OptionsRaw) (*Middleware, E.Error) {
 	mr := new(modifyRequest)
 	var mrFunc RewriteFunc
 	if common.IsDebug {

internal/net/http/middleware/modify_response.go

@@ -24,7 +24,7 @@ var ModifyResponse = &modifyResponse{
 	m: &Middleware{withOptions: NewModifyResponse},
 }
 
-func NewModifyResponse(optsRaw OptionsRaw) (*Middleware, E.NestedError) {
+func NewModifyResponse(optsRaw OptionsRaw) (*Middleware, E.Error) {
 	mr := new(modifyResponse)
 	mr.m = &Middleware{impl: mr}
 	if common.IsDebug {

internal/net/http/middleware/oauth2.go

@@ -26,7 +26,7 @@ var OAuth2 = &oAuth2{
 	m: &Middleware{withOptions: NewAuthentikOAuth2},
 }
 
-func NewAuthentikOAuth2(opts OptionsRaw) (*Middleware, E.NestedError) {
+func NewAuthentikOAuth2(opts OptionsRaw) (*Middleware, E.Error) {
 	oauth := new(oAuth2)
 	oauth.m = &Middleware{
 		impl:   oauth,

internal/net/http/middleware/real_ip.go

@@ -41,7 +41,7 @@ var realIPOptsDefault = func() *realIPOpts {
 	}
 }
 
-func NewRealIP(opts OptionsRaw) (*Middleware, E.NestedError) {
+func NewRealIP(opts OptionsRaw) (*Middleware, E.Error) {
 	riWithOpts := new(realIP)
 	riWithOpts.m = &Middleware{
 		impl:   riWithOpts,

internal/net/http/middleware/test_utils.go

@@ -72,7 +72,7 @@ type testArgs struct {
 	scheme        string
 }
 
-func newMiddlewareTest(middleware *Middleware, args *testArgs) (*TestResult, E.NestedError) {
+func newMiddlewareTest(middleware *Middleware, args *testArgs) (*TestResult, E.Error) {
 	var body io.Reader
 	var rr requestRecorder
 	var proxyURL *url.URL

internal/net/types/cidr.go

@@ -9,7 +9,7 @@ import (
 
 type CIDR net.IPNet
 
-func (cidr *CIDR) ConvertFrom(val any) E.NestedError {
+func (cidr *CIDR) ConvertFrom(val any) E.Error {
 	cidrStr, ok := val.(string)
 	if !ok {
 		return E.TypeMismatch[string](val)

internal/net/types/stream.go

@@ -7,13 +7,7 @@ import (
 
 type Stream interface {
 	fmt.Stringer
+	net.Listener
 	Setup() error
-	Accept() (conn StreamConn, err error)
-	Handle(conn StreamConn) error
-	CloseListeners()
-}
-
-type StreamConn interface {
-	RemoteAddr() net.Addr
-	Close() error
+	Handle(conn net.Conn) error
 }