<s> from threading import Event <EOL> import zmq <EOL> from onitu . referee import UP , DEL , MOV <EOL> from onitu . utils import get_brocker_uri , log_traceback <EOL> from onitu . brocker . commands import GET_CHUNK , GET_FILE <EOL> from onitu . brocker . responses import ERROR <EOL> from onitu . escalator . client import EscalatorClosed <EOL> from . metadata import Metadata <EOL> from . exceptions import AbortOperation <EOL> class Worker ( object ) : <EOL> def __init__ ( self , dealer , fid ) : <EOL> super ( Worker , self ) . __init__ ( ) <EOL> self . dealer = dealer <EOL> self . fid = fid <EOL> self . logger = dealer . logger <EOL> self . call = dealer . plug . call <EOL> self . session = dealer . plug . session <EOL> self . _stop = Event ( ) <EOL> self . context = zmq . Context ( ) <EOL> self . escalator = self . dealer . escalator . clone ( context = self . context ) <EOL> def __call__ ( self ) : <EOL> try : <EOL> self . metadata = Metadata . get_by_id ( self . dealer . plug , self . fid ) <EOL> if not self . metadata : <EOL> return <EOL> self . filename = self . metadata . filename <EOL> self . do ( ) <EOL> except Exception : <EOL> log_traceback ( self . logger ) <EOL> except EscalatorClosed : <EOL> pass <EOL> finally : <EOL> self . escalator . close ( ) <EOL> self . context . destroy ( ) <EOL> if self . fid in self . dealer . in_progress : <EOL> self . dealer . in_progress . pop ( self . fid ) <EOL> def do ( self ) : <EOL> raise NotImplementedError ( ) <EOL> def stop ( self ) : <EOL> self . _stop . set ( ) <EOL> class TransferWorker ( Worker ) : <EOL> def __init__ ( self , dealer , fid , offset = <NUM_LIT:0> , restart = False ) : <EOL> super ( TransferWorker , self ) . __init__ ( dealer , fid ) <EOL> self . offset = offset <EOL> self . restart = restart <EOL> self . chunk_size = self . dealer . plug . options [ '<STR_LIT>' ] <EOL> driver_chunk_size = self . call ( '<STR_LIT>' , self . chunk_size ) <EOL> if driver_chunk_size is not None : <EOL> self . chunk_size = driver_chunk_size <EOL> self . transfer_key = ( u'<STR_LIT>' <EOL> . format ( self . dealer . name , self . fid ) ) <EOL> def do ( self ) : <EOL> success = False <EOL> dealer = None <EOL> try : <EOL> self . start_transfer ( ) <EOL> dealer = self . context . socket ( zmq . DEALER ) <EOL> dealer . connect ( get_brocker_uri ( self . session ) ) <EOL> if self . metadata . size < self . chunk_size * <NUM_LIT:2> : <EOL> self . get_file_oneshot ( dealer ) <EOL> else : <EOL> if self . dealer . plug . has_handler ( '<STR_LIT>' ) : <EOL> self . get_file_multipart ( dealer ) <EOL> else : <EOL> self . get_file_oneshot ( dealer ) <EOL> except AbortOperation : <EOL> pass <EOL> else : <EOL> success = True <EOL> finally : <EOL> if dealer : <EOL> dealer . close ( ) <EOL> self . end_transfer ( success ) <EOL> def start_transfer ( self ) : <EOL> if self . restart : <EOL> self . call ( '<STR_LIT>' , self . metadata , self . offset ) <EOL> self . logger . info ( "<STR_LIT>" , self . filename ) <EOL> else : <EOL> self . escalator . put ( self . transfer_key , self . offset ) <EOL> self . call ( '<STR_LIT>' , self . metadata ) <EOL> self . logger . info ( "<STR_LIT>" , self . filename ) <EOL> def get_file_oneshot ( self , dealer ) : <EOL> self . logger . debug ( "<STR_LIT>" , self . filename ) <EOL> dealer . send_multipart ( ( GET_FILE , str ( self . fid ) . encode ( ) ) ) <EOL> resp = dealer . recv_multipart ( ) <EOL> if resp [ <NUM_LIT:0> ] == ERROR : <EOL> raise AbortOperation ( ) <EOL> self . logger . debug ( "<STR_LIT>" , self . filename ) <EOL> if self . dealer . plug . has_handler ( '<STR_LIT>' ) : <EOL> self . call ( '<STR_LIT>' , self . metadata , resp [ <NUM_LIT:1> ] ) <EOL> else : <EOL> self . call ( '<STR_LIT>' , self . metadata , <NUM_LIT:0> , resp [ <NUM_LIT:1> ] ) <EOL> def get_file_multipart ( self , dealer ) : <EOL> while self . offset < self . metadata . size : <EOL> if self . _stop . is_set ( ) : <EOL> raise AbortOperation ( ) <EOL> dealer . send_multipart ( ( <EOL> GET_CHUNK , <EOL> str ( self . fid ) . encode ( ) , <EOL> str ( self . offset ) . encode ( ) , <EOL> str ( self . chunk_size ) . encode ( ) <EOL> ) ) <EOL> resp = dealer . recv_multipart ( ) <EOL> if len ( resp ) < <NUM_LIT:2> or resp [ <NUM_LIT:0> ] == ERROR : <EOL> raise AbortOperation ( ) <EOL> chunk = resp [ <NUM_LIT:1> ] <EOL> if not chunk or len ( chunk ) == <NUM_LIT:0> : <EOL> raise AbortOperation ( ) <EOL> self . call ( '<STR_LIT>' , self . metadata , self . offset , chunk ) <EOL> self . offset += len ( chunk ) <EOL> self . escalator . put ( self . transfer_key , self . offset ) <EOL> def end_transfer ( self , success ) : <EOL> if self . _stop . is_set ( ) : <EOL> success = False <EOL> if success : <EOL> try : <EOL> self . call ( '<STR_LIT>' , self . metadata ) <EOL> except AbortOperation : <EOL> success = False <EOL> if not success : <EOL> try : <EOL> self . call ( '<STR_LIT>' , self . metadata ) <EOL> except AbortOperation : <EOL> pass <EOL> if success : <EOL> self . escalator . delete ( self . transfer_key ) <EOL> self . metadata . set_uptodate ( ) <EOL> self . metadata . write ( ) <EOL> self . logger . info ( "<STR_LIT>" , self . filename ) <EOL> else : <EOL> self . logger . info ( "<STR_LIT>" , self . filename ) <EOL> class DeletionWorker ( Worker ) : <EOL> def __init__ ( self , dealer , fid ) : <EOL> super ( DeletionWorker , self ) . __init__ ( dealer , fid ) <EOL> def do ( self ) : <EOL> self . logger . debug ( "<STR_LIT>" , self . filename ) <EOL> try : <EOL> self . call ( '<STR_LIT>' , self . metadata ) <EOL> except AbortOperation : <EOL> pass <EOL> self . metadata . delete ( ) <EOL> self . logger . info ( "<STR_LIT>" , self . filename ) <EOL> class MoveWorker ( Worker ) : <EOL> def __init__ ( self , dealer , fid , new_fid ) : <EOL> super ( MoveWorker , self ) . __init__ ( dealer , fid ) <EOL> self . new_fid = new_fid <EOL> def do ( self ) : <EOL> new_metadata = Metadata . get_by_id ( self . dealer . plug , self . new_fid ) <EOL> new_filename = new_metadata . filename <EOL> self . logger . debug ( "<STR_LIT>" , self . filename , new_filename ) <EOL> has_move_file = self . dealer . plug . has_handler ( '<STR_LIT>' ) <EOL> try : <EOL> if has_move_file and self . metadata . is_uptodate : <EOL> self . call ( '<STR_LIT>' , self . metadata , new_metadata ) <EOL> self . metadata . delete ( ) <EOL> new_metadata . set_uptodate ( ) <EOL> else : <EOL> self . call ( '<STR_LIT>' , self . metadata ) <EOL> self . metadata . delete ( ) <EOL> transfer = TransferWorker ( self . dealer , self . new_fid ) <EOL> transfer ( ) <EOL> except AbortOperation : <EOL> pass <EOL> self . logger . info ( "<STR_LIT>" , self . filename , new_filename ) <EOL> WORKERS = { <EOL> UP : TransferWorker , <EOL> DEL : DeletionWorker , <EOL> MOV : MoveWorker , <EOL> } </s>
