gearman-go/client/client.go

// Copyright 2011 Xing Xing <mikespook@gmail.com>.
// All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package client

import (
	"io"
	"net"
	"sync"
	"github.com/mikespook/golib/idgen"
)

/*
The client side api for gearman

usage:
c := client.New("tcp4", "127.0.0.1:4730")
handle := c.Do("foobar", []byte("data here"), JOB_LOW | JOB_BG)

*/
type Client struct {
	net, addr, lastcall    string
	respHandler  map[string]ResponseHandler
	innerHandler map[string]ResponseHandler
	in           chan []byte
	isConn       bool
	conn         net.Conn
	mutex        sync.RWMutex

	ErrorHandler ErrorHandler
	IdGen idgen.IdGen
}

// Create a new client.
// Connect to "addr" through "network"
// Eg.
//      client, err := client.New("127.0.0.1:4730")
func New(net, addr string) (client *Client, err error) {
	client = &Client{
		net:          net,
		addr:         addr,
		respHandler:  make(map[string]ResponseHandler, QUEUE_SIZE),
		innerHandler: make(map[string]ResponseHandler, QUEUE_SIZE),
		in:           make(chan []byte, QUEUE_SIZE),
		IdGen:        idgen.NewObjectId(),
	}
	if err = client.connect(); err != nil {
		return
	}
	client.isConn = true
	go client.readLoop()
	go client.processLoop()
	return
}

// {{{ private functions

//
func (client *Client) connect() (err error) {
	client.conn, err = net.Dial(client.net, client.addr)
	return
}

// Internal write
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.conn.Write(buf[i:])
		if err != nil {
			return
		}
	}
	return
}

// read length bytes from the socket
func (client *Client) read(length int) (data []byte, err error) {
	n := 0
	buf := getBuffer(BUFFER_SIZE)
	// read until data can be unpacked
	for i := length; i > 0 || len(data) < MIN_PACKET_LEN; i -= n {
		if n, err = client.conn.Read(buf); err != nil {
			if !client.isConn {
				err = ErrConnClosed
				return
			}
			if err == io.EOF && n == 0 {
				if data == nil {
					err = ErrConnection
				}
			}
			return
		}
		data = append(data, buf[0:n]...)
		if n < BUFFER_SIZE {
			break
		}
	}
	return
}

// read data from socket
func (client *Client) readLoop() {
	var data []byte
	var err error
	for client.isConn {
		if data, err = client.read(BUFFER_SIZE); err != nil {
			if err == ErrConnClosed {
				break
			}
			client.err(err)
			continue
		}
		client.in <- data
	}
	close(client.in)
}

// decode data & process it
func (client *Client) processLoop() {
	var resp *Response
	var l int
	var err error
	var data, leftdata []byte
	for data = range client.in {
		if len(leftdata) > 0 { // some data left for processing
			data = append(leftdata, data...)
		}
		l = len(data)
		if l < MIN_PACKET_LEN { // not enough data
			leftdata = data
			continue
		}
		if resp, l, err = decodeResponse(data); err != nil {
			client.err(err)
			continue
		}
		leftdata = nil
		switch resp.DataType {
		case ERROR:
			if client.lastcall != "" {
				client.handleInner(client.lastcall, resp)
				client.lastcall = ""
			} else {
				client.err(GetError(resp.Data))
			}
		case STATUS_RES:
			client.handleInner("s"+resp.Handle, resp)
		case JOB_CREATED:
			client.handleInner("c", resp)
		case ECHO_RES:
			client.handleInner("e", resp)
		case WORK_DATA, WORK_WARNING, WORK_STATUS, WORK_COMPLETE,
			WORK_FAIL, WORK_EXCEPTION:
			client.handleResponse(resp.Handle, resp)
		}
		if len(data) > l {
			leftdata = data[l:]
		}
	}
}

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

// job handler
func (client *Client) handleResponse(key string, resp *Response) {
	client.mutex.RLock()
	defer client.mutex.RUnlock()
	if h, ok := client.respHandler[key]; ok {
		h(resp)
		delete(client.respHandler, key)
	}
}

// job handler
func (client *Client) handleInner(key string, resp *Response) {
	if h, ok := client.innerHandler[key]; ok {
		h(resp)
		delete(client.innerHandler, key)
	}
}

// Internal do
func (client *Client) do(funcname string, data []byte,
	flag uint32) (handle string, err error) {
	id := client.IdGen.Id().(string)
	req := getJob(id, []byte(funcname), data)
	req.DataType = flag
	client.write(req)
	var wg sync.WaitGroup
	wg.Add(1)
	client.mutex.RLock()
	client.lastcall = "c"
	client.innerHandler["c"] = ResponseHandler(func(resp *Response) {
		defer wg.Done()
		defer client.mutex.RUnlock()
		if resp.DataType == ERROR {
			err = GetError(resp.Data)
			return
		}
		handle = resp.Handle
	})
	wg.Wait()
	return
}

// }}}

// Do the function.
// funcname is a string with function name.
// data is encoding to byte array.
// flag set the job type, include running level: JOB_LOW, JOB_NORMAL, JOB_HIGH
func (client *Client) Do(funcname string, data []byte,
	flag byte, h ResponseHandler) (handle string, err error) {
	var datatype uint32
	switch flag {
	case JOB_LOW:
		datatype = SUBMIT_JOB_LOW
	case JOB_HIGH:
		datatype = SUBMIT_JOB_HIGH
	default:
		datatype = SUBMIT_JOB
	}
	handle, err = client.do(funcname, data, datatype)
	client.mutex.Lock()
	defer client.mutex.Unlock()
	if h != nil {
		client.respHandler[handle] = h
	}
	return
}

// Do the function at background.
// funcname is a string with function name.
// data is encoding to byte array.
// flag set the job type, include running level: JOB_LOW, JOB_NORMAL, JOB_HIGH
func (client *Client) DoBg(funcname string, data []byte,
	flag byte) (handle string, err error) {
	var datatype uint32
	switch flag {
	case JOB_LOW:
		datatype = SUBMIT_JOB_LOW_BG
	case JOB_HIGH:
		datatype = SUBMIT_JOB_HIGH_BG
	default:
		datatype = SUBMIT_JOB_BG
	}
	handle, err = client.do(funcname, data, datatype)
	return
}

// Get job status from job server.
// !!!Not fully tested.!!!
func (client *Client) Status(handle string) (status *Status, err error) {
	req := getRequest()
	req.DataType = GET_STATUS
	req.Data = []byte(handle)
	client.write(req)
	var wg sync.WaitGroup
	wg.Add(1)
	client.mutex.Lock()
	client.lastcall = "s" + handle
	client.innerHandler["s" + handle] = ResponseHandler(func(resp *Response) {
		defer wg.Done()
		defer client.mutex.Unlock()
		var err error
		status, err = resp.Status()
		if err != nil {
			client.err(err)
		}
	})
	wg.Wait()
	return
}

// Send a something out, get the samething back.
func (client *Client) Echo(data []byte) (echo []byte, err error) {
	req := getRequest()
	req.DataType = ECHO_REQ
	req.Data = data
	client.write(req)
	var wg sync.WaitGroup
	wg.Add(1)
	client.mutex.Lock()
	client.lastcall = "e"
	client.innerHandler["e"] = ResponseHandler(func(resp *Response) {
		defer wg.Done()
		defer client.mutex.Unlock()
		echo = resp.Data
	})
	wg.Wait()
	return
}

// Close
func (client *Client) Close() (err error) {
	client.isConn = false
	return client.conn.Close()
}