Refactor SqliteEventBus for async event handling

- Updated SqliteEventBus to use asyncio for cross-process event listening, replacing the previous threading implementation.
- Changed the processed field in the database from CharField to BooleanField for better data integrity.
- Adjusted event processing logic to accommodate the new async structure, ensuring events are handled correctly across processes.
- Enhanced test cases to support async operations and validate cross-process event communication.

Author: Dylan Huang

eval_protocol/event_bus/sqlite_event_bus.py

@@ -1,3 +1,5 @@
+import asyncio
+import os
 import threading
 import time
 from typing import Any, Optional
@@ -16,17 +18,14 @@ def __init__(self, db_path: Optional[str] = None):
 
         # Use the same database as the evaluation row store
         if db_path is None:
-            import os
-
             from eval_protocol.directory_utils import find_eval_protocol_dir
 
             eval_protocol_dir = find_eval_protocol_dir()
             db_path = os.path.join(eval_protocol_dir, "logs.db")
 
-        self._db = SqliteEventBusDatabase(db_path)
+        self._db: SqliteEventBusDatabase = SqliteEventBusDatabase(db_path)
         self._running = False
-        self._listener_thread: Optional[threading.Thread] = None
-        self._process_id = str(uuid4())
+        self._process_id = str(os.getpid())
 
     def emit(self, event_type: str, data: Any) -> None:
         """Emit an event to all subscribers.
@@ -64,67 +63,54 @@ def start_listening(self) -> None:
 
         logger.debug("[CROSS_PROCESS_LISTEN] Starting cross-process event listening")
         self._running = True
-        self._start_database_listener()
-        logger.debug("[CROSS_PROCESS_LISTEN] Started database listener thread")
+        loop = asyncio.get_running_loop()
+        loop.create_task(self._database_listener_task())
+        logger.debug("[CROSS_PROCESS_LISTEN] Started async database listener task")
 
     def stop_listening(self) -> None:
         """Stop listening for cross-process events."""
         logger.debug("[CROSS_PROCESS_LISTEN] Stopping cross-process event listening")
         self._running = False
-        if self._listener_thread and self._listener_thread.is_alive():
-            logger.debug("[CROSS_PROCESS_LISTEN] Waiting for listener thread to stop")
-            self._listener_thread.join(timeout=1)
-            logger.debug("[CROSS_PROCESS_LISTEN] Listener thread stopped")
-
-    def _start_database_listener(self) -> None:
-        """Start database-based event listener."""
-
-        def database_listener():
-            logger.debug("[CROSS_PROCESS_LISTENER] Starting database listener loop")
-            last_cleanup = time.time()
-
-            while self._running:
-                try:
-                    # Get unprocessed events from other processes
-                    events = self._db.get_unprocessed_events(self._process_id)
-                    if events:
-                        logger.debug(f"[CROSS_PROCESS_LISTENER] Found {len(events)} unprocessed events")
-
-                    for event in events:
-                        if not self._running:
-                            break
-
-                        try:
-                            logger.debug(
-                                f"[CROSS_PROCESS_LISTENER] Processing event {event['event_id']} of type {event['event_type']}"
-                            )
-                            # Handle the event
-                            self._handle_cross_process_event(event["event_type"], event["data"])
-                            logger.debug(f"[CROSS_PROCESS_LISTENER] Successfully processed event {event['event_id']}")
-
-                            # Mark as processed
-                            self._db.mark_event_processed(event["event_id"])
-                            logger.debug(f"[CROSS_PROCESS_LISTENER] Marked event {event['event_id']} as processed")
-
-                        except Exception as e:
-                            logger.debug(f"[CROSS_PROCESS_LISTENER] Failed to process event {event['event_id']}: {e}")
-
-                    # Clean up old events every hour
-                    current_time = time.time()
-                    if current_time - last_cleanup >= 3600:
-                        logger.debug("[CROSS_PROCESS_LISTENER] Cleaning up old events")
-                        self._db.cleanup_old_events()
-                        last_cleanup = current_time
-
-                    # Small sleep to prevent busy waiting
-                    time.sleep(0.1)
-
-                except Exception as e:
-                    logger.debug(f"[CROSS_PROCESS_LISTENER] Database listener error: {e}")
-                    time.sleep(1)
-
-        self._listener_thread = threading.Thread(target=database_listener, daemon=True)
-        self._listener_thread.start()
+
+    async def _database_listener_task(self) -> None:
+        """Single database listener task that processes events and recreates itself."""
+        if not self._running:
+            # this should end the task loop
+            logger.debug("[CROSS_PROCESS_LISTENER] Stopping database listener task")
+            return
+
+        # Get unprocessed events from other processes
+        events = self._db.get_unprocessed_events(str(self._process_id))
+        if events:
+            logger.debug(f"[CROSS_PROCESS_LISTENER] Found {len(events)} unprocessed events")
+        else:
+            logger.debug(f"[CROSS_PROCESS_LISTENER] No unprocessed events found for process {self._process_id}")
+
+        for event in events:
+            logger.debug(
+                f"[CROSS_PROCESS_LISTENER] Processing event {event['event_id']} of type {event['event_type']}"
+            )
+            # Handle the event
+            self._handle_cross_process_event(event["event_type"], event["data"])
+            logger.debug(f"[CROSS_PROCESS_LISTENER] Successfully processed event {event['event_id']}")
+
+            # Mark as processed
+            self._db.mark_event_processed(event["event_id"])
+            logger.debug(f"[CROSS_PROCESS_LISTENER] Marked event {event['event_id']} as processed")
+
+        # Clean up old events every hour
+        current_time = time.time()
+        if not hasattr(self, "_last_cleanup"):
+            self._last_cleanup = current_time
+        elif current_time - self._last_cleanup >= 3600:
+            logger.debug("[CROSS_PROCESS_LISTENER] Cleaning up old events")
+            self._db.cleanup_old_events()
+            self._last_cleanup = current_time
+
+        # Schedule the next task if still running
+        await asyncio.sleep(1.0)
+        loop = asyncio.get_running_loop()
+        loop.create_task(self._database_listener_task())
 
     def _handle_cross_process_event(self, event_type: str, data: Any) -> None:
         """Handle events received from other processes."""

eval_protocol/event_bus/sqlite_event_bus_database.py

@@ -2,7 +2,7 @@
 from typing import Any, List
 from uuid import uuid4
 
-from peewee import CharField, DateTimeField, Model, SqliteDatabase
+from peewee import BooleanField, CharField, DateTimeField, Model, SqliteDatabase
 from playhouse.sqlite_ext import JSONField
 
 from eval_protocol.event_bus.logger import logger
@@ -25,7 +25,7 @@ class Event(BaseModel):  # type: ignore
             data = JSONField()
             timestamp = DateTimeField()
             process_id = CharField()
-            processed = CharField(default="false")  # Track if event has been processed
+            processed = BooleanField(default=False)  # Track if event has been processed
 
         self._Event = Event
         self._db.connect()
@@ -46,7 +46,7 @@ def publish_event(self, event_type: str, data: Any, process_id: str) -> None:
                 data=serialized_data,
                 timestamp=time.time(),
                 process_id=process_id,
-                processed="false",
+                processed=False,
             )
         except Exception as e:
             logger.warning(f"Failed to publish event to database: {e}")
@@ -56,7 +56,7 @@ def get_unprocessed_events(self, process_id: str) -> List[dict]:
         try:
             query = (
                 self._Event.select()
-                .where((self._Event.process_id != process_id) & (self._Event.processed == "false"))
+                .where((self._Event.process_id != process_id) & (~self._Event.processed))
                 .order_by(self._Event.timestamp)
             )
 
@@ -80,16 +80,14 @@ def get_unprocessed_events(self, process_id: str) -> List[dict]:
     def mark_event_processed(self, event_id: str) -> None:
         """Mark an event as processed."""
         try:
-            self._Event.update(processed="true").where(self._Event.event_id == event_id).execute()
+            self._Event.update(processed=True).where(self._Event.event_id == event_id).execute()
         except Exception as e:
             logger.debug(f"Failed to mark event as processed: {e}")
 
     def cleanup_old_events(self, max_age_hours: int = 24) -> None:
         """Clean up old processed events."""
         try:
             cutoff_time = time.time() - (max_age_hours * 3600)
-            self._Event.delete().where(
-                (self._Event.processed == "true") & (self._Event.timestamp < cutoff_time)
-            ).execute()
+            self._Event.delete().where((self._Event.processed) & (self._Event.timestamp < cutoff_time)).execute()
         except Exception as e:
             logger.debug(f"Failed to cleanup old events: {e}")

test_event_bus_helper.py

@@ -0,0 +1,74 @@
+#!/usr/bin/env python3
+"""Helper script for testing event bus cross-process communication."""
+
+import asyncio
+import sys
+import json
+from eval_protocol.event_bus import SqliteEventBus
+from eval_protocol.models import EvaluationRow, InputMetadata
+
+
+async def listener_process(db_path: str):
+    """Run an event bus listener in a separate process."""
+    try:
+        event_bus = SqliteEventBus(db_path=db_path)
+
+        received_events = []
+
+        def test_listener(event_type: str, data):
+            received_events.append((event_type, data))
+
+        event_bus.subscribe(test_listener)
+        event_bus.start_listening()
+
+        # Wait for events for up to 5 seconds
+        start_time = asyncio.get_event_loop().time()
+        while asyncio.get_event_loop().time() - start_time < 5.0:
+            await asyncio.sleep(0.1)
+            if received_events:
+                break
+
+        # Output results to stdout
+        print(json.dumps(received_events))
+        event_bus.stop_listening()
+
+    except Exception as e:
+        print(f"Error in listener process: {e}", file=sys.stderr)
+        sys.exit(1)
+
+
+async def emitter_process(db_path: str, event_type: str, data_json: str):
+    """Emit an event from a separate process."""
+    try:
+        event_bus = SqliteEventBus(db_path=db_path)
+
+        # Parse the data
+        if data_json:
+            data = json.loads(data_json)
+        else:
+            data = None
+
+        event_bus.emit(event_type, data)
+
+    except Exception as e:
+        print(f"Error in emitter process: {e}", file=sys.stderr)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    if len(sys.argv) < 3:
+        print("Usage: python test_event_bus_helper.py <mode> <db_path> [event_type] [data_json]", file=sys.stderr)
+        sys.exit(1)
+
+    mode = sys.argv[1]
+    db_path = sys.argv[2]
+
+    if mode == "listener":
+        asyncio.run(listener_process(db_path))
+    elif mode == "emitter":
+        event_type = sys.argv[3]
+        data_json = sys.argv[4] if len(sys.argv) > 4 else ""
+        asyncio.run(emitter_process(db_path, event_type, data_json))
+    else:
+        print(f"Unknown mode: {mode}", file=sys.stderr)
+        sys.exit(1)