Merge pull request #62 from kujohn/pool-custom-handler
Export SelectionHandler
This commit is contained in:
commit
6d9b2fba51
@ -16,15 +16,15 @@ var (
|
|||||||
SelectRandom = selectRandom
|
SelectRandom = selectRandom
|
||||||
)
|
)
|
||||||
|
|
||||||
type poolClient struct {
|
type PoolClient struct {
|
||||||
*Client
|
*Client
|
||||||
Rate int
|
Rate int
|
||||||
mutex sync.Mutex
|
mutex sync.Mutex
|
||||||
}
|
}
|
||||||
|
|
||||||
type SelectionHandler func(map[string]*poolClient, string) string
|
type SelectionHandler func(map[string]*PoolClient, string) string
|
||||||
|
|
||||||
func selectWithRate(pool map[string]*poolClient,
|
func selectWithRate(pool map[string]*PoolClient,
|
||||||
last string) (addr string) {
|
last string) (addr string) {
|
||||||
total := 0
|
total := 0
|
||||||
for _, item := range pool {
|
for _, item := range pool {
|
||||||
@ -36,7 +36,7 @@ func selectWithRate(pool map[string]*poolClient,
|
|||||||
return last
|
return last
|
||||||
}
|
}
|
||||||
|
|
||||||
func selectRandom(pool map[string]*poolClient,
|
func selectRandom(pool map[string]*PoolClient,
|
||||||
last string) (addr string) {
|
last string) (addr string) {
|
||||||
r := rand.Intn(len(pool))
|
r := rand.Intn(len(pool))
|
||||||
i := 0
|
i := 0
|
||||||
@ -52,9 +52,9 @@ func selectRandom(pool map[string]*poolClient,
|
|||||||
type Pool struct {
|
type Pool struct {
|
||||||
SelectionHandler SelectionHandler
|
SelectionHandler SelectionHandler
|
||||||
ErrorHandler ErrorHandler
|
ErrorHandler ErrorHandler
|
||||||
|
Clients map[string]*PoolClient
|
||||||
|
|
||||||
clients map[string]*poolClient
|
last string
|
||||||
last string
|
|
||||||
|
|
||||||
mutex sync.Mutex
|
mutex sync.Mutex
|
||||||
}
|
}
|
||||||
@ -62,7 +62,7 @@ type Pool struct {
|
|||||||
// Return a new pool.
|
// Return a new pool.
|
||||||
func NewPool() (pool *Pool) {
|
func NewPool() (pool *Pool) {
|
||||||
return &Pool{
|
return &Pool{
|
||||||
clients: make(map[string]*poolClient, poolSize),
|
Clients: make(map[string]*PoolClient, poolSize),
|
||||||
SelectionHandler: SelectWithRate,
|
SelectionHandler: SelectWithRate,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@ -71,16 +71,16 @@ func NewPool() (pool *Pool) {
|
|||||||
func (pool *Pool) Add(net, addr string, rate int) (err error) {
|
func (pool *Pool) Add(net, addr string, rate int) (err error) {
|
||||||
pool.mutex.Lock()
|
pool.mutex.Lock()
|
||||||
defer pool.mutex.Unlock()
|
defer pool.mutex.Unlock()
|
||||||
var item *poolClient
|
var item *PoolClient
|
||||||
var ok bool
|
var ok bool
|
||||||
if item, ok = pool.clients[addr]; ok {
|
if item, ok = pool.Clients[addr]; ok {
|
||||||
item.Rate = rate
|
item.Rate = rate
|
||||||
} else {
|
} else {
|
||||||
var client *Client
|
var client *Client
|
||||||
client, err = New(net, addr)
|
client, err = New(net, addr)
|
||||||
if err == nil {
|
if err == nil {
|
||||||
item = &poolClient{Client: client, Rate: rate}
|
item = &PoolClient{Client: client, Rate: rate}
|
||||||
pool.clients[addr] = item
|
pool.Clients[addr] = item
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
return
|
return
|
||||||
@ -90,7 +90,7 @@ func (pool *Pool) Add(net, addr string, rate int) (err error) {
|
|||||||
func (pool *Pool) Remove(addr string) {
|
func (pool *Pool) Remove(addr string) {
|
||||||
pool.mutex.Lock()
|
pool.mutex.Lock()
|
||||||
defer pool.mutex.Unlock()
|
defer pool.mutex.Unlock()
|
||||||
delete(pool.clients, addr)
|
delete(pool.Clients, addr)
|
||||||
}
|
}
|
||||||
|
|
||||||
func (pool *Pool) Do(funcname string, data []byte,
|
func (pool *Pool) Do(funcname string, data []byte,
|
||||||
@ -116,7 +116,7 @@ func (pool *Pool) DoBg(funcname string, data []byte,
|
|||||||
// Get job status from job server.
|
// Get job status from job server.
|
||||||
// !!!Not fully tested.!!!
|
// !!!Not fully tested.!!!
|
||||||
func (pool *Pool) Status(addr, handle string) (status *Status, err error) {
|
func (pool *Pool) Status(addr, handle string) (status *Status, err error) {
|
||||||
if client, ok := pool.clients[addr]; ok {
|
if client, ok := pool.Clients[addr]; ok {
|
||||||
client.Lock()
|
client.Lock()
|
||||||
defer client.Unlock()
|
defer client.Unlock()
|
||||||
status, err = client.Status(handle)
|
status, err = client.Status(handle)
|
||||||
@ -128,12 +128,12 @@ func (pool *Pool) Status(addr, handle string) (status *Status, err error) {
|
|||||||
|
|
||||||
// Send a something out, get the samething back.
|
// Send a something out, get the samething back.
|
||||||
func (pool *Pool) Echo(addr string, data []byte) (echo []byte, err error) {
|
func (pool *Pool) Echo(addr string, data []byte) (echo []byte, err error) {
|
||||||
var client *poolClient
|
var client *PoolClient
|
||||||
if addr == "" {
|
if addr == "" {
|
||||||
client = pool.selectServer()
|
client = pool.selectServer()
|
||||||
} else {
|
} else {
|
||||||
var ok bool
|
var ok bool
|
||||||
if client, ok = pool.clients[addr]; !ok {
|
if client, ok = pool.Clients[addr]; !ok {
|
||||||
err = ErrNotFound
|
err = ErrNotFound
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
@ -147,18 +147,18 @@ func (pool *Pool) Echo(addr string, data []byte) (echo []byte, err error) {
|
|||||||
// Close
|
// Close
|
||||||
func (pool *Pool) Close() (err map[string]error) {
|
func (pool *Pool) Close() (err map[string]error) {
|
||||||
err = make(map[string]error)
|
err = make(map[string]error)
|
||||||
for _, c := range pool.clients {
|
for _, c := range pool.Clients {
|
||||||
err[c.addr] = c.Close()
|
err[c.addr] = c.Close()
|
||||||
}
|
}
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
// selecting server
|
// selecting server
|
||||||
func (pool *Pool) selectServer() (client *poolClient) {
|
func (pool *Pool) selectServer() (client *PoolClient) {
|
||||||
for client == nil {
|
for client == nil {
|
||||||
addr := pool.SelectionHandler(pool.clients, pool.last)
|
addr := pool.SelectionHandler(pool.Clients, pool.last)
|
||||||
var ok bool
|
var ok bool
|
||||||
if client, ok = pool.clients[addr]; ok {
|
if client, ok = pool.Clients[addr]; ok {
|
||||||
pool.last = addr
|
pool.last = addr
|
||||||
break
|
break
|
||||||
}
|
}
|
||||||
|
@ -18,8 +18,8 @@ func TestPoolAdd(t *testing.T) {
|
|||||||
t.Log(err)
|
t.Log(err)
|
||||||
c -= 1
|
c -= 1
|
||||||
}
|
}
|
||||||
if len(pool.clients) != c {
|
if len(pool.Clients) != c {
|
||||||
t.Errorf("%d servers expected, %d got.", c, len(pool.clients))
|
t.Errorf("%d servers expected, %d got.", c, len(pool.Clients))
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user