gearman-go/client/pool.go

168 lines
3.2 KiB
Go
Raw Normal View History

2013-01-04 21:12:49 +08:00
package client
import (
2013-08-29 16:51:23 +08:00
"errors"
"math/rand"
"sync"
)
const (
poolSize = 10
)
2013-01-25 15:16:11 +08:00
var (
2014-01-09 17:58:02 +08:00
ErrNotFound = errors.New("Server Not Found")
SelectWithRate = selectWithRate
2014-01-09 17:58:02 +08:00
SelectRandom = selectRandom
2013-01-25 15:16:11 +08:00
)
2013-01-24 18:13:02 +08:00
type poolClient struct {
2013-08-29 16:51:23 +08:00
*Client
Rate int
mutex sync.Mutex
}
2013-01-24 18:13:02 +08:00
type SelectionHandler func(map[string]*poolClient, string) string
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
}
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
2013-01-25 15:16:11 +08:00
2013-08-29 16:51:23 +08:00
clients map[string]*poolClient
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
}
// 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{
clients: make(map[string]*poolClient, poolSize),
2013-08-29 16:51:23 +08:00
SelectionHandler: SelectWithRate,
}
2013-01-04 21:12:49 +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()
var item *poolClient
var ok bool
if item, ok = pool.clients[addr]; ok {
item.Rate = rate
} else {
var client *Client
client, err = New(net, addr)
2013-09-22 22:02:05 +08:00
if err == nil {
item = &poolClient{Client: client, Rate: rate}
pool.clients[addr] = item
}
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()
delete(pool.clients, addr)
2013-01-24 18:13:02 +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()
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()
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) {
2013-08-29 16:51:23 +08:00
if client, ok := pool.clients[addr]; ok {
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) {
2013-08-29 16:51:23 +08:00
var client *poolClient
if addr == "" {
client = pool.selectServer()
} else {
var ok bool
if client, ok = pool.clients[addr]; !ok {
err = ErrNotFound
return
}
}
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)
for _, c := range pool.clients {
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
2013-01-25 15:16:11 +08:00
func (pool *Pool) selectServer() (client *poolClient) {
2013-08-29 16:51:23 +08:00
for client == nil {
addr := pool.SelectionHandler(pool.clients, pool.last)
var ok bool
if client, ok = pool.clients[addr]; ok {
pool.last = addr
break
}
}
return
2013-01-25 15:16:11 +08:00
}