<s> import os , sys <EOL> import zookeeper <EOL> import threading <EOL> import functools <EOL> import time <EOL> zookeeper . set_log_stream ( sys . stdout ) <EOL> ZOO_OPEN_ACL_UNSAFE = { "<STR_LIT>" : <NUM_LIT> , "<STR_LIT>" : "<STR_LIT>" , "<STR_LIT:id>" : "<STR_LIT>" } <EOL> def retry_on ( * excepts ) : <EOL> """<STR_LIT>""" <EOL> def _decorator ( fn ) : <EOL> @ functools . wraps ( fn ) <EOL> def wrapper ( * args , ** kwargs ) : <EOL> while True : <EOL> try : <EOL> return fn ( * args , ** kwargs ) <EOL> except Exception , e : <EOL> if not any ( [ isinstance ( e , _ ) for _ in excepts ] ) : <EOL> raise <EOL> time . sleep ( <NUM_LIT:5> ) <EOL> return wrapper <EOL> return _decorator <EOL> class ZooKeeper ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , quorum , timeout = <NUM_LIT> ) : <EOL> """<STR_LIT>""" <EOL> self . _quorum = quorum <EOL> self . _timeout = timeout <EOL> self . _handle = None <EOL> self . _connected = False <EOL> self . _cv = threading . Condition ( ) <EOL> self . _connect ( ) <EOL> def _connect ( self ) : <EOL> """<STR_LIT>""" <EOL> def watcher ( handle , type , state , path ) : <EOL> self . _cv . acquire ( ) <EOL> try : <EOL> self . _connected = True <EOL> self . _cv . notify ( ) <EOL> finally : <EOL> self . _cv . release ( ) <EOL> self . _cv . acquire ( ) <EOL> try : <EOL> self . _handle = zookeeper . init ( self . _quorum , watcher , self . _timeout ) <EOL> self . _cv . wait ( self . _timeout / <NUM_LIT:1000> ) <EOL> if not self . _connected : <EOL> print >> sys . stderr , '<STR_LIT>' <EOL> sys . exit ( - <NUM_LIT:1> ) <EOL> finally : <EOL> self . _cv . release ( ) <EOL> def __getattr__ ( self , name ) : <EOL> """<STR_LIT>""" <EOL> value = getattr ( zookeeper , name ) <EOL> if callable ( value ) : <EOL> return functools . partial ( <EOL> retry_on ( zookeeper . ConnectionLossException , zookeeper . OperationTimeoutException ) ( value ) , <EOL> self . _handle <EOL> ) <EOL> else : <EOL> return name <EOL> def ensure_exists ( self , name , data = '<STR_LIT>' ) : <EOL> try : <EOL> self . create ( name , data , [ ZOO_OPEN_ACL_UNSAFE ] , <NUM_LIT:0> ) <EOL> except zookeeper . NodeExistsException : <EOL> pass <EOL> @ retry_on ( zookeeper . NoNodeException ) <EOL> def create_sequence ( self , name , data ) : <EOL> """<STR_LIT>""" <EOL> name = self . create ( name , '<STR_LIT>' , <EOL> [ ZOO_OPEN_ACL_UNSAFE ] , <EOL> zookeeper . SEQUENCE <EOL> ) <EOL> try : <EOL> self . set ( name , data , <NUM_LIT:0> ) <EOL> except zookeeper . BadVersionException : <EOL> pass <EOL> class Producer ( object ) : <EOL> def __init__ ( self , zk ) : <EOL> self . _zk = zk <EOL> self . _zk . ensure_exists ( '<STR_LIT>' ) <EOL> self . _zk . ensure_exists ( '<STR_LIT>' ) <EOL> def put ( self , data , priority = <NUM_LIT:100> ) : <EOL> return self . _zk . create_sequence ( '<STR_LIT>' % priority , str ( data ) ) <EOL> class Consumer ( object ) : <EOL> def __init__ ( self , zk , active = None ) : <EOL> self . _zk = zk <EOL> self . _id = None <EOL> self . active = active <EOL> map ( self . _zk . ensure_exists , ( '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' , '<STR_LIT>' ) ) <EOL> self . _register ( ) <EOL> def _register ( self ) : <EOL> self . _id = self . _zk . create ( "<STR_LIT>" , '<STR_LIT>' , <EOL> [ ZOO_OPEN_ACL_UNSAFE ] , zookeeper . SEQUENCE ) <EOL> if self . active is None : <EOL> self . _zk . create ( self . _fullpath ( '<STR_LIT>' ) , '<STR_LIT>' , <EOL> [ ZOO_OPEN_ACL_UNSAFE ] , zookeeper . EPHEMERAL ) <EOL> elif self . active : <EOL> self . _zk . create ( self . _fullpath ( '<STR_LIT>' ) , '<STR_LIT>' , <EOL> [ ZOO_OPEN_ACL_UNSAFE ] , <NUM_LIT:0> ) <EOL> self . _zk . create ( self . _fullpath ( '<STR_LIT>' ) , '<STR_LIT>' , <EOL> [ ZOO_OPEN_ACL_UNSAFE ] , <NUM_LIT:0> ) <EOL> def _fullpath ( self , suffix ) : <EOL> return self . _id + suffix <EOL> @ retry_on ( zookeeper . BadVersionException ) <EOL> def _move ( self , src , dest ) : <EOL> try : <EOL> ( data , stat ) = self . _zk . get ( src , None ) <EOL> if data : <EOL> self . _zk . set ( src , data , stat [ '<STR_LIT:version>' ] ) <EOL> self . _zk . set ( dest , data ) <EOL> self . _zk . delete ( src , stat [ '<STR_LIT:version>' ] + <NUM_LIT:1> ) <EOL> return data <EOL> elif stat [ '<STR_LIT>' ] < <NUM_LIT:1000> * ( time . time ( ) - <NUM_LIT> ) : <EOL> try : <EOL> self . _zk . delete ( src ) <EOL> except zookeeper . NoNodeException : <EOL> pass <EOL> return None <EOL> except zookeeper . NoNodeException : <EOL> self . _zk . set ( dest , '<STR_LIT>' ) <EOL> return None <EOL> def reserve ( self , block = False ) : <EOL> if block : <EOL> return self . _blocking_reserve ( ) <EOL> else : <EOL> return self . _simple_reserve ( ) <EOL> def _blocking_reserve ( self ) : <EOL> def queue_watcher ( * args , ** kwargs ) : <EOL> self . _zk . _cv . acquire ( ) <EOL> try : <EOL> self . _zk . _cv . notify ( ) <EOL> finally : <EOL> self . _zk . _cv . release ( ) <EOL> while True : <EOL> self . _zk . _cv . acquire ( ) <EOL> try : <EOL> children = sorted ( self . _zk . get_children ( '<STR_LIT>' , queue_watcher ) ) <EOL> for child in children : <EOL> data = self . _move ( '<STR_LIT>' + child , self . _fullpath ( '<STR_LIT>' ) ) <EOL> if data : <EOL> return data <EOL> self . _zk . _cv . wait ( ) <EOL> finally : <EOL> self . _zk . _cv . release ( ) <EOL> def _simple_reserve ( self ) : <EOL> while True : <EOL> children = sorted ( self . _zk . get_children ( '<STR_LIT>' , None ) ) <EOL> if len ( children ) == <NUM_LIT:0> : <EOL> return None <EOL> for child in children : <EOL> data = self . _move ( '<STR_LIT>' + child , self . _fullpath ( '<STR_LIT>' ) ) <EOL> if data : return data <EOL> def done ( self ) : <EOL> self . _zk . set ( self . _fullpath ( '<STR_LIT>' ) , '<STR_LIT>' ) <EOL> def close ( self ) : <EOL> map ( self . _zk . delete , ( self . _fullpath ( '<STR_LIT>' ) , <EOL> self . _fullpath ( '<STR_LIT>' ) , self . _id ) ) <EOL> self . _id = None <EOL> def __repr__ ( self ) : <EOL> return '<STR_LIT>' % self . _id <EOL> class GarbageCollector ( object ) : <EOL> def __init__ ( self , zk ) : <EOL> self . _zk = zk <EOL> map ( self . _zk . ensure_exists , ( '<STR_LIT>' , <EOL> '<STR_LIT>' , '<STR_LIT>' ) ) <EOL> def _get_dlock ( self ) : <EOL> try : <EOL> self . _zk . create ( '<STR_LIT>' , '<STR_LIT>' , [ ZOO_OPEN_ACL_UNSAFE ] , zookeeper . EPHEMERAL ) <EOL> return True <EOL> except zookeeper . NodeExistsException : <EOL> return False <EOL> def _remove_dlock ( self ) : <EOL> try : <EOL> self . _zk . delete ( '<STR_LIT>' ) <EOL> except zookeeper . NoNodeException : <EOL> pass <EOL> def collect ( self ) : <EOL> if self . _get_dlock ( ) : <EOL> try : <EOL> children = self . _zk . get_children ( '<STR_LIT>' , None ) <EOL> for child in children : <EOL> cpath = '<STR_LIT>' + child <EOL> if not self . _zk . exists ( cpath + '<STR_LIT>' ) : <EOL> ( data , stat ) = self . _zk . get ( cpath , None ) <EOL> if stat [ '<STR_LIT>' ] < <NUM_LIT:1000> * ( time . time ( ) - <NUM_LIT> ) : <EOL> ( data , stat ) = self . _zk . get ( cpath + '<STR_LIT>' , None ) <EOL> if data : <EOL> self . _zk . create_sequence ( '<STR_LIT>' , data ) <EOL> try : <EOL> self . _zk . delete ( cpath + '<STR_LIT>' ) <EOL> except zookeeper . NoNodeException : pass <EOL> try : <EOL> self . _zk . delete ( cpath ) <EOL> except zookeeper . NoNodeException : pass <EOL> finally : <EOL> self . _remove_dlock ( ) <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> zk = ZooKeeper ( "<STR_LIT>" "<STR_LIT>" ) <EOL> p = Producer ( zk ) <EOL> map ( p . put , map ( str , range ( <NUM_LIT:1> , <NUM_LIT:10> ) ) ) <EOL> c = Consumer ( zk ) <EOL> print c <EOL> print c . reserve ( block = True ) <EOL> c . done ( ) <EOL> while True : <EOL> data = c . reserve ( ) <EOL> if not data : break <EOL> print data <EOL> c . done ( ) <EOL> c . close ( ) <EOL> gc = GarbageCollector ( zk ) <EOL> gc . collect ( ) <EOL> zk . close ( ) </s>
