github.com/imannamdari/v2ray-core/v5@v5.0.5/app/router/balancing.go (about)

     1  //go:build !confonly
     2  // +build !confonly
     3  
     4  package router
     5  
     6  import (
     7  	"context"
     8  
     9  	"github.com/imannamdari/v2ray-core/v5/features/extension"
    10  	"github.com/imannamdari/v2ray-core/v5/features/outbound"
    11  )
    12  
    13  type BalancingStrategy interface {
    14  	PickOutbound([]string) string
    15  }
    16  
    17  type BalancingPrincipleTarget interface {
    18  	GetPrincipleTarget([]string) []string
    19  }
    20  
    21  type Balancer struct {
    22  	selectors   []string
    23  	strategy    BalancingStrategy
    24  	ohm         outbound.Manager
    25  	fallbackTag string
    26  
    27  	override override
    28  }
    29  
    30  // PickOutbound picks the tag of an outbound
    31  func (b *Balancer) PickOutbound() (string, error) {
    32  	candidates, err := b.SelectOutbounds()
    33  	if err != nil {
    34  		if b.fallbackTag != "" {
    35  			newError("fallback to [", b.fallbackTag, "], due to error: ", err).AtInfo().WriteToLog()
    36  			return b.fallbackTag, nil
    37  		}
    38  		return "", err
    39  	}
    40  	var tag string
    41  	if o := b.override.Get(); o != "" {
    42  		tag = o
    43  	} else {
    44  		tag = b.strategy.PickOutbound(candidates)
    45  	}
    46  	if tag == "" {
    47  		if b.fallbackTag != "" {
    48  			newError("fallback to [", b.fallbackTag, "], due to empty tag returned").AtInfo().WriteToLog()
    49  			return b.fallbackTag, nil
    50  		}
    51  		// will use default handler
    52  		return "", newError("balancing strategy returns empty tag")
    53  	}
    54  	return tag, nil
    55  }
    56  
    57  func (b *Balancer) InjectContext(ctx context.Context) {
    58  	if contextReceiver, ok := b.strategy.(extension.ContextReceiver); ok {
    59  		contextReceiver.InjectContext(ctx)
    60  	}
    61  }
    62  
    63  // SelectOutbounds select outbounds with selectors of the Balancer
    64  func (b *Balancer) SelectOutbounds() ([]string, error) {
    65  	hs, ok := b.ohm.(outbound.HandlerSelector)
    66  	if !ok {
    67  		return nil, newError("outbound.Manager is not a HandlerSelector")
    68  	}
    69  	tags := hs.Select(b.selectors)
    70  	return tags, nil
    71  }
    72  
    73  // GetPrincipleTarget implements routing.BalancerPrincipleTarget
    74  func (r *Router) GetPrincipleTarget(tag string) ([]string, error) {
    75  	if b, ok := r.balancers[tag]; ok {
    76  		if s, ok := b.strategy.(BalancingPrincipleTarget); ok {
    77  			candidates, err := b.SelectOutbounds()
    78  			if err != nil {
    79  				return nil, newError("unable to select outbounds").Base(err)
    80  			}
    81  			return s.GetPrincipleTarget(candidates), nil
    82  		}
    83  		return nil, newError("unsupported GetPrincipleTarget")
    84  	}
    85  	return nil, newError("cannot find tag")
    86  }
    87  
    88  // SetOverrideTarget implements routing.BalancerOverrider
    89  func (r *Router) SetOverrideTarget(tag, target string) error {
    90  	if b, ok := r.balancers[tag]; ok {
    91  		b.override.Put(target)
    92  		return nil
    93  	}
    94  	return newError("cannot find tag")
    95  }
    96  
    97  // GetOverrideTarget implements routing.BalancerOverrider
    98  func (r *Router) GetOverrideTarget(tag string) (string, error) {
    99  	if b, ok := r.balancers[tag]; ok {
   100  		return b.override.Get(), nil
   101  	}
   102  	return "", newError("cannot find tag")
   103  }