<s> import re <EOL> import time <EOL> import requests <EOL> _drain_methods = { } <EOL> def register_drain_method ( name ) : <EOL> """<STR_LIT>""" <EOL> def outer ( drain_method ) : <EOL> _drain_methods [ name ] = drain_method <EOL> return drain_method <EOL> return outer <EOL> def get_drain_method ( name , service , instance , nerve_ns , drain_method_params = None ) : <EOL> return _drain_methods [ name ] ( service , instance , nerve_ns , ** ( drain_method_params or { } ) ) <EOL> def list_drain_methods ( ) : <EOL> return sorted ( _drain_methods . keys ( ) ) <EOL> class DrainMethod ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , service , instance , nerve_ns , ** kwargs ) : <EOL> self . service = service <EOL> self . instance = instance <EOL> self . nerve_ns = nerve_ns <EOL> def drain ( self , task ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError ( ) <EOL> def stop_draining ( self , task ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError ( ) <EOL> def is_draining ( self , task ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError ( ) <EOL> def is_safe_to_kill ( self , task ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError ( ) <EOL> @ register_drain_method ( '<STR_LIT>' ) <EOL> class NoopDrainMethod ( DrainMethod ) : <EOL> """<STR_LIT>""" <EOL> def drain ( self , task ) : <EOL> pass <EOL> def stop_draining ( self , task ) : <EOL> pass <EOL> def is_draining ( self , task ) : <EOL> return False <EOL> def is_safe_to_kill ( self , task ) : <EOL> return True <EOL> @ register_drain_method ( '<STR_LIT:test>' ) <EOL> class TestDrainMethod ( DrainMethod ) : <EOL> """<STR_LIT>""" <EOL> downed_task_ids = set ( ) <EOL> safe_to_kill_task_ids = set ( ) <EOL> def drain ( self , task ) : <EOL> self . downed_task_ids . add ( task . id ) <EOL> def stop_draining ( self , task ) : <EOL> self . downed_task_ids -= set ( [ task . id ] ) <EOL> self . safe_to_kill_task_ids -= set ( [ task . id ] ) <EOL> def is_draining ( self , task ) : <EOL> return task . id in ( self . downed_task_ids | self . safe_to_kill_task_ids ) <EOL> def is_safe_to_kill ( self , task ) : <EOL> return task . id in self . safe_to_kill_task_ids <EOL> @ classmethod <EOL> def mark_arbitrary_task_as_safe_to_kill ( cls ) : <EOL> cls . safe_to_kill_task_ids . add ( cls . downed_task_ids . pop ( ) ) <EOL> @ register_drain_method ( '<STR_LIT>' ) <EOL> class HacheckDrainMethod ( DrainMethod ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , service , instance , nerve_ns , delay = <NUM_LIT> , hacheck_port = <NUM_LIT> , expiration = <NUM_LIT:0> , ** kwargs ) : <EOL> super ( HacheckDrainMethod , self ) . __init__ ( service , instance , nerve_ns ) <EOL> self . delay = float ( delay ) <EOL> self . hacheck_port = hacheck_port <EOL> self . expiration = float ( expiration ) or float ( delay ) * <NUM_LIT:10> <EOL> def spool_url ( self , task ) : <EOL> return '<STR_LIT>' % { <EOL> '<STR_LIT>' : task . host , <EOL> '<STR_LIT>' : task . ports [ <NUM_LIT:0> ] , <EOL> '<STR_LIT>' : self . hacheck_port , <EOL> '<STR_LIT>' : self . service , <EOL> '<STR_LIT>' : self . nerve_ns , <EOL> } <EOL> def post_spool ( self , task , status ) : <EOL> resp = requests . post ( <EOL> self . spool_url ( task ) , <EOL> data = { <EOL> '<STR_LIT:status>' : status , <EOL> '<STR_LIT>' : time . time ( ) + self . expiration , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> } , <EOL> ) <EOL> resp . raise_for_status ( ) <EOL> def get_spool ( self , task ) : <EOL> """<STR_LIT>""" <EOL> response = requests . get ( self . spool_url ( task ) ) <EOL> if response . status_code == <NUM_LIT:200> : <EOL> return { <EOL> '<STR_LIT:state>' : '<STR_LIT>' , <EOL> } <EOL> regex = '<STR_LIT>' . join ( [ <EOL> "<STR_LIT>" , <EOL> r"<STR_LIT>" , <EOL> r"<STR_LIT>" , <EOL> r"<STR_LIT>" , <EOL> r"<STR_LIT>" , <EOL> r"<STR_LIT>" , <EOL> "<STR_LIT:$>" <EOL> ] ) <EOL> match = re . match ( regex , response . text ) <EOL> groupdict = match . groupdict ( ) <EOL> info = { } <EOL> info [ '<STR_LIT>' ] = groupdict [ '<STR_LIT>' ] <EOL> info [ '<STR_LIT:state>' ] = groupdict [ '<STR_LIT:state>' ] <EOL> if '<STR_LIT>' in groupdict : <EOL> info [ '<STR_LIT>' ] = float ( groupdict [ '<STR_LIT>' ] or <NUM_LIT:0> ) <EOL> if '<STR_LIT>' in groupdict : <EOL> info [ '<STR_LIT>' ] = float ( groupdict [ '<STR_LIT>' ] or <NUM_LIT:0> ) <EOL> if '<STR_LIT>' in groupdict : <EOL> info [ '<STR_LIT>' ] = groupdict [ '<STR_LIT>' ] <EOL> return info <EOL> def drain ( self , task ) : <EOL> self . post_spool ( task , '<STR_LIT>' ) <EOL> def stop_draining ( self , task ) : <EOL> self . post_spool ( task , '<STR_LIT>' ) <EOL> def is_draining ( self , task ) : <EOL> info = self . get_spool ( task ) <EOL> if info [ "<STR_LIT:state>" ] == "<STR_LIT>" : <EOL> return False <EOL> else : <EOL> return True <EOL> def is_safe_to_kill ( self , task ) : <EOL> info = self . get_spool ( task ) <EOL> if info [ "<STR_LIT:state>" ] == "<STR_LIT>" : <EOL> return False <EOL> else : <EOL> return info . get ( "<STR_LIT>" , <NUM_LIT:0> ) < ( time . time ( ) - self . delay ) </s>
