TaskTable.py 1.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556
  1. import threading
  2. from google.protobuf import message as _message
  3. class TaskTable(object):
  4. def __init__(self):
  5. self.taskMap = {}
  6. self.lock = threading.Lock()
  7. def PushTask(self, primary_key, msg: _message.Message):
  8. if len(self.taskMap.items()) > 0:
  9. return False
  10. else:
  11. self.lock.acquire()
  12. self.taskMap[primary_key] = msg
  13. self.lock.release()
  14. return True
  15. def UpdateResult(self, primary_key, ret):
  16. if primary_key in self.taskMap:
  17. taskTable = self.taskMap[primary_key]
  18. self.lock.acquire()
  19. self.taskMap[primary_key] = [taskTable, ret]
  20. self.lock.release()
  21. return True
  22. else:
  23. return False
  24. def GetTask(self):
  25. if len(self.taskMap.items()) == 0:
  26. return None
  27. else:
  28. for item in self.taskMap.items():
  29. key, task = item
  30. if isinstance(task, (_message.Message)):
  31. return task
  32. else:
  33. return None
  34. def GetResult(self, primary_key):
  35. if len(self.taskMap.items()) == 0:
  36. return None
  37. else:
  38. if primary_key in self.taskMap:
  39. result = self.taskMap[primary_key]
  40. if isinstance(result, (list)):
  41. ret = result[1]
  42. else:
  43. return None
  44. self.lock.acquire()
  45. self.taskMap.clear()
  46. self.lock.release()
  47. return ret
  48. else:
  49. return None