You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
104 lines
2.0 KiB
Go
104 lines
2.0 KiB
Go
package internal
|
|
|
|
import "sync"
|
|
|
|
type (
|
|
DialFn func(server string) (interface{}, error)
|
|
CloseFn func(server string, conn interface{}) error
|
|
|
|
Balancer interface {
|
|
AddConn(kv KV) error
|
|
IsEmpty() bool
|
|
Next(key ...string) (interface{}, bool)
|
|
RemoveKey(key string)
|
|
initialize()
|
|
setListener(listener Listener)
|
|
}
|
|
|
|
serverConn struct {
|
|
key string
|
|
conn interface{}
|
|
}
|
|
|
|
baseBalancer struct {
|
|
exclusive bool
|
|
servers map[string][]string
|
|
mapping map[string]string
|
|
lock sync.Mutex
|
|
dialFn DialFn
|
|
closeFn CloseFn
|
|
listener Listener
|
|
}
|
|
)
|
|
|
|
func newBaseBalancer(dialFn DialFn, closeFn CloseFn, exclusive bool) *baseBalancer {
|
|
return &baseBalancer{
|
|
exclusive: exclusive,
|
|
servers: make(map[string][]string),
|
|
mapping: make(map[string]string),
|
|
dialFn: dialFn,
|
|
closeFn: closeFn,
|
|
}
|
|
}
|
|
|
|
// addKv adds the kv, returns if there are already other keys associate with the server
|
|
func (b *baseBalancer) addKv(key, value string) ([]string, bool) {
|
|
b.lock.Lock()
|
|
defer b.lock.Unlock()
|
|
|
|
keys := b.servers[value]
|
|
previous := append([]string(nil), keys...)
|
|
early := len(keys) > 0
|
|
if b.exclusive && early {
|
|
for _, each := range keys {
|
|
b.doRemoveKv(each)
|
|
}
|
|
}
|
|
b.servers[value] = append(b.servers[value], key)
|
|
b.mapping[key] = value
|
|
|
|
if early {
|
|
return previous, true
|
|
} else {
|
|
return nil, false
|
|
}
|
|
}
|
|
|
|
func (b *baseBalancer) doRemoveKv(key string) (server string, keepConn bool) {
|
|
server, ok := b.mapping[key]
|
|
if !ok {
|
|
return "", true
|
|
}
|
|
|
|
delete(b.mapping, key)
|
|
keys := b.servers[server]
|
|
remain := keys[:0]
|
|
|
|
for _, k := range keys {
|
|
if k != key {
|
|
remain = append(remain, k)
|
|
}
|
|
}
|
|
|
|
if len(remain) > 0 {
|
|
b.servers[server] = remain
|
|
return server, true
|
|
} else {
|
|
delete(b.servers, server)
|
|
return server, false
|
|
}
|
|
}
|
|
|
|
func (b *baseBalancer) removeKv(key string) (server string, keepConn bool) {
|
|
b.lock.Lock()
|
|
defer b.lock.Unlock()
|
|
|
|
return b.doRemoveKv(key)
|
|
}
|
|
|
|
func (b *baseBalancer) setListener(listener Listener) {
|
|
b.lock.Lock()
|
|
b.listener = listener
|
|
b.lock.Unlock()
|
|
}
|