github.com/xtls/xray-core@v1.8.12-0.20240518155711-3168d27b0bdb/app/router/balancing.go (about)

     1  package router
     2  
     3  import (
     4  	"context"
     5  	sync "sync"
     6  
     7  	"github.com/xtls/xray-core/app/observatory"
     8  	"github.com/xtls/xray-core/common"
     9  	"github.com/xtls/xray-core/core"
    10  	"github.com/xtls/xray-core/features/extension"
    11  	"github.com/xtls/xray-core/features/outbound"
    12  )
    13  
    14  type BalancingStrategy interface {
    15  	PickOutbound([]string) string
    16  }
    17  
    18  type BalancingPrincipleTarget interface {
    19  	GetPrincipleTarget([]string) []string
    20  }
    21  
    22  type RoundRobinStrategy struct {
    23  	FallbackTag string
    24  
    25  	ctx         context.Context
    26  	observatory extension.Observatory
    27  	mu    sync.Mutex
    28  	index int
    29  }
    30  
    31  func (s *RoundRobinStrategy) InjectContext(ctx context.Context) {
    32  	s.ctx = ctx
    33  }
    34  
    35  func (s *RoundRobinStrategy) GetPrincipleTarget(strings []string) []string {
    36  	return strings
    37  }
    38  
    39  func (s *RoundRobinStrategy) PickOutbound(tags []string) string {
    40  	if len(s.FallbackTag) > 0 && s.observatory == nil {
    41  		common.Must(core.RequireFeatures(s.ctx, func(observatory extension.Observatory) error {
    42  			s.observatory = observatory
    43  			return nil
    44  		}))
    45  	}
    46  	if s.observatory != nil {
    47  		observeReport, err := s.observatory.GetObservation(s.ctx)
    48  		if err == nil {
    49  			aliveTags := make([]string, 0)
    50  			if result, ok := observeReport.(*observatory.ObservationResult); ok {
    51  				status := result.Status
    52  				statusMap := make(map[string]*observatory.OutboundStatus)
    53  				for _, outboundStatus := range status {
    54  					statusMap[outboundStatus.OutboundTag] = outboundStatus
    55  				}
    56  				for _, candidate := range tags {
    57  					if outboundStatus, found := statusMap[candidate]; found {
    58  						if outboundStatus.Alive {
    59  							aliveTags = append(aliveTags, candidate)
    60  						}
    61  					} else {
    62  						// unfound candidate is considered alive
    63  						aliveTags = append(aliveTags, candidate)
    64  					}
    65  				}
    66  				tags = aliveTags
    67  			}
    68  		}
    69  	}
    70  
    71  	n := len(tags)
    72  	if n == 0 {
    73  		// goes to fallbackTag
    74  		return ""
    75  	}
    76  
    77  	s.mu.Lock()
    78  	defer s.mu.Unlock()
    79  	tag := tags[s.index%n]
    80  	s.index = (s.index + 1) % n
    81  	return tag
    82  }
    83  
    84  type Balancer struct {
    85  	selectors   []string
    86  	strategy    BalancingStrategy
    87  	ohm         outbound.Manager
    88  	fallbackTag string
    89  
    90  	override override
    91  }
    92  
    93  // PickOutbound picks the tag of a outbound
    94  func (b *Balancer) PickOutbound() (string, error) {
    95  	candidates, err := b.SelectOutbounds()
    96  	if err != nil {
    97  		if b.fallbackTag != "" {
    98  			newError("fallback to [", b.fallbackTag, "], due to error: ", err).AtInfo().WriteToLog()
    99  			return b.fallbackTag, nil
   100  		}
   101  		return "", err
   102  	}
   103  	var tag string
   104  	if o := b.override.Get(); o != "" {
   105  		tag = o
   106  	} else {
   107  		tag = b.strategy.PickOutbound(candidates)
   108  	}
   109  	if tag == "" {
   110  		if b.fallbackTag != "" {
   111  			newError("fallback to [", b.fallbackTag, "], due to empty tag returned").AtInfo().WriteToLog()
   112  			return b.fallbackTag, nil
   113  		}
   114  		// will use default handler
   115  		return "", newError("balancing strategy returns empty tag")
   116  	}
   117  	return tag, nil
   118  }
   119  
   120  func (b *Balancer) InjectContext(ctx context.Context) {
   121  	if contextReceiver, ok := b.strategy.(extension.ContextReceiver); ok {
   122  		contextReceiver.InjectContext(ctx)
   123  	}
   124  }
   125  
   126  // SelectOutbounds select outbounds with selectors of the Balancer
   127  func (b *Balancer) SelectOutbounds() ([]string, error) {
   128  	hs, ok := b.ohm.(outbound.HandlerSelector)
   129  	if !ok {
   130  		return nil, newError("outbound.Manager is not a HandlerSelector")
   131  	}
   132  	tags := hs.Select(b.selectors)
   133  	return tags, nil
   134  }
   135  
   136  // GetPrincipleTarget implements routing.BalancerPrincipleTarget
   137  func (r *Router) GetPrincipleTarget(tag string) ([]string, error) {
   138  	if b, ok := r.balancers[tag]; ok {
   139  		if s, ok := b.strategy.(BalancingPrincipleTarget); ok {
   140  			candidates, err := b.SelectOutbounds()
   141  			if err != nil {
   142  				return nil, newError("unable to select outbounds").Base(err)
   143  			}
   144  			return s.GetPrincipleTarget(candidates), nil
   145  		}
   146  		return nil, newError("unsupported GetPrincipleTarget")
   147  	}
   148  	return nil, newError("cannot find tag")
   149  }
   150  
   151  // SetOverrideTarget implements routing.BalancerOverrider
   152  func (r *Router) SetOverrideTarget(tag, target string) error {
   153  	if b, ok := r.balancers[tag]; ok {
   154  		b.override.Put(target)
   155  		return nil
   156  	}
   157  	return newError("cannot find tag")
   158  }
   159  
   160  // GetOverrideTarget implements routing.BalancerOverrider
   161  func (r *Router) GetOverrideTarget(tag string) (string, error) {
   162  	if b, ok := r.balancers[tag]; ok {
   163  		return b.override.Get(), nil
   164  	}
   165  	return "", newError("cannot find tag")
   166  }