forked from yuxh/gearman-go
		
	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