Add tags cache to app.proxyman.ohm.Select() (#2927)

* Add tags cache to ohm.Select().

* Refactor round-robin.

* Fix a bug.

---------

Co-authored-by: nobody <nobody@nowhere.mars>
This commit is contained in:
nobody 2024-01-12 23:36:48 +08:00 committed by GitHub
parent 0ea2a50264
commit 7f7f57d3b6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 22 additions and 24 deletions

View File

@ -22,12 +22,14 @@ type Manager struct {
taggedHandler map[string]outbound.Handler taggedHandler map[string]outbound.Handler
untaggedHandlers []outbound.Handler untaggedHandlers []outbound.Handler
running bool running bool
tagsCache map[string][]string
} }
// New creates a new Manager. // New creates a new Manager.
func New(ctx context.Context, config *proxyman.OutboundConfig) (*Manager, error) { func New(ctx context.Context, config *proxyman.OutboundConfig) (*Manager, error) {
m := &Manager{ m := &Manager{
taggedHandler: make(map[string]outbound.Handler), taggedHandler: make(map[string]outbound.Handler),
tagsCache: make(map[string][]string),
} }
return m, nil return m, nil
} }
@ -104,6 +106,8 @@ func (m *Manager) AddHandler(ctx context.Context, handler outbound.Handler) erro
m.access.Lock() m.access.Lock()
defer m.access.Unlock() defer m.access.Unlock()
m.tagsCache = make(map[string][]string)
if m.defaultHandler == nil { if m.defaultHandler == nil {
m.defaultHandler = handler m.defaultHandler = handler
} }
@ -133,6 +137,8 @@ func (m *Manager) RemoveHandler(ctx context.Context, tag string) error {
m.access.Lock() m.access.Lock()
defer m.access.Unlock() defer m.access.Unlock()
m.tagsCache = make(map[string][]string)
delete(m.taggedHandler, tag) delete(m.taggedHandler, tag)
if m.defaultHandler != nil && m.defaultHandler.Tag() == tag { if m.defaultHandler != nil && m.defaultHandler.Tag() == tag {
m.defaultHandler = nil m.defaultHandler = nil
@ -146,6 +152,11 @@ func (m *Manager) Select(selectors []string) []string {
m.access.RLock() m.access.RLock()
defer m.access.RUnlock() defer m.access.RUnlock()
key := strings.Join(selectors, ",")
if cache, ok := m.tagsCache[key]; ok {
return cache
}
tags := make([]string, 0, len(selectors)) tags := make([]string, 0, len(selectors))
for tag := range m.taggedHandler { for tag := range m.taggedHandler {
@ -156,7 +167,10 @@ func (m *Manager) Select(selectors []string) []string {
} }
} }
} }
sort.Strings(tags) sort.Strings(tags)
m.tagsCache[key] = tags
return tags return tags
} }

View File

@ -2,7 +2,6 @@ package router
import ( import (
"context" "context"
reflect "reflect"
sync "sync" sync "sync"
"github.com/xtls/xray-core/common/dice" "github.com/xtls/xray-core/common/dice"
@ -27,34 +26,19 @@ func (s *RandomStrategy) PickOutbound(tags []string) string {
type RoundRobinStrategy struct { type RoundRobinStrategy struct {
mu sync.Mutex mu sync.Mutex
tags []string
index int index int
roundRobin *RoundRobinStrategy
}
func NewRoundRobin(tags []string) *RoundRobinStrategy {
return &RoundRobinStrategy{
tags: tags,
}
}
func (r *RoundRobinStrategy) NextTag() string {
r.mu.Lock()
defer r.mu.Unlock()
tags := r.tags[r.index]
r.index = (r.index + 1) % len(r.tags)
return tags
} }
func (s *RoundRobinStrategy) PickOutbound(tags []string) string { func (s *RoundRobinStrategy) PickOutbound(tags []string) string {
if len(tags) == 0 { n := len(tags)
if n == 0 {
panic("0 tags") panic("0 tags")
} }
if s.roundRobin == nil || !reflect.DeepEqual(s.roundRobin.tags, tags) {
s.roundRobin = NewRoundRobin(tags)
}
tag := s.roundRobin.NextTag()
s.mu.Lock()
defer s.mu.Unlock()
tag := tags[s.index%n]
s.index = (s.index + 1) % n
return tag return tag
} }