engine_group_policy.go 2.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116
  1. // Copyright 2017 The Xorm Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style
  3. // license that can be found in the LICENSE file.
  4. package xorm
  5. import (
  6. "math/rand"
  7. "sync"
  8. "time"
  9. )
  10. // GroupPolicy is be used by chosing the current slave from slaves
  11. type GroupPolicy interface {
  12. Slave(*EngineGroup) *Engine
  13. }
  14. // GroupPolicyHandler should be used when a function is a GroupPolicy
  15. type GroupPolicyHandler func(*EngineGroup) *Engine
  16. // Slave implements the chosen of slaves
  17. func (h GroupPolicyHandler) Slave(eg *EngineGroup) *Engine {
  18. return h(eg)
  19. }
  20. // RandomPolicy implmentes randomly chose the slave of slaves
  21. func RandomPolicy() GroupPolicyHandler {
  22. var r = rand.New(rand.NewSource(time.Now().UnixNano()))
  23. return func(g *EngineGroup) *Engine {
  24. return g.Slaves()[r.Intn(len(g.Slaves()))]
  25. }
  26. }
  27. // WeightRandomPolicy implmentes randomly chose the slave of slaves
  28. func WeightRandomPolicy(weights []int) GroupPolicyHandler {
  29. var rands = make([]int, 0, len(weights))
  30. for i := 0; i < len(weights); i++ {
  31. for n := 0; n < weights[i]; n++ {
  32. rands = append(rands, i)
  33. }
  34. }
  35. var r = rand.New(rand.NewSource(time.Now().UnixNano()))
  36. return func(g *EngineGroup) *Engine {
  37. var slaves = g.Slaves()
  38. idx := rands[r.Intn(len(rands))]
  39. if idx >= len(slaves) {
  40. idx = len(slaves) - 1
  41. }
  42. return slaves[idx]
  43. }
  44. }
  45. func RoundRobinPolicy() GroupPolicyHandler {
  46. var pos = -1
  47. var lock sync.Mutex
  48. return func(g *EngineGroup) *Engine {
  49. var slaves = g.Slaves()
  50. lock.Lock()
  51. defer lock.Unlock()
  52. pos++
  53. if pos >= len(slaves) {
  54. pos = 0
  55. }
  56. return slaves[pos]
  57. }
  58. }
  59. func WeightRoundRobinPolicy(weights []int) GroupPolicyHandler {
  60. var rands = make([]int, 0, len(weights))
  61. for i := 0; i < len(weights); i++ {
  62. for n := 0; n < weights[i]; n++ {
  63. rands = append(rands, i)
  64. }
  65. }
  66. var pos = -1
  67. var lock sync.Mutex
  68. return func(g *EngineGroup) *Engine {
  69. var slaves = g.Slaves()
  70. lock.Lock()
  71. defer lock.Unlock()
  72. pos++
  73. if pos >= len(rands) {
  74. pos = 0
  75. }
  76. idx := rands[pos]
  77. if idx >= len(slaves) {
  78. idx = len(slaves) - 1
  79. }
  80. return slaves[idx]
  81. }
  82. }
  83. // LeastConnPolicy implements GroupPolicy, every time will get the least connections slave
  84. func LeastConnPolicy() GroupPolicyHandler {
  85. return func(g *EngineGroup) *Engine {
  86. var slaves = g.Slaves()
  87. connections := 0
  88. idx := 0
  89. for i := 0; i < len(slaves); i++ {
  90. openConnections := slaves[i].DB().Stats().OpenConnections
  91. if i == 0 {
  92. connections = openConnections
  93. idx = i
  94. } else if openConnections <= connections {
  95. connections = openConnections
  96. idx = i
  97. }
  98. }
  99. return slaves[idx]
  100. }
  101. }