2013-01-04 21:12:49 +08:00
|
|
|
package client
|
|
|
|
|
|
|
|
import (
|
2013-08-29 16:51:23 +08:00
|
|
|
"errors"
|
|
|
|
"math/rand"
|
|
|
|
"sync"
|
2013-01-05 15:24:06 +08:00
|
|
|
)
|
|
|
|
|
|
|
|
const (
|
2013-12-26 15:28:42 +08:00
|
|
|
poolSize = 10
|
2013-01-05 15:24:06 +08:00
|
|
|
)
|
|
|
|
|
2013-01-25 15:16:11 +08:00
|
|
|
var (
|
2014-01-09 17:58:02 +08:00
|
|
|
ErrNotFound = errors.New("Server Not Found")
|
2013-12-26 15:28:42 +08:00
|
|
|
SelectWithRate = selectWithRate
|
2014-01-09 17:58:02 +08:00
|
|
|
SelectRandom = selectRandom
|
2013-01-25 15:16:11 +08:00
|
|
|
)
|
|
|
|
|
2015-12-11 03:19:04 +08:00
|
|
|
type PoolClient struct {
|
2013-08-29 16:51:23 +08:00
|
|
|
*Client
|
2014-05-16 22:23:44 +08:00
|
|
|
Rate int
|
|
|
|
mutex sync.Mutex
|
2013-01-05 15:24:06 +08:00
|
|
|
}
|
|
|
|
|
2015-12-11 03:19:04 +08:00
|
|
|
type SelectionHandler func(map[string]*PoolClient, string) string
|
2013-01-05 15:24:06 +08:00
|
|
|
|
2015-12-11 03:19:04 +08:00
|
|
|
func selectWithRate(pool map[string]*PoolClient,
|
2013-08-29 16:51:23 +08:00
|
|
|
last string) (addr string) {
|
|
|
|
total := 0
|
|
|
|
for _, item := range pool {
|
|
|
|
total += item.Rate
|
|
|
|
if rand.Intn(total) < item.Rate {
|
|
|
|
return item.addr
|
|
|
|
}
|
|
|
|
}
|
|
|
|
return last
|
2013-01-05 15:24:06 +08:00
|
|
|
}
|
|
|
|
|
2015-12-11 03:19:04 +08:00
|
|
|
func selectRandom(pool map[string]*PoolClient,
|
2013-08-29 16:51:23 +08:00
|
|
|
last string) (addr string) {
|
|
|
|
r := rand.Intn(len(pool))
|
|
|
|
i := 0
|
|
|
|
for k, _ := range pool {
|
|
|
|
if r == i {
|
|
|
|
return k
|
|
|
|
}
|
|
|
|
i++
|
|
|
|
}
|
|
|
|
return last
|
2013-01-04 21:12:49 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
type Pool struct {
|
2013-08-29 16:51:23 +08:00
|
|
|
SelectionHandler SelectionHandler
|
|
|
|
ErrorHandler ErrorHandler
|
2015-12-11 05:54:48 +08:00
|
|
|
Clients map[string]*PoolClient
|
2013-01-25 15:16:11 +08:00
|
|
|
|
2015-12-11 05:54:48 +08:00
|
|
|
last string
|
2013-01-25 15:16:11 +08:00
|
|
|
|
2013-08-29 16:51:23 +08:00
|
|
|
mutex sync.Mutex
|
2013-01-04 21:12:49 +08:00
|
|
|
}
|
|
|
|
|
2013-12-26 15:28:42 +08:00
|
|
|
// Return a new pool.
|
2013-01-04 21:12:49 +08:00
|
|
|
func NewPool() (pool *Pool) {
|
2013-08-29 16:51:23 +08:00
|
|
|
return &Pool{
|
2015-12-11 05:54:48 +08:00
|
|
|
Clients: make(map[string]*PoolClient, poolSize),
|
2013-08-29 16:51:23 +08:00
|
|
|
SelectionHandler: SelectWithRate,
|
|
|
|
}
|
2013-01-04 21:12:49 +08:00
|
|
|
}
|
|
|
|
|
2013-01-05 15:24:06 +08:00
|
|
|
// Add a server with rate.
|
2013-08-29 16:51:23 +08:00
|
|
|
func (pool *Pool) Add(net, addr string, rate int) (err error) {
|
|
|
|
pool.mutex.Lock()
|
|
|
|
defer pool.mutex.Unlock()
|
2015-12-11 03:19:04 +08:00
|
|
|
var item *PoolClient
|
2013-08-29 16:51:23 +08:00
|
|
|
var ok bool
|
2015-12-11 05:54:48 +08:00
|
|
|
if item, ok = pool.Clients[addr]; ok {
|
2013-08-29 16:51:23 +08:00
|
|
|
item.Rate = rate
|
|
|
|
} else {
|
|
|
|
var client *Client
|
|
|
|
client, err = New(net, addr)
|
2013-09-22 22:02:05 +08:00
|
|
|
if err == nil {
|
2015-12-11 03:19:04 +08:00
|
|
|
item = &PoolClient{Client: client, Rate: rate}
|
2015-12-11 05:54:48 +08:00
|
|
|
pool.Clients[addr] = item
|
2013-09-22 22:02:05 +08:00
|
|
|
}
|
2013-08-29 16:51:23 +08:00
|
|
|
}
|
|
|
|
return
|
2013-01-04 21:12:49 +08:00
|
|
|
}
|
|
|
|
|
2013-01-24 18:13:02 +08:00
|
|
|
// Remove a server.
|
|
|
|
func (pool *Pool) Remove(addr string) {
|
2013-08-29 16:51:23 +08:00
|
|
|
pool.mutex.Lock()
|
|
|
|
defer pool.mutex.Unlock()
|
2015-12-11 05:54:48 +08:00
|
|
|
delete(pool.Clients, addr)
|
2013-01-24 18:13:02 +08:00
|
|
|
}
|
|
|
|
|
2013-01-05 15:24:06 +08:00
|
|
|
func (pool *Pool) Do(funcname string, data []byte,
|
2013-08-30 11:20:51 +08:00
|
|
|
flag byte, h ResponseHandler) (addr, handle string, err error) {
|
2013-08-29 16:51:23 +08:00
|
|
|
client := pool.selectServer()
|
2014-05-16 22:23:44 +08:00
|
|
|
client.Lock()
|
|
|
|
defer client.Unlock()
|
2013-08-30 11:20:51 +08:00
|
|
|
handle, err = client.Do(funcname, data, flag, h)
|
2013-08-29 16:51:23 +08:00
|
|
|
addr = client.addr
|
|
|
|
return
|
2013-01-18 17:03:45 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
func (pool *Pool) DoBg(funcname string, data []byte,
|
2013-08-30 11:20:51 +08:00
|
|
|
flag byte) (addr, handle string, err error) {
|
2013-08-29 16:51:23 +08:00
|
|
|
client := pool.selectServer()
|
2014-05-16 22:23:44 +08:00
|
|
|
client.Lock()
|
|
|
|
defer client.Unlock()
|
2013-08-30 11:20:51 +08:00
|
|
|
handle, err = client.DoBg(funcname, data, flag)
|
2013-08-29 16:51:23 +08:00
|
|
|
addr = client.addr
|
|
|
|
return
|
2013-01-04 21:12:49 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
// Get job status from job server.
|
|
|
|
// !!!Not fully tested.!!!
|
2013-08-30 11:20:51 +08:00
|
|
|
func (pool *Pool) Status(addr, handle string) (status *Status, err error) {
|
2015-12-11 05:54:48 +08:00
|
|
|
if client, ok := pool.Clients[addr]; ok {
|
2014-05-16 22:23:44 +08:00
|
|
|
client.Lock()
|
|
|
|
defer client.Unlock()
|
2013-08-30 11:20:51 +08:00
|
|
|
status, err = client.Status(handle)
|
2013-08-29 16:51:23 +08:00
|
|
|
} else {
|
|
|
|
err = ErrNotFound
|
|
|
|
}
|
|
|
|
return
|
2013-01-04 21:12:49 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
// Send a something out, get the samething back.
|
2013-08-30 11:20:51 +08:00
|
|
|
func (pool *Pool) Echo(addr string, data []byte) (echo []byte, err error) {
|
2015-12-11 03:19:04 +08:00
|
|
|
var client *PoolClient
|
2013-08-29 16:51:23 +08:00
|
|
|
if addr == "" {
|
|
|
|
client = pool.selectServer()
|
|
|
|
} else {
|
|
|
|
var ok bool
|
2015-12-11 05:54:48 +08:00
|
|
|
if client, ok = pool.Clients[addr]; !ok {
|
2013-08-29 16:51:23 +08:00
|
|
|
err = ErrNotFound
|
|
|
|
return
|
|
|
|
}
|
|
|
|
}
|
2014-05-16 22:23:44 +08:00
|
|
|
client.Lock()
|
|
|
|
defer client.Unlock()
|
2013-08-29 18:08:05 +08:00
|
|
|
echo, err = client.Echo(data)
|
2013-08-29 16:51:23 +08:00
|
|
|
return
|
2013-01-04 21:12:49 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
// Close
|
2013-01-08 17:23:10 +08:00
|
|
|
func (pool *Pool) Close() (err map[string]error) {
|
2013-08-29 16:51:23 +08:00
|
|
|
err = make(map[string]error)
|
2015-12-11 05:54:48 +08:00
|
|
|
for _, c := range pool.Clients {
|
2013-08-29 16:51:23 +08:00
|
|
|
err[c.addr] = c.Close()
|
|
|
|
}
|
|
|
|
return
|
2013-01-04 21:12:49 +08:00
|
|
|
}
|
2013-01-25 15:16:11 +08:00
|
|
|
|
2013-08-29 16:51:23 +08:00
|
|
|
// selecting server
|
2015-12-11 03:19:04 +08:00
|
|
|
func (pool *Pool) selectServer() (client *PoolClient) {
|
2013-08-29 16:51:23 +08:00
|
|
|
for client == nil {
|
2015-12-11 05:54:48 +08:00
|
|
|
addr := pool.SelectionHandler(pool.Clients, pool.last)
|
2013-08-29 16:51:23 +08:00
|
|
|
var ok bool
|
2015-12-11 05:54:48 +08:00
|
|
|
if client, ok = pool.Clients[addr]; ok {
|
2013-08-29 16:51:23 +08:00
|
|
|
pool.last = addr
|
|
|
|
break
|
|
|
|
}
|
|
|
|
}
|
|
|
|
return
|
2013-01-25 15:16:11 +08:00
|
|
|
}
|