<s> """<STR_LIT>""" <EOL> from collections import defaultdict <EOL> from nydus . utils import ThreadPool <EOL> from nydus . db . exceptions import CommandError <EOL> from nydus . db . promise import EventualCommand , change_resolution <EOL> class BaseDistributedConnection ( object ) : <EOL> def __init__ ( self , cluster , workers = None , fail_silently = False ) : <EOL> self . _commands = [ ] <EOL> self . _complete = False <EOL> self . _errors = [ ] <EOL> self . _resolved = False <EOL> self . _cluster = cluster <EOL> self . _fail_silently = fail_silently <EOL> self . _workers = min ( workers or len ( cluster ) , <NUM_LIT:16> ) <EOL> def __getattr__ ( self , attr ) : <EOL> command = EventualCommand ( attr ) <EOL> self . _commands . append ( command ) <EOL> return command <EOL> def _build_pending_commands ( self ) : <EOL> pending_commands = defaultdict ( list ) <EOL> for command in self . _commands : <EOL> if not command . was_called ( ) : <EOL> continue <EOL> if self . _cluster . router : <EOL> name , args , kwargs = command . get_command ( ) <EOL> db_nums = self . _cluster . router . get_dbs ( <EOL> cluster = self . _cluster , <EOL> attr = name , <EOL> args = args , <EOL> kwargs = kwargs , <EOL> ) <EOL> else : <EOL> db_nums = self . _cluster . keys ( ) <EOL> for db_num in db_nums : <EOL> pending_commands [ db_num ] . append ( command ) <EOL> return pending_commands <EOL> def get_pool ( self , commands ) : <EOL> return ThreadPool ( min ( self . _workers , len ( commands ) ) ) <EOL> def resolve ( self ) : <EOL> pending_commands = self . _build_pending_commands ( ) <EOL> num_commands = sum ( len ( v ) for v in pending_commands . itervalues ( ) ) <EOL> if num_commands == <NUM_LIT:1> : <EOL> db_num , ( command , ) = pending_commands . items ( ) [ <NUM_LIT:0> ] <EOL> self . _commands = [ command . resolve ( self . _cluster [ db_num ] ) ] <EOL> elif num_commands > <NUM_LIT:1> : <EOL> results = self . execute ( self . _cluster , pending_commands ) <EOL> for command in self . _commands : <EOL> result = results . get ( command ) <EOL> if result : <EOL> for value in result : <EOL> if isinstance ( value , Exception ) : <EOL> self . _errors . append ( ( command . get_name ( ) , value ) ) <EOL> if len ( result ) == <NUM_LIT:1> : <EOL> result = result [ <NUM_LIT:0> ] <EOL> change_resolution ( command , result ) <EOL> self . _resolved = True <EOL> if not self . _fail_silently and self . _errors : <EOL> raise CommandError ( self . _errors ) <EOL> def execute ( self , cluster , commands ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError <EOL> def get_results ( self ) : <EOL> """<STR_LIT>""" <EOL> assert self . _resolved , '<STR_LIT>' <EOL> return self . _commands <EOL> def get_errors ( self ) : <EOL> """<STR_LIT>""" <EOL> assert self . _resolved , '<STR_LIT>' <EOL> return self . _errors <EOL> class DistributedConnection ( BaseDistributedConnection ) : <EOL> """<STR_LIT>""" <EOL> def execute ( self , cluster , commands ) : <EOL> pool = self . get_pool ( commands ) <EOL> for db_num , command_list in commands . iteritems ( ) : <EOL> for command in command_list : <EOL> pool . add ( command , command . clone ( ) . resolve , [ cluster [ db_num ] ] ) <EOL> return dict ( pool . join ( ) ) <EOL> class PipelinedDistributedConnection ( BaseDistributedConnection ) : <EOL> """<STR_LIT>""" <EOL> def execute ( self , cluster , commands ) : <EOL> pipes = { } <EOL> pool = self . get_pool ( commands ) <EOL> for db_num , command_list in commands . iteritems ( ) : <EOL> pipes [ db_num ] = cluster [ db_num ] . get_pipeline ( ) <EOL> for command in command_list : <EOL> pipes [ db_num ] . add ( command . clone ( ) ) <EOL> for db_num , pipe in pipes . iteritems ( ) : <EOL> pool . add ( db_num , pipe . execute , ( ) , { } ) <EOL> db_result_map = pool . join ( ) <EOL> results = defaultdict ( list ) <EOL> for db_num , db_results in db_result_map . iteritems ( ) : <EOL> assert len ( db_results ) == <NUM_LIT:1> <EOL> db_results = db_results [ <NUM_LIT:0> ] <EOL> if isinstance ( db_results , Exception ) : <EOL> for command in commands [ db_num ] : <EOL> results [ command ] . append ( db_results ) <EOL> continue <EOL> for command , result in db_results . iteritems ( ) : <EOL> results [ command ] . append ( result ) <EOL> return results <EOL> class DistributedContextManager ( object ) : <EOL> def __init__ ( self , cluster , workers = None , ** kwargs ) : <EOL> if self . can_pipeline ( cluster ) : <EOL> cls = PipelinedDistributedConnection <EOL> else : <EOL> cls = DistributedConnection <EOL> self . connection = cls ( cluster , workers , ** kwargs ) <EOL> def __enter__ ( self ) : <EOL> return self . connection <EOL> def __exit__ ( self , exc_type , exc_value , tb ) : <EOL> self . connection . resolve ( ) <EOL> def can_pipeline ( self , cluster ) : <EOL> return all ( cluster [ n ] . supports_pipelines for n in cluster ) </s>
