From c618aab9e127f56ec229d136a16016390c070e8b Mon Sep 17 00:00:00 2001
From: WhiteBatman <xiaolipeng.icewater@gmail.com>
Date: Tue, 26 Sep 2017 10:22:13 +0800
Subject: [PATCH] modify policy interface

---
 engine_group.go        | 10 +++++-----
 engine_group_policy.go | 12 ++++++++----
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/engine_group.go b/engine_group.go
index 341de2c7..9bbe854a 100644
--- a/engine_group.go
+++ b/engine_group.go
@@ -62,7 +62,7 @@ func newGroup1(driverName string, dataSourceNames string, policy ...Policy) (*En
 			s_count: len(engines[1:]),
 			policy:  policy[0],
 		}
-		eg.policy.SetEngineGroup(eg)
+		eg.policy.Init()
 		return eg, nil
 	} else {
 		xPolicy := new(XormEngineGroupPolicy)
@@ -73,7 +73,7 @@ func newGroup1(driverName string, dataSourceNames string, policy ...Policy) (*En
 			s_count: len(engines[1:]),
 			policy:  xPolicy,
 		}
-		xPolicy.SetEngineGroup(eg)
+		xPolicy.Init()
 		return eg, nil
 	}
 
@@ -91,7 +91,7 @@ func newGroup2(Master *Engine, Slaves []*Engine, policy ...Policy) (*EngineGroup
 			s_count: len(Slaves),
 			policy:  policy[0],
 		}
-		eg.policy.SetEngineGroup(eg)
+		eg.policy.Init()
 		return eg, nil
 	} else {
 		xPolicy := new(XormEngineGroupPolicy)
@@ -102,7 +102,7 @@ func newGroup2(Master *Engine, Slaves []*Engine, policy ...Policy) (*EngineGroup
 			s_count: len(Slaves),
 			policy:  xPolicy,
 		}
-		xPolicy.SetEngineGroup(eg)
+		xPolicy.Init()
 		return eg, nil
 	}
 }
@@ -149,7 +149,7 @@ func (eg *EngineGroup) Slave() *Engine {
 	if eg.count == 1 {
 		return eg.master
 	}
-	return eg.slaves[eg.policy.Slave()]
+	return eg.policy.Slave(eg)
 }
 
 func (eg *EngineGroup) Slaves() []*Engine {
diff --git a/engine_group_policy.go b/engine_group_policy.go
index 1d55bde1..2219d9f5 100644
--- a/engine_group_policy.go
+++ b/engine_group_policy.go
@@ -19,8 +19,8 @@ const (
 )
 
 type Policy interface {
-	Slave() int
-	SetEngineGroup(*EngineGroup)
+	Init()
+	Slave(*EngineGroup) *Engine
 }
 
 type XormEngineGroupPolicy struct {
@@ -30,9 +30,13 @@ type XormEngineGroupPolicy struct {
 	r      *rand.Rand
 }
 
-func (xgep *XormEngineGroupPolicy) SetEngineGroup(eg *EngineGroup) {
+func (xgep *XormEngineGroupPolicy) Init() {
 	xgep.r = rand.New(rand.NewSource(time.Now().UnixNano()))
+}
+
+func (xgep *XormEngineGroupPolicy) Slave(eg *EngineGroup) *Engine {
 	xgep.eg = eg
+	return eg.slaves[xgep.slave()]
 }
 
 func (xgep *XormEngineGroupPolicy) SetWeight() {
@@ -44,7 +48,7 @@ func (xgep *XormEngineGroupPolicy) SetWeight() {
 	}
 }
 
-func (xgep *XormEngineGroupPolicy) Slave() int {
+func (xgep *XormEngineGroupPolicy) slave() int {
 	switch xgep.eg.p {
 	case ENGINE_GROUP_POLICY_RANDOM:
 		return xgep.Random()