152 lines
5.2 KiB
Python
152 lines
5.2 KiB
Python
# This file is part of Tryton. The COPYRIGHT file at the top level of
|
|
# this repository contains the full copyright notices and license terms.
|
|
import datetime as dt
|
|
import logging
|
|
import random
|
|
import select
|
|
import signal
|
|
import time
|
|
from multiprocessing import Pool as MPool
|
|
from multiprocessing import cpu_count
|
|
|
|
from sql import Flavor
|
|
|
|
from trytond import backend
|
|
from trytond.config import config
|
|
from trytond.exceptions import UserError, UserWarning
|
|
from trytond.pool import Pool
|
|
from trytond.status import processing
|
|
from trytond.transaction import Transaction
|
|
|
|
__all__ = ['work']
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class Queue(object):
|
|
def __init__(self, pool, mpool):
|
|
self.database = backend.Database(pool.database_name).connect()
|
|
self.connection = self.database.get_connection(autocommit=True)
|
|
self.pool = pool
|
|
self.mpool = mpool
|
|
|
|
def pull(self, name=None):
|
|
Queue = self.pool.get('ir.queue')
|
|
return Queue.pull(self.database, self.connection, name=name)
|
|
|
|
def run(self, task_id):
|
|
return self.mpool.apply_async(
|
|
run_task, (self.pool.database_name, task_id))
|
|
|
|
|
|
class TaskList(list):
|
|
def filter(self):
|
|
for t in list(self):
|
|
if t.ready():
|
|
self.remove(t)
|
|
return self
|
|
|
|
|
|
def work(options):
|
|
Flavor.set(backend.Database.flavor)
|
|
if not config.getboolean('queue', 'worker', default=False):
|
|
return
|
|
try:
|
|
processes = options.processes or cpu_count()
|
|
except NotImplementedError:
|
|
processes = 1
|
|
logger.info("start %d workers", processes)
|
|
mpool = MPool(
|
|
processes, initializer, (options,), options.maxtasksperchild)
|
|
queues = [Queue(pool, mpool) for pool in initializer(options, False)]
|
|
|
|
tasks = TaskList()
|
|
timeout = options.timeout
|
|
try:
|
|
while True:
|
|
while len(tasks.filter()) >= processes:
|
|
time.sleep(0.1)
|
|
for queue in queues:
|
|
task_id, next_ = queue.pull(options.name)
|
|
timeout = min(
|
|
next_ or options.timeout, timeout, options.timeout)
|
|
if task_id:
|
|
tasks.append(queue.run(task_id))
|
|
break
|
|
else:
|
|
connections = [q.connection for q in queues]
|
|
connections, _, _ = select.select(connections, [], [], timeout)
|
|
for connection in connections:
|
|
connection.poll()
|
|
while connection.notifies:
|
|
connection.notifies.pop(0)
|
|
except KeyboardInterrupt:
|
|
mpool.close()
|
|
|
|
|
|
def initializer(options, worker=True):
|
|
if worker:
|
|
signal.signal(signal.SIGINT, signal.SIG_IGN)
|
|
pools = []
|
|
database_list = Pool.database_list()
|
|
for database_name in options.database_names:
|
|
pool = Pool(database_name)
|
|
if database_name not in database_list:
|
|
with Transaction().start(database_name, 0, readonly=True):
|
|
pool.init()
|
|
pools.append(pool)
|
|
return pools
|
|
|
|
|
|
def run_task(pool, task_id):
|
|
if not isinstance(pool, Pool):
|
|
pool = Pool(pool)
|
|
Queue = pool.get('ir.queue')
|
|
Error = pool.get('ir.error')
|
|
name = '<Task %s@%s>' % (task_id, pool.database_name)
|
|
logger.info('%s started', name)
|
|
retry = config.getint('database', 'retry')
|
|
try:
|
|
for count in range(retry, -1, -1):
|
|
if count != retry:
|
|
time.sleep(0.02 * (retry - count))
|
|
with Transaction().start(pool.database_name, 0) as transaction:
|
|
try:
|
|
try:
|
|
task, = Queue.search([('id', '=', task_id)])
|
|
except ValueError:
|
|
# the task was rollbacked, nothing to do
|
|
break
|
|
with processing(name):
|
|
task.run()
|
|
break
|
|
except backend.DatabaseOperationalError:
|
|
if count:
|
|
transaction.rollback()
|
|
continue
|
|
raise
|
|
except (UserError, UserWarning) as e:
|
|
Error.log(task, e)
|
|
raise
|
|
logger.info('%s done', name)
|
|
except backend.DatabaseOperationalError:
|
|
logger.info('%s failed, retrying', name, exc_info=True)
|
|
if not config.getboolean('queue', 'worker', default=False):
|
|
time.sleep(0.02 * retry)
|
|
try:
|
|
with Transaction().start(pool.database_name, 0) as transaction:
|
|
if not transaction.database.has_channel():
|
|
logger.critical('%s failed', name, exc_info=True)
|
|
return
|
|
task = Queue(task_id)
|
|
scheduled_at = dt.datetime.now()
|
|
scheduled_at += dt.timedelta(
|
|
seconds=random.randint(0, 2 * retry))
|
|
Queue.push(task.name, task.data, scheduled_at=scheduled_at)
|
|
except Exception:
|
|
logger.critical(
|
|
'rescheduling %s failed', name, exc_info=True)
|
|
except (UserError, UserWarning):
|
|
logger.info('%s failed', name)
|
|
except Exception:
|
|
logger.critical('%s failed', name, exc_info=True)
|