v5: Health Check & LeastLoad Strategy

app/dispatcher/default.go

@@ -189,6 +189,7 @@ func (d *DefaultDispatcher) Dispatch(ctx context.Context, destination net.Destin
 	if !destination.IsValid() {
 		panic("Dispatcher: Invalid destination.")
 	}
+
 	ob := &session.Outbound{
 		Target: destination,
 	}
@@ -200,9 +201,15 @@ func (d *DefaultDispatcher) Dispatch(ctx context.Context, destination net.Destin
 		content = new(session.Content)
 		ctx = session.ContextWithContent(ctx, content)
 	}
+
+	handler := session.HandlerFromContext(ctx)
 	sniffingRequest := content.SniffingRequest
 	if destination.Network != net.Network_TCP || !sniffingRequest.Enabled {
-		go d.routedDispatch(ctx, outbound, destination)
+		if handler != nil {
+			go d.targetedDispatch(ctx, outbound, handler.Tag)
+		} else {
+			go d.routedDispatch(ctx, outbound, destination)
+		}
 	} else {
 		go func() {
 			cReader := &cachedReader{
@@ -219,7 +226,11 @@ func (d *DefaultDispatcher) Dispatch(ctx context.Context, destination net.Destin
 				destination.Address = net.ParseAddress(domain)
 				ob.Target = destination
 			}
-			d.routedDispatch(ctx, outbound, destination)
+			if handler != nil {
+				d.targetedDispatch(ctx, outbound, handler.Tag)
+			} else {
+				d.routedDispatch(ctx, outbound, destination)
+			}
 		}()
 	}
 	return inbound, nil
@@ -255,6 +266,24 @@ func sniffer(ctx context.Context, cReader *cachedReader) (SniffResult, error) {
 	}
 }
 
+func (d *DefaultDispatcher) targetedDispatch(ctx context.Context, link *transport.Link, tag string) {
+	handler := d.ohm.GetHandler(tag)
+	if handler == nil {
+		newError("outbound handler [", tag, "] not exist").AtError().WriteToLog(session.ExportIDToError(ctx))
+		common.Close(link.Writer)
+		common.Interrupt(link.Reader)
+		return
+	}
+	if accessMessage := log.AccessMessageFromContext(ctx); accessMessage != nil {
+		if tag := handler.Tag(); tag != "" {
+			accessMessage.Detour = tag
+		}
+		log.Record(accessMessage)
+	}
+
+	handler.Dispatch(ctx, link)
+}
+
 func (d *DefaultDispatcher) routedDispatch(ctx context.Context, link *transport.Link, destination net.Destination) {
 	var handler outbound.Handler
 

app/router/balancing.go

@@ -1,44 +1,53 @@
 package router
 
 import (
-	"v2ray.com/core/common/dice"
 	"v2ray.com/core/features/outbound"
+	"v2ray.com/core/features/routing"
 )
 
-type BalancingStrategy interface {
-	PickOutbound([]string) string
-}
-
-type RandomStrategy struct {
-}
-
-func (s *RandomStrategy) PickOutbound(tags []string) string {
-	n := len(tags)
-	if n == 0 {
-		panic("0 tags")
-	}
-
-	return tags[dice.Roll(n)]
-}
-
+// Balancer represents a balancer
 type Balancer struct {
-	selectors []string
-	strategy  BalancingStrategy
-	ohm       outbound.Manager
+	selectors   []string
+	strategy    routing.BalancingStrategy
+	ohm         outbound.Manager
+	fallbackTag string
+
+	override overridden
 }
 
+// PickOutbound picks the tag of a outbound
 func (b *Balancer) PickOutbound() (string, error) {
-	hs, ok := b.ohm.(outbound.HandlerSelector)
-	if !ok {
-		return "", newError("outbound.Manager is not a HandlerSelector")
+	candidates, err := b.SelectOutbounds()
+	if err != nil {
+		if b.fallbackTag != "" {
+			newError("fallback to [", b.fallbackTag, "], due to error: ", err).AtInfo().WriteToLog()
+			return b.fallbackTag, nil
+		}
+		return "", err
 	}
-	tags := hs.Select(b.selectors)
-	if len(tags) == 0 {
-		return "", newError("no available outbounds selected")
+	var tag string
+	if o := b.override.Get(); o != nil {
+		tag = b.strategy.Pick(o.selects)
+	} else {
+		tag = b.strategy.SelectAndPick(candidates)
 	}
-	tag := b.strategy.PickOutbound(tags)
 	if tag == "" {
+		if b.fallbackTag != "" {
+			newError("fallback to [", b.fallbackTag, "], due to empty tag returned").AtInfo().WriteToLog()
+			return b.fallbackTag, nil
+		}
+		// will use default handler
 		return "", newError("balancing strategy returns empty tag")
 	}
 	return tag, nil
 }
+
+// SelectOutbounds select outbounds with selectors of the Balancer
+func (b *Balancer) SelectOutbounds() ([]string, error) {
+	hs, ok := b.ohm.(outbound.HandlerSelector)
+	if !ok {
+		return nil, newError("outbound.Manager is not a HandlerSelector")
+	}
+	tags := hs.Select(b.selectors)
+	return tags, nil
+}

app/router/balancing_override.go

@@ -0,0 +1,83 @@
+package router
+
+import (
+	sync "sync"
+	"time"
+
+	"v2ray.com/core/features/outbound"
+)
+
+func (b *Balancer) overrideSelecting(selects []string, validity time.Duration) error {
+	if validity <= 0 {
+		b.override.Clear()
+		return nil
+	}
+	hs, ok := b.ohm.(outbound.HandlerSelector)
+	if !ok {
+		return newError("outbound.Manager is not a HandlerSelector")
+	}
+	tags := hs.Select(selects)
+	if len(tags) == 0 {
+		return newError("no outbound selected")
+	}
+	b.override.Put(tags, time.Now().Add(validity))
+	return nil
+}
+
+// OverrideSelecting implements routing.BalancingOverrider
+func (r *Router) OverrideSelecting(balancer string, selects []string, validity time.Duration) error {
+	var b *Balancer
+	for tag, bl := range r.balancers {
+		if tag == balancer {
+			b = bl
+			break
+		}
+	}
+	if b == nil {
+		return newError("balancer '", balancer, "' not found")
+	}
+	err := b.overrideSelecting(selects, validity)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+type overriddenSettings struct {
+	selects []string
+	until   time.Time
+}
+
+type overridden struct {
+	access   sync.RWMutex
+	settings overriddenSettings
+}
+
+// Get gets the overridden settings
+func (o *overridden) Get() *overriddenSettings {
+	o.access.RLock()
+	defer o.access.RUnlock()
+	if len(o.settings.selects) == 0 || time.Now().After(o.settings.until) {
+		return nil
+	}
+	return &overriddenSettings{
+		selects: o.settings.selects,
+		until:   o.settings.until,
+	}
+}
+
+// Put updates the overridden settings
+func (o *overridden) Put(selects []string, until time.Time) {
+	o.access.Lock()
+	defer o.access.Unlock()
+	o.settings.selects = selects
+	o.settings.until = until
+}
+
+// Clear clears the overridden settings
+func (o *overridden) Clear() {
+	o.access.Lock()
+	defer o.access.Unlock()
+	o.settings.selects = nil
+	o.settings.until = time.Time{}
+}

app/router/command/command.go

@@ -7,6 +7,8 @@ import (
 	"time"
 
 	"google.golang.org/grpc"
+	codes "google.golang.org/grpc/codes"
+	status "google.golang.org/grpc/status"
 	"v2ray.com/core"
 	"v2ray.com/core/common"
 	"v2ray.com/core/features/routing"
@@ -72,6 +74,80 @@ func (s *routingServer) SubscribeRoutingStats(request *SubscribeRoutingStatsRequ
 	}
 }
 
+func (s *routingServer) GetBalancers(ctx context.Context, request *GetBalancersRequest) (*GetBalancersResponse, error) {
+	h, ok := s.router.(routing.RouterChecker)
+	if !ok {
+		return nil, status.Errorf(codes.Unavailable, "current router is not a health checker")
+	}
+	results, err := h.GetBalancersInfo(request.BalancerTags)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, err.Error())
+	}
+	rsp := &GetBalancersResponse{
+		Balancers: make([]*BalancerMsg, 0),
+	}
+	for _, result := range results {
+		var override *OverrideSelectingMsg
+		if result.Override != nil {
+			override = &OverrideSelectingMsg{
+				Until:   result.Override.Until.Local().String(),
+				Selects: result.Override.Selects,
+			}
+		}
+		stat := &BalancerMsg{
+			Tag:              result.Tag,
+			StrategySettings: result.Strategy.Settings,
+			Titles:           result.Strategy.ValueTitles,
+			Override:         override,
+			Selects:          make([]*OutboundMsg, 0),
+			Others:           make([]*OutboundMsg, 0),
+		}
+		for _, item := range result.Strategy.Selects {
+			stat.Selects = append(stat.Selects, &OutboundMsg{
+				Tag:    item.Tag,
+				Values: item.Values,
+			})
+		}
+		for _, item := range result.Strategy.Others {
+			stat.Others = append(stat.Others, &OutboundMsg{
+				Tag:    item.Tag,
+				Values: item.Values,
+			})
+		}
+		rsp.Balancers = append(rsp.Balancers, stat)
+	}
+	return rsp, nil
+}
+func (s *routingServer) CheckBalancers(ctx context.Context, request *CheckBalancersRequest) (*CheckBalancersResponse, error) {
+	h, ok := s.router.(routing.RouterChecker)
+	if !ok {
+		return nil, status.Errorf(codes.Unavailable, "current router is not a health checker")
+	}
+	go func() {
+		err := h.CheckBalancers(request.BalancerTags)
+		if err != nil {
+			newError("CheckBalancers error:", err).AtInfo().WriteToLog()
+		}
+	}()
+	return &CheckBalancersResponse{}, nil
+}
+
+func (s *routingServer) OverrideSelecting(ctx context.Context, request *OverrideSelectingRequest) (*OverrideSelectingResponse, error) {
+	bo, ok := s.router.(routing.BalancingOverrider)
+	if !ok {
+		return nil, status.Errorf(codes.Unavailable, "current router doesn't support balancing override")
+	}
+	err := bo.OverrideSelecting(
+		request.BalancerTag,
+		request.Selectors,
+		time.Duration(request.Validity),
+	)
+	if err != nil {
+		return nil, status.Errorf(codes.InvalidArgument, err.Error())
+	}
+	return &OverrideSelectingResponse{}, nil
+}
+
 func (s *routingServer) mustEmbedUnimplementedRoutingServiceServer() {}
 
 type service struct {