Merge pull request #62 from kujohn/pool-custom-handler

Export SelectionHandler
This commit is contained in:
Xing 2015-12-11 12:44:35 +08:00
commit 6d9b2fba51
2 changed files with 21 additions and 21 deletions

View File

@ -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
} }

View File

@ -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))
} }
} }