jrpc_py.py

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import json
import random
import time
from builtins import *

import zmq

try:
    import queue
except ImportError:
    import Queue as queue
import threading
import msgpack
import snappy
import copy

qEmpty = copy.copy(queue.Empty)


def _unpack_msgpack_snappy(str):
    if str.startswith(b'S'):
        tmp = snappy.uncompress(str[1:])
        # print "SNAPPY: ", len(str), len(tmp)
        obj = msgpack.loads(tmp, encoding='utf-8')
    elif str.startswith(b'\0'):
        obj = msgpack.loads(str[1:], encoding='utf-8')
    else:
        return None
    
    return obj


def _pack_msgpack_snappy(obj):
    # print "pack", obj
    tmp = msgpack.dumps(obj, encoding='utf-8')
    if len(tmp) > 1000:
        return b'S' + snappy.compress(tmp)
    else:
        return b'\0' + tmp


def _unpack_msgpack(str):
    return msgpack.loads(str, encoding='utf-8')


def _pack_msgpack(obj):
    return msgpack.dumps(obj, encoding='utf-8')


def _unpack_json(str):
    return json.loads(str, encoding='utf-8')


def _pack_json(obj):
    return json.dumps(obj, encoding='utf-8')


class JRpcClient(object):
    def __init__(self, data_format="msgpack_snappy"):
        self._waiter_lock = threading.Lock()
        self._waiter_map = {}

        self._should_close = False
        self._next_callid = 0
        self._send_lock = threading.Lock()
        self._callid_lock = threading.Lock()

        self._last_heartbeat_rsp_time = 0
        self._connected = False

        self.on_disconnected = None
        self.on_rpc_callback = None
        self._callback_queue = queue.Queue()
        self._call_wait_queue = queue.Queue()

        self._ctx = zmq.Context()
        self._pull_sock = self._ctx.socket(zmq.PULL)
        self._pull_sock.bind("inproc://pull_sock")
        self._push_sock = self._ctx.socket(zmq.PUSH)
        self._push_sock.connect("inproc://pull_sock")

        self._heartbeat_interval = 1
        self._heartbeat_timeout = 3

        self._addr = None
        
        if data_format == "msgpack_snappy":
            self._pack = _pack_msgpack_snappy
            self._unpack = _unpack_msgpack_snappy

        elif data_format == "msgpack":
            self._pack = _pack_msgpack
            self._unpack = _unpack_msgpack
        
        elif data_format == "json":
            self._pack = _pack_json
            self._unpack = _unpack_json
        
        else:
            assert False, "unknown data_format " + data_format
        
        t = threading.Thread(target=self._recv_run)
        t.setDaemon(True)
        t.start()
        self._recv_thread = t

        t = threading.Thread(target=self._callback_run)
        t.setDaemon(True)
        t.start()
        self._callback_thread = t

    def __del__(self):
        self.close()

    def next_callid(self):
        self._callid_lock.acquire()
        self._next_callid += 1
        callid = self._next_callid
        self._callid_lock.release()
        return callid

    def set_heartbeat_options(self, interval, timeout):
        self._heartbeat_interval = interval
        self._heartbeat_timeout = timeout

    def _recv_run(self):

        heartbeat_time = 0

        poller = zmq.Poller()
        poller.register(self._pull_sock, zmq.POLLIN)

        remote_sock = None

        while not self._should_close:

            try:
                if self._connected and time.time() - self._last_heartbeat_rsp_time > self._heartbeat_timeout:
                    self._connected = False
                    if self.on_disconnected: self._async_call(self.on_disconnected)

                if remote_sock and time.time() - heartbeat_time > self._heartbeat_interval:
                    self._send_hearbeat()
                    heartbeat_time = time.time()

                socks = dict(poller.poll(500))
                if self._pull_sock in socks and socks[self._pull_sock] == zmq.POLLIN:
                    cmd = self._pull_sock.recv()
                    if cmd == b"CONNECT":
                        # print time.ctime(), "CONNECT " + self._addr
                        if remote_sock:
                            poller.unregister(remote_sock)
                            remote_sock.close()
                            remote_sock = None

                        remote_sock = self._do_connect()

                        if remote_sock:
                            poller.register(remote_sock, zmq.POLLIN)

                    elif cmd.startswith(b"SEND:") and remote_sock:
                        # print time.ctime(), "SEND " + cmd[5:]
                        remote_sock.send(cmd[5:])

                if remote_sock and remote_sock in socks and socks[remote_sock] == zmq.POLLIN:
                    data = remote_sock.recv()
                    if data:
                        # if not data.find("heartbeat"):
                        #    print time.ctime(), "RECV", data
                        self._on_data_arrived(data)

            except zmq.error.Again as e:
                # print "RECV timeout: ", e
                pass
            except Exception as e:
                print("_recv_run:", e)

    def _callback_run(self):
        while not self._should_close:
            try:
                r = self._callback_queue.get(timeout=1)
                if r:
                    r()
            except qEmpty as e:
                pass
            except TypeError as e:
                if str(e) == "'NoneType' object is not callable":
                    pass
                else:
                    print("_callback_run {}".format(r), type(e), e)
            except Exception as e:
                print("_callback_run {}".format(r), type(e), e)

    def _async_call(self, func):
        self._callback_queue.put(func)

    def _send_request(self, json):

        try:
            self._send_lock.acquire()
            self._push_sock.send(b"SEND:" + json)

        finally:
            self._send_lock.release()

    def connect(self, addr):
        self._addr = addr
        self._push_sock.send_string('CONNECT', encoding='utf-8')

    def _do_connect(self):

        client_id = str(random.randint(1000000, 100000000))

        socket = self._ctx.socket(zmq.DEALER)
        identity = (client_id) + '$' + str(random.randint(1000000, 1000000000))
        identity = identity.encode('utf-8')
        socket.setsockopt(zmq.IDENTITY, identity)
        socket.setsockopt(zmq.RCVTIMEO, 500)
        socket.setsockopt(zmq.SNDTIMEO, 500)
        socket.setsockopt(zmq.LINGER, 0)
        socket.connect(self._addr)

        return socket

    def close(self):
        self._should_close = True
        self._callback_thread.join()
        self._recv_thread.join()

    def _on_data_arrived(self, str):
        try:
            msg = self._unpack(str)
            # print "RECV", msg

            if not msg:
                print("wrong message format")
                return

            if 'method' in msg and msg['method'] == '.sys.heartbeat':
                self._last_heartbeat_rsp_time = time.time()
                if not self._connected:
                    self._connected = True
                    if self.on_connected:
                        self._async_call(self.on_connected)
                
                # Let user has a chance to check message in .sys.heartbeat
                if 'result' in msg and self.on_rpc_callback:
                    self._async_call(lambda: self.on_rpc_callback(msg['method'], msg['result']))

            elif 'id' in msg and msg['id']:
                
                # Call result
                id = int(msg['id'])
                
                if self._waiter_lock.acquire():
                    if id in self._waiter_map:
                        q = self._waiter_map[id]
                        if q: q.put(msg)
                    self._waiter_lock.release()
            else:
                # Notification message
                if 'method' in msg and 'result' in msg and self.on_rpc_callback:
                    self._async_call(lambda: self.on_rpc_callback(msg['method'], msg['result']))
        
        except Exception as e:
            print("_on_data_arrived:", e)
            pass
    
    def _send_hearbeat(self):
        msg = {'jsonrpc' : '2.0',
               'method'  : '.sys.heartbeat',
               'params'  : { 'time': time.time()},
               'id'      : str(self.next_callid())}
        json_str = self._pack(msg)
        self._send_request(json_str)

    def _alloc_wait_queue(self):
        self._waiter_lock.acquire()
        if self._call_wait_queue:
            q = self._call_wait_queue
            self._call_wait_queue = None
        else:
            q = queue.Queue()
        self._waiter_lock.release()
        return q

    def _free_wait_queue(self, q):
        self._waiter_lock.acquire()
        if not self._call_wait_queue:
            self._call_wait_queue = q
        else:
            del q
        self._waiter_lock.release()

    def call(self, method, params, timeout=6):
        # print "call", method, params, timeout
        callid = self.next_callid()
        if timeout:
            q = self._alloc_wait_queue()

            self._waiter_lock.acquire()
            self._waiter_map[callid] = q
            self._waiter_lock.release()
        
        msg = {'jsonrpc' : '2.0',
               'method'  : method,
               'params'  : params,
               'id'      : str(callid) }

        # print "SEND", msg
        json_str = self._pack(msg)
        self._send_request(json_str)
        
        if timeout:
            ret = {}
            try:
                r = q.get(timeout=timeout)
                q.task_done()
            except qEmpty:
                r = None

            self._waiter_lock.acquire()
            self._waiter_map[callid] = None
            self._waiter_lock.release()
            self._free_wait_queue(q)

            if r:
                if 'result' in r:
                    ret['result'] = r['result']

                if 'error' in r:
                    ret['error'] = r['error']

            return ret if ret else {'error': {'error': -1, 'message': "timeout"}}
        else:
            return {'result': True}