# Path: cc/message.py
# class CCMessage(object):
#     """CC multipart message.
# 
#     Format is similar to usual REQ+REP sockets:
# 
#     - id_hopX
#     - id_hop2
#     - id_hop1
#     - id_req (optional)
#     - empty part
#     - cc dest
#     - cc payload (json)
#     - cc signature
#     - cc blob data
#     """
#     __slots__ = ('zmsg', 'rpos', 'parsed', 'signature')
# 
#     def __init__(self, zmsg):
#         assert isinstance(zmsg, list)
#         self.zmsg = zmsg
#         self.rpos = zmsg.index('')
#         self.parsed = None
#         self.signature = None
#         assert_msg_req (self.get_dest())
# 
#     def get_route(self):
#         """Route parts"""
#         return self.zmsg[ : self.rpos]
# 
#     def get_non_route(self):
#         """Payload parts"""
#         return self.zmsg[ self.rpos + 1 : ]
# 
#     def get_dest(self):
#         """Return destination part"""
#         return self.zmsg[self.rpos + 1]
# 
#     def get_part1(self):
#         """Return body (json) as string"""
#         return self.zmsg[self.rpos + 2]
# 
#     def get_part2(self):
#         """Return signature"""
#         if self.rpos + 3 >= len(self.zmsg):
#             return ''
#         return self.zmsg[self.rpos + 3]
# 
#     def get_part3(self):
#         """Return blob"""
#         if self.rpos + 4 >= len(self.zmsg):
#             return None
#         return self.zmsg[self.rpos + 4]
# 
#     def get_size(self):
#         return zmsg_size (self.zmsg)
# 
#     def __repr__(self):
#         x = repr(self.zmsg)
#         if len(x) > 300:
#             x = x[:300] + '...'
#         return 'CCMessage(%s)' % x
# 
#     def __str__(self):
#         x = repr(self.get_non_route())
#         if len(x) > 300:
#             x = x[:300] + '...'
#         return 'CCMessage(%s)' % x
# 
#     def set_route (self, route):
#         """Fill local route with another route."""
#         self.zmsg[:self.rpos] = route
#         self.rpos = len(route)
# 
#     def take_route(self, cmsg):
#         """Fill local route with route from another message."""
#         r = cmsg.get_route()
#         self.set_route(r)
# 
#     def get_payload(self, xtx):
#         if self.parsed:
#             return self.parsed
#         msg, sgn = xtx.parse_cmsg(self)
#         if not msg:
#             return None
#         self.parsed = msg
#         self.signature = sgn
#         return msg
# 
#     def get_signature(self, xtx):
#         self.get_payload(xtx)
#         return self.signature
# 
#     def send_to (self, sock):
#         sock.send_multipart (self.zmsg)
# 
# def zmsg_size (zmsg):
#     n = 0
#     for p in zmsg:
#         n += len(p)
#     return n
# 
# Path: cc/util.py
# def stat_inc (key, increase = 1):
#     """ Increase a stat value. """
#     global stat_dict
#     try:
#         stat_dict[key] += increase
#     except KeyError:
#         stat_dict[key] = increase


import sys
import time
import zmq
import skytools
from zmq.eventloop import IOLoop
from zmq.eventloop.zmqstream import ZMQStream
from cc.message import CCMessage, zmsg_size
from cc.util import stat_inc





__all__ = ['CCStream', 'CCReqStream']

#
# simple wrapper around ZMQStream
#

class CCStream (ZMQStream):
    """
    Adds CCMessage methods to ZMQStream as well as protection (on by default)
    against unlimited memory (send queue) growth.
    """

    def __init__ (self, *args, **kwargs):
        self.qmaxsize = kwargs.pop ('qmaxsize', None)
        if self.qmaxsize is None:
            self.qmaxsize = 1000
        elif self.qmaxsize <= 0:
            self.qmaxsize = sys.maxsize
        super(CCStream, self).__init__(*args, **kwargs)

    def send_multipart (self, msg, *args, **kwargs):
        if self._send_queue.qsize() < self.qmaxsize:
            super(CCStream, self).send_multipart (msg, *args, **kwargs)
        else:
            stat_inc ('count.dropped', 1)

# Path: cc/message.py
# class CCMessage(object):
#     """CC multipart message.
# 
#     Format is similar to usual REQ+REP sockets:
# 
#     - id_hopX
#     - id_hop2
#     - id_hop1
#     - id_req (optional)
#     - empty part
#     - cc dest
#     - cc payload (json)
#     - cc signature
#     - cc blob data
#     """
#     __slots__ = ('zmsg', 'rpos', 'parsed', 'signature')
# 
#     def __init__(self, zmsg):
#         assert isinstance(zmsg, list)
#         self.zmsg = zmsg
#         self.rpos = zmsg.index('')
#         self.parsed = None
#         self.signature = None
#         assert_msg_req (self.get_dest())
# 
#     def get_route(self):
#         """Route parts"""
#         return self.zmsg[ : self.rpos]
# 
#     def get_non_route(self):
#         """Payload parts"""
#         return self.zmsg[ self.rpos + 1 : ]
# 
#     def get_dest(self):
#         """Return destination part"""
#         return self.zmsg[self.rpos + 1]
# 
#     def get_part1(self):
#         """Return body (json) as string"""
#         return self.zmsg[self.rpos + 2]
# 
#     def get_part2(self):
#         """Return signature"""
#         if self.rpos + 3 >= len(self.zmsg):
#             return ''
#         return self.zmsg[self.rpos + 3]
# 
#     def get_part3(self):
#         """Return blob"""
#         if self.rpos + 4 >= len(self.zmsg):
#             return None
#         return self.zmsg[self.rpos + 4]
# 
#     def get_size(self):
#         return zmsg_size (self.zmsg)
# 
#     def __repr__(self):
#         x = repr(self.zmsg)
#         if len(x) > 300:
#             x = x[:300] + '...'
#         return 'CCMessage(%s)' % x
# 
#     def __str__(self):
#         x = repr(self.get_non_route())
#         if len(x) > 300:
#             x = x[:300] + '...'
#         return 'CCMessage(%s)' % x
# 
#     def set_route (self, route):
#         """Fill local route with another route."""
#         self.zmsg[:self.rpos] = route
#         self.rpos = len(route)
# 
#     def take_route(self, cmsg):
#         """Fill local route with route from another message."""
#         r = cmsg.get_route()
#         self.set_route(r)
# 
#     def get_payload(self, xtx):
#         if self.parsed:
#             return self.parsed
#         msg, sgn = xtx.parse_cmsg(self)
#         if not msg:
#             return None
#         self.parsed = msg
#         self.signature = sgn
#         return msg
# 
#     def get_signature(self, xtx):
#         self.get_payload(xtx)
#         return self.signature
# 
#     def send_to (self, sock):
#         sock.send_multipart (self.zmsg)
# 
# def zmsg_size (zmsg):
#     n = 0
#     for p in zmsg:
#         n += len(p)
#     return n
# 
# Path: cc/util.py
# def stat_inc (key, increase = 1):
#     """ Increase a stat value. """
#     global stat_dict
#     try:
#         stat_dict[key] += increase
#     except KeyError:
#         stat_dict[key] = increase


# Path: cc/stream.py
import sys
import time
import zmq
import skytools
from zmq.eventloop import IOLoop
from zmq.eventloop.zmqstream import ZMQStream
from cc.message import CCMessage, zmsg_size
from cc.util import stat_inc






__all__ = ['CCStream', 'CCReqStream']

#
# simple wrapper around ZMQStream
#

class CCStream (ZMQStream):
    """
    Adds CCMessage methods to ZMQStream as well as protection (on by default)
    against unlimited memory (send queue) growth.
    """

    def __init__ (self, *args, **kwargs):
        self.qmaxsize = kwargs.pop ('qmaxsize', None)
        if self.qmaxsize is None:
            self.qmaxsize = 1000
        elif self.qmaxsize <= 0:
            self.qmaxsize = sys.maxsize
        super(CCStream, self).__init__(*args, **kwargs)

    def send_multipart (self, msg, *args, **kwargs):
        if self._send_queue.qsize() < self.qmaxsize:
            super(CCStream, self).send_multipart (msg, *args, **kwargs)
        else:
            stat_inc ('count.dropped', 1)

            stat_inc ('bytes.dropped', zmsg_size (msg))