<s> import time <EOL> import heapq <EOL> import logging <EOL> import threading <EOL> try : <EOL> from UserDict import DictMixin <EOL> except ImportError : <EOL> from collections import Mapping as DictMixin <EOL> from . token_bucket import Bucket <EOL> from six . moves import queue as Queue <EOL> logger = logging . getLogger ( '<STR_LIT>' ) <EOL> try : <EOL> cmp <EOL> except NameError : <EOL> cmp = lambda x , y : ( x > y ) - ( x < y ) <EOL> class InQueueTask ( DictMixin ) : <EOL> __slots__ = ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ) <EOL> __getitem__ = lambda * x : getattr ( * x ) <EOL> __setitem__ = lambda * x : setattr ( * x ) <EOL> __iter__ = lambda self : iter ( self . __slots__ ) <EOL> __len__ = lambda self : len ( self . __slots__ ) <EOL> keys = lambda self : self . __slots__ <EOL> def __init__ ( self , taskid , priority = <NUM_LIT:0> , exetime = <NUM_LIT:0> ) : <EOL> self . taskid = taskid <EOL> self . priority = priority <EOL> self . exetime = exetime <EOL> def __cmp__ ( self , other ) : <EOL> if self . exetime == <NUM_LIT:0> and other . exetime == <NUM_LIT:0> : <EOL> return - cmp ( self . priority , other . priority ) <EOL> else : <EOL> return cmp ( self . exetime , other . exetime ) <EOL> def __lt__ ( self , other ) : <EOL> return self . __cmp__ ( other ) < <NUM_LIT:0> <EOL> class PriorityTaskQueue ( Queue . Queue ) : <EOL> '''<STR_LIT>''' <EOL> def _init ( self , maxsize ) : <EOL> self . queue = [ ] <EOL> self . queue_dict = dict ( ) <EOL> def _qsize ( self , len = len ) : <EOL> return len ( self . queue_dict ) <EOL> def _put ( self , item , heappush = heapq . heappush ) : <EOL> if item . taskid in self . queue_dict : <EOL> task = self . queue_dict [ item . taskid ] <EOL> changed = False <EOL> if item . priority > task . priority : <EOL> task . priority = item . priority <EOL> changed = True <EOL> if item . exetime < task . exetime : <EOL> task . exetime = item . exetime <EOL> changed = True <EOL> if changed : <EOL> self . _resort ( ) <EOL> else : <EOL> heappush ( self . queue , item ) <EOL> self . queue_dict [ item . taskid ] = item <EOL> def _get ( self , heappop = heapq . heappop ) : <EOL> while self . queue : <EOL> item = heappop ( self . queue ) <EOL> if item . taskid is None : <EOL> continue <EOL> self . queue_dict . pop ( item . taskid , None ) <EOL> return item <EOL> return None <EOL> @ property <EOL> def top ( self ) : <EOL> while self . queue and self . queue [ <NUM_LIT:0> ] . taskid is None : <EOL> heapq . heappop ( self . queue ) <EOL> if self . queue : <EOL> return self . queue [ <NUM_LIT:0> ] <EOL> return None <EOL> def _resort ( self ) : <EOL> heapq . heapify ( self . queue ) <EOL> def __contains__ ( self , taskid ) : <EOL> return taskid in self . queue_dict <EOL> def __getitem__ ( self , taskid ) : <EOL> return self . queue_dict [ taskid ] <EOL> def __setitem__ ( self , taskid , item ) : <EOL> assert item . taskid == taskid <EOL> self . put ( item ) <EOL> def __delitem__ ( self , taskid ) : <EOL> self . queue_dict . pop ( taskid ) . taskid = None <EOL> class TaskQueue ( object ) : <EOL> '''<STR_LIT>''' <EOL> processing_timeout = <NUM_LIT:10> * <NUM_LIT> <EOL> def __init__ ( self , rate = <NUM_LIT:0> , burst = <NUM_LIT:0> ) : <EOL> self . mutex = threading . RLock ( ) <EOL> self . priority_queue = PriorityTaskQueue ( ) <EOL> self . time_queue = PriorityTaskQueue ( ) <EOL> self . processing = PriorityTaskQueue ( ) <EOL> self . bucket = Bucket ( rate = rate , burst = burst ) <EOL> @ property <EOL> def rate ( self ) : <EOL> return self . bucket . rate <EOL> @ rate . setter <EOL> def rate ( self , value ) : <EOL> self . bucket . rate = value <EOL> @ property <EOL> def burst ( self ) : <EOL> return self . burst . burst <EOL> @ burst . setter <EOL> def burst ( self , value ) : <EOL> self . bucket . burst = value <EOL> def check_update ( self ) : <EOL> '''<STR_LIT>''' <EOL> self . _check_time_queue ( ) <EOL> self . _check_processing ( ) <EOL> def _check_time_queue ( self ) : <EOL> now = time . time ( ) <EOL> self . mutex . acquire ( ) <EOL> while self . time_queue . qsize ( ) and self . time_queue . top and self . time_queue . top . exetime < now : <EOL> task = self . time_queue . get_nowait ( ) <EOL> task . exetime = <NUM_LIT:0> <EOL> self . priority_queue . put ( task ) <EOL> self . mutex . release ( ) <EOL> def _check_processing ( self ) : <EOL> now = time . time ( ) <EOL> self . mutex . acquire ( ) <EOL> while self . processing . qsize ( ) and self . processing . top and self . processing . top . exetime < now : <EOL> task = self . processing . get_nowait ( ) <EOL> if task . taskid is None : <EOL> continue <EOL> task . exetime = <NUM_LIT:0> <EOL> self . priority_queue . put ( task ) <EOL> logger . info ( "<STR_LIT>" , task . taskid ) <EOL> self . mutex . release ( ) <EOL> def put ( self , taskid , priority = <NUM_LIT:0> , exetime = <NUM_LIT:0> ) : <EOL> '''<STR_LIT>''' <EOL> now = time . time ( ) <EOL> task = InQueueTask ( taskid , priority , exetime ) <EOL> self . mutex . acquire ( ) <EOL> if taskid in self . priority_queue : <EOL> self . priority_queue . put ( task ) <EOL> elif taskid in self . time_queue : <EOL> self . time_queue . put ( task ) <EOL> elif taskid in self . processing and self . processing [ taskid ] . taskid : <EOL> pass <EOL> else : <EOL> if exetime and exetime > now : <EOL> self . time_queue . put ( task ) <EOL> else : <EOL> self . priority_queue . put ( task ) <EOL> self . mutex . release ( ) <EOL> def get ( self ) : <EOL> '''<STR_LIT>''' <EOL> if self . bucket . get ( ) < <NUM_LIT:1> : <EOL> return None <EOL> now = time . time ( ) <EOL> self . mutex . acquire ( ) <EOL> try : <EOL> task = self . priority_queue . get_nowait ( ) <EOL> self . bucket . desc ( ) <EOL> except Queue . Empty : <EOL> self . mutex . release ( ) <EOL> return None <EOL> task . exetime = now + self . processing_timeout <EOL> self . processing . put ( task ) <EOL> self . mutex . release ( ) <EOL> return task . taskid <EOL> def done ( self , taskid ) : <EOL> '''<STR_LIT>''' <EOL> if taskid in self . processing : <EOL> self . mutex . acquire ( ) <EOL> if taskid in self . processing : <EOL> del self . processing [ taskid ] <EOL> self . mutex . release ( ) <EOL> return True <EOL> return False <EOL> def size ( self ) : <EOL> return self . priority_queue . qsize ( ) + self . time_queue . qsize ( ) + self . processing . qsize ( ) <EOL> def __len__ ( self ) : <EOL> return self . size ( ) <EOL> def __contains__ ( self , taskid ) : <EOL> if taskid in self . priority_queue or taskid in self . time_queue : <EOL> return True <EOL> if taskid in self . processing and self . processing [ taskid ] . taskid : <EOL> return True <EOL> return False <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> task_queue = TaskQueue ( ) <EOL> task_queue . processing_timeout = <NUM_LIT:0.1> <EOL> task_queue . put ( '<STR_LIT>' , <NUM_LIT:3> , time . time ( ) + <NUM_LIT:0.1> ) <EOL> task_queue . put ( '<STR_LIT>' , <NUM_LIT:1> ) <EOL> task_queue . put ( '<STR_LIT>' , <NUM_LIT:2> ) <EOL> assert task_queue . get ( ) == '<STR_LIT>' <EOL> time . sleep ( <NUM_LIT:0.1> ) <EOL> task_queue . _check_time_queue ( ) <EOL> assert task_queue . get ( ) == '<STR_LIT>' <EOL> assert task_queue . get ( ) == '<STR_LIT>' <EOL> task_queue . _check_processing ( ) <EOL> assert task_queue . get ( ) == '<STR_LIT>' <EOL> assert len ( task_queue ) == <NUM_LIT:0> </s>
