delayed_queue.py 2.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283
  1. # -*- coding: utf-8 -*-
  2. #
  3. # Copyright 2014 Thomas Amland <thomas.amland@gmail.com>
  4. #
  5. # Licensed under the Apache License, Version 2.0 (the "License");
  6. # you may not use this file except in compliance with the License.
  7. # You may obtain a copy of the License at
  8. #
  9. # http://www.apache.org/licenses/LICENSE-2.0
  10. #
  11. # Unless required by applicable law or agreed to in writing, software
  12. # distributed under the License is distributed on an "AS IS" BASIS,
  13. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14. # See the License for the specific language governing permissions and
  15. # limitations under the License.
  16. import time
  17. import threading
  18. from collections import deque
  19. class DelayedQueue(object):
  20. def __init__(self, delay):
  21. self.delay_sec = delay
  22. self._lock = threading.Lock()
  23. self._not_empty = threading.Condition(self._lock)
  24. self._queue = deque()
  25. self._closed = False
  26. def put(self, element, delay=False):
  27. """Add element to queue."""
  28. self._lock.acquire()
  29. self._queue.append((element, time.time(), delay))
  30. self._not_empty.notify()
  31. self._lock.release()
  32. def close(self):
  33. """Close queue, indicating no more items will be added."""
  34. self._closed = True
  35. # Interrupt the blocking _not_empty.wait() call in get
  36. self._not_empty.acquire()
  37. self._not_empty.notify()
  38. self._not_empty.release()
  39. def get(self):
  40. """Remove and return an element from the queue, or this queue has been
  41. closed raise the Closed exception.
  42. """
  43. while True:
  44. # wait for element to be added to queue
  45. self._not_empty.acquire()
  46. while len(self._queue) == 0 and not self._closed:
  47. self._not_empty.wait()
  48. if self._closed:
  49. self._not_empty.release()
  50. return None
  51. head, insert_time, delay = self._queue[0]
  52. self._not_empty.release()
  53. # wait for delay if required
  54. if delay:
  55. time_left = insert_time + self.delay_sec - time.time()
  56. while time_left > 0:
  57. time.sleep(time_left)
  58. time_left = insert_time + self.delay_sec - time.time()
  59. # return element if it's still in the queue
  60. with self._lock:
  61. if len(self._queue) > 0 and self._queue[0][0] is head:
  62. self._queue.popleft()
  63. return head
  64. def remove(self, predicate):
  65. """Remove and return the first items for which predicate is True,
  66. ignoring delay."""
  67. with self._lock:
  68. for i, (elem, t, delay) in enumerate(self._queue):
  69. if predicate(elem):
  70. del self._queue[i]
  71. return elem
  72. return None