asynpool.py 46 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256
  1. # -*- coding: utf-8 -*-
  2. """
  3. celery.concurrency.asynpool
  4. ~~~~~~~~~~~~~~~~~~~~~~~~~~~
  5. .. note::
  6. This module will be moved soon, so don't use it directly.
  7. Non-blocking version of :class:`multiprocessing.Pool`.
  8. This code deals with three major challenges:
  9. 1) Starting up child processes and keeping them running.
  10. 2) Sending jobs to the processes and receiving results back.
  11. 3) Safely shutting down this system.
  12. """
  13. from __future__ import absolute_import
  14. import errno
  15. import os
  16. import select
  17. import socket
  18. import struct
  19. import sys
  20. import time
  21. from collections import deque, namedtuple
  22. from io import BytesIO
  23. from numbers import Integral
  24. from pickle import HIGHEST_PROTOCOL
  25. from time import sleep
  26. from weakref import WeakValueDictionary, ref
  27. from amqp.utils import promise
  28. from billiard.pool import RUN, TERMINATE, ACK, NACK, WorkersJoined
  29. from billiard import pool as _pool
  30. from billiard.compat import buf_t, setblocking, isblocking
  31. from billiard.queues import _SimpleQueue
  32. from kombu.async import READ, WRITE, ERR
  33. from kombu.serialization import pickle as _pickle
  34. from kombu.utils import fxrange
  35. from kombu.utils.eventio import SELECT_BAD_FD
  36. from celery.five import Counter, items, values
  37. from celery.utils.log import get_logger
  38. from celery.worker import state as worker_state
  39. try:
  40. from _billiard import read as __read__
  41. from struct import unpack_from as _unpack_from
  42. memoryview = memoryview
  43. readcanbuf = True
  44. if sys.version_info[0] == 2 and sys.version_info < (2, 7, 6):
  45. def unpack_from(fmt, view, _unpack_from=_unpack_from): # noqa
  46. return _unpack_from(fmt, view.tobytes()) # <- memoryview
  47. else:
  48. # unpack_from supports memoryview in 2.7.6 and 3.3+
  49. unpack_from = _unpack_from # noqa
  50. except (ImportError, NameError): # pragma: no cover
  51. def __read__(fd, buf, size, read=os.read): # noqa
  52. chunk = read(fd, size)
  53. n = len(chunk)
  54. if n != 0:
  55. buf.write(chunk)
  56. return n
  57. readcanbuf = False # noqa
  58. def unpack_from(fmt, iobuf, unpack=struct.unpack): # noqa
  59. return unpack(fmt, iobuf.getvalue()) # <-- BytesIO
  60. logger = get_logger(__name__)
  61. error, debug = logger.error, logger.debug
  62. UNAVAIL = frozenset({errno.EAGAIN, errno.EINTR})
  63. #: Constant sent by child process when started (ready to accept work)
  64. WORKER_UP = 15
  65. #: A process must have started before this timeout (in secs.) expires.
  66. PROC_ALIVE_TIMEOUT = 4.0
  67. SCHED_STRATEGY_PREFETCH = 1
  68. SCHED_STRATEGY_FAIR = 4
  69. SCHED_STRATEGIES = {
  70. None: SCHED_STRATEGY_PREFETCH,
  71. 'fair': SCHED_STRATEGY_FAIR,
  72. }
  73. Ack = namedtuple('Ack', ('id', 'fd', 'payload'))
  74. def gen_not_started(gen):
  75. # gi_frame is None when generator stopped.
  76. return gen.gi_frame and gen.gi_frame.f_lasti == -1
  77. def _get_job_writer(job):
  78. try:
  79. writer = job._writer
  80. except AttributeError:
  81. pass
  82. else:
  83. return writer() # is a weakref
  84. if hasattr(select, 'poll'):
  85. def _select_imp(readers=None, writers=None, err=None, timeout=0,
  86. poll=select.poll, POLLIN=select.POLLIN,
  87. POLLOUT=select.POLLOUT, POLLERR=select.POLLERR):
  88. poller = poll()
  89. register = poller.register
  90. if readers:
  91. [register(fd, POLLIN) for fd in readers]
  92. if writers:
  93. [register(fd, POLLOUT) for fd in writers]
  94. if err:
  95. [register(fd, POLLERR) for fd in err]
  96. R, W = set(), set()
  97. timeout = 0 if timeout and timeout < 0 else round(timeout * 1e3)
  98. events = poller.poll(timeout)
  99. for fd, event in events:
  100. if not isinstance(fd, Integral):
  101. fd = fd.fileno()
  102. if event & POLLIN:
  103. R.add(fd)
  104. if event & POLLOUT:
  105. W.add(fd)
  106. if event & POLLERR:
  107. R.add(fd)
  108. return R, W, 0
  109. else:
  110. def _select_imp(readers=None, writers=None, err=None, timeout=0):
  111. r, w, e = select.select(readers, writers, err, timeout)
  112. if e:
  113. r = list(set(r) | set(e))
  114. return r, w, 0
  115. def _select(readers=None, writers=None, err=None, timeout=0,
  116. poll=_select_imp):
  117. """Simple wrapper to :class:`~select.select`, using :`~select.poll`
  118. as the implementation.
  119. :param readers: Set of reader fds to test if readable.
  120. :param writers: Set of writer fds to test if writable.
  121. :param err: Set of fds to test for error condition.
  122. All fd sets passed must be mutable as this function
  123. will remove non-working fds from them, this also means
  124. the caller must make sure there are still fds in the sets
  125. before calling us again.
  126. :returns: tuple of ``(readable, writable, again)``, where
  127. ``readable`` is a set of fds that have data available for read,
  128. ``writable`` is a set of fds that is ready to be written to
  129. and ``again`` is a flag that if set means the caller must
  130. throw away the result and call us again.
  131. """
  132. readers = set() if readers is None else readers
  133. writers = set() if writers is None else writers
  134. err = set() if err is None else err
  135. try:
  136. return poll(readers, writers, err, timeout)
  137. except (select.error, socket.error) as exc:
  138. if exc.errno == errno.EINTR:
  139. return set(), set(), 1
  140. elif exc.errno in SELECT_BAD_FD:
  141. for fd in readers | writers | err:
  142. try:
  143. select.select([fd], [], [], 0)
  144. except (select.error, socket.error) as exc:
  145. if getattr(exc, 'errno', None) not in SELECT_BAD_FD:
  146. raise
  147. readers.discard(fd)
  148. writers.discard(fd)
  149. err.discard(fd)
  150. return set(), set(), 1
  151. else:
  152. raise
  153. class Worker(_pool.Worker):
  154. """Pool worker process."""
  155. def on_loop_start(self, pid):
  156. # our version sends a WORKER_UP message when the process is ready
  157. # to accept work, this will tell the parent that the inqueue fd
  158. # is writable.
  159. self.outq.put((WORKER_UP, (pid,)))
  160. class ResultHandler(_pool.ResultHandler):
  161. """Handles messages from the pool processes."""
  162. def __init__(self, *args, **kwargs):
  163. self.fileno_to_outq = kwargs.pop('fileno_to_outq')
  164. self.on_process_alive = kwargs.pop('on_process_alive')
  165. super(ResultHandler, self).__init__(*args, **kwargs)
  166. # add our custom message handler
  167. self.state_handlers[WORKER_UP] = self.on_process_alive
  168. def _recv_message(self, add_reader, fd, callback,
  169. __read__=__read__, readcanbuf=readcanbuf,
  170. BytesIO=BytesIO, unpack_from=unpack_from,
  171. load=_pickle.load):
  172. Hr = Br = 0
  173. if readcanbuf:
  174. buf = bytearray(4)
  175. bufv = memoryview(buf)
  176. else:
  177. buf = bufv = BytesIO()
  178. # header
  179. while Hr < 4:
  180. try:
  181. n = __read__(
  182. fd, bufv[Hr:] if readcanbuf else bufv, 4 - Hr,
  183. )
  184. except OSError as exc:
  185. if exc.errno not in UNAVAIL:
  186. raise
  187. yield
  188. else:
  189. if n == 0:
  190. raise (OSError('End of file during message') if Hr
  191. else EOFError())
  192. Hr += n
  193. body_size, = unpack_from('>i', bufv)
  194. if readcanbuf:
  195. buf = bytearray(body_size)
  196. bufv = memoryview(buf)
  197. else:
  198. buf = bufv = BytesIO()
  199. while Br < body_size:
  200. try:
  201. n = __read__(
  202. fd, bufv[Br:] if readcanbuf else bufv, body_size - Br,
  203. )
  204. except OSError as exc:
  205. if exc.errno not in UNAVAIL:
  206. raise
  207. yield
  208. else:
  209. if n == 0:
  210. raise (OSError('End of file during message') if Br
  211. else EOFError())
  212. Br += n
  213. add_reader(fd, self.handle_event, fd)
  214. if readcanbuf:
  215. message = load(BytesIO(bufv))
  216. else:
  217. bufv.seek(0)
  218. message = load(bufv)
  219. if message:
  220. callback(message)
  221. def _make_process_result(self, hub):
  222. """Coroutine that reads messages from the pool processes
  223. and calls the appropriate handler."""
  224. fileno_to_outq = self.fileno_to_outq
  225. on_state_change = self.on_state_change
  226. add_reader = hub.add_reader
  227. remove_reader = hub.remove_reader
  228. recv_message = self._recv_message
  229. def on_result_readable(fileno):
  230. try:
  231. fileno_to_outq[fileno]
  232. except KeyError: # process gone
  233. return remove_reader(fileno)
  234. it = recv_message(add_reader, fileno, on_state_change)
  235. try:
  236. next(it)
  237. except StopIteration:
  238. pass
  239. except (IOError, OSError, EOFError):
  240. remove_reader(fileno)
  241. else:
  242. add_reader(fileno, it)
  243. return on_result_readable
  244. def register_with_event_loop(self, hub):
  245. self.handle_event = self._make_process_result(hub)
  246. def handle_event(self, fileno):
  247. raise RuntimeError('Not registered with event loop')
  248. def on_stop_not_started(self):
  249. """This method is always used to stop when the helper thread is not
  250. started."""
  251. cache = self.cache
  252. check_timeouts = self.check_timeouts
  253. fileno_to_outq = self.fileno_to_outq
  254. on_state_change = self.on_state_change
  255. join_exited_workers = self.join_exited_workers
  256. # flush the processes outqueues until they have all terminated.
  257. outqueues = set(fileno_to_outq)
  258. while cache and outqueues and self._state != TERMINATE:
  259. if check_timeouts is not None:
  260. # make sure tasks with a time limit will time out.
  261. check_timeouts()
  262. # cannot iterate and remove at the same time
  263. pending_remove_fd = set()
  264. for fd in outqueues:
  265. self._flush_outqueue(
  266. fd, pending_remove_fd.discard, fileno_to_outq,
  267. on_state_change,
  268. )
  269. try:
  270. join_exited_workers(shutdown=True)
  271. except WorkersJoined:
  272. return debug('result handler: all workers terminated')
  273. outqueues.difference_update(pending_remove_fd)
  274. def _flush_outqueue(self, fd, remove, process_index, on_state_change):
  275. try:
  276. proc = process_index[fd]
  277. except KeyError:
  278. # process already found terminated
  279. # which means its outqueue has already been processed
  280. # by the worker lost handler.
  281. return remove(fd)
  282. reader = proc.outq._reader
  283. try:
  284. setblocking(reader, 1)
  285. except (OSError, IOError):
  286. return remove(fd)
  287. try:
  288. if reader.poll(0):
  289. task = reader.recv()
  290. else:
  291. task = None
  292. sleep(0.5)
  293. except (IOError, EOFError):
  294. return remove(fd)
  295. else:
  296. if task:
  297. on_state_change(task)
  298. finally:
  299. try:
  300. setblocking(reader, 0)
  301. except (OSError, IOError):
  302. return remove(fd)
  303. class AsynPool(_pool.Pool):
  304. """Pool version that uses AIO instead of helper threads."""
  305. ResultHandler = ResultHandler
  306. Worker = Worker
  307. def WorkerProcess(self, worker):
  308. worker = super(AsynPool, self).WorkerProcess(worker)
  309. worker.dead = False
  310. return worker
  311. def __init__(self, processes=None, synack=False,
  312. sched_strategy=None, *args, **kwargs):
  313. self.sched_strategy = SCHED_STRATEGIES.get(sched_strategy,
  314. sched_strategy)
  315. processes = self.cpu_count() if processes is None else processes
  316. self.synack = synack
  317. # create queue-pairs for all our processes in advance.
  318. self._queues = {
  319. self.create_process_queues(): None for _ in range(processes)
  320. }
  321. # inqueue fileno -> process mapping
  322. self._fileno_to_inq = {}
  323. # outqueue fileno -> process mapping
  324. self._fileno_to_outq = {}
  325. # synqueue fileno -> process mapping
  326. self._fileno_to_synq = {}
  327. # We keep track of processes that have not yet
  328. # sent a WORKER_UP message. If a process fails to send
  329. # this message within proc_up_timeout we terminate it
  330. # and hope the next process will recover.
  331. self._proc_alive_timeout = PROC_ALIVE_TIMEOUT
  332. self._waiting_to_start = set()
  333. # denormalized set of all inqueues.
  334. self._all_inqueues = set()
  335. # Set of fds being written to (busy)
  336. self._active_writes = set()
  337. # Set of active co-routines currently writing jobs.
  338. self._active_writers = set()
  339. # Set of fds that are busy (executing task)
  340. self._busy_workers = set()
  341. self._mark_worker_as_available = self._busy_workers.discard
  342. # Holds jobs waiting to be written to child processes.
  343. self.outbound_buffer = deque()
  344. self.write_stats = Counter()
  345. super(AsynPool, self).__init__(processes, *args, **kwargs)
  346. for proc in self._pool:
  347. # create initial mappings, these will be updated
  348. # as processes are recycled, or found lost elsewhere.
  349. self._fileno_to_outq[proc.outqR_fd] = proc
  350. self._fileno_to_synq[proc.synqW_fd] = proc
  351. self.on_soft_timeout = self._timeout_handler.on_soft_timeout
  352. self.on_hard_timeout = self._timeout_handler.on_hard_timeout
  353. def _event_process_exit(self, hub, fd):
  354. # This method is called whenever the process sentinel is readable.
  355. hub.remove(fd)
  356. self.maintain_pool()
  357. def register_with_event_loop(self, hub):
  358. """Registers the async pool with the current event loop."""
  359. self._result_handler.register_with_event_loop(hub)
  360. self.handle_result_event = self._result_handler.handle_event
  361. self._create_timelimit_handlers(hub)
  362. self._create_process_handlers(hub)
  363. self._create_write_handlers(hub)
  364. # Add handler for when a process exits (calls maintain_pool)
  365. [hub.add_reader(fd, self._event_process_exit, hub, fd)
  366. for fd in self.process_sentinels]
  367. # Handle_result_event is called whenever one of the
  368. # result queues are readable.
  369. [hub.add_reader(fd, self.handle_result_event, fd)
  370. for fd in self._fileno_to_outq]
  371. # Timers include calling maintain_pool at a regular interval
  372. # to be certain processes are restarted.
  373. for handler, interval in items(self.timers):
  374. hub.call_repeatedly(interval, handler)
  375. hub.on_tick.add(self.on_poll_start)
  376. def _create_timelimit_handlers(self, hub, now=time.time):
  377. """For async pool this sets up the handlers used
  378. to implement time limits."""
  379. call_later = hub.call_later
  380. trefs = self._tref_for_id = WeakValueDictionary()
  381. def on_timeout_set(R, soft, hard):
  382. if soft:
  383. trefs[R._job] = call_later(
  384. soft, self._on_soft_timeout, R._job, soft, hard, hub,
  385. )
  386. elif hard:
  387. trefs[R._job] = call_later(
  388. hard, self._on_hard_timeout, R._job,
  389. )
  390. self.on_timeout_set = on_timeout_set
  391. def _discard_tref(job):
  392. try:
  393. tref = trefs.pop(job)
  394. tref.cancel()
  395. del(tref)
  396. except (KeyError, AttributeError):
  397. pass # out of scope
  398. self._discard_tref = _discard_tref
  399. def on_timeout_cancel(R):
  400. _discard_tref(R._job)
  401. self.on_timeout_cancel = on_timeout_cancel
  402. def _on_soft_timeout(self, job, soft, hard, hub, now=time.time):
  403. # only used by async pool.
  404. if hard:
  405. self._tref_for_id[job] = hub.call_at(
  406. now() + (hard - soft), self._on_hard_timeout, job,
  407. )
  408. try:
  409. result = self._cache[job]
  410. except KeyError:
  411. pass # job ready
  412. else:
  413. self.on_soft_timeout(result)
  414. finally:
  415. if not hard:
  416. # remove tref
  417. self._discard_tref(job)
  418. def _on_hard_timeout(self, job):
  419. # only used by async pool.
  420. try:
  421. result = self._cache[job]
  422. except KeyError:
  423. pass # job ready
  424. else:
  425. self.on_hard_timeout(result)
  426. finally:
  427. # remove tref
  428. self._discard_tref(job)
  429. def on_job_ready(self, job, i, obj, inqW_fd):
  430. self._mark_worker_as_available(inqW_fd)
  431. def _create_process_handlers(self, hub, READ=READ, ERR=ERR):
  432. """For async pool this will create the handlers called
  433. when a process is up/down and etc."""
  434. add_reader, remove_reader, remove_writer = (
  435. hub.add_reader, hub.remove_reader, hub.remove_writer,
  436. )
  437. cache = self._cache
  438. all_inqueues = self._all_inqueues
  439. fileno_to_inq = self._fileno_to_inq
  440. fileno_to_outq = self._fileno_to_outq
  441. fileno_to_synq = self._fileno_to_synq
  442. busy_workers = self._busy_workers
  443. event_process_exit = self._event_process_exit
  444. handle_result_event = self.handle_result_event
  445. process_flush_queues = self.process_flush_queues
  446. waiting_to_start = self._waiting_to_start
  447. def verify_process_alive(proc):
  448. if proc._is_alive() and proc in waiting_to_start:
  449. assert proc.outqR_fd in fileno_to_outq
  450. assert fileno_to_outq[proc.outqR_fd] is proc
  451. assert proc.outqR_fd in hub.readers
  452. error('Timed out waiting for UP message from %r', proc)
  453. os.kill(proc.pid, 9)
  454. def on_process_up(proc):
  455. """Called when a process has started."""
  456. # If we got the same fd as a previous process then we will also
  457. # receive jobs in the old buffer, so we need to reset the
  458. # job._write_to and job._scheduled_for attributes used to recover
  459. # message boundaries when processes exit.
  460. infd = proc.inqW_fd
  461. for job in values(cache):
  462. if job._write_to and job._write_to.inqW_fd == infd:
  463. job._write_to = proc
  464. if job._scheduled_for and job._scheduled_for.inqW_fd == infd:
  465. job._scheduled_for = proc
  466. fileno_to_outq[proc.outqR_fd] = proc
  467. # maintain_pool is called whenever a process exits.
  468. add_reader(
  469. proc.sentinel, event_process_exit, hub, proc.sentinel,
  470. )
  471. assert not isblocking(proc.outq._reader)
  472. # handle_result_event is called when the processes outqueue is
  473. # readable.
  474. add_reader(proc.outqR_fd, handle_result_event, proc.outqR_fd)
  475. waiting_to_start.add(proc)
  476. hub.call_later(
  477. self._proc_alive_timeout, verify_process_alive, proc,
  478. )
  479. self.on_process_up = on_process_up
  480. def _remove_from_index(obj, proc, index, remove_fun, callback=None):
  481. # this remove the file descriptors for a process from
  482. # the indices. we have to make sure we don't overwrite
  483. # another processes fds, as the fds may be reused.
  484. try:
  485. fd = obj.fileno()
  486. except (IOError, OSError):
  487. return
  488. try:
  489. if index[fd] is proc:
  490. # fd has not been reused so we can remove it from index.
  491. index.pop(fd, None)
  492. except KeyError:
  493. pass
  494. else:
  495. remove_fun(fd)
  496. if callback is not None:
  497. callback(fd)
  498. return fd
  499. def on_process_down(proc):
  500. """Called when a worker process exits."""
  501. if proc.dead:
  502. return
  503. process_flush_queues(proc)
  504. _remove_from_index(
  505. proc.outq._reader, proc, fileno_to_outq, remove_reader,
  506. )
  507. if proc.synq:
  508. _remove_from_index(
  509. proc.synq._writer, proc, fileno_to_synq, remove_writer,
  510. )
  511. inq = _remove_from_index(
  512. proc.inq._writer, proc, fileno_to_inq, remove_writer,
  513. callback=all_inqueues.discard,
  514. )
  515. if inq:
  516. busy_workers.discard(inq)
  517. remove_reader(proc.sentinel)
  518. waiting_to_start.discard(proc)
  519. self._active_writes.discard(proc.inqW_fd)
  520. remove_writer(proc.inqW_fd)
  521. remove_reader(proc.outqR_fd)
  522. if proc.synqR_fd:
  523. remove_reader(proc.synqR_fd)
  524. if proc.synqW_fd:
  525. self._active_writes.discard(proc.synqW_fd)
  526. remove_reader(proc.synqW_fd)
  527. self.on_process_down = on_process_down
  528. def _create_write_handlers(self, hub,
  529. pack=struct.pack, dumps=_pickle.dumps,
  530. protocol=HIGHEST_PROTOCOL):
  531. """For async pool this creates the handlers used to write data to
  532. child processes."""
  533. fileno_to_inq = self._fileno_to_inq
  534. fileno_to_synq = self._fileno_to_synq
  535. outbound = self.outbound_buffer
  536. pop_message = outbound.popleft
  537. put_message = outbound.append
  538. all_inqueues = self._all_inqueues
  539. active_writes = self._active_writes
  540. active_writers = self._active_writers
  541. busy_workers = self._busy_workers
  542. diff = all_inqueues.difference
  543. add_writer = hub.add_writer
  544. hub_add, hub_remove = hub.add, hub.remove
  545. mark_write_fd_as_active = active_writes.add
  546. mark_write_gen_as_active = active_writers.add
  547. mark_worker_as_busy = busy_workers.add
  548. write_generator_done = active_writers.discard
  549. get_job = self._cache.__getitem__
  550. write_stats = self.write_stats
  551. is_fair_strategy = self.sched_strategy == SCHED_STRATEGY_FAIR
  552. revoked_tasks = worker_state.revoked
  553. getpid = os.getpid
  554. precalc = {ACK: self._create_payload(ACK, (0,)),
  555. NACK: self._create_payload(NACK, (0,))}
  556. def _put_back(job, _time=time.time):
  557. # puts back at the end of the queue
  558. if job._terminated is not None or \
  559. job.correlation_id in revoked_tasks:
  560. if not job._accepted:
  561. job._ack(None, _time(), getpid(), None)
  562. job._set_terminated(job._terminated)
  563. else:
  564. # XXX linear lookup, should find a better way,
  565. # but this happens rarely and is here to protect against races.
  566. if job not in outbound:
  567. outbound.appendleft(job)
  568. self._put_back = _put_back
  569. # called for every event loop iteration, and if there
  570. # are messages pending this will schedule writing one message
  571. # by registering the 'schedule_writes' function for all currently
  572. # inactive inqueues (not already being written to)
  573. # consolidate means the event loop will merge them
  574. # and call the callback once with the list writable fds as
  575. # argument. Using this means we minimize the risk of having
  576. # the same fd receive every task if the pipe read buffer is not
  577. # full.
  578. if is_fair_strategy:
  579. def on_poll_start():
  580. if outbound and len(busy_workers) < len(all_inqueues):
  581. # print('ALL: %r ACTIVE: %r' % (len(all_inqueues),
  582. # len(active_writes)))
  583. inactive = diff(active_writes)
  584. [hub_add(fd, None, WRITE | ERR, consolidate=True)
  585. for fd in inactive]
  586. else:
  587. [hub_remove(fd) for fd in diff(active_writes)]
  588. else:
  589. def on_poll_start(): # noqa
  590. if outbound:
  591. [hub_add(fd, None, WRITE | ERR, consolidate=True)
  592. for fd in diff(active_writes)]
  593. else:
  594. [hub_remove(fd) for fd in diff(active_writes)]
  595. self.on_poll_start = on_poll_start
  596. def on_inqueue_close(fd, proc):
  597. # Makes sure the fd is removed from tracking when
  598. # the connection is closed, this is essential as fds may be reused.
  599. busy_workers.discard(fd)
  600. try:
  601. if fileno_to_inq[fd] is proc:
  602. fileno_to_inq.pop(fd, None)
  603. active_writes.discard(fd)
  604. all_inqueues.discard(fd)
  605. hub_remove(fd)
  606. except KeyError:
  607. pass
  608. self.on_inqueue_close = on_inqueue_close
  609. def schedule_writes(ready_fds, total_write_count=[0]):
  610. # Schedule write operation to ready file descriptor.
  611. # The file descriptor is writeable, but that does not
  612. # mean the process is currently reading from the socket.
  613. # The socket is buffered so writeable simply means that
  614. # the buffer can accept at least 1 byte of data.
  615. # This means we have to cycle between the ready fds.
  616. # the first version used shuffle, but this version
  617. # using `total_writes % ready_fds` is about 30% faster
  618. # with many processes, and also leans more towards fairness
  619. # in write stats when used with many processes
  620. # [XXX On OS X, this may vary depending
  621. # on event loop implementation (i.e select vs epoll), so
  622. # have to test further]
  623. num_ready = len(ready_fds)
  624. for i in range(num_ready):
  625. ready_fd = ready_fds[total_write_count[0] % num_ready]
  626. total_write_count[0] += 1
  627. if ready_fd in active_writes:
  628. # already writing to this fd
  629. continue
  630. if is_fair_strategy and ready_fd in busy_workers:
  631. # worker is already busy with another task
  632. continue
  633. if ready_fd not in all_inqueues:
  634. hub_remove(ready_fd)
  635. continue
  636. try:
  637. job = pop_message()
  638. except IndexError:
  639. # no more messages, remove all inactive fds from the hub.
  640. # this is important since the fds are always writeable
  641. # as long as there's 1 byte left in the buffer, and so
  642. # this may create a spinloop where the event loop
  643. # always wakes up.
  644. for inqfd in diff(active_writes):
  645. hub_remove(inqfd)
  646. break
  647. else:
  648. if not job._accepted: # job not accepted by another worker
  649. try:
  650. # keep track of what process the write operation
  651. # was scheduled for.
  652. proc = job._scheduled_for = fileno_to_inq[ready_fd]
  653. except KeyError:
  654. # write was scheduled for this fd but the process
  655. # has since exited and the message must be sent to
  656. # another process.
  657. put_message(job)
  658. continue
  659. cor = _write_job(proc, ready_fd, job)
  660. job._writer = ref(cor)
  661. mark_write_gen_as_active(cor)
  662. mark_write_fd_as_active(ready_fd)
  663. mark_worker_as_busy(ready_fd)
  664. # Try to write immediately, in case there's an error.
  665. try:
  666. next(cor)
  667. except StopIteration:
  668. pass
  669. except OSError as exc:
  670. if exc.errno != errno.EBADF:
  671. raise
  672. else:
  673. add_writer(ready_fd, cor)
  674. hub.consolidate_callback = schedule_writes
  675. def send_job(tup):
  676. # Schedule writing job request for when one of the process
  677. # inqueues are writable.
  678. body = dumps(tup, protocol=protocol)
  679. body_size = len(body)
  680. header = pack('>I', body_size)
  681. # index 1,0 is the job ID.
  682. job = get_job(tup[1][0])
  683. job._payload = buf_t(header), buf_t(body), body_size
  684. put_message(job)
  685. self._quick_put = send_job
  686. def on_not_recovering(proc, fd, job):
  687. error('Process inqueue damaged: %r %r' % (proc, proc.exitcode))
  688. if proc._is_alive():
  689. proc.terminate()
  690. hub.remove(fd)
  691. self._put_back(job)
  692. def _write_job(proc, fd, job):
  693. # writes job to the worker process.
  694. # Operation must complete if more than one byte of data
  695. # was written. If the broker connection is lost
  696. # and no data was written the operation shall be cancelled.
  697. header, body, body_size = job._payload
  698. errors = 0
  699. try:
  700. # job result keeps track of what process the job is sent to.
  701. job._write_to = proc
  702. send = proc.send_job_offset
  703. Hw = Bw = 0
  704. # write header
  705. while Hw < 4:
  706. try:
  707. Hw += send(header, Hw)
  708. except Exception as exc:
  709. if getattr(exc, 'errno', None) not in UNAVAIL:
  710. raise
  711. # suspend until more data
  712. errors += 1
  713. if errors > 100:
  714. on_not_recovering(proc, fd, job)
  715. raise StopIteration()
  716. yield
  717. else:
  718. errors = 0
  719. # write body
  720. while Bw < body_size:
  721. try:
  722. Bw += send(body, Bw)
  723. except Exception as exc:
  724. if getattr(exc, 'errno', None) not in UNAVAIL:
  725. raise
  726. # suspend until more data
  727. errors += 1
  728. if errors > 100:
  729. on_not_recovering(proc, fd, job)
  730. raise StopIteration()
  731. yield
  732. else:
  733. errors = 0
  734. finally:
  735. hub_remove(fd)
  736. write_stats[proc.index] += 1
  737. # message written, so this fd is now available
  738. active_writes.discard(fd)
  739. write_generator_done(job._writer()) # is a weakref
  740. def send_ack(response, pid, job, fd, WRITE=WRITE, ERR=ERR):
  741. # Only used when synack is enabled.
  742. # Schedule writing ack response for when the fd is writeable.
  743. msg = Ack(job, fd, precalc[response])
  744. callback = promise(write_generator_done)
  745. cor = _write_ack(fd, msg, callback=callback)
  746. mark_write_gen_as_active(cor)
  747. mark_write_fd_as_active(fd)
  748. callback.args = (cor,)
  749. add_writer(fd, cor)
  750. self.send_ack = send_ack
  751. def _write_ack(fd, ack, callback=None):
  752. # writes ack back to the worker if synack enabled.
  753. # this operation *MUST* complete, otherwise
  754. # the worker process will hang waiting for the ack.
  755. header, body, body_size = ack[2]
  756. try:
  757. try:
  758. proc = fileno_to_synq[fd]
  759. except KeyError:
  760. # process died, we can safely discard the ack at this
  761. # point.
  762. raise StopIteration()
  763. send = proc.send_syn_offset
  764. Hw = Bw = 0
  765. # write header
  766. while Hw < 4:
  767. try:
  768. Hw += send(header, Hw)
  769. except Exception as exc:
  770. if getattr(exc, 'errno', None) not in UNAVAIL:
  771. raise
  772. yield
  773. # write body
  774. while Bw < body_size:
  775. try:
  776. Bw += send(body, Bw)
  777. except Exception as exc:
  778. if getattr(exc, 'errno', None) not in UNAVAIL:
  779. raise
  780. # suspend until more data
  781. yield
  782. finally:
  783. if callback:
  784. callback()
  785. # message written, so this fd is now available
  786. active_writes.discard(fd)
  787. def flush(self):
  788. if self._state == TERMINATE:
  789. return
  790. # cancel all tasks that have not been accepted so that NACK is sent.
  791. for job in values(self._cache):
  792. if not job._accepted:
  793. job._cancel()
  794. # clear the outgoing buffer as the tasks will be redelivered by
  795. # the broker anyway.
  796. if self.outbound_buffer:
  797. self.outbound_buffer.clear()
  798. self.maintain_pool()
  799. try:
  800. # ...but we must continue writing the payloads we already started
  801. # to keep message boundaries.
  802. # The messages may be NACK'ed later if synack is enabled.
  803. if self._state == RUN:
  804. # flush outgoing buffers
  805. intervals = fxrange(0.01, 0.1, 0.01, repeatlast=True)
  806. owned_by = {}
  807. for job in values(self._cache):
  808. writer = _get_job_writer(job)
  809. if writer is not None:
  810. owned_by[writer] = job
  811. while self._active_writers:
  812. writers = list(self._active_writers)
  813. for gen in writers:
  814. if (gen.__name__ == '_write_job' and
  815. gen_not_started(gen)):
  816. # has not started writing the job so can
  817. # discard the task, but we must also remove
  818. # it from the Pool._cache.
  819. try:
  820. job = owned_by[gen]
  821. except KeyError:
  822. pass
  823. else:
  824. # removes from Pool._cache
  825. job.discard()
  826. self._active_writers.discard(gen)
  827. else:
  828. try:
  829. job = owned_by[gen]
  830. except KeyError:
  831. pass
  832. else:
  833. job_proc = job._write_to
  834. if job_proc._is_alive():
  835. self._flush_writer(job_proc, gen)
  836. # workers may have exited in the meantime.
  837. self.maintain_pool()
  838. sleep(next(intervals)) # don't busyloop
  839. finally:
  840. self.outbound_buffer.clear()
  841. self._active_writers.clear()
  842. self._active_writes.clear()
  843. self._busy_workers.clear()
  844. def _flush_writer(self, proc, writer):
  845. fds = {proc.inq._writer}
  846. try:
  847. while fds:
  848. if not proc._is_alive():
  849. break # process exited
  850. readable, writable, again = _select(
  851. writers=fds, err=fds, timeout=0.5,
  852. )
  853. if not again and (writable or readable):
  854. try:
  855. next(writer)
  856. except (StopIteration, OSError, IOError, EOFError):
  857. break
  858. finally:
  859. self._active_writers.discard(writer)
  860. def get_process_queues(self):
  861. """Get queues for a new process.
  862. Here we will find an unused slot, as there should always
  863. be one available when we start a new process.
  864. """
  865. return next(q for q, owner in items(self._queues)
  866. if owner is None)
  867. def on_grow(self, n):
  868. """Grow the pool by ``n`` proceses."""
  869. diff = max(self._processes - len(self._queues), 0)
  870. if diff:
  871. self._queues.update({
  872. self.create_process_queues(): None for _ in range(diff)
  873. })
  874. def on_shrink(self, n):
  875. """Shrink the pool by ``n`` processes."""
  876. pass
  877. def create_process_queues(self):
  878. """Creates new in, out (and optionally syn) queues,
  879. returned as a tuple."""
  880. # NOTE: Pipes must be set O_NONBLOCK at creation time (the original
  881. # fd), otherwise it will not be possible to change the flags until
  882. # there is an actual reader/writer on the other side.
  883. inq = _SimpleQueue(wnonblock=True)
  884. outq = _SimpleQueue(rnonblock=True)
  885. synq = None
  886. assert isblocking(inq._reader)
  887. assert not isblocking(inq._writer)
  888. assert not isblocking(outq._reader)
  889. assert isblocking(outq._writer)
  890. if self.synack:
  891. synq = _SimpleQueue(wnonblock=True)
  892. assert isblocking(synq._reader)
  893. assert not isblocking(synq._writer)
  894. return inq, outq, synq
  895. def on_process_alive(self, pid):
  896. """Handler called when the :const:`WORKER_UP` message is received
  897. from a child process, which marks the process as ready
  898. to receive work."""
  899. try:
  900. proc = next(w for w in self._pool if w.pid == pid)
  901. except StopIteration:
  902. return logger.warning('process with pid=%s already exited', pid)
  903. assert proc.inqW_fd not in self._fileno_to_inq
  904. assert proc.inqW_fd not in self._all_inqueues
  905. self._waiting_to_start.discard(proc)
  906. self._fileno_to_inq[proc.inqW_fd] = proc
  907. self._fileno_to_synq[proc.synqW_fd] = proc
  908. self._all_inqueues.add(proc.inqW_fd)
  909. def on_job_process_down(self, job, pid_gone):
  910. """Handler called for each job when the process it was assigned to
  911. exits."""
  912. if job._write_to and not job._write_to._is_alive():
  913. # job was partially written
  914. self.on_partial_read(job, job._write_to)
  915. elif job._scheduled_for and not job._scheduled_for._is_alive():
  916. # job was only scheduled to be written to this process,
  917. # but no data was sent so put it back on the outbound_buffer.
  918. self._put_back(job)
  919. def on_job_process_lost(self, job, pid, exitcode):
  920. """Handler called for each *started* job when the process it
  921. was assigned to exited by mysterious means (error exitcodes and
  922. signals)"""
  923. self.mark_as_worker_lost(job, exitcode)
  924. def human_write_stats(self):
  925. if self.write_stats is None:
  926. return 'N/A'
  927. vals = list(values(self.write_stats))
  928. total = sum(vals)
  929. def per(v, total):
  930. return '{0:.2f}%'.format((float(v) / total) * 100.0 if v else 0)
  931. return {
  932. 'total': total,
  933. 'avg': per(total / len(self.write_stats) if total else 0, total),
  934. 'all': ', '.join(per(v, total) for v in vals),
  935. 'raw': ', '.join(map(str, vals)),
  936. 'inqueues': {
  937. 'total': len(self._all_inqueues),
  938. 'active': len(self._active_writes),
  939. }
  940. }
  941. def _process_cleanup_queues(self, proc):
  942. """Handler called to clean up a processes queues after process
  943. exit."""
  944. if not proc.dead:
  945. try:
  946. self._queues[self._find_worker_queues(proc)] = None
  947. except (KeyError, ValueError):
  948. pass
  949. @staticmethod
  950. def _stop_task_handler(task_handler):
  951. """Called at shutdown to tell processes that we are shutting down."""
  952. for proc in task_handler.pool:
  953. try:
  954. setblocking(proc.inq._writer, 1)
  955. except (OSError, IOError):
  956. pass
  957. else:
  958. try:
  959. proc.inq.put(None)
  960. except OSError as exc:
  961. if exc.errno != errno.EBADF:
  962. raise
  963. def create_result_handler(self):
  964. return super(AsynPool, self).create_result_handler(
  965. fileno_to_outq=self._fileno_to_outq,
  966. on_process_alive=self.on_process_alive,
  967. )
  968. def _process_register_queues(self, proc, queues):
  969. """Marks new ownership for ``queues`` so that the fileno indices are
  970. updated."""
  971. assert queues in self._queues
  972. b = len(self._queues)
  973. self._queues[queues] = proc
  974. assert b == len(self._queues)
  975. def _find_worker_queues(self, proc):
  976. """Find the queues owned by ``proc``."""
  977. try:
  978. return next(q for q, owner in items(self._queues)
  979. if owner == proc)
  980. except StopIteration:
  981. raise ValueError(proc)
  982. def _setup_queues(self):
  983. # this is only used by the original pool which uses a shared
  984. # queue for all processes.
  985. # these attributes makes no sense for us, but we will still
  986. # have to initialize them.
  987. self._inqueue = self._outqueue = \
  988. self._quick_put = self._quick_get = self._poll_result = None
  989. def process_flush_queues(self, proc):
  990. """Flushes all queues, including the outbound buffer, so that
  991. all tasks that have not been started will be discarded.
  992. In Celery this is called whenever the transport connection is lost
  993. (consumer restart), and when a process is terminated.
  994. """
  995. resq = proc.outq._reader
  996. on_state_change = self._result_handler.on_state_change
  997. fds = {resq}
  998. while fds and not resq.closed and self._state != TERMINATE:
  999. readable, _, again = _select(fds, None, fds, timeout=0.01)
  1000. if readable:
  1001. try:
  1002. task = resq.recv()
  1003. except (OSError, IOError, EOFError) as exc:
  1004. _errno = getattr(exc, 'errno', None)
  1005. if _errno == errno.EINTR:
  1006. continue
  1007. elif _errno == errno.EAGAIN:
  1008. break
  1009. elif _errno not in UNAVAIL:
  1010. debug('got %r while flushing process %r',
  1011. exc, proc, exc_info=1)
  1012. break
  1013. else:
  1014. if task is None:
  1015. debug('got sentinel while flushing process %r', proc)
  1016. break
  1017. else:
  1018. on_state_change(task)
  1019. else:
  1020. break
  1021. def on_partial_read(self, job, proc):
  1022. """Called when a job was only partially written to a child process
  1023. and it exited."""
  1024. # worker terminated by signal:
  1025. # we cannot reuse the sockets again, because we don't know if
  1026. # the process wrote/read anything frmo them, and if so we cannot
  1027. # restore the message boundaries.
  1028. if not job._accepted:
  1029. # job was not acked, so find another worker to send it to.
  1030. self._put_back(job)
  1031. writer = _get_job_writer(job)
  1032. if writer:
  1033. self._active_writers.discard(writer)
  1034. del(writer)
  1035. if not proc.dead:
  1036. proc.dead = True
  1037. # Replace queues to avoid reuse
  1038. before = len(self._queues)
  1039. try:
  1040. queues = self._find_worker_queues(proc)
  1041. if self.destroy_queues(queues, proc):
  1042. self._queues[self.create_process_queues()] = None
  1043. except ValueError:
  1044. pass
  1045. assert len(self._queues) == before
  1046. def destroy_queues(self, queues, proc):
  1047. """Destroy queues that can no longer be used, so that they
  1048. be replaced by new sockets."""
  1049. assert not proc._is_alive()
  1050. self._waiting_to_start.discard(proc)
  1051. removed = 1
  1052. try:
  1053. self._queues.pop(queues)
  1054. except KeyError:
  1055. removed = 0
  1056. try:
  1057. self.on_inqueue_close(queues[0]._writer.fileno(), proc)
  1058. except IOError:
  1059. pass
  1060. for queue in queues:
  1061. if queue:
  1062. for sock in (queue._reader, queue._writer):
  1063. if not sock.closed:
  1064. try:
  1065. sock.close()
  1066. except (IOError, OSError):
  1067. pass
  1068. return removed
  1069. def _create_payload(self, type_, args,
  1070. dumps=_pickle.dumps, pack=struct.pack,
  1071. protocol=HIGHEST_PROTOCOL):
  1072. body = dumps((type_, args), protocol=protocol)
  1073. size = len(body)
  1074. header = pack('>I', size)
  1075. return header, body, size
  1076. @classmethod
  1077. def _set_result_sentinel(cls, _outqueue, _pool):
  1078. # unused
  1079. pass
  1080. def _help_stuff_finish_args(self):
  1081. # Pool._help_stuff_finished is a classmethod so we have to use this
  1082. # trick to modify the arguments passed to it.
  1083. return (self._pool,)
  1084. @classmethod
  1085. def _help_stuff_finish(cls, pool):
  1086. debug(
  1087. 'removing tasks from inqueue until task handler finished',
  1088. )
  1089. fileno_to_proc = {}
  1090. inqR = set()
  1091. for w in pool:
  1092. try:
  1093. fd = w.inq._reader.fileno()
  1094. inqR.add(fd)
  1095. fileno_to_proc[fd] = w
  1096. except IOError:
  1097. pass
  1098. while inqR:
  1099. readable, _, again = _select(inqR, timeout=0.5)
  1100. if again:
  1101. continue
  1102. if not readable:
  1103. break
  1104. for fd in readable:
  1105. fileno_to_proc[fd].inq._reader.recv()
  1106. sleep(0)
  1107. @property
  1108. def timers(self):
  1109. return {self.maintain_pool: 5.0}