Select alive only node when fallbackTag is given

- Apply to random and roundrobin strategy
- Require observatory config

Co-authored-by: Mark Ma <38940419+mkmark@users.noreply.github.com>
This commit is contained in:
yuhan6665 2024-05-05 10:07:40 -04:00
parent eba2906d3a
commit 84eeb56ae4
4 changed files with 100 additions and 5 deletions

View File

@ -4,6 +4,9 @@ import (
"context" "context"
sync "sync" sync "sync"
"github.com/xtls/xray-core/app/observatory"
"github.com/xtls/xray-core/common"
"github.com/xtls/xray-core/core"
"github.com/xtls/xray-core/features/extension" "github.com/xtls/xray-core/features/extension"
"github.com/xtls/xray-core/features/outbound" "github.com/xtls/xray-core/features/outbound"
) )
@ -17,14 +20,58 @@ type BalancingPrincipleTarget interface {
} }
type RoundRobinStrategy struct { type RoundRobinStrategy struct {
FallbackTag string
ctx context.Context
observatory extension.Observatory
mu sync.Mutex mu sync.Mutex
index int index int
} }
func (s *RoundRobinStrategy) InjectContext(ctx context.Context) {
s.ctx = ctx
}
func (s *RoundRobinStrategy) GetPrincipleTarget(strings []string) []string {
return strings
}
func (s *RoundRobinStrategy) PickOutbound(tags []string) string { func (s *RoundRobinStrategy) PickOutbound(tags []string) string {
if len(s.FallbackTag) > 0 && s.observatory == nil {
common.Must(core.RequireFeatures(s.ctx, func(observatory extension.Observatory) error {
s.observatory = observatory
return nil
}))
}
if s.observatory != nil {
observeReport, err := s.observatory.GetObservation(s.ctx)
if err == nil {
aliveTags := make([]string, 0)
if result, ok := observeReport.(*observatory.ObservationResult); ok {
status := result.Status
statusMap := make(map[string]*observatory.OutboundStatus)
for _, outboundStatus := range status {
statusMap[outboundStatus.OutboundTag] = outboundStatus
}
for _, candidate := range tags {
if outboundStatus, found := statusMap[candidate]; found {
if outboundStatus.Alive {
aliveTags = append(aliveTags, candidate)
}
} else {
// unfound candidate is considered alive
aliveTags = append(aliveTags, candidate)
}
}
tags = aliveTags
}
}
}
n := len(tags) n := len(tags)
if n == 0 { if n == 0 {
panic("0 tags") // goes to fallbackTag
return ""
} }
s.mu.Lock() s.mu.Lock()

View File

@ -135,7 +135,7 @@ func (br *BalancingRule) Build(ohm outbound.Manager, dispatcher routing.Dispatch
case "roundrobin": case "roundrobin":
return &Balancer{ return &Balancer{
selectors: br.OutboundSelector, selectors: br.OutboundSelector,
strategy: &RoundRobinStrategy{}, strategy: &RoundRobinStrategy{FallbackTag: br.FallbackTag},
fallbackTag: br.FallbackTag, fallbackTag: br.FallbackTag,
ohm: ohm, ohm: ohm,
}, nil }, nil
@ -162,7 +162,7 @@ func (br *BalancingRule) Build(ohm outbound.Manager, dispatcher routing.Dispatch
selectors: br.OutboundSelector, selectors: br.OutboundSelector,
ohm: ohm, ohm: ohm,
fallbackTag: br.FallbackTag, fallbackTag: br.FallbackTag,
strategy: &RandomStrategy{}, strategy: &RandomStrategy{FallbackTag: br.FallbackTag},
}, nil }, nil
default: default:
return nil, newError("unrecognized balancer type") return nil, newError("unrecognized balancer type")

View File

@ -1,17 +1,63 @@
package router package router
import ( import (
"context"
"github.com/xtls/xray-core/app/observatory"
"github.com/xtls/xray-core/common"
"github.com/xtls/xray-core/common/dice" "github.com/xtls/xray-core/common/dice"
"github.com/xtls/xray-core/core"
"github.com/xtls/xray-core/features/extension"
) )
// RandomStrategy represents a random balancing strategy // RandomStrategy represents a random balancing strategy
type RandomStrategy struct{} type RandomStrategy struct{
FallbackTag string
ctx context.Context
observatory extension.Observatory
}
func (s *RandomStrategy) InjectContext(ctx context.Context) {
s.ctx = ctx
}
func (s *RandomStrategy) GetPrincipleTarget(strings []string) []string { func (s *RandomStrategy) GetPrincipleTarget(strings []string) []string {
return strings return strings
} }
func (s *RandomStrategy) PickOutbound(candidates []string) string { func (s *RandomStrategy) PickOutbound(candidates []string) string {
if len(s.FallbackTag) > 0 && s.observatory == nil {
common.Must(core.RequireFeatures(s.ctx, func(observatory extension.Observatory) error {
s.observatory = observatory
return nil
}))
}
if s.observatory != nil {
observeReport, err := s.observatory.GetObservation(s.ctx)
if err == nil {
aliveTags := make([]string, 0)
if result, ok := observeReport.(*observatory.ObservationResult); ok {
status := result.Status
statusMap := make(map[string]*observatory.OutboundStatus)
for _, outboundStatus := range status {
statusMap[outboundStatus.OutboundTag] = outboundStatus
}
for _, candidate := range candidates {
if outboundStatus, found := statusMap[candidate]; found {
if outboundStatus.Alive {
aliveTags = append(aliveTags, candidate)
}
} else {
// unfound candidate is considered alive
aliveTags = append(aliveTags, candidate)
}
}
candidates = aliveTags
}
}
}
count := len(candidates) count := len(candidates)
if count == 0 { if count == 0 {
// goes to fallbackTag // goes to fallbackTag

View File

@ -97,7 +97,8 @@ func TestRouterConfig(t *testing.T) {
"balancers": [ "balancers": [
{ {
"tag": "b1", "tag": "b1",
"selector": ["test"] "selector": ["test"],
"fallbackTag": "fall"
}, },
{ {
"tag": "b2", "tag": "b2",
@ -137,6 +138,7 @@ func TestRouterConfig(t *testing.T) {
Tag: "b1", Tag: "b1",
OutboundSelector: []string{"test"}, OutboundSelector: []string{"test"},
Strategy: "random", Strategy: "random",
FallbackTag: "fall",
}, },
{ {
Tag: "b2", Tag: "b2",