Repository: proxystore/aeris

Issue:
Create special purpose `@loop` decorators
A common type of loop is one that runs at regular intervals and one that waits for an event to be triggered.
```python
class Example(Behavior):
    def __init__(self) -> None:
        self.foo = threading.Event  # maybe a custom aeris provided event that has a lock?

    @timer_loop(interval=60)
    def run_every_minute(self) -> None:
        # The decorator handles the `shutdown` event.
        ...
    
    @event_loop(name='foo')
    def run_when_foo_is_set(self) -> None:
        # Maybe decorator holds a lock on the event? (optional?)
        ...
        # Decorator clears foo when function done
```

Hints from issue comments:
@gpauloski: Closed by #72 

Solution patch:
diff --git a/aeris/behavior.py b/aeris/behavior.py
index 2f71248..70f18cb 100644
--- a/aeris/behavior.py
+++ b/aeris/behavior.py
@@ -5,6 +5,7 @@
 import logging
 import sys
 import threading
+from datetime import timedelta
 from typing import Any
 from typing import Callable
 from typing import Generic
@@ -24,6 +25,7 @@
 else:  # pragma: <3.11 cover
     from typing_extensions import Self
 
+from aeris.event import or_event
 from aeris.handle import Handle
 from aeris.handle import HandleDict
 from aeris.handle import HandleList
@@ -81,7 +83,7 @@ def behavior_actions(self) -> dict[str, Action[Any, Any]]:
         actions: dict[str, Action[Any, Any]] = {}
         for name in dir(self):
             attr = getattr(self, name)
-            if _is_actor_method_type(attr, 'action'):
+            if _is_agent_method_type(attr, 'action'):
                 actions[name] = attr
         return actions
 
@@ -94,7 +96,7 @@ def behavior_loops(self) -> dict[str, ControlLoop]:
         loops: dict[str, ControlLoop] = {}
         for name in dir(self):
             attr = getattr(self, name)
-            if _is_actor_method_type(attr, 'loop'):
+            if _is_agent_method_type(attr, 'loop'):
                 loops[name] = attr
         return loops
 
@@ -211,7 +213,7 @@ def perform(self):
                 ...
         ```
     """
-    method._actor_method_type = 'action'  # type: ignore[attr-defined]
+    method._agent_method_type = 'action'  # type: ignore[attr-defined]
     return method
 
 
@@ -241,7 +243,7 @@ def listen(self, shutdown: threading.Event) -> None:
         TypeError: if the method signature does not conform to the
             [`ControlLoop`][aeris.behavior.ControlLoop] protocol.
     """
-    method._actor_method_type = 'loop'  # type: ignore[attr-defined]
+    method._agent_method_type = 'loop'  # type: ignore[attr-defined]
 
     if sys.version_info >= (3, 10):  # pragma: >=3.10 cover
         found_sig = inspect.signature(method, eval_str=True)
@@ -269,9 +271,141 @@ def _wrapped(self: BehaviorT, *args: P.args, **kwargs: P.kwargs) -> R:
     return _wrapped
 
 
-def _is_actor_method_type(obj: Any, kind: str) -> bool:
+def event(
+    name: str,
+) -> Callable[
+    [Callable[[BehaviorT], None]],
+    Callable[[BehaviorT, threading.Event], None],
+]:
+    """Decorator that annotates a method of a behavior as an event loop.
+
+    An event loop is a special type of control loop that runs when a
+    [`threading.Event`][threading.Event] is set. The event is cleared
+    after the loop runs.
+
+    Example:
+        ```python
+        import threading
+        from aeris.behavior import Behavior, timer
+
+        class Example(Behavior):
+            def __init__(self) -> None:
+                self.alert = threading.Event()
+
+            @event('alert')
+            def handle(self) -> None:
+                # Runs every time alter is set
+                ...
+        ```
+
+    Args:
+        name: Attribute name of the [`threading.Event`][theading.Event]
+            to wait on.
+
+    Raises:
+        AttributeError: Raised at runtime if no attribute named `name`
+            exists on the behavior.
+        TypeError: Raised at runtime if the attribute named `name` is not
+            a [`threading.Event`][threading.Event].
+    """
+
+    def decorator(
+        method: Callable[[BehaviorT], None],
+    ) -> Callable[[BehaviorT, threading.Event], None]:
+        method._agent_method_type = 'loop'  # type: ignore[attr-defined]
+
+        @functools.wraps(method)
+        def _wrapped(self: BehaviorT, shutdown: threading.Event) -> None:
+            event = getattr(self, name)
+            if not isinstance(event, threading.Event):
+                raise TypeError(
+                    f'Attribute {name} of {type(self).__class__} has type '
+                    f'{type(event).__class__}. Expected threading.Event.',
+                )
+
+            logger.debug(
+                'Started %r event loop for %s (event: %r)',
+                method.__name__,
+                self,
+                name,
+            )
+            combined = or_event(shutdown, event)
+            while True:
+                combined.wait()
+                if shutdown.is_set():
+                    break
+                elif event.is_set():
+                    try:
+                        method(self)
+                    finally:
+                        event.clear()
+                else:
+                    raise AssertionError('Unreachable.')
+            logger.debug('Exited %r event loop for %s', method.__name__, self)
+
+        return _wrapped
+
+    return decorator
+
+
+def timer(
+    interval: float | timedelta,
+) -> Callable[
+    [Callable[[BehaviorT], None]],
+    Callable[[BehaviorT, threading.Event], None],
+]:
+    """Decorator that annotates a method of a behavior as a timer loop.
+
+    A timer loop is a special type of control loop that runs at a set
+    interval. The method will always be called once before the first
+    sleep.
+
+    Example:
+        ```python
+        from aeris.behavior import Behavior, timer
+
+        class Example(Behavior):
+            @timer(interval=1)
+            def listen(self) -> None:
+                # Runs every 1 second
+                ...
+        ```
+
+    Args:
+        interval: Seconds or a [`timedelta`][datetime.timedelta] to wait
+            between invoking the method.
+    """
+    interval = (
+        interval.total_seconds()
+        if isinstance(interval, timedelta)
+        else interval
+    )
+
+    def decorator(
+        method: Callable[[BehaviorT], None],
+    ) -> Callable[[BehaviorT, threading.Event], None]:
+        method._agent_method_type = 'loop'  # type: ignore[attr-defined]
+
+        @functools.wraps(method)
+        def _wrapped(self: BehaviorT, shutdown: threading.Event) -> None:
+            logger.debug(
+                'Started %r timer loop for %s (interval: %fs)',
+                method.__name__,
+                self,
+                interval,
+            )
+            while not shutdown.wait(interval):
+                method(self)
+            logger.debug('Exited %r timer loop for %s', method.__name__, self)
+
+        return _wrapped
+
+    return decorator
+
+
+def _is_agent_method_type(obj: Any, kind: str) -> bool:
     return (
         callable(obj)
-        and hasattr(obj, '_actor_method_type')
-        and obj._actor_method_type == kind
+        and hasattr(obj, '_agent_method_type')
+        and obj._agent_method_type == kind
     )
diff --git a/aeris/event.py b/aeris/event.py
new file mode 100644
index 0000000..ca33d00
--- /dev/null
+++ b/aeris/event.py
@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+import threading
+from typing import Any
+from typing import Callable
+
+
+def _or_set(event: Any) -> None:
+    event._set()
+    event.changed()
+
+
+def _or_clear(event: Any) -> None:
+    event._clear()
+    event.changed()
+
+
+def _orify(
+    event: threading.Event,
+    changed_callback: Callable[[], None],
+) -> None:
+    event._set = event.set  # type: ignore[attr-defined]
+    event._clear = event.clear  # type: ignore[attr-defined]
+    event.changed = changed_callback  # type: ignore[attr-defined]
+    event.set = lambda: _or_set(event)  # type: ignore[method-assign]
+    event.clear = lambda: _or_clear(event)  # type: ignore[method-assign]
+
+
+def or_event(*events: threading.Event) -> threading.Event:
+    """Create a combined event that is set when any input events are set.
+
+    Note:
+        The creator can wait on the combined event, but must still check
+        each individual event to see which was set.
+
+    Warning:
+        This works by dynamically replacing methods on the inputs events
+        with custom methods that trigger callbacks.
+
+    Note:
+        Based on the
+        [Stack Overflow answer][https://stackoverflow.com/a/12320352].
+
+    Args:
+        events: One or more events to combine.
+
+    Returns:
+        A single event that is set when any of the input events is set.
+    """
+    combined = threading.Event()
+
+    def changed() -> None:
+        bools = [e.is_set() for e in events]
+        if any(bools):
+            combined.set()
+        else:
+            combined.clear()
+
+    for e in events:
+        _orify(e, changed)
+
+    changed()
+    return combined


Test patch:
diff --git a/tests/behavior_test.py b/tests/behavior_test.py
index cc16503..df44bf7 100644
--- a/tests/behavior_test.py
+++ b/tests/behavior_test.py
@@ -1,11 +1,14 @@
 from __future__ import annotations
 
 import threading
+import time
 
 import pytest
 
 from aeris.behavior import Behavior
+from aeris.behavior import event
 from aeris.behavior import loop
+from aeris.behavior import timer
 from aeris.handle import Handle
 from aeris.handle import HandleDict
 from aeris.handle import HandleList
@@ -14,6 +17,8 @@
 from testing.behavior import HandleBehavior
 from testing.behavior import IdentityBehavior
 from testing.behavior import WaitBehavior
+from testing.constant import TEST_LOOP_SLEEP
+from testing.constant import TEST_THREAD_JOIN_TIMEOUT
 
 
 def test_initialize_base_type_error() -> None:
@@ -63,6 +68,72 @@ def test_behavior_loops() -> None:
     behavior.on_shutdown()
 
 
+def test_behavior_event() -> None:
+    class _Event(Behavior):
+        def __init__(self) -> None:
+            self.event = threading.Event()
+            self.ran = threading.Event()
+            self.bad = 42
+
+        @event('event')
+        def run(self) -> None:
+            self.ran.set()
+
+        @event('missing')
+        def missing_event(self) -> None: ...
+
+        @event('bad')
+        def bad_event(self) -> None: ...
+
+    behavior = _Event()
+
+    loops = behavior.behavior_loops()
+    assert set(loops) == {'bad_event', 'missing_event', 'run'}
+
+    shutdown = threading.Event()
+
+    with pytest.raises(AttributeError, match='missing'):
+        behavior.missing_event(shutdown)
+    with pytest.raises(TypeError, match='bad'):
+        behavior.bad_event(shutdown)
+
+    handle = threading.Thread(target=behavior.run, args=(shutdown,))
+    handle.start()
+
+    for _ in range(5):
+        assert not behavior.ran.is_set()
+        behavior.event.set()
+        assert behavior.ran.wait(1)
+        behavior.ran.clear()
+
+    shutdown.set()
+    handle.join(TEST_THREAD_JOIN_TIMEOUT)
+
+
+def test_behavior_timer() -> None:
+    class _Timer(Behavior):
+        def __init__(self) -> None:
+            self.count = 0
+
+        @timer(TEST_LOOP_SLEEP)
+        def counter(self) -> None:
+            self.count += 1
+
+    behavior = _Timer()
+
+    loops = behavior.behavior_loops()
+    assert set(loops) == {'counter'}
+
+    shutdown = threading.Event()
+    handle = threading.Thread(target=behavior.counter, args=(shutdown,))
+    handle.start()
+
+    time.sleep(TEST_LOOP_SLEEP * 10)
+    shutdown.set()
+
+    handle.join(TEST_THREAD_JOIN_TIMEOUT)
+
+
 def test_behavior_handles() -> None:
     handle = ProxyHandle(EmptyBehavior())
     behavior = HandleBehavior(handle)
diff --git a/tests/event_test.py b/tests/event_test.py
new file mode 100644
index 0000000..ae7943f
--- /dev/null
+++ b/tests/event_test.py
@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import threading
+
+from aeris.event import or_event
+
+
+def test_or_event() -> None:
+    a = threading.Event()
+    b = threading.Event()
+
+    c = or_event(a, b)
+
+    assert not c.is_set()
+    a.set()
+    assert c.is_set()
+    a.clear()
+    assert not c.is_set()
+
+    a.set()
+    b.set()
+    assert c.is_set()
+
+    # Both events must be cleared
+    a.clear()
+    assert c.is_set()
+    b.clear()
+    assert not c.is_set()