<s> '''<STR_LIT>''' <EOL> import os <EOL> import threading <EOL> try : <EOL> import cPickle as pickle <EOL> except ImportError : <EOL> import pickle <EOL> from collections import defaultdict <EOL> import socket <EOL> from cola . core . rpc import client_call <EOL> from cola . core . utils import get_rpc_prefix <EOL> from cola . core . mq . store import Store <EOL> from cola . core . mq . distributor import Distributor <EOL> MQ_STATUS_FILENAME = '<STR_LIT>' <EOL> PRIORITY_STORE_FN = '<STR_LIT:store>' <EOL> BACKUP_STORE_FN = '<STR_LIT>' <EOL> INCR_STORE_FN = '<STR_LIT>' <EOL> CACHE_SIZE = <NUM_LIT:20> <EOL> class LocalMessageQueueNode ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , base_dir , rpc_server , addr , addrs , <EOL> copies = <NUM_LIT:1> , n_priorities = <NUM_LIT:3> , deduper = None , <EOL> app_name = None ) : <EOL> self . dir_ = base_dir <EOL> self . rpc_server = rpc_server <EOL> assert addr in addrs <EOL> self . addr = addr <EOL> self . addrs = addrs <EOL> self . other_addrs = [ n for n in self . addrs if n != self . addr ] <EOL> self . copies = max ( min ( len ( self . addrs ) - <NUM_LIT:1> , copies ) , <NUM_LIT:0> ) <EOL> self . n_priorities = max ( n_priorities , <NUM_LIT:1> ) <EOL> self . deduper = deduper <EOL> self . app_name = app_name <EOL> self . _lock = threading . Lock ( ) <EOL> self . _register_rpc ( ) <EOL> self . inited = False <EOL> def init ( self ) : <EOL> with self . _lock : <EOL> if self . inited : return <EOL> get_priority_store_dir = lambda priority : os . path . join ( self . dir_ , <EOL> PRIORITY_STORE_FN , str ( priority ) ) <EOL> self . priority_stores = [ Store ( get_priority_store_dir ( i ) , <EOL> deduper = self . deduper , <EOL> mkdirs = True ) for i in range ( self . n_priorities ) ] <EOL> backup_store_dir = os . path . join ( self . dir_ , BACKUP_STORE_FN ) <EOL> self . backup_stores = { } <EOL> for backup_addr in self . other_addrs : <EOL> backup_node_dir = backup_addr . replace ( '<STR_LIT::>' , '<STR_LIT:_>' ) <EOL> backup_path = os . path . join ( backup_store_dir , backup_node_dir ) <EOL> self . backup_stores [ backup_addr ] = Store ( backup_path , <EOL> size = <NUM_LIT> * <NUM_LIT> , mkdirs = True ) <EOL> inc_store_dir = os . path . join ( self . dir_ , INCR_STORE_FN ) <EOL> self . inc_store = Store ( inc_store_dir , mkdirs = True ) <EOL> self . inited = True <EOL> def _register_rpc ( self ) : <EOL> if self . rpc_server : <EOL> self . register_rpc ( self , self . rpc_server , app_name = self . app_name ) <EOL> @ classmethod <EOL> def register_rpc ( cls , node , rpc_server , app_name = None ) : <EOL> prefix = get_rpc_prefix ( app_name , '<STR_LIT>' ) <EOL> rpc_server . register_function ( node . put_proxy , name = '<STR_LIT>' , <EOL> prefix = prefix ) <EOL> rpc_server . register_function ( node . batch_put_proxy , name = '<STR_LIT>' , <EOL> prefix = prefix ) <EOL> rpc_server . register_function ( node . put_backup_proxy , name = '<STR_LIT>' , <EOL> prefix = prefix ) <EOL> rpc_server . register_function ( node . get_proxy , name = '<STR_LIT>' , <EOL> prefix = prefix ) <EOL> rpc_server . register_function ( node . exist , name = '<STR_LIT>' , <EOL> prefix = prefix ) <EOL> def put ( self , objs , force = False , priority = <NUM_LIT:0> ) : <EOL> self . init ( ) <EOL> priority = max ( min ( priority , self . n_priorities - <NUM_LIT:1> ) , <NUM_LIT:0> ) <EOL> priority_store = self . priority_stores [ priority ] <EOL> priority_store . put ( objs , force = force ) <EOL> def put_proxy ( self , pickled_objs , force = False , priority = <NUM_LIT:0> ) : <EOL> """<STR_LIT>""" <EOL> objs = pickle . loads ( pickled_objs ) <EOL> self . put ( objs , force = force , priority = priority ) <EOL> def batch_put ( self , objs ) : <EOL> self . init ( ) <EOL> puts = defaultdict ( lambda : defaultdict ( list ) ) <EOL> for obj in objs : <EOL> priority = getattr ( obj , '<STR_LIT>' , <NUM_LIT:0> ) <EOL> force = getattr ( obj , '<STR_LIT>' , False ) <EOL> puts [ priority ] [ force ] . append ( obj ) <EOL> for priority , m in puts . iteritems ( ) : <EOL> for force , obs in m . iteritems ( ) : <EOL> self . put ( obs , force = force , priority = priority ) <EOL> def batch_put_proxy ( self , pickled_objs ) : <EOL> """<STR_LIT>""" <EOL> objs = pickle . loads ( pickled_objs ) <EOL> self . batch_put ( objs ) <EOL> def put_backup ( self , addr , objs , force = False ) : <EOL> self . init ( ) <EOL> backup_store = self . backup_stores [ addr ] <EOL> backup_store . put ( objs , force = force ) <EOL> def put_backup_proxy ( self , addr , pickled_objs , force = False ) : <EOL> """<STR_LIT>""" <EOL> objs = pickle . loads ( pickled_objs ) <EOL> self . put_backup ( addr , objs , force = force ) <EOL> def put_inc ( self , objs , force = True ) : <EOL> self . init ( ) <EOL> self . inc_store . put ( objs , force = force ) <EOL> def get ( self , size = <NUM_LIT:1> , priority = <NUM_LIT:0> ) : <EOL> self . init ( ) <EOL> priority = max ( min ( priority , self . n_priorities - <NUM_LIT:1> ) , <NUM_LIT:0> ) <EOL> priority_store = self . priority_stores [ priority ] <EOL> return priority_store . get ( size = size ) <EOL> def get_proxy ( self , size = <NUM_LIT:1> , priority = <NUM_LIT:0> ) : <EOL> """<STR_LIT>""" <EOL> return pickle . dumps ( self . get ( size = size , priority = priority ) ) <EOL> def get_backup ( self , addr , size = <NUM_LIT:1> ) : <EOL> self . init ( ) <EOL> backup_store = self . backup_stores [ addr ] <EOL> return backup_store . get ( size = size ) <EOL> def get_inc ( self , size = <NUM_LIT:1> ) : <EOL> self . init ( ) <EOL> return self . inc_store . get ( size = size ) <EOL> def add_node ( self , addr ) : <EOL> """<STR_LIT>""" <EOL> if addr in self . addrs : return <EOL> self . addrs . append ( addr ) <EOL> backup_store_dir = os . path . join ( self . dir_ , BACKUP_STORE_FN ) <EOL> backup_node_dir = addr . replace ( '<STR_LIT::>' , '<STR_LIT:_>' ) <EOL> backup_path = os . path . join ( backup_store_dir , backup_node_dir ) <EOL> self . backup_stores [ addr ] = Store ( backup_path , <EOL> size = <NUM_LIT> * <NUM_LIT> , mkdirs = True ) <EOL> def remove_node ( self , addr ) : <EOL> """<STR_LIT>""" <EOL> if addr not in self . addrs : return <EOL> self . addrs . remove ( addr ) <EOL> self . backup_stores [ addr ] . shutdown ( ) <EOL> del self . backup_stores [ addr ] <EOL> def exist ( self , obj ) : <EOL> if self . deduper : <EOL> return self . deduper . exist ( str ( obj ) ) <EOL> return False <EOL> def shutdown ( self ) : <EOL> if not self . inited : return <EOL> [ store . shutdown ( ) for store in self . priority_stores ] <EOL> for backup_store in self . backup_stores . values ( ) : <EOL> backup_store . shutdown ( ) <EOL> self . inc_store . shutdown ( ) <EOL> class MessageQueueNodeProxy ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , base_dir , rpc_server , addr , addrs , <EOL> copies = <NUM_LIT:1> , n_priorities = <NUM_LIT:3> , deduper = None , <EOL> app_name = None , logger = None ) : <EOL> self . dir_ = base_dir <EOL> self . addr_ = addr <EOL> self . addrs = list ( addrs ) <EOL> self . mq_node = LocalMessageQueueNode ( <EOL> base_dir , rpc_server , addr , addrs , <EOL> copies = copies , n_priorities = n_priorities , deduper = deduper , <EOL> app_name = app_name ) <EOL> self . distributor = Distributor ( addrs , copies = copies ) <EOL> self . logger = logger <EOL> self . prefix = get_rpc_prefix ( app_name , '<STR_LIT>' ) <EOL> self . _lock = threading . Lock ( ) <EOL> self . inited = False <EOL> @ classmethod <EOL> def register_rpc ( cls , node , rpc_server , app_name = None ) : <EOL> LocalMessageQueueNode . register_rpc ( node . mq_node , rpc_server , <EOL> app_name = app_name ) <EOL> def init ( self ) : <EOL> with self . _lock : <EOL> if self . inited : return <EOL> self . load ( ) <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> self . caches = dict ( ( addr , [ ] ) for addr in self . addrs ) <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> self . caches_inited = dict ( ( addr , False ) for addr in self . addrs ) <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> self . backup_caches = dict ( ( addr , { } ) for addr in self . addrs ) <EOL> for addr in self . addrs : <EOL> for other_addr in [ n for n in self . addrs if addr != n ] : <EOL> self . backup_caches [ addr ] [ other_addr ] = [ ] <EOL> self . mq_node . init ( ) <EOL> self . inited = True <EOL> def load ( self ) : <EOL> save_file = os . path . join ( self . dir_ , MQ_STATUS_FILENAME ) <EOL> if not os . path . exists ( save_file ) : <EOL> return <EOL> with open ( save_file , '<STR_LIT:r>' ) as f : <EOL> self . caches , self . caches_inited , self . backup_caches = pickle . load ( f ) <EOL> def save ( self ) : <EOL> if not self . inited : <EOL> return <EOL> save_file = os . path . join ( self . dir_ , MQ_STATUS_FILENAME ) <EOL> with open ( save_file , '<STR_LIT:w>' ) as f : <EOL> t = ( self . caches , self . caches_inited , self . backup_caches ) <EOL> pickle . dump ( t , f ) <EOL> def _check_empty ( self , objs ) : <EOL> if objs is None : <EOL> return True <EOL> elif isinstance ( objs , list ) and len ( objs ) == <NUM_LIT:0> : <EOL> return True <EOL> return False <EOL> def _remote_or_local_put ( self , addr , objs , force = False , priority = <NUM_LIT:0> ) : <EOL> if self . _check_empty ( objs ) : <EOL> return <EOL> if addr == self . addr_ : <EOL> self . mq_node . put ( objs , force = force , priority = priority ) <EOL> else : <EOL> client_call ( addr , self . prefix + '<STR_LIT>' , pickle . dumps ( objs ) , <EOL> force , priority ) <EOL> def _remote_or_local_batch_put ( self , addr , objs ) : <EOL> if self . _check_empty ( objs ) : <EOL> return <EOL> if addr == self . addr_ : <EOL> self . mq_node . batch_put ( objs ) <EOL> else : <EOL> client_call ( addr , self . prefix + '<STR_LIT>' , pickle . dumps ( objs ) ) <EOL> def _remote_or_local_get ( self , addr , size = <NUM_LIT:1> , priority = <NUM_LIT:0> ) : <EOL> objs = None <EOL> if addr == self . addr_ : <EOL> objs = self . mq_node . get ( size = size , priority = priority ) <EOL> else : <EOL> objs = pickle . loads ( client_call ( addr , self . prefix + '<STR_LIT>' , <EOL> size , priority ) ) <EOL> addr_caches = self . caches . get ( addr , [ ] ) <EOL> if size == <NUM_LIT:1> and objs is None and len ( addr_caches ) > <NUM_LIT:0> : <EOL> return addr_caches . pop ( <NUM_LIT:0> ) <EOL> elif size > <NUM_LIT:1> and len ( objs ) == <NUM_LIT:0> and len ( addr_caches ) > <NUM_LIT:0> : <EOL> return addr_caches [ : size ] <EOL> return objs <EOL> def _remote_or_local_put_backup ( self , addr , backup_addr , objs , <EOL> force = False ) : <EOL> if self . _check_empty ( objs ) : <EOL> return <EOL> if addr == self . addr_ : <EOL> self . mq_node . put_backup ( backup_addr , objs , force = force ) <EOL> else : <EOL> client_call ( addr , self . prefix + '<STR_LIT>' , backup_addr , <EOL> pickle . dumps ( objs ) , force ) <EOL> def put ( self , objects , flush = False ) : <EOL> """<STR_LIT>""" <EOL> self . init ( ) <EOL> addrs_objs , backup_addrs_objs = self . distributor . distribute ( objects ) <EOL> if flush is True : <EOL> for addr in self . addrs : <EOL> if addr not in addrs_objs : <EOL> addrs_objs [ addr ] = [ ] <EOL> if addr not in backup_addrs_objs : <EOL> backup_addrs_objs [ addr ] = { } <EOL> for addr , objs in addrs_objs . iteritems ( ) : <EOL> self . caches [ addr ] . extend ( objs ) <EOL> if not self . caches_inited [ addr ] or len ( self . caches [ addr ] ) >= CACHE_SIZE or flush : <EOL> try : <EOL> self . _remote_or_local_batch_put ( addr , self . caches [ addr ] ) <EOL> except socket . error , e : <EOL> if self . logger : <EOL> self . logger . exception ( e ) <EOL> else : <EOL> self . caches [ addr ] = [ ] <EOL> if not self . caches_inited [ addr ] : <EOL> self . caches_inited [ addr ] = True <EOL> for addr , m in backup_addrs_objs . iteritems ( ) : <EOL> for backup_addr , objs in m . iteritems ( ) : <EOL> self . backup_caches [ addr ] [ backup_addr ] . extend ( objs ) <EOL> size = sum ( [ len ( obs ) for obs in self . backup_caches [ addr ] . values ( ) ] ) <EOL> if size >= CACHE_SIZE or flush : <EOL> for backup_addr , objs in self . backup_caches [ addr ] . iteritems ( ) : <EOL> try : <EOL> self . _remote_or_local_put_backup ( <EOL> addr , backup_addr , objs ) <EOL> except socket . error , e : <EOL> if self . logger : <EOL> self . logger . exception ( e ) <EOL> else : <EOL> self . backup_caches [ addr ] [ backup_addr ] = [ ] <EOL> def get ( self , size = <NUM_LIT:1> , priority = <NUM_LIT:0> ) : <EOL> """<STR_LIT>""" <EOL> self . init ( ) <EOL> if size < <NUM_LIT:1> : size = <NUM_LIT:1> <EOL> results = [ ] <EOL> _addrs = sorted ( self . addrs , key = lambda k : k == self . addr_ , <EOL> reverse = True ) <EOL> for addr in _addrs : <EOL> left = size - len ( results ) <EOL> if left <= <NUM_LIT:0> : <EOL> break <EOL> objs = None <EOL> try : <EOL> objs = self . _remote_or_local_get ( addr , size = left , <EOL> priority = priority ) <EOL> except socket . error , e : <EOL> if self . logger : <EOL> self . logger . exception ( e ) <EOL> if objs is None : <EOL> continue <EOL> if not isinstance ( objs , list ) : <EOL> objs = [ objs , ] <EOL> results . extend ( objs ) <EOL> if size == <NUM_LIT:1> : <EOL> if len ( results ) == <NUM_LIT:0> : <EOL> return <EOL> return results [ <NUM_LIT:0> ] <EOL> return results <EOL> def put_inc ( self , objs ) : <EOL> self . mq_node . put_inc ( objs ) <EOL> def get_inc ( self , size = <NUM_LIT:1> ) : <EOL> return self . mq_node . get_inc ( size = size ) <EOL> def flush ( self ) : <EOL> self . put ( [ ] , flush = True ) <EOL> def add_node ( self , addr ) : <EOL> if addr in self . addrs : return <EOL> self . init ( ) <EOL> self . distributor . add_node ( addr ) <EOL> self . addrs . append ( addr ) <EOL> self . caches [ addr ] = [ ] <EOL> self . caches_inited [ addr ] = False <EOL> self . backup_caches [ addr ] = { } <EOL> for o_addr in self . addrs : <EOL> if o_addr != addr : <EOL> self . backup_caches [ addr ] [ o_addr ] = [ ] <EOL> self . backup_caches [ o_addr ] [ addr ] = [ ] <EOL> self . mq_node . add_node ( addr ) <EOL> def remove_node ( self , addr ) : <EOL> if addr not in self . addrs : return <EOL> self . init ( ) <EOL> self . distributor . remove_node ( addr ) <EOL> self . addrs . remove ( addr ) <EOL> self . mq_node . batch_put ( self . caches [ addr ] ) <EOL> del self . caches [ addr ] <EOL> del self . caches_inited [ addr ] <EOL> del self . backup_caches [ addr ] <EOL> for o_addr in self . addrs : <EOL> if o_addr != addr : <EOL> del self . backup_caches [ o_addr ] [ addr ] <EOL> self . flush ( ) <EOL> BATCH_SIZE = <NUM_LIT:10> <EOL> objs = self . mq_node . get_backup ( addr , size = BATCH_SIZE ) <EOL> while len ( objs ) > <NUM_LIT:0> : <EOL> self . mq_node . batch_put ( objs ) <EOL> objs = self . mq_node . get_backup ( addr , size = BATCH_SIZE ) <EOL> self . mq_node . remove_node ( addr ) <EOL> def exist ( self , obj ) : <EOL> return self . mq_node . exist ( obj ) <EOL> def shutdown ( self ) : <EOL> if not self . inited : return <EOL> self . mq_node . shutdown ( ) <EOL> self . save ( ) <EOL> def __enter__ ( self ) : <EOL> return self <EOL> def __exit__ ( self , type_ , value , traceback ) : <EOL> self . shutdown ( ) </s>
