github.com/xraypb/Xray-core@v1.8.1/app/router/balancing.go (about)

     1  package router
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/xraypb/Xray-core/common/dice"
     7  	"github.com/xraypb/Xray-core/features/extension"
     8  	"github.com/xraypb/Xray-core/features/outbound"
     9  )
    10  
    11  type BalancingStrategy interface {
    12  	PickOutbound([]string) string
    13  }
    14  
    15  type RandomStrategy struct{}
    16  
    17  func (s *RandomStrategy) PickOutbound(tags []string) string {
    18  	n := len(tags)
    19  	if n == 0 {
    20  		panic("0 tags")
    21  	}
    22  
    23  	return tags[dice.Roll(n)]
    24  }
    25  
    26  type Balancer struct {
    27  	selectors []string
    28  	strategy  BalancingStrategy
    29  	ohm       outbound.Manager
    30  }
    31  
    32  func (b *Balancer) PickOutbound() (string, error) {
    33  	hs, ok := b.ohm.(outbound.HandlerSelector)
    34  	if !ok {
    35  		return "", newError("outbound.Manager is not a HandlerSelector")
    36  	}
    37  	tags := hs.Select(b.selectors)
    38  	if len(tags) == 0 {
    39  		return "", newError("no available outbounds selected")
    40  	}
    41  	tag := b.strategy.PickOutbound(tags)
    42  	if tag == "" {
    43  		return "", newError("balancing strategy returns empty tag")
    44  	}
    45  	return tag, nil
    46  }
    47  
    48  func (b *Balancer) InjectContext(ctx context.Context) {
    49  	if contextReceiver, ok := b.strategy.(extension.ContextReceiver); ok {
    50  		contextReceiver.InjectContext(ctx)
    51  	}
    52  }