From 01c14a599450ffd576cced599098c895d5790c2d Mon Sep 17 00:00:00 2001
From: Hossin Asaadi <hossin.asaadi@gmail.com>
Date: Sun, 24 Dec 2023 23:29:10 +0300
Subject: [PATCH] add Round-Robin Strategy to balancer (#2844)

* add Round-Robin Strategy

* clean up
---
 app/router/balancing.go       | 34 ++++++++++++++++++++++++++++++++++
 app/router/config.go          |  6 ++++++
 infra/conf/router.go          |  2 ++
 infra/conf/router_strategy.go |  5 +++--
 4 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/app/router/balancing.go b/app/router/balancing.go
index 50b84388..ad0a3b36 100644
--- a/app/router/balancing.go
+++ b/app/router/balancing.go
@@ -2,6 +2,7 @@ package router
 
 import (
 	"context"
+	sync "sync"
 
 	"github.com/xtls/xray-core/common/dice"
 	"github.com/xtls/xray-core/features/extension"
@@ -23,6 +24,39 @@ func (s *RandomStrategy) PickOutbound(tags []string) string {
 	return tags[dice.Roll(n)]
 }
 
+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
+}
+
+func (s *RoundRobinStrategy) PickOutbound(tags []string) string {
+	if len(tags) == 0 {
+		panic("0 tags")
+	}
+	if s.roundRobin == nil {
+		s.roundRobin = NewRoundRobin(tags)
+	}
+	tag := s.roundRobin.NextTag()
+
+	return tag
+}
+
 type Balancer struct {
 	selectors []string
 	strategy  BalancingStrategy
diff --git a/app/router/config.go b/app/router/config.go
index f50f02a1..5dc32fa8 100644
--- a/app/router/config.go
+++ b/app/router/config.go
@@ -129,6 +129,12 @@ func (br *BalancingRule) Build(ohm outbound.Manager) (*Balancer, error) {
 			strategy:  &LeastPingStrategy{},
 			ohm:       ohm,
 		}, nil
+	case "roundRobin":
+		return &Balancer{
+			selectors: br.OutboundSelector,
+			strategy:  &RoundRobinStrategy{},
+			ohm:       ohm,
+		}, nil
 	case "random":
 		fallthrough
 	default:
diff --git a/infra/conf/router.go b/infra/conf/router.go
index 79176fa5..a3285e85 100644
--- a/infra/conf/router.go
+++ b/infra/conf/router.go
@@ -43,6 +43,8 @@ func (r *BalancingRule) Build() (*router.BalancingRule, error) {
 		strategy = strategyRandom
 	case strategyLeastPing:
 		strategy = "leastPing"
+	case strategyRoundRobin:
+		strategy = "roundRobin"
 	default:
 		return nil, newError("unknown balancing strategy: " + r.Strategy.Type)
 	}
diff --git a/infra/conf/router_strategy.go b/infra/conf/router_strategy.go
index b8536330..ef2abc26 100644
--- a/infra/conf/router_strategy.go
+++ b/infra/conf/router_strategy.go
@@ -1,6 +1,7 @@
 package conf
 
 const (
-	strategyRandom    string = "random"
-	strategyLeastPing string = "leastping"
+	strategyRandom     string = "random"
+	strategyLeastPing  string = "leastping"
+	strategyRoundRobin string = "roundrobin"
 )