From 7f7f57d3b6a3fbce3cd2ab3d34abe27e6f6a3829 Mon Sep 17 00:00:00 2001
From: nobody <59990325+vrnobody@users.noreply.github.com>
Date: Fri, 12 Jan 2024 23:36:48 +0800
Subject: [PATCH] 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>
---
 app/proxyman/outbound/outbound.go | 14 ++++++++++++++
 app/router/balancing.go           | 32 ++++++++-----------------------
 2 files changed, 22 insertions(+), 24 deletions(-)

diff --git a/app/proxyman/outbound/outbound.go b/app/proxyman/outbound/outbound.go
index 1bc4c403..3bd0d85c 100644
--- a/app/proxyman/outbound/outbound.go
+++ b/app/proxyman/outbound/outbound.go
@@ -22,12 +22,14 @@ type Manager struct {
 	taggedHandler    map[string]outbound.Handler
 	untaggedHandlers []outbound.Handler
 	running          bool
+	tagsCache        map[string][]string
 }
 
 // New creates a new Manager.
 func New(ctx context.Context, config *proxyman.OutboundConfig) (*Manager, error) {
 	m := &Manager{
 		taggedHandler: make(map[string]outbound.Handler),
+		tagsCache:     make(map[string][]string),
 	}
 	return m, nil
 }
@@ -104,6 +106,8 @@ func (m *Manager) AddHandler(ctx context.Context, handler outbound.Handler) erro
 	m.access.Lock()
 	defer m.access.Unlock()
 
+	m.tagsCache = make(map[string][]string)
+
 	if m.defaultHandler == nil {
 		m.defaultHandler = handler
 	}
@@ -133,6 +137,8 @@ func (m *Manager) RemoveHandler(ctx context.Context, tag string) error {
 	m.access.Lock()
 	defer m.access.Unlock()
 
+	m.tagsCache = make(map[string][]string)
+
 	delete(m.taggedHandler, tag)
 	if m.defaultHandler != nil && m.defaultHandler.Tag() == tag {
 		m.defaultHandler = nil
@@ -146,6 +152,11 @@ func (m *Manager) Select(selectors []string) []string {
 	m.access.RLock()
 	defer m.access.RUnlock()
 
+	key := strings.Join(selectors, ",")
+	if cache, ok := m.tagsCache[key]; ok {
+		return cache
+	}
+
 	tags := make([]string, 0, len(selectors))
 
 	for tag := range m.taggedHandler {
@@ -156,7 +167,10 @@ func (m *Manager) Select(selectors []string) []string {
 			}
 		}
 	}
+
 	sort.Strings(tags)
+	m.tagsCache[key] = tags
+
 	return tags
 }
 
diff --git a/app/router/balancing.go b/app/router/balancing.go
index e6f29a67..458e6838 100644
--- a/app/router/balancing.go
+++ b/app/router/balancing.go
@@ -2,7 +2,6 @@ package router
 
 import (
 	"context"
-	reflect "reflect"
 	sync "sync"
 
 	"github.com/xtls/xray-core/common/dice"
@@ -26,35 +25,20 @@ func (s *RandomStrategy) PickOutbound(tags []string) string {
 }
 
 type RoundRobinStrategy struct {
-	mu         sync.Mutex
-	tags       []string
-	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
+	mu    sync.Mutex
+	index int
 }
 
 func (s *RoundRobinStrategy) PickOutbound(tags []string) string {
-	if len(tags) == 0 {
+	n := len(tags)
+	if n == 0 {
 		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
 }