gearman-go/client/client.go

// The client package helps developers connect to Gearmand, send
// jobs and fetch result.
package client

import (
	"bufio"
	"net"
	"sync"
	"time"
)

var (
	DefaultTimeout time.Duration = time.Second
)

// One client connect to one server.
// Use Pool for multi-connections.
type Client struct {
	sync.Mutex

	net, addr    string
	innerHandler *responseHandlerMap
	in           chan *Response
	conn         net.Conn
	rw           *bufio.ReadWriter

	ResponseTimeout time.Duration // response timeout for do()

	ErrorHandler ErrorHandler
}

type responseHandlerMap struct {
	sync.Mutex
	holder map[string]handledResponse
}

type handledResponse struct {
	internal ResponseHandler // internal handler, always non-nil
	external ResponseHandler // handler passed in from (*Client).Do, sometimes nil
}

func newResponseHandlerMap() *responseHandlerMap {
	return &responseHandlerMap{holder: make(map[string]handledResponse, queueSize)}
}

func (r *responseHandlerMap) remove(key string) {
	r.Lock()
	delete(r.holder, key)
	r.Unlock()
}

func (r *responseHandlerMap) getAndRemove(key string) (handledResponse, bool) {
	r.Lock()
	rh, b := r.holder[key]
	delete(r.holder, key)
	r.Unlock()
	return rh, b
}

func (r *responseHandlerMap) putWithExternalHandler(key string, internal, external ResponseHandler) {
	r.Lock()
	r.holder[key] = handledResponse{internal: internal, external: external}
	r.Unlock()
}

func (r *responseHandlerMap) put(key string, rh ResponseHandler) {
	r.putWithExternalHandler(key, rh, nil)
}

// New returns a client.
func New(network, addr string) (client *Client, err error) {
	client = &Client{
		net:             network,
		addr:            addr,
		innerHandler:    newResponseHandlerMap(),
		in:              make(chan *Response, queueSize),
		ResponseTimeout: DefaultTimeout,
	}
	client.conn, err = net.Dial(client.net, client.addr)
	if err != nil {
		return
	}
	client.rw = bufio.NewReadWriter(bufio.NewReader(client.conn),
		bufio.NewWriter(client.conn))
	go client.readLoop()
	go client.processLoop()
	return
}

func (client *Client) write(req *request) (err error) {
	var n int
	buf := req.Encode()
	for i := 0; i < len(buf); i += n {
		n, err = client.rw.Write(buf[i:])
		if err != nil {
			return
		}
	}
	return client.rw.Flush()
}

func (client *Client) read(length int) (data []byte, err error) {
	n := 0
	buf := getBuffer(bufferSize)
	// read until data can be unpacked
	for i := length; i > 0 || len(data) < minPacketLength; i -= n {
		if n, err = client.rw.Read(buf); err != nil {
			return
		}
		data = append(data, buf[0:n]...)
		if n < bufferSize {
			break
		}
	}
	return
}

func (client *Client) readLoop() {
	defer close(client.in)
	var data, leftdata []byte
	var err error
	var resp *Response
ReadLoop:
	for client.conn != nil {
		if data, err = client.read(bufferSize); err != nil {
			if opErr, ok := err.(*net.OpError); ok {
				if opErr.Timeout() {
					client.err(err)
				}
				if opErr.Temporary() {
					continue
				}
				break
			}
			client.err(err)
			// If it is unexpected error and the connection wasn't
			// closed by Gearmand, the client should close the conection
			// and reconnect to job server.
			client.Close()
			client.conn, err = net.Dial(client.net, client.addr)
			if err != nil {
				client.err(err)
				break
			}
			client.rw = bufio.NewReadWriter(bufio.NewReader(client.conn),
				bufio.NewWriter(client.conn))
			continue
		}
		if len(leftdata) > 0 { // some data left for processing
			data = append(leftdata, data...)
			leftdata = nil
		}
		for {
			l := len(data)
			if l < minPacketLength { // not enough data
				leftdata = data
				continue ReadLoop
			}
			if resp, l, err = decodeResponse(data); err != nil {
				leftdata = data[l:]
				continue ReadLoop
			} else {
				client.in <- resp
			}
			data = data[l:]
			if len(data) > 0 {
				continue
			}
			break
		}
	}
}

func (client *Client) processLoop() {
	rhandlers := map[string]ResponseHandler{}
	for resp := range client.in {
		switch resp.DataType {
		case dtError:
			client.err(getError(resp.Data))
		case dtStatusRes:
			client.handleInner("s"+resp.Handle, resp, nil)
		case dtJobCreated:
			client.handleInner("c", resp, rhandlers)
		case dtEchoRes:
			client.handleInner("e", resp, nil)
		case dtWorkData, dtWorkWarning, dtWorkStatus:
			if cb := rhandlers[resp.Handle]; cb != nil {
				cb(resp)
			}
		case dtWorkComplete, dtWorkFail, dtWorkException:
			if cb := rhandlers[resp.Handle]; cb != nil {
				cb(resp)
				delete(rhandlers, resp.Handle)
			}
		}
	}
}

func (client *Client) err(e error) {
	if client.ErrorHandler != nil {
		client.ErrorHandler(e)
	}
}

func (client *Client) handleInner(key string, resp *Response, rhandlers map[string]ResponseHandler) {
	if h, ok := client.innerHandler.getAndRemove(key); ok {
		if h.external != nil && resp.Handle != "" {
			rhandlers[resp.Handle] = h.external
		}
		h.internal(resp)
	}
}

type handleOrError struct {
	handle string
	err    error
}

func (client *Client) do(funcname string, data []byte,
	flag uint32, h ResponseHandler) (handle string, err error) {
	if client.conn == nil {
		return "", ErrLostConn
	}
	var result = make(chan handleOrError, 1)
	client.Lock()
	defer client.Unlock()
	client.innerHandler.putWithExternalHandler("c", func(resp *Response) {
		if resp.DataType == dtError {
			err = getError(resp.Data)
			result <- handleOrError{"", err}
			return
		}
		handle = resp.Handle
		result <- handleOrError{handle, nil}
	}, h)
	id := IdGen.Id()
	req := getJob(id, []byte(funcname), data)
	req.DataType = flag
	if err = client.write(req); err != nil {
		client.innerHandler.remove("c")
		return
	}
	var timer = time.After(client.ResponseTimeout)
	select {
	case ret := <-result:
		return ret.handle, ret.err
	case <-timer:
		client.innerHandler.remove("c")
		return "", ErrLostConn
	}
	return
}

// Call the function and get a response.
// flag can be set to: JobLow, JobNormal and JobHigh
func (client *Client) Do(funcname string, data []byte,
	flag byte, h ResponseHandler) (handle string, err error) {
	var datatype uint32
	switch flag {
	case JobLow:
		datatype = dtSubmitJobLow
	case JobHigh:
		datatype = dtSubmitJobHigh
	default:
		datatype = dtSubmitJob
	}

	handle, err = client.do(funcname, data, datatype, h)
	return
}

// Call the function in background, no response needed.
// flag can be set to: JobLow, JobNormal and JobHigh
func (client *Client) DoBg(funcname string, data []byte,
	flag byte) (handle string, err error) {
	if client.conn == nil {
		return "", ErrLostConn
	}
	var datatype uint32
	switch flag {
	case JobLow:
		datatype = dtSubmitJobLowBg
	case JobHigh:
		datatype = dtSubmitJobHighBg
	default:
		datatype = dtSubmitJobBg
	}
	handle, err = client.do(funcname, data, datatype, nil)
	return
}

// Status gets job status from job server.
func (client *Client) Status(handle string) (status *Status, err error) {
	if client.conn == nil {
		return nil, ErrLostConn
	}
	var mutex sync.Mutex
	mutex.Lock()
	client.innerHandler.put("s"+handle, func(resp *Response) {
		defer mutex.Unlock()
		var err error
		status, err = resp._status()
		if err != nil {
			client.err(err)
		}
	})
	req := getRequest()
	req.DataType = dtGetStatus
	req.Data = []byte(handle)
	client.write(req)
	mutex.Lock()
	return
}

// Echo.
func (client *Client) Echo(data []byte) (echo []byte, err error) {
	if client.conn == nil {
		return nil, ErrLostConn
	}
	var mutex sync.Mutex
	mutex.Lock()
	client.innerHandler.put("e", func(resp *Response) {
		echo = resp.Data
		mutex.Unlock()
	})
	req := getRequest()
	req.DataType = dtEchoReq
	req.Data = data
	client.write(req)
	mutex.Lock()
	return
}

// Close connection
func (client *Client) Close() (err error) {
	client.Lock()
	defer client.Unlock()
	if client.conn != nil {
		err = client.conn.Close()
		client.conn = nil
	}
	return
}