1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
|
// Copyright (c) 2012 The gocql Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package gocql
import (
"sync"
"sync/atomic"
)
type Node interface {
Pick(qry *Query) *Conn
Close()
}
type RoundRobin struct {
pool []Node
pos uint32
mu sync.RWMutex
}
func NewRoundRobin() *RoundRobin {
return &RoundRobin{}
}
func (r *RoundRobin) AddNode(node Node) {
r.mu.Lock()
r.pool = append(r.pool, node)
r.mu.Unlock()
}
func (r *RoundRobin) RemoveNode(node Node) {
r.mu.Lock()
n := len(r.pool)
for i := 0; i < n; i++ {
if r.pool[i] == node {
r.pool[i], r.pool[n-1] = r.pool[n-1], r.pool[i]
r.pool = r.pool[:n-1]
break
}
}
r.mu.Unlock()
}
func (r *RoundRobin) Size() int {
r.mu.RLock()
n := len(r.pool)
r.mu.RUnlock()
return n
}
func (r *RoundRobin) Pick(qry *Query) *Conn {
pos := atomic.AddUint32(&r.pos, 1)
var node Node
r.mu.RLock()
if len(r.pool) > 0 {
node = r.pool[pos%uint32(len(r.pool))]
}
r.mu.RUnlock()
if node == nil {
return nil
}
return node.Pick(qry)
}
func (r *RoundRobin) Close() {
r.mu.Lock()
for i := 0; i < len(r.pool); i++ {
r.pool[i].Close()
}
r.pool = nil
r.mu.Unlock()
}
|
