100行代码实现任务队列

#! -*- coding: utf-8 -*-try:    import cPickle as pickleexcept ImportError:    import pickletry:    import simplejson as jsonexcept ImportError:    import jsonimport osimport errnoimport Queueimport randomimport loggingfrom functools import wrapsfrom threading import Timer, RLock, Threadfrom time import sleep, timefrom base64 import b64encode, b64decode# json 的数据结构# tasks = {#     index: {#         cycle_num: [(func, bargs)]#     }# }logging.basicConfig(level=logging.DEBUG,
                   format='(%(asctime)-15s) %(message)s',)
tasks_file = 'tasks.json'flags = os.O_CREAT | os.O_EXCL | os.O_WRONLY# 为了防止任务太多需要生成过多的线程，我们使用Queue 来限制生成的线程数量WORKER_NUMS = 2q = Queue.Queue(WORKER_NUMS)lock = RLock()def check_file():
   try:
       file_handle = os.open(tasks_file, flags)    except OSError as e:        if e.errno == errno.EEXIST:  # Failed as the file already exists.
           pass
       else:            raise
   else:        with os.fdopen(file_handle, 'w') as file_obj:
           file_obj.write("{}")def set_delay_task(func_name, *args, **kwargs):
   # 使用锁来保证每次只要一个线程写入文件，防止数据出错
   with lock:        with open(tasks_file, 'r+') as json_file:
           count_down = kwargs.pop('count_down', 0)
           tasks = json.load(json_file)            # 执行时间
           exec_time = int(time()) + count_down            # 循环索引
           index = str(exec_time % 3600)            # 圈数
           cycle_num = str(exec_time / 3600 + 1)
           dargs = pickle.dumps((args, kwargs))
           bargs = b64encode(dargs)
           index_data = tasks.get(index, {})
           index_data.setdefault(cycle_num, []).append((func_name, bargs))
           tasks[index] = index_data
           json_file.seek(0)
           json.dump(tasks, json_file)
           logging.debug('Received task: %s' % func_name)def get_delay_tasks():
   with open(tasks_file, 'r+') as json_file:
       tasks = json.load(json_file)        # 执行时间
       current_time = int(time())        # 循环索引
       index = str(current_time % 3600)        # 圈数
       cycle_num = str(current_time / 3600 + 1)
       current_tasks = tasks.get(index, {}).get(cycle_num, [])
   tasks = []    for func, bargs in current_tasks:
       dargs = b64decode(bargs)
       args, kwargs = pickle.loads(dargs)
       tasks.append((func, (args, kwargs)))    return tasksdef get_method_by_name(method_name):
   possibles = globals().copy()
   possibles.update(locals())
   method = possibles.get(method_name)    return methoddef create_task(task_class, func, task_name=None, **kwargs):   def execute(self):
       args, kwargs = self.data or ((), {})        return func(*args, **kwargs)   attrs = {        'execute': execute,        'func_name': func.__name__,        '__module__': func.__module__,        '__doc__': func.__doc__
   }
   attrs.update(kwargs)   klass = type(
       task_name or func.__name__,
       (task_class,),
       attrs
   )    return klassclass Hu(object):   def __init__(self, func_name=None):
       self.func_name = func_name
       check_file()    def task(self):
       def deco(func):
           self.func_name = func.__name__
           klass = create_task(Hu, func, self.func_name)
           func.delay = klass(func_name=klass.func_name).delay            @wraps(func)
           def wrapper(*args, **kwargs):
               return func(*args, **kwargs)            return wrapper        return deco    def delay(self, *args, **kwargs):
       _args = [self.func_name]
       _args.extend(args)
       Timer(0, set_delay_task, _args, kwargs).start()        return Truedef boss():
   while True:
       current_tasks = get_delay_tasks()        for func, params in current_tasks:            # Task accepted: auth.tasks.send_msg
           logging.debug('Task accepted: %s' % func)
           q.put((func, params))
       sleep(1)def worker():
   while True:
       func, params = q.get()        print 'get task: %s\n' % func
       method = get_method_by_name(func)
       args, kwargs = params        # Task auth.tasks.send_msgsucceeded in
       start_time = time()
       method(*args, **kwargs)
       end_time = time()
       logging.debug('Task %s succeeded in %s' % (str(func), end_time - start_time))
       q.task_done()def main():
   check_file()
   print('starting at:', time())    for target in (boss, worker):
       t = Thread(target=target)
       t.start()
   print('all DONE at:', time())hu = Hu()# 使用方式如下：@hu.task()def test(num):
   sleep(2)    print 'test: %s' % numif __name__ == '__main__':    for i in range(10):
       test.delay(i, count_down=random.randint(1, 10))
   main()# output(2017-03-21 15:59:20,394) Received task: test
(2017-03-21 15:59:20,396) Received task: test
(2017-03-21 15:59:20,397) Received task: test
(2017-03-21 15:59:20,398) Received task: test
(2017-03-21 15:59:20,400) Received task: test
(2017-03-21 15:59:20,401) Received task: test
(2017-03-21 15:59:20,403) Received task: test
(2017-03-21 15:59:20,404) Received task: test
(2017-03-21 15:59:20,406) Received task: test
(2017-03-21 15:59:20,408) Received task: test
get task: test(2017-03-21 15:59:21,395) Task accepted: test
(2017-03-21 15:59:22,397) Task accepted: test
(2017-03-21 15:59:22,397) Task accepted: test
(2017-03-21 15:59:22,397) Task accepted: test
test: 2get task: test(2017-03-21 15:59:23,399) Task test succeeded in 2.0037419796(2017-03-21 15:59:24,404) Task accepted: test
test: 1get task: test

{
   index: {
       cycle_num: [(func, bargs)]
   }
}