<s> import socket <EOL> import zmq <EOL> import logging <EOL> import time <EOL> from dpark . env import env <EOL> from dpark . util import spawn <EOL> logger = logging . getLogger ( __name__ ) <EOL> class TrackerMessage ( object ) : <EOL> pass <EOL> class StopTrackerMessage ( TrackerMessage ) : <EOL> pass <EOL> class SetValueMessage ( TrackerMessage ) : <EOL> def __init__ ( self , key , value ) : <EOL> self . key = key <EOL> self . value = value <EOL> class AddItemMessage ( TrackerMessage ) : <EOL> def __init__ ( self , key , item ) : <EOL> self . key = key <EOL> self . item = item <EOL> class RemoveItemMessage ( TrackerMessage ) : <EOL> def __init__ ( self , key , item ) : <EOL> self . key = key <EOL> self . item = item <EOL> class GetValueMessage ( TrackerMessage ) : <EOL> def __init__ ( self , key ) : <EOL> self . key = key <EOL> class TrackerServer ( object ) : <EOL> locs = { } <EOL> def __init__ ( self ) : <EOL> self . addr = None <EOL> self . thread = None <EOL> def start ( self ) : <EOL> self . thread = spawn ( self . run ) <EOL> while self . addr is None : <EOL> time . sleep ( <NUM_LIT> ) <EOL> def stop ( self ) : <EOL> sock = env . ctx . socket ( zmq . REQ ) <EOL> sock . connect ( self . addr ) <EOL> sock . send_pyobj ( StopTrackerMessage ( ) ) <EOL> confirm_msg = sock . recv_pyobj ( ) <EOL> sock . close ( ) <EOL> self . thread . join ( ) <EOL> return confirm_msg <EOL> def get ( self , key ) : <EOL> return self . locs . get ( key , [ ] ) <EOL> def set ( self , key , value ) : <EOL> if not isinstance ( value , list ) : <EOL> value = [ value ] <EOL> self . locs [ key ] = value <EOL> def add ( self , key , item ) : <EOL> if key not in self . locs : <EOL> self . locs [ key ] = [ ] <EOL> self . locs [ key ] . append ( item ) <EOL> def remove ( self , key , item ) : <EOL> if item in self . locs [ key ] : <EOL> self . locs [ key ] . remove ( item ) <EOL> def run ( self ) : <EOL> locs = self . locs <EOL> sock = env . ctx . socket ( zmq . REP ) <EOL> port = sock . bind_to_random_port ( "<STR_LIT>" ) <EOL> self . addr = "<STR_LIT>" % ( socket . gethostname ( ) , port ) <EOL> logger . debug ( "<STR_LIT>" , self . addr ) <EOL> def reply ( msg ) : <EOL> sock . send_pyobj ( msg ) <EOL> while True : <EOL> msg = sock . recv_pyobj ( ) <EOL> if isinstance ( msg , SetValueMessage ) : <EOL> self . set ( msg . key , msg . value ) <EOL> reply ( '<STR_LIT:OK>' ) <EOL> elif isinstance ( msg , AddItemMessage ) : <EOL> self . add ( msg . key , msg . item ) <EOL> reply ( '<STR_LIT:OK>' ) <EOL> elif isinstance ( msg , RemoveItemMessage ) : <EOL> self . remove ( msg . key , msg . item ) <EOL> reply ( '<STR_LIT:OK>' ) <EOL> elif isinstance ( msg , GetValueMessage ) : <EOL> reply ( self . get ( msg . key ) ) <EOL> elif isinstance ( msg , StopTrackerMessage ) : <EOL> reply ( '<STR_LIT:OK>' ) <EOL> break <EOL> else : <EOL> logger . error ( "<STR_LIT>" , msg , type ( msg ) ) <EOL> reply ( '<STR_LIT>' ) <EOL> sock . close ( ) <EOL> logger . debug ( "<STR_LIT>" , self . addr ) <EOL> class TrackerClient ( object ) : <EOL> def __init__ ( self , addr ) : <EOL> self . addr = addr <EOL> def call ( self , msg ) : <EOL> try : <EOL> sock = env . ctx . socket ( zmq . REQ ) <EOL> sock . connect ( self . addr ) <EOL> sock . send_pyobj ( msg ) <EOL> return sock . recv_pyobj ( ) <EOL> finally : <EOL> sock . close ( ) </s>
