def __init__(self, channel, prefetch_count=0): self.channel = channel self.prefetch_count = prefetch_count or 0 self._delivered = OrderedDict() self._delivered.restored = False self._dirty = set() self._quick_ack = self._dirty.add self._quick_append = self._delivered.__setitem__ self._on_collect = Finalize(self, self.restore_unacked_once, exitpriority=1)
def __init__(self, address, backlog=None): self._address = address self._handle_queue = [self._new_handle(first=True)] self._last_accepted = None sub_debug('listener created with address=%r', self._address) self.close = Finalize(self, PipeListener._finalize_pipe_listener, args=(self._handle_queue, self._address), exitpriority=0)
def __init__(self, address, backlog=None): self._address = address handle = win32.CreateNamedPipe(address, win32.PIPE_ACCESS_DUPLEX, win32.PIPE_TYPE_MESSAGE | win32.PIPE_READMODE_MESSAGE | win32.PIPE_WAIT, win32.PIPE_UNLIMITED_INSTANCES, BUFSIZE, BUFSIZE, win32.NMPWAIT_WAIT_FOREVER, win32.NULL) self._handle_queue = [ handle] self._last_accepted = None sub_debug('listener created with address=%r', self._address) self.close = Finalize(self, PipeListener._finalize_pipe_listener, args=( self._handle_queue, self._address), exitpriority=0) return
def build_pool(nproc, log_file_name=None): """ Make a multiprocessing pool object and a shared log (either a file or stdout) object. nproc - number of processes log_file_name - name/None returns pool, log """ print('Building pool with', nproc, 'processors') lock = mp.Lock() # a lock we can prevent collisions with proc_counter = mp.Value('i', 0) if log_file_name is None: buf = sys.stdout else: if os.path.exists(log_file_name): os.remove(log_file_name) buf = open(log_file_name, 'a') # Acquire a lock before writing a line buf = LockedLineIO([buf], lock) log = PrefixedIO(buf, str(' ' * 1)) log = VerboseTimingLog(bufs=[log], also_stdout=False, insert_timings=True) # make as many pipes as processors (pass all-to-all, will dynamically use) connections = [mp.Pipe() for i in range(nproc)] slave_connections = [c[1] for c in connections] my_connections = [c[0] for c in connections] t = time() pool = mp.Pool(nproc, initializer=init_worker, initargs=(lock, proc_counter, t, log_file_name, slave_connections)) for i in range(nproc): obj = my_connections[0].recv() print(obj, file=log) print('Built worker pool with {:,} processors'.format(proc_counter.value), file=log) global _myproc _myproc = MyProc(log, lock, t, my_connections) global _free_connections _free_connections = list(range(nproc)) # prevent exceptions on teardown by cleaning up logs def cleanup_log(): """ Cleanup the logs """ get_log().close() Finalize(None, cleanup_log, exitpriority=16) return pool, log
def __init__(self, processes=None, initializer=None, initargs=(), maxtasksperchild=None): self._setup_queues() self._taskqueue = Queue.Queue() self._cache = {} self._state = RUN self._maxtasksperchild = maxtasksperchild self._initializer = initializer self._initargs = initargs if processes is None: try: processes = cpu_count() except NotImplementedError: processes = 1 if initializer is not None and not hasattr(initializer, '__call__'): raise TypeError('initializer must be a callable') self._processes = processes self._pool = [] self._repopulate_pool() self._worker_handler = threading.Thread( target=Pool._handle_workers, args=(self, ) ) self._worker_handler.daemon = True self._worker_handler._state = RUN self._worker_handler.start() self._task_handler = threading.Thread( target=Pool._handle_tasks, args=(self._taskqueue, self._quick_put, self._outqueue, self._pool) ) self._task_handler.daemon = True self._task_handler._state = RUN self._task_handler.start() self._result_handler = threading.Thread( target=Pool._handle_results, args=(self._outqueue, self._quick_get, self._cache) ) self._result_handler.daemon = True self._result_handler._state = RUN self._result_handler.start() self._terminate = Finalize( self, self._terminate_pool, args=(self._taskqueue, self._inqueue, self._outqueue, self._pool, self._worker_handler, self._task_handler, self._result_handler, self._cache), exitpriority=15 )
def __init__(self, processes=None, initializer=None, initargs=()): self._setup_queues() self._taskqueue = Queue.Queue() self._cache = {} self._state = RUN self._initializer = initializer self._initargs = initargs if processes is None: try: processes = cpu_count() except NotImplementedError: processes = 1 self._size = processes if initializer is not None and not hasattr(initializer, '__call__'): raise TypeError('initializer must be a callable') self._pool = [] for i in range(processes): self._create_worker_process() self._task_handler = threading.Thread( target=Pool._handle_tasks, args=(self._taskqueue, self._quick_put, self._outqueue, self._pool)) self._task_handler.daemon = True self._task_handler._state = RUN self._task_handler.start() # Thread processing acknowledgements form the ackqueue. self._ack_handler = threading.Thread(target=Pool._handle_ack, args=(self._ackqueue, self._quick_get_ack, self._cache)) self._ack_handler.daemon = True self._ack_handler._state = RUN self._ack_handler.start() # Thread processing results in the outqueue. self._result_handler = threading.Thread(target=Pool._handle_results, args=(self._outqueue, self._quick_get, self._cache)) self._result_handler.daemon = True self._result_handler._state = RUN self._result_handler.start() self._terminate = Finalize( self, self._terminate_pool, args=(self._taskqueue, self._inqueue, self._outqueue, self._ackqueue, self._pool, self._ack_handler, self._task_handler, self._result_handler, self._cache), exitpriority=15)
def __init__(self, file_, callbackfn=None, desc='daemon lock', debug=False): self.pidfile = file_ self.callbackfn = callbackfn self.desc = desc self.debug = debug self.held = False # run the lock automatically! self.lock() self._finalize = Finalize(self, DaemonLock._on_finalize, args=(self, debug), exitpriority=10)
def pool_init(): global conn global api conn = pymysql.connect(host='localhost', user='******', password='******', db='vk17', charset='utf8', cursorclass=pymysql.cursors.DictCursor) session = vk.Session() api = vk.API(session, v='5.62', lang='ru') Finalize(conn, pool_close, exitpriority=0)
def get_temp_dir(): ''' Function returning a temp directory which will be removed on exit ''' # get name of a temp directory which will be automatically cleaned up if current_process()._tempdir is None: tempdir = tempfile.mkdtemp(prefix='pymp-') info('created temp directory %s', tempdir) Finalize(None, shutil.rmtree, args=[tempdir], exitpriority=-100) current_process()._tempdir = tempdir return current_process()._tempdir
def _start_thread( self): # Start thread which transfers data from buffer to pipe self._buffer.clear() self._thread = threading.Thread( target=Queue._feed, args=(self._buffer, self._notempty, self._send_bytes, self._wlock, self._writer.close, self._ignore_epipe, self._sem), name='QueueFeederThread', daemon=True, ) self._thread.start() if not self._joincancelled: self._jointhread = Finalize(self._thread, Queue._finalize_join, [weakref.ref(self._thread)], exitpriority=-5) self._close = Finalize( self, Queue._finalize_close, [self._buffer, self._notempty], exitpriority=10 ) # Send sentinel to the thread queue object when garbage collected
def _start_thread(self): debug('Queue._start_thread()') self._buffer.clear() self._thread = threading.Thread(target=Queue._feed, args=(self._buffer, self._notempty, self._send, self._wlock, self._writer.close), name='QueueFeederThread') self._thread.daemon = True debug('doing self._thread.start()') self._thread.start() debug('... done self._thread.start()') if not self._joincancelled: self._jointhread = Finalize(self._thread, Queue._finalize_join, [weakref.ref(self._thread)], exitpriority=-5) self._close = Finalize(self, Queue._finalize_close, [self._buffer, self._notempty], exitpriority=10)
def func_after_fork(tracer): def exit_routine(): self.exit_routine() from multiprocessing.util import Finalize import signal Finalize(tracer, exit_routine, exitpriority=32) def term_handler(signalnum, frame): self.exit_routine() signal.signal(signal.SIGTERM, term_handler)
def make_instance(stages, context, store): SequentialWorker.instance = SequentialWorker(stages, context, global_store=store) SequentialWorker.instance.__enter__() # https://stackoverflow.com/a/24724452 Finalize( SequentialWorker.instance, SequentialWorker.instance.__exit__, args=(None, None, None), exitpriority=10, )
def __init__(self, address, family, backlog=1): self._socket = socket.socket(getattr(socket, family)) self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) self._socket.bind(address) self._socket.listen(backlog) self._address = self._socket.getsockname() self._family = family self._last_accepted = None if family == 'AF_UNIX': self._unlink = Finalize(self, os.unlink, args=(address,), exitpriority=0) else: self._unlink = None return
def run(self): ''' Bind the pub and pull sockets for events ''' salt.utils.appendproctitle(self.__class__.__name__) self.io_loop = tornado.ioloop.IOLoop() with salt.utils. async .current_ioloop(self.io_loop): if self.opts['ipc_mode'] == 'tcp': epub_uri = int(self.opts['tcp_master_pub_port']) epull_uri = int(self.opts['tcp_master_pull_port']) else: epub_uri = os.path.join(self.opts['sock_dir'], 'master_event_pub.ipc') epull_uri = os.path.join(self.opts['sock_dir'], 'master_event_pull.ipc') self.publisher = salt.transport.ipc.IPCMessagePublisher( epub_uri, io_loop=self.io_loop) self.puller = salt.transport.ipc.IPCMessageServer( epull_uri, io_loop=self.io_loop, payload_handler=self.handle_publish, ) # Start the master event publisher old_umask = os.umask(0o177) try: self.publisher.start() self.puller.start() if self.opts['client_acl'] or self.opts['client_acl_blacklist']: salt.utils.warn_until( 'Nitrogen', 'ACL rules should be configured with \'publisher_acl\' and ' '\'publisher_acl_blacklist\' not \'client_acl\' and ' '\'client_acl_blacklist\'. This functionality will be removed in Salt ' 'Nitrogen.') if (self.opts['ipc_mode'] != 'tcp' and (self.opts['publisher_acl'] or self.opts['client_acl'] or self.opts['external_auth'])): os.chmod( os.path.join(self.opts['sock_dir'], 'master_event_pub.ipc'), 0o666) finally: os.umask(old_umask) # Make sure the IO loop and respective sockets are closed and # destroyed Finalize(self, self.close, exitpriority=15) self.io_loop.start()
def _init_worker(*args, **kwargs): """Initialize a :class:`multiprocessing.Pool` worker. Call the Django's ``ParallelTestSuite.init_worker`` and then also start the manager infrastructure. """ result = django_init_worker(*args, **kwargs) # Further patch channel names and the like with our current pid, # so that parallel managers and executors don't clash on the # same channels and directories. resolwe_settings.FLOW_MANAGER_SETTINGS['REDIS_PREFIX'] += '-parallel-pid{}'.format(os.getpid()) _create_test_dirs() overrides = _prepare_settings() overrides.__enter__() Finalize(overrides, lambda: overrides.__exit__(None, None, None), exitpriority=16) _manager_setup() state_cleanup = AtScopeExit(manager.state.destroy_channels) state_cleanup.__enter__() Finalize(state_cleanup, state_cleanup.__exit__, exitpriority=16) listener = CommandContext('runlistener', '--clear-queue') listener.__enter__() Finalize(listener, listener.__exit__, exitpriority=16) workers = CommandContext('runworker', only_channels=[state.MANAGER_CONTROL_CHANNEL]) workers.__enter__() Finalize(workers, workers.__exit__, exitpriority=16) signal_override = override_settings(FLOW_MANAGER_SYNC_AUTO_CALLS=True) signal_override.__enter__() Finalize(signal_override, lambda: signal_override.__exit__(None, None, None), exitpriority=16) return result
def __init__(self, *args, **kwargs): """Initialize the database scheduler.""" self.app = kwargs['app'] self.dburi = kwargs.get('dburi') or self.app.conf.get( 'beat_dburi') or DEFAULT_BEAT_DBURI self.engine, self.Session = session_manager.create_session(self.dburi) session_manager.prepare_models(self.engine) self._dirty = set() Scheduler.__init__(self, *args, **kwargs) self._finalize = Finalize(self, self.sync, exitpriority=5) self.max_interval = (kwargs.get('max_interval') or self.app.conf.beat_max_loop_interval or DEFAULT_MAX_INTERVAL)
def sqp_initializer(*args): """The initializer used by the Scrape Quotes Pool""" # Preclude the need for worker processes to even care about KeyboardInterrupt in the first place. signal.signal(signal.SIGINT, signal.SIG_IGN) global driver_manager driver_creator = args[0] scraping_config = args[1] driver_manager = DriverManager(driver_creator, scraping_config) driver_manager.__enter__() Finalize(driver_manager, driver_manager.__exit__, exitpriority=16) log_queue = args[2] configure_log_dispatcher(log_queue)
def run(self): ''' Bind the pub and pull sockets for events ''' salt.utils.process.appendproctitle(self.__class__.__name__) self.io_loop = tornado.ioloop.IOLoop() with salt.utils.asynchronous.current_ioloop(self.io_loop): if self.opts['ipc_mode'] == 'tcp': epub_uri = int(self.opts['tcp_master_pub_port']) epull_uri = int(self.opts['tcp_master_pull_port']) else: epub_uri = os.path.join( self.opts['sock_dir'], 'master_event_pub.ipc' ) epull_uri = os.path.join( self.opts['sock_dir'], 'master_event_pull.ipc' ) self.publisher = salt.transport.ipc.IPCMessagePublisher( self.opts, epub_uri, io_loop=self.io_loop ) self.puller = salt.transport.ipc.IPCMessageServer( epull_uri, io_loop=self.io_loop, payload_handler=self.handle_publish, ) # Start the master event publisher old_umask = os.umask(0o177) try: self.publisher.start() self.puller.start() if (self.opts['ipc_mode'] != 'tcp' and ( self.opts['publisher_acl'] or self.opts['external_auth'])): os.chmod(os.path.join( self.opts['sock_dir'], 'master_event_pub.ipc'), 0o666) finally: os.umask(old_umask) # Make sure the IO loop and respective sockets are closed and # destroyed Finalize(self, self.close, exitpriority=15) self.io_loop.start()
def __init__(self, file_=None, callbackfn=None, desc='daemon lock', debug=False): self.pidfile = file_ if file_ else os.path.join( os.path.dirname(__file__), 'running.lock') self.callbackfn = callbackfn self.desc = desc self.debug = debug self.held = False #run the lock automatically ! self.lock() self._finalize = Finalize(self, DaemonLock._on_finalize, args=(self, debug), exitpriority=10)
def _start_thread(self): debug('Queue._start_thread()') # Start thread which transfers data from buffer to pipe self._buffer.clear() self._thread = threading.Thread( target=Queue._feed, args=(self._buffer, self._notempty, self._send, self._wlock, self._writer.close), name='QueueFeederThread' ) self._thread.daemon = True debug('doing self._thread.start()') self._thread.start() debug('... done self._thread.start()') # On process exit we will wait for data to be flushed to pipe. # # However, if this process created the queue then all # processes which use the queue will be descendants of this # process. Therefore waiting for the queue to be flushed # is pointless once all the child processes have been joined. created_by_this_process = (self._opid == os.getpid()) if not self._joincancelled and not created_by_this_process: self._jointhread = Finalize( self._thread, Queue._finalize_join, [weakref.ref(self._thread)], exitpriority=-5 ) # Send sentinel to the thread queue object when garbage collected self._close = Finalize( self, Queue._finalize_close, [self._buffer, self._notempty], exitpriority=10 )
def extractor_pool(self, func, iterable): ''' Extract items (billard multiprocessing use) :param func: function :param iterable: list ''' _finalizers = list() p = Pool(processes=cpu_count()) _finalizers.append(Finalize(p, p.terminate)) try: p.map_async(func, iterable) p.close() p.join() finally: p.terminate()
def _start_thread(self): debug("Queue._start_thread()") # Start thread which transfers data from buffer to pipe self._buffer.clear() self._thread = threading.Thread( target=MPQueueFixed._feed, args=( self._buffer, self._notempty, self._send, self._wlock, self._writer.close, ), name="QueueFeederThread", ) self._thread.daemon = True debug("doing self._thread.start()") self._thread.start() debug("... done self._thread.start()") # On process exit we will wait for data to be flushed to pipe. if not self._joincancelled: self._jointhread = Finalize( self._thread, Queue._finalize_join, [weakref.ref(self._thread)], exitpriority=-5, ) # Send sentinel to the thread queue object when garbage collected self._close = Finalize(self, Queue._finalize_close, [self._buffer, self._notempty], exitpriority=10)
def __init__(self, *args, **kwargs): database_proxy.initialize(kwargs['app'].database) self.database_proxy = database_proxy """Initialize the database scheduler.""" self._dirty = set() Scheduler.__init__(self, *args, **kwargs) self._finalize = Finalize(self, self.sync, exitpriority=5) self.max_interval = (kwargs.get('max_interval') or self.app.conf.beat_max_loop_interval or DEFAULT_MAX_INTERVAL) self.database_proxy.create_tables([ PeriodicTask, PeriodicTasks, CrontabSchedule, IntervalSchedule, SolarSchedule ], safe=True)
def __init__(self, processes=None, initializer=None, initargs=()): self._setup_queues() self._taskqueue = queue.Queue() self._cache = {} self._state = RUN if processes is None: try: processes = cpu_count() except NotImplementedError: processes = 1 if initializer is not None and not hasattr(initializer, '__call__'): raise TypeError('initializer must be a callable') self._pool = [] for i in range(processes): w = self.Process( target=worker, args=(self._inqueue, self._outqueue, initializer, initargs) ) self._pool.append(w) w.name = w.name.replace('Process', 'PoolWorker') w.daemon = True w.start() self._task_handler = threading.Thread( target=Pool._handle_tasks, args=(self._taskqueue, self._quick_put, self._outqueue, self._pool) ) self._task_handler.daemon = True self._task_handler._state = RUN self._task_handler.start() self._result_handler = threading.Thread( target=Pool._handle_results, args=(self._outqueue, self._quick_get, self._cache) ) self._result_handler.daemon = True self._result_handler._state = RUN self._result_handler.start() self._terminate = Finalize( self, self._terminate_pool, args=(self._taskqueue, self._inqueue, self._outqueue, self._pool, self._task_handler, self._result_handler, self._cache), exitpriority=15 )
def start_monitor_thread(self): "Run check_subprocesses in a seperate thread. Kill it gracefully." from threading import Thread t = Thread(target=self.check_subprocesses) t.setDaemon(True) t.start() def kill_thread(self): "Let the child process checker finish" self.__stopping = True t.join() self.__term = Finalize(self, kill_thread, args=(self, ), exitpriority=20)
def func_after_fork(tracer): def exit_routine(): self.exit_routine() from multiprocessing.util import Finalize import signal Finalize(tracer, exit_routine, exitpriority=32) def term_handler(signalnum, frame): self.exit_routine() signal.signal(signal.SIGTERM, term_handler) tracer.clear() tracer._set_curr_stack_depth(1) if tracer._afterfork_cb: tracer._afterfork_cb(tracer, *tracer._afterfork_args, **tracer._afterfork_kwargs)
def worker_init_entry_point( init_function, arg_queue=None, backup_arg_queue=None): kwargs = {} if arg_queue: try: kwargs = arg_queue.get(block=False) except queue.Empty: print("Argument queue empty. Using round robin arg queue.") kwargs = backup_arg_queue.get(block=True) backup_arg_queue.put(kwargs) # On exit we add the init args back to the queue so restarted workers # (e.g. when when running with maxtasksperchild) will pickup init # arguments from a previously exited worker. Finalize(None, arg_queue.put, (kwargs,), exitpriority=1) print("Initializing worker: %s" % str(kwargs)) init_function(**kwargs)
def init_parallel_runner(*args): """ This initializes a new runner in each child process that calls it. We need to make the runner global so we can access it from compute_for_sticker_mp / compute_for_event_mp Clean up from: https://stackoverflow.com/questions/24717468/context-managers-and-multiprocessing-pools/24724452#24724452 :param args: :return: """ global parallel_runner parallel_runner = args[0](extra_feature_classes=args[1], **args[2]) try: current_niceness = os.nice(0) os.nice(current_niceness - 10) except: pass Finalize(object, clean_up_mp, exitpriority=16)
def pool_init(access_tokens, shared_countries, shared_cities, shared_count): global conn global api global countries global cities global count conn = pymysql.connect(host='localhost', user='******', password='******', db='vk17', charset='utf8') pname, _, pid = current_process().name.partition('-') access_token = access_tokens[int(pid) % len(access_tokens)] session = vk.Session(access_token=access_token) api = vk.API(session, v='5.62', lang='ru') countries = shared_countries cities = shared_cities count = shared_count Finalize(conn, conn.close, exitpriority=0)
def setup_worker(f, *args): global resource url, schema = args f.schema = schema Session = sessionmaker() f.session = Session() resource_cm = open_resource(url, schema) E.debug(f"{os.getpid()}: setting up worker for resource={id(resource)}") old_resource = resource resource = resource_cm.__enter__() E.debug(f"{os.getpid()}: new worker for resource={id(resource)} (old_resource={id(old_resource)})") # Register a finalizer to flush table cache Finalize(resource, resource.__exit__, exitpriority=16) E.debug(f"{os.getpid()}: adding cache={id(resource.table_cache)} from resource={id(resource)} " f"to worker={id(f)}, session={id(f.session)}") f.table_cache = resource.table_cache
class QoS(object): """Quality of Service guarantees. Only supports `prefetch_count` at this point. Arguments: channel (ChannelT): Connection channel. prefetch_count (int): Initial prefetch count (defaults to 0). """ #: current prefetch count value prefetch_count = 0 #: :class:`~collections.OrderedDict` of active messages. #: *NOTE*: Can only be modified by the consuming thread. _delivered = None #: acks can be done by other threads than the consuming thread. #: Instead of a mutex, which doesn't perform well here, we mark #: the delivery tags as dirty, so subsequent calls to append() can remove #: them. _dirty = None #: If disabled, unacked messages won't be restored at shutdown. restore_at_shutdown = True def __init__(self, channel, prefetch_count=0): self.channel = channel self.prefetch_count = prefetch_count or 0 self._delivered = OrderedDict() self._delivered.restored = False self._dirty = set() self._quick_ack = self._dirty.add self._quick_append = self._delivered.__setitem__ self._on_collect = Finalize( self, self.restore_unacked_once, exitpriority=1, ) def can_consume(self): """Return true if the channel can be consumed from. Used to ensure the client adhers to currently active prefetch limits. """ pcount = self.prefetch_count return not pcount or len(self._delivered) - len(self._dirty) < pcount def can_consume_max_estimate(self): """Return the maximum number of messages allowed to be returned. Returns an estimated number of messages that a consumer may be allowed to consume at once from the broker. This is used for services where bulk 'get message' calls are preferred to many individual 'get message' calls - like SQS. Returns: int: greater than zero. """ pcount = self.prefetch_count if pcount: return max(pcount - (len(self._delivered) - len(self._dirty)), 0) def append(self, message, delivery_tag): """Append message to transactional state.""" if self._dirty: self._flush() self._quick_append(delivery_tag, message) def get(self, delivery_tag): return self._delivered[delivery_tag] def _flush(self): """Flush dirty (acked/rejected) tags from.""" dirty = self._dirty delivered = self._delivered while 1: try: dirty_tag = dirty.pop() except KeyError: break delivered.pop(dirty_tag, None) def ack(self, delivery_tag): """Acknowledge message and remove from transactional state.""" self._quick_ack(delivery_tag) def reject(self, delivery_tag, requeue=False): """Remove from transactional state and requeue message.""" if requeue: self.channel._restore_at_beginning(self._delivered[delivery_tag]) self._quick_ack(delivery_tag) def restore_unacked(self): """Restore all unacknowledged messages.""" self._flush() delivered = self._delivered errors = [] restore = self.channel._restore pop_message = delivered.popitem while delivered: try: _, message = pop_message() except KeyError: # pragma: no cover break try: restore(message) except BaseException as exc: errors.append((exc, message)) delivered.clear() return errors def restore_unacked_once(self, stderr=None): """Restore all unacknowledged messages at shutdown/gc collect. Note: Can only be called once for each instance, subsequent calls will be ignored. """ self._on_collect.cancel() self._flush() stderr = sys.stderr if stderr is None else stderr state = self._delivered if not self.restore_at_shutdown or not self.channel.do_restore: return if getattr(state, 'restored', None): assert not state return try: if state: print(RESTORING_FMT.format(len(self._delivered)), file=stderr) unrestored = self.restore_unacked() if unrestored: errors, messages = list(zip(*unrestored)) print(RESTORE_PANIC_FMT.format(len(errors), errors), file=stderr) emergency_dump_state(messages, stderr=stderr) finally: state.restored = True def restore_visible(self, *args, **kwargs): """Restore any pending unackwnowledged messages. To be filled in for visibility_timeout style implementations. Note: This is implementation optional, and currently only used by the Redis transport. """ pass
class QoS(object): """Quality of Service guarantees. Only supports `prefetch_count` at this point. :param channel: AMQ Channel. :keyword prefetch_count: Initial prefetch count (defaults to 0). """ #: current prefetch count value prefetch_count = 0 #: :class:`~collections.OrderedDict` of active messages. #: *NOTE*: Can only be modified by the consuming thread. _delivered = None #: acks can be done by other threads than the consuming thread. #: Instead of a mutex, which doesn't perform well here, we mark #: the delivery tags as dirty, so subsequent calls to append() can remove #: them. _dirty = None #: If disabled, unacked messages won't be restored at shutdown. restore_at_shutdown = True def __init__(self, channel, prefetch_count=0): self.channel = channel self.prefetch_count = prefetch_count or 0 self._delivered = OrderedDict() self._delivered.restored = False self._dirty = set() self._quick_ack = self._dirty.add self._quick_append = self._delivered.__setitem__ self._on_collect = Finalize( self, self.restore_unacked_once, exitpriority=1, ) def can_consume(self): """Return true if the channel can be consumed from. Used to ensure the client adhers to currently active prefetch limits. """ pcount = self.prefetch_count return not pcount or len(self._delivered) - len(self._dirty) < pcount def append(self, message, delivery_tag): """Append message to transactional state.""" if self._dirty: self._flush() self._quick_append(delivery_tag, message) def get(self, delivery_tag): return self._delivered[delivery_tag] def _flush(self): """Flush dirty (acked/rejected) tags from.""" dirty = self._dirty delivered = self._delivered while 1: try: dirty_tag = dirty.pop() except KeyError: break delivered.pop(dirty_tag, None) def ack(self, delivery_tag): """Acknowledge message and remove from transactional state.""" self._quick_ack(delivery_tag) def reject(self, delivery_tag, requeue=False): """Remove from transactional state and requeue message.""" if requeue: self.channel._restore_at_beginning(self._delivered[delivery_tag]) self._quick_ack(delivery_tag) def restore_unacked(self): """Restore all unacknowledged messages.""" self._flush() delivered = self._delivered errors = [] restore = self.channel._restore pop_message = delivered.popitem while delivered: try: _, message = pop_message() except KeyError: # pragma: no cover break try: restore(message) except BaseException as exc: errors.append((exc, message)) delivered.clear() return errors def restore_unacked_once(self): """Restores all unacknowledged messages at shutdown/gc collect. Will only be done once for each instance. """ self._on_collect.cancel() self._flush() state = self._delivered if not self.restore_at_shutdown or not self.channel.do_restore: return if getattr(state, 'restored', None): assert not state return try: if state: say('Restoring {0!r} unacknowledged message(s).', len(self._delivered)) unrestored = self.restore_unacked() if unrestored: errors, messages = list(zip(*unrestored)) say('UNABLE TO RESTORE {0} MESSAGES: {1}', len(errors), errors) emergency_dump_state(messages) finally: state.restored = True