From d41264d1f7e8ce676a89f147be1cfef50ba3375f Mon Sep 17 00:00:00 2001
From: Kirill Lepikhin <lepihin.k@gmail.com>
Date: Mon, 30 Sep 2019 15:09:24 +0800
Subject: [PATCH] Renamed old RoundRobin strategy to Random. Added RoundRobin
 strategy. Added StrategyFactory to config.

---
 broker/broker.go                |  3 +++
 moleculer.go                    |  2 ++
 registry/actionCatalog_test.go  |  2 +-
 registry/registry.go            |  6 +++++-
 strategy/random.go              | 18 ++++++++++++++++++
 strategy/round-robin.go         | 22 ++++++++++++++--------
 strategy/strategy_suite_test.go |  4 ++--
 7 files changed, 45 insertions(+), 12 deletions(-)
 create mode 100644 strategy/random.go

diff --git a/broker/broker.go b/broker/broker.go
index 96a13d0b..c6e03531 100644
--- a/broker/broker.go
+++ b/broker/broker.go
@@ -37,6 +37,9 @@ func mergeConfigs(baseConfig moleculer.Config, userConfig []*moleculer.Config) m
 			if config.TransporterFactory != nil {
 				baseConfig.TransporterFactory = config.TransporterFactory
 			}
+			if config.StrategyFactory != nil {
+				baseConfig.StrategyFactory = config.StrategyFactory
+			}
 			if config.DisableInternalMiddlewares {
 				baseConfig.DisableInternalMiddlewares = config.DisableInternalMiddlewares
 			}
diff --git a/moleculer.go b/moleculer.go
index 3571a663..fe9bb9fb 100644
--- a/moleculer.go
+++ b/moleculer.go
@@ -111,6 +111,7 @@ type Mixin struct {
 }
 
 type TransporterFactoryFunc func() interface{}
+type StrategyFactoryFunc func() interface{}
 
 type Config struct {
 	LogLevel                   string
@@ -118,6 +119,7 @@ type Config struct {
 	DiscoverNodeID             func() string
 	Transporter                string
 	TransporterFactory         TransporterFactoryFunc
+	StrategyFactory            StrategyFactoryFunc
 	HeartbeatFrequency         time.Duration
 	HeartbeatTimeout           time.Duration
 	OfflineCheckFrequency      time.Duration
diff --git a/registry/actionCatalog_test.go b/registry/actionCatalog_test.go
index 623e824a..4d4b284a 100644
--- a/registry/actionCatalog_test.go
+++ b/registry/actionCatalog_test.go
@@ -12,7 +12,7 @@ import (
 
 var _ = Describe("Actions Catalog", func() {
 	logger := log.WithField("unit test pkg", "registry_test")
-	strategy := strategy.RoundRobinStrategy{}
+	strategy := strategy.RandomStrategy{}
 	params := moleculer.ObjectSchema{nil}
 	node1 := registry.CreateNode("node-test-1", true, logger)
 	node2 := registry.CreateNode("node-test-2", true, logger)
diff --git a/registry/registry.go b/registry/registry.go
index 92c3e20a..ae84a572 100644
--- a/registry/registry.go
+++ b/registry/registry.go
@@ -50,7 +50,11 @@ func createTransit(broker *moleculer.BrokerDelegates) transit.Transit {
 // createStrategy create a strsategy instance based on the config.
 func createStrategy(broker *moleculer.BrokerDelegates) strategy.Strategy {
 	//TODO: when new strategies are addes.. adde config check here to load the right one.
-	return strategy.RoundRobinStrategy{}
+	if broker.Config.StrategyFactory != nil {
+		return broker.Config.StrategyFactory().(strategy.Strategy)
+	}
+
+	return strategy.RandomStrategy{}
 }
 
 func CreateRegistry(nodeID string, broker *moleculer.BrokerDelegates) *ServiceRegistry {
diff --git a/strategy/random.go b/strategy/random.go
new file mode 100644
index 00000000..c73a5c69
--- /dev/null
+++ b/strategy/random.go
@@ -0,0 +1,18 @@
+package strategy
+
+import (
+	"math/rand"
+)
+
+// RoundRobinStrategy exposes the type as a strategy option
+type RandomStrategy struct {
+}
+
+func (randomStrategy RandomStrategy) Select(nodes []Selector) *Selector {
+	if len(nodes) == 0 {
+		return nil
+	}
+	// Returns a number among the indexes up to the length
+	// of the slice
+	return &nodes[rand.Intn(len(nodes))]
+}
diff --git a/strategy/round-robin.go b/strategy/round-robin.go
index f05f4739..8e6c2633 100644
--- a/strategy/round-robin.go
+++ b/strategy/round-robin.go
@@ -1,18 +1,24 @@
 package strategy
 
-import (
-	"math/rand"
-)
-
 // RoundRobinStrategy exposes the type as a strategy option
 type RoundRobinStrategy struct {
+	counter int
+}
+
+func NewRoundRobinStrategy() Strategy {
+	return &RoundRobinStrategy{counter: -1}
 }
 
-func (roundRobinStrategy RoundRobinStrategy) Select(nodes []Selector) *Selector {
+func (roundRobinStrategy *RoundRobinStrategy) Select(nodes []Selector) *Selector {
 	if len(nodes) == 0 {
 		return nil
 	}
-	// Returns a number among the indexes up to the length
-	// of the slice
-	return &nodes[rand.Intn(len(nodes))]
+
+	roundRobinStrategy.counter++
+
+	if roundRobinStrategy.counter >= len(nodes) {
+		roundRobinStrategy.counter = 0
+	}
+
+	return &nodes[roundRobinStrategy.counter]
 }
diff --git a/strategy/strategy_suite_test.go b/strategy/strategy_suite_test.go
index d19807d9..adbadfc2 100644
--- a/strategy/strategy_suite_test.go
+++ b/strategy/strategy_suite_test.go
@@ -19,13 +19,13 @@ func (sel SelectorImpl) TargetNodeID() string {
 }
 
 var _ = Describe("Strategy", func() {
-	thisStrategy := strategy.RoundRobinStrategy{}
+	thisStrategy := strategy.RandomStrategy{}
 
 	list := []strategy.Selector{SelectorImpl{"alpha"}, SelectorImpl{"beta"}, SelectorImpl{"gamma"}, SelectorImpl{"delta"}}
 
 	emptyList := []strategy.Selector{}
 
-	It("Should return a ramdon target node according to strategy", func() {
+	It("Should return a random target node according to strategy", func() {
 		thisNode := thisStrategy.Select(list)
 		Expect(thisNode).Should(Not(BeNil()))
 	})