<s> import gevent <EOL> import gevent . pool <EOL> import os <EOL> import signal <EOL> import datetime <EOL> import time <EOL> import socket <EOL> import traceback <EOL> import psutil <EOL> import sys <EOL> import json as json_stdlib <EOL> import ujson as json <EOL> import BaseHTTPServer <EOL> from bson import ObjectId <EOL> from collections import defaultdict <EOL> from . job import Job <EOL> from . exceptions import ( TimeoutInterrupt , StopRequested , JobInterrupt , AbortInterrupt , <EOL> RetryInterrupt , MaxRetriesInterrupt ) <EOL> from . context import ( set_current_worker , set_current_job , get_current_job , get_current_config , <EOL> connections , enable_greenlet_tracing ) <EOL> from . queue import Queue <EOL> from . utils import MongoJSONEncoder <EOL> class Worker ( object ) : <EOL> """<STR_LIT>""" <EOL> job_class = Job <EOL> status = "<STR_LIT>" <EOL> mongodb_jobs = None <EOL> mongodb_logs = None <EOL> redis = None <EOL> def __init__ ( self ) : <EOL> set_current_worker ( self ) <EOL> if self . config . get ( "<STR_LIT>" ) : <EOL> enable_greenlet_tracing ( ) <EOL> self . datestarted = datetime . datetime . utcnow ( ) <EOL> self . done_jobs = <NUM_LIT:0> <EOL> self . max_jobs = self . config [ "<STR_LIT>" ] <EOL> self . connected = False <EOL> self . exitcode = <NUM_LIT:0> <EOL> self . process = psutil . Process ( os . getpid ( ) ) <EOL> self . greenlet = gevent . getcurrent ( ) <EOL> self . graceful_stop = None <EOL> self . idle_event = gevent . event . Event ( ) <EOL> self . idle_wait_count = <NUM_LIT:0> <EOL> self . id = ObjectId ( ) <EOL> if self . config . get ( "<STR_LIT:name>" ) : <EOL> self . name = self . config [ "<STR_LIT:name>" ] <EOL> else : <EOL> self . name = "<STR_LIT>" % ( socket . gethostname ( ) . split ( "<STR_LIT:.>" ) [ <NUM_LIT:0> ] , os . getpid ( ) ) <EOL> self . pool_size = self . config [ "<STR_LIT>" ] <EOL> from . logger import LogHandler <EOL> self . log_handler = LogHandler ( quiet = self . config [ "<STR_LIT>" ] ) <EOL> self . log = self . log_handler . get_logger ( worker = self . id ) <EOL> self . queues = [ Queue ( x ) for x in self . config [ "<STR_LIT>" ] if x ] <EOL> self . log . info ( <EOL> "<STR_LIT>" % <EOL> self . pool_size ) <EOL> self . gevent_pool = gevent . pool . Pool ( self . pool_size ) <EOL> self . greenlets = { } <EOL> self . _traced_io = { <EOL> "<STR_LIT>" : defaultdict ( float ) , <EOL> "<STR_LIT>" : defaultdict ( float ) , <EOL> "<STR_LIT>" : <NUM_LIT:0> <EOL> } <EOL> @ property <EOL> def config ( self ) : <EOL> return get_current_config ( ) <EOL> def connect ( self , force = False ) : <EOL> if self . connected and not force : <EOL> return <EOL> self . redis = connections . redis <EOL> self . mongodb_jobs = connections . mongodb_jobs <EOL> self . mongodb_logs = connections . mongodb_logs <EOL> if self . mongodb_logs : <EOL> self . log_handler . set_collection ( self . mongodb_logs . mrq_logs ) <EOL> self . connected = True <EOL> if not self . config [ "<STR_LIT>" ] : <EOL> self . ensure_indexes ( ) <EOL> def ensure_indexes ( self ) : <EOL> if self . mongodb_logs : <EOL> self . mongodb_logs . mrq_logs . ensure_index ( <EOL> [ ( "<STR_LIT>" , <NUM_LIT:1> ) ] , background = False ) <EOL> self . mongodb_logs . mrq_logs . ensure_index ( <EOL> [ ( "<STR_LIT>" , <NUM_LIT:1> ) ] , background = False , sparse = True ) <EOL> self . mongodb_jobs . mrq_workers . ensure_index ( <EOL> [ ( "<STR_LIT:status>" , <NUM_LIT:1> ) ] , background = False ) <EOL> self . mongodb_jobs . mrq_workers . ensure_index ( <EOL> [ ( "<STR_LIT>" , <NUM_LIT:1> ) ] , background = False , expireAfterSeconds = <NUM_LIT> ) <EOL> self . mongodb_jobs . mrq_jobs . ensure_index ( <EOL> [ ( "<STR_LIT:status>" , <NUM_LIT:1> ) ] , background = False ) <EOL> self . mongodb_jobs . mrq_jobs . ensure_index ( <EOL> [ ( "<STR_LIT:path>" , <NUM_LIT:1> ) , ( "<STR_LIT:status>" , <NUM_LIT:1> ) ] , background = False ) <EOL> self . mongodb_jobs . mrq_jobs . ensure_index ( <EOL> [ ( "<STR_LIT>" , <NUM_LIT:1> ) , ( "<STR_LIT:status>" , <NUM_LIT:1> ) ] , background = False , sparse = True ) <EOL> self . mongodb_jobs . mrq_jobs . ensure_index ( <EOL> [ ( "<STR_LIT>" , <NUM_LIT:1> ) , ( "<STR_LIT:status>" , <NUM_LIT:1> ) ] , background = False ) <EOL> self . mongodb_jobs . mrq_jobs . ensure_index ( <EOL> [ ( "<STR_LIT>" , <NUM_LIT:1> ) ] , sparse = True , background = False , expireAfterSeconds = <NUM_LIT:0> ) <EOL> self . mongodb_jobs . mrq_jobs . ensure_index ( <EOL> [ ( "<STR_LIT>" , <NUM_LIT:1> ) ] , sparse = True , background = False ) <EOL> self . mongodb_jobs . mrq_scheduled_jobs . ensure_index ( <EOL> [ ( "<STR_LIT>" , <NUM_LIT:1> ) ] , unique = True , background = False , drop_dups = True ) <EOL> try : <EOL> self . mongodb_jobs . command ( <EOL> { "<STR_LIT>" : "<STR_LIT>" , "<STR_LIT>" : True } ) <EOL> self . mongodb_jobs . command ( <EOL> { "<STR_LIT>" : "<STR_LIT>" , "<STR_LIT>" : True } ) <EOL> except : <EOL> pass <EOL> def greenlet_scheduler ( self ) : <EOL> from . scheduler import Scheduler <EOL> scheduler = Scheduler ( self . mongodb_jobs . mrq_scheduled_jobs ) <EOL> scheduler . sync_tasks ( self . config . get ( "<STR_LIT>" ) or [ ] ) <EOL> while True : <EOL> scheduler . check ( ) <EOL> time . sleep ( int ( self . config [ "<STR_LIT>" ] ) ) <EOL> def greenlet_report ( self ) : <EOL> """<STR_LIT>""" <EOL> while True : <EOL> try : <EOL> self . report_worker ( ) <EOL> except Exception as e : <EOL> self . log . error ( "<STR_LIT>" % e ) <EOL> finally : <EOL> time . sleep ( self . config [ "<STR_LIT>" ] ) <EOL> def greenlet_logs ( self ) : <EOL> """<STR_LIT>""" <EOL> while True : <EOL> try : <EOL> self . flush_logs ( w = <NUM_LIT:0> ) <EOL> except Exception as e : <EOL> self . log . error ( "<STR_LIT>" % e ) <EOL> finally : <EOL> time . sleep ( self . config [ "<STR_LIT>" ] ) <EOL> def greenlet_subqueues ( self ) : <EOL> while True : <EOL> queues = [ ] <EOL> try : <EOL> for queue in self . config [ "<STR_LIT>" ] : <EOL> if queue . endswith ( get_current_config ( ) . get ( "<STR_LIT>" ) ) : <EOL> queues . append ( Queue ( queue [ : - <NUM_LIT:1> ] ) ) <EOL> queues . append ( Queue ( queue ) ) <EOL> queues += Queue . redis_known_subqueues ( queue ) <EOL> except Exception as e : <EOL> self . log . error ( "<STR_LIT>" , e ) <EOL> else : <EOL> self . queues = queues <EOL> finally : <EOL> time . sleep ( self . config [ "<STR_LIT>" ] ) <EOL> def get_memory ( self ) : <EOL> mmaps = self . process . get_memory_maps ( ) <EOL> mem = { <EOL> "<STR_LIT>" : sum ( [ x . rss for x in mmaps ] ) , <EOL> "<STR_LIT>" : sum ( [ getattr ( x , '<STR_LIT>' , getattr ( x , '<STR_LIT>' , <NUM_LIT:0> ) ) for x in mmaps ] ) <EOL> } <EOL> mem [ "<STR_LIT>" ] = mem [ "<STR_LIT>" ] + mem [ "<STR_LIT>" ] <EOL> return mem <EOL> def get_worker_report ( self , with_memory = False ) : <EOL> """<STR_LIT>""" <EOL> greenlets = [ ] <EOL> for greenlet in self . gevent_pool : <EOL> g = { } <EOL> short_stack = [ ] <EOL> stack = traceback . format_stack ( greenlet . gr_frame ) <EOL> for s in stack [ <NUM_LIT:1> : ] : <EOL> if "<STR_LIT>" in s : <EOL> break <EOL> short_stack . append ( s ) <EOL> g [ "<STR_LIT>" ] = short_stack <EOL> job = get_current_job ( id ( greenlet ) ) <EOL> if job : <EOL> job . save ( ) <EOL> if job . data : <EOL> g [ "<STR_LIT:path>" ] = job . data [ "<STR_LIT:path>" ] <EOL> g [ "<STR_LIT>" ] = job . datestarted <EOL> g [ "<STR_LIT:id>" ] = str ( job . id ) <EOL> g [ "<STR_LIT:time>" ] = getattr ( greenlet , "<STR_LIT>" , <NUM_LIT:0> ) <EOL> g [ "<STR_LIT>" ] = getattr ( greenlet , "<STR_LIT>" , None ) <EOL> if job . _current_io is not None : <EOL> g [ "<STR_LIT>" ] = job . _current_io <EOL> greenlets . append ( g ) <EOL> if ( not with_memory ) or ( self . config [ "<STR_LIT>" ] != "<STR_LIT:0>" and self . config [ "<STR_LIT>" ] ) : <EOL> cpu = { <EOL> "<STR_LIT:user>" : <NUM_LIT:0> , <EOL> "<STR_LIT>" : <NUM_LIT:0> , <EOL> "<STR_LIT>" : <NUM_LIT:0> <EOL> } <EOL> mem = { "<STR_LIT>" : <NUM_LIT:0> , "<STR_LIT>" : <NUM_LIT:0> , "<STR_LIT>" : <NUM_LIT:0> } <EOL> else : <EOL> cpu_times = self . process . get_cpu_times ( ) <EOL> cpu = { <EOL> "<STR_LIT:user>" : cpu_times . user , <EOL> "<STR_LIT>" : cpu_times . system , <EOL> "<STR_LIT>" : self . process . get_cpu_percent ( <NUM_LIT:0> ) <EOL> } <EOL> mem = self . get_memory ( ) <EOL> whitelisted_config = [ <EOL> "<STR_LIT>" , <EOL> "<STR_LIT>" , <EOL> "<STR_LIT>" , <EOL> "<STR_LIT>" , <EOL> "<STR_LIT>" , <EOL> "<STR_LIT:name>" , <EOL> "<STR_LIT>" <EOL> ] <EOL> io = None <EOL> if self . _traced_io : <EOL> io = { } <EOL> for k , v in self . _traced_io . items ( ) : <EOL> if k == "<STR_LIT>" : <EOL> io [ k ] = v <EOL> else : <EOL> io [ k ] = sorted ( v . items ( ) , reverse = True , key = lambda x : x [ <NUM_LIT:1> ] ) <EOL> return { <EOL> "<STR_LIT:status>" : self . status , <EOL> "<STR_LIT>" : { k : v for k , v in self . config . iteritems ( ) if k in whitelisted_config } , <EOL> "<STR_LIT>" : self . done_jobs , <EOL> "<STR_LIT>" : self . datestarted , <EOL> "<STR_LIT>" : datetime . datetime . utcnow ( ) , <EOL> "<STR_LIT:name>" : self . name , <EOL> "<STR_LIT>" : io , <EOL> "<STR_LIT>" : str ( self . id ) , <EOL> "<STR_LIT>" : { <EOL> "<STR_LIT>" : self . process . pid , <EOL> "<STR_LIT>" : cpu , <EOL> "<STR_LIT>" : mem <EOL> } , <EOL> "<STR_LIT>" : greenlets <EOL> } <EOL> def report_worker ( self , w = <NUM_LIT:0> ) : <EOL> report = self . get_worker_report ( with_memory = True ) <EOL> if self . config [ "<STR_LIT>" ] > <NUM_LIT:0> : <EOL> if report [ "<STR_LIT>" ] [ "<STR_LIT>" ] [ "<STR_LIT>" ] > ( self . config [ "<STR_LIT>" ] * <NUM_LIT> * <NUM_LIT> ) : <EOL> self . shutdown_max_memory ( ) <EOL> if self . config [ "<STR_LIT>" ] : <EOL> with open ( self . config [ "<STR_LIT>" ] , "<STR_LIT:wb>" ) as f : <EOL> f . write ( json . dumps ( report , ensure_ascii = False ) ) <EOL> if "<STR_LIT>" in report : <EOL> del report [ "<STR_LIT>" ] <EOL> try : <EOL> self . mongodb_jobs . mrq_workers . update ( { <EOL> "<STR_LIT>" : ObjectId ( self . id ) <EOL> } , { "<STR_LIT>" : report } , upsert = True , w = w ) <EOL> except Exception as e : <EOL> self . log . debug ( "<STR_LIT>" % e ) <EOL> def greenlet_admin ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . config [ "<STR_LIT>" ] > <NUM_LIT:1> : <EOL> self . log . debug ( <EOL> "<STR_LIT>" ) <EOL> return <EOL> class Devnull ( object ) : <EOL> def write ( self , * _ ) : <EOL> pass <EOL> from gevent import wsgi <EOL> def admin_routes ( env , start_response ) : <EOL> path = env [ "<STR_LIT>" ] <EOL> status = "<STR_LIT>" <EOL> res = "<STR_LIT>" <EOL> if path in [ "<STR_LIT:/>" , "<STR_LIT>" , "<STR_LIT>" ] : <EOL> report = self . get_worker_report ( with_memory = ( path == "<STR_LIT>" ) ) <EOL> res = json_stdlib . dumps ( report , cls = MongoJSONEncoder ) <EOL> elif path == "<STR_LIT>" : <EOL> self . idle_wait_count = <NUM_LIT:0> <EOL> self . idle_event . clear ( ) <EOL> self . idle_event . wait ( ) <EOL> res = "<STR_LIT>" <EOL> else : <EOL> status = "<STR_LIT>" <EOL> start_response ( status , [ ( '<STR_LIT:Content-Type>' , '<STR_LIT:application/json>' ) ] ) <EOL> return [ res ] <EOL> server = wsgi . WSGIServer ( ( self . config [ "<STR_LIT>" ] , self . config [ "<STR_LIT>" ] ) , admin_routes , log = Devnull ( ) ) <EOL> try : <EOL> self . log . debug ( "<STR_LIT>" % self . config [ "<STR_LIT>" ] ) <EOL> server . serve_forever ( ) <EOL> except Exception as e : <EOL> self . log . debug ( "<STR_LIT>" % e ) <EOL> def flush_logs ( self , w = <NUM_LIT:0> ) : <EOL> self . log_handler . flush ( w = w ) <EOL> def work ( self ) : <EOL> """<STR_LIT>""" <EOL> self . work_init ( ) <EOL> self . work_loop ( max_jobs = self . max_jobs ) <EOL> return self . work_stop ( ) <EOL> def work_init ( self ) : <EOL> self . connect ( ) <EOL> self . status = "<STR_LIT>" <EOL> self . greenlets [ "<STR_LIT>" ] = gevent . spawn ( self . greenlet_subqueues ) <EOL> self . greenlets [ "<STR_LIT>" ] = gevent . spawn ( self . greenlet_report ) <EOL> self . greenlets [ "<STR_LIT>" ] = gevent . spawn ( self . greenlet_logs ) <EOL> if self . config [ "<STR_LIT>" ] : <EOL> self . greenlets [ "<STR_LIT>" ] = gevent . spawn ( self . greenlet_scheduler ) <EOL> if self . config [ "<STR_LIT>" ] : <EOL> self . greenlets [ "<STR_LIT>" ] = gevent . spawn ( self . greenlet_admin ) <EOL> self . install_signal_handlers ( ) <EOL> def work_loop ( self , max_jobs = None ) : <EOL> self . done_jobs = <NUM_LIT:0> <EOL> self . idle_wait_count = <NUM_LIT:0> <EOL> try : <EOL> while True : <EOL> if self . graceful_stop : <EOL> break <EOL> while True : <EOL> free_pool_slots = self . gevent_pool . free_count ( ) <EOL> if max_jobs : <EOL> total_started = ( self . pool_size - free_pool_slots ) + self . done_jobs <EOL> free_pool_slots = min ( free_pool_slots , max_jobs - total_started ) <EOL> if free_pool_slots == <NUM_LIT:0> : <EOL> break <EOL> if free_pool_slots > <NUM_LIT:0> : <EOL> self . status = "<STR_LIT>" <EOL> break <EOL> self . status = "<STR_LIT>" <EOL> gevent . sleep ( <NUM_LIT> ) <EOL> jobs = [ ] <EOL> for queue in self . queues : <EOL> max_jobs_per_queue = free_pool_slots - len ( jobs ) <EOL> if max_jobs_per_queue <= <NUM_LIT:0> : <EOL> break <EOL> if self . config [ "<STR_LIT>" ] == "<STR_LIT>" : <EOL> max_jobs_per_queue = <NUM_LIT:1> <EOL> jobs += queue . dequeue_jobs ( <EOL> max_jobs = max_jobs_per_queue , <EOL> job_class = self . job_class , <EOL> worker = self <EOL> ) <EOL> for job in jobs : <EOL> self . gevent_pool . spawn ( self . perform_job , job ) <EOL> if max_jobs and self . done_jobs >= max_jobs : <EOL> self . log . info ( "<STR_LIT>" % self . done_jobs ) <EOL> break <EOL> if len ( jobs ) == <NUM_LIT:0> : <EOL> if self . config [ "<STR_LIT>" ] == "<STR_LIT>" : <EOL> self . log . info ( "<STR_LIT>" ) <EOL> break <EOL> if ( <EOL> not self . idle_event . is_set ( ) and <EOL> len ( jobs ) == <NUM_LIT:0> and <EOL> free_pool_slots == self . pool_size and <EOL> self . idle_wait_count > <NUM_LIT:0> <EOL> ) : <EOL> self . idle_event . set ( ) <EOL> self . status = "<STR_LIT>" <EOL> self . idle_wait_count += <NUM_LIT:1> <EOL> gevent . sleep ( min ( self . config [ "<STR_LIT>" ] , <NUM_LIT> * self . idle_wait_count ) ) <EOL> else : <EOL> self . idle_wait_count = <NUM_LIT:0> <EOL> except StopRequested : <EOL> pass <EOL> finally : <EOL> try : <EOL> self . log . debug ( "<STR_LIT>" ) <EOL> self . status = "<STR_LIT>" <EOL> self . gevent_pool . join ( timeout = None , raise_error = False ) <EOL> self . log . debug ( "<STR_LIT>" ) <EOL> except StopRequested : <EOL> pass <EOL> def work_stop ( self ) : <EOL> self . status = "<STR_LIT>" <EOL> self . gevent_pool . kill ( exception = JobInterrupt , block = True ) <EOL> for g in self . greenlets : <EOL> g_time = getattr ( self . greenlets [ g ] , "<STR_LIT>" , <NUM_LIT:0> ) <EOL> g_switches = getattr ( self . greenlets [ g ] , "<STR_LIT>" , None ) <EOL> self . greenlets [ g ] . kill ( block = True ) <EOL> self . log . debug ( <EOL> "<STR_LIT>" % <EOL> ( g , g_time , g_switches ) ) <EOL> self . status = "<STR_LIT>" <EOL> self . report_worker ( w = <NUM_LIT:1> ) <EOL> self . flush_logs ( w = <NUM_LIT:1> ) <EOL> g_time = getattr ( self . greenlet , "<STR_LIT>" , <NUM_LIT:0> ) <EOL> g_switches = getattr ( self . greenlet , "<STR_LIT>" , None ) <EOL> self . log . debug ( <EOL> "<STR_LIT>" % <EOL> ( g_time , g_switches ) ) <EOL> return self . exitcode <EOL> def perform_job ( self , job ) : <EOL> """<STR_LIT>""" <EOL> if self . config [ "<STR_LIT>" ] : <EOL> job . trace_memory_start ( ) <EOL> set_current_job ( job ) <EOL> gevent_timeout = None <EOL> if job . timeout : <EOL> gevent_timeout = gevent . Timeout ( <EOL> job . timeout , <EOL> TimeoutInterrupt ( <EOL> '<STR_LIT>' % <EOL> job . timeout <EOL> ) <EOL> ) <EOL> gevent_timeout . start ( ) <EOL> try : <EOL> job . perform ( ) <EOL> except RetryInterrupt : <EOL> self . log . error ( "<STR_LIT>" ) <EOL> job . save_retry ( sys . exc_info ( ) [ <NUM_LIT:1> ] ) <EOL> except MaxRetriesInterrupt : <EOL> self . log . error ( "<STR_LIT>" ) <EOL> job . _save_status ( "<STR_LIT>" , exception = True ) <EOL> except AbortInterrupt : <EOL> self . log . error ( "<STR_LIT>" ) <EOL> job . save_abort ( ) <EOL> except TimeoutInterrupt : <EOL> self . log . error ( "<STR_LIT>" % job . timeout ) <EOL> job . _save_status ( "<STR_LIT>" , exception = True ) <EOL> except JobInterrupt : <EOL> self . log . error ( "<STR_LIT>" ) <EOL> job . _save_status ( "<STR_LIT>" , exception = True ) <EOL> except Exception : <EOL> self . log . error ( "<STR_LIT>" ) <EOL> job . _save_status ( "<STR_LIT>" , exception = True ) <EOL> finally : <EOL> if gevent_timeout : <EOL> gevent_timeout . cancel ( ) <EOL> set_current_job ( None ) <EOL> self . done_jobs += <NUM_LIT:1> <EOL> if self . config [ "<STR_LIT>" ] : <EOL> job . trace_memory_stop ( ) <EOL> def shutdown_graceful ( self ) : <EOL> """<STR_LIT>""" <EOL> self . exitcode = <NUM_LIT:2> <EOL> self . log . info ( "<STR_LIT>" ) <EOL> raise StopRequested ( ) <EOL> def shutdown_max_memory ( self ) : <EOL> self . exitcode = <NUM_LIT:4> <EOL> self . log . info ( "<STR_LIT>" ) <EOL> self . graceful_stop = True <EOL> def shutdown_now ( self ) : <EOL> """<STR_LIT>""" <EOL> self . exitcode = <NUM_LIT:3> <EOL> self . log . info ( "<STR_LIT>" ) <EOL> self . status = "<STR_LIT>" <EOL> self . gevent_pool . kill ( exception = JobInterrupt , block = False ) <EOL> raise StopRequested ( ) <EOL> def install_signal_handlers ( self ) : <EOL> """<STR_LIT>""" <EOL> self . graceful_stop = False <EOL> def request_shutdown_now ( ) : <EOL> self . shutdown_now ( ) <EOL> def request_shutdown_graceful ( ) : <EOL> if self . graceful_stop : <EOL> request_shutdown_now ( ) <EOL> else : <EOL> self . graceful_stop = True <EOL> self . shutdown_graceful ( ) <EOL> gevent . signal ( signal . SIGINT , request_shutdown_graceful ) <EOL> gevent . signal ( signal . SIGTERM , request_shutdown_now ) </s>
