refac: Attempt to use PriorityQueue

Author: fgmacedo

statemachine/engines/async_.py

@@ -1,5 +1,4 @@
 import asyncio
-import heapq
 from time import time
 from typing import TYPE_CHECKING
 
@@ -63,8 +62,8 @@ async def processing_loop(self):
         first_result = self._sentinel
         try:
             # Execute the triggers in the queue in FIFO order until the queue is empty
-            while self._external_queue:
-                trigger_data = heapq.heappop(self._external_queue)
+            while self._running and not self.empty():
+                trigger_data = self.pop()
                 current_time = time()
                 if trigger_data.execution_time > current_time:
                     self.put(trigger_data)
@@ -77,7 +76,7 @@ async def processing_loop(self):
                 except Exception:
                     # Whe clear the queue as we don't have an expected behavior
                     # and cannot keep processing
-                    self._external_queue.clear()
+                    self.clear()
                     raise
         finally:
             self._processing.release()

statemachine/engines/base.py

@@ -1,4 +1,5 @@
-import heapq
+from queue import PriorityQueue
+from queue import Queue
 from threading import Lock
 from typing import TYPE_CHECKING
 from weakref import proxy
@@ -15,19 +16,46 @@
 
 class BaseEngine:
     def __init__(self, sm: "StateMachine", rtc: bool = True):
-        self.sm: StateMachine = proxy(sm)
-        self._external_queue: list = []
         self._sentinel = object()
         self._rtc = rtc
-        self._processing = Lock()
         self._running = True
+        self._init(sm)
+
+    def _init(self, sm: "StateMachine"):
+        self.sm: StateMachine = proxy(sm)
+        self._external_queue: Queue = PriorityQueue()
+        self._processing = Lock()
+
+    def __getstate__(self) -> dict:
+        state = self.__dict__.copy()
+        del state["_external_queue"]
+        del state["_processing"]
+        del state["sm"]
+        return state
+
+    def __setstate__(self, state: dict) -> None:
+        for attr, value in state.items():
+            setattr(self, attr, value)
+
+    def empty(self):
+        return self._external_queue.qsize() == 0
 
     def put(self, trigger_data: TriggerData):
         """Put the trigger on the queue without blocking the caller."""
         if not self._running and not self.sm.allow_event_without_transition:
             raise TransitionNotAllowed(trigger_data.event, self.sm.current_state)
 
-        heapq.heappush(self._external_queue, trigger_data)
+        self._external_queue.put(trigger_data)
+
+    def pop(self):
+        try:
+            return self._external_queue.get(block=False)
+        except Exception:
+            return None
+
+    def clear(self):
+        with self._external_queue.mutex:
+            self._external_queue.queue.clear()
 
     def start(self):
         if self.sm.current_state_value is not None:

statemachine/engines/sync.py

@@ -1,4 +1,3 @@
-import heapq
 from time import sleep
 from time import time
 from typing import TYPE_CHECKING
@@ -50,7 +49,7 @@ def processing_loop(self):
         """
         if not self._rtc:
             # The machine is in "synchronous" mode
-            trigger_data = heapq.heappop(self._external_queue)
+            trigger_data = self.pop()
             return self._trigger(trigger_data)
 
         # We make sure that only the first event enters the processing critical section,
@@ -64,8 +63,8 @@ def processing_loop(self):
         first_result = self._sentinel
         try:
             # Execute the triggers in the queue in FIFO order until the queue is empty
-            while self._running and self._external_queue:
-                trigger_data = heapq.heappop(self._external_queue)
+            while self._running and not self.empty():
+                trigger_data = self.pop()
                 current_time = time()
                 if trigger_data.execution_time > current_time:
                     self.put(trigger_data)
@@ -78,7 +77,7 @@ def processing_loop(self):
                 except Exception:
                     # Whe clear the queue as we don't have an expected behavior
                     # and cannot keep processing
-                    self._external_queue.clear()
+                    self.clear()
                     raise
         finally:
             self._processing.release()
@@ -137,7 +136,7 @@ def _activate(self, trigger_data: TriggerData, transition: "Transition"):  # noq
         self.sm._callbacks.call(transition.after.key, *args, **kwargs)
 
         if target.final:
-            self._external_queue.clear()
+            self.clear()
             self._running = False
 
         if len(result) == 0:

statemachine/statemachine.py

@@ -148,7 +148,7 @@ def __setstate__(self, state):
         self._listeners: Dict[Any, Any] = {}
 
         self._register_callbacks([])
-        self.add_listener(*listeners.keys())
+        self.add_listener(*listeners.values())
         self._engine = self._get_engine(rtc)
 
     def _get_initial_state(self):