任务分发系统设计（二）

class FunctionLock(object):
   TIMEOUT = 60  # 锁过期超时


   @classmethod
   def gen_name(cls, key):
       return "{prefix}:{key}".format(prefix=cls.__name__, key=key)


   @classmethod
   def get_lock(cls, name, timeout=0):
       return cod.set(name, "1", ex=timeout or cls.TIMEOUT, nx=True)


   @classmethod
   def acquire(cls, key, timeout=0):
       name = cls.gen_name(key)
       return cls.get_lock(name, timeout=timeout)


   @classmethod
   def release(cls, key):
       name = cls.gen_name(key)
       cod.delete(name)




def synchronized(prefix, local=0, timeout=60):
   def decorator(func):
       @functools.wraps(func)
       def wrapper(*args, **kwargs):
           key = prefix + str(args[local])
           if not FunctionLock.acquire(key, timeout=timeout):
               raise RuntimeError("PROCESSING LOCKED")
           try:
               return func(*args, **kwargs)
           finally:
               FunctionLock.release(id)
       return wrapper
   return decorator
复制

class TaskDispenser(object):
 key = "_dispenser_"


  def __str__(self):
      return self.pool_key
     
  def __init__(self, pool_key, router_key, bucket_prefix):
      self.pool_key = pool_key
      self.task_pool = TaskPool(pool_key)
      self.task_router = TaskRouter(router_key, bucket_prefix)
  @synchronized(key)
  def recreate_task_pool(self, mapping):
      """构建任务池
      :param mapping: dict, {value: score}
      """
      buckets = self.task_router.buckets()


      # 收集已经分配的任务
      assigned_flow_ids = []
      for name in buckets:
          tb = TaskBucket(name)
          assigned_flow_ids.extend(tb.peek_all())


      # 移除已经分配的任务
      for fid in assigned_flow_ids:
          mapping.pop(fid, None)


      # 重构任务池
      self.task_pool.clear()
      self.task_pool.add_many(mapping)
 
  @synchronized(key)
  def assign_tasks(self, limit, worker):
      """任务分配
      :param limit: int, 任务数量
      :param worker: str, 任务处理者
      :return: list[int]
      """
      bucket_name = self.task_router.get(worker)
      if not bucket_name:
          bucket_name = self.task_router.add(worker)


      tb = TaskBucket(bucket_name)


      tasks = tb.peek_all()
      size = len(tasks)
      if size == limit:
          return tasks


      # 分配任务时需要先检查自己的桶里是否有未处理的任务
      # 并且确认数量是否满足要求
      if size < limit:
          # 少补
          new_tasks = self.task_pool.assign(limit-size)
          if new_tasks:
              tb.add(new_tasks)
              tasks.extend(new_tasks)
      else:
          # 多不退，通过下次重建任务池的方式重新分配
          for _ in range(size-limit):
              tb.pop(reverse=True)
          tasks = tb.peek_all()


      return tasks


  def check_tasks(self, worker, tasks):
      """校验任务
      :param worker: str, 任务处理者
      :param tasks: list, 需要校验的任务
      :return: bool
      """
      bucket_name = self.task_router.get(worker)
      if not bucket_name:
          return False


      tb = TaskBucket(bucket_name)
      bucket_tasks = tb.peek_all()


      return set(bucket_tasks) == set(tasks)


  @synchronized(key)
  def tasks_done(self, worker, tasks):
      """任务完成"""
      bucket_name = self.task_router.get(worker)
      if not bucket_name:
          return False
      tb = TaskBucket(bucket_name)
      tb.clear()
      return True
复制