Use environment variables to ensure debug mode is on for event loops

1 files changed, 26 insertions, 5 deletions

diff --git a/aergia b/aergia
index 2918ef6..5fca019 100755
--- a/aergia
+++ b/aergia
@@ -45,10 +45,11 @@ from collections import defaultdict
 from typing import Optional
 import argparse
 import asyncio
+import os
 import signal
 import sys
-import time
 import threading
+import time
 import traceback
 
 
@@ -78,6 +79,20 @@ def thread_join_replacement(
 threading.Thread.join = thread_join_replacement
 
 
+class DebugEventLoopPolicy(asyncio.DefaultEventLoopPolicy):
+    '''
+    An alternative method of ensuring all new event loops are in debug mode.
+    '''
+
+    def new_event_loop(self):
+        loop = super().new_event_loop()
+        loop.set_debug(True)
+        return loop
+
+
+# asyncio.set_event_loop_policy(DebugEventLoopPolicy())
+
+
 class Aergia(object):
 
     # a key-value pair where keys represent frame metadata (see
@@ -97,6 +112,15 @@ class Aergia(object):
 
     @staticmethod
     def start():
+        '''
+        Turns on asyncio debug mode and sets up our signals.
+
+        Debug mode must be on by default to avoid losing samples.
+        Debug mode is required to view the current future being waited on
+        in `Aergia.get_idle_task_frames'. The TimerHandler object otherwise
+        does not keep track of a _source_traceback.
+        '''
+        os.environ["PYTHONASYNCIODEBUG"] = "1"
         signal.signal(signal.SIGALRM,
                       Aergia.cpu_signal_handler)
         signal.setitimer(signal.ITIMER_REAL,
@@ -105,6 +129,7 @@ class Aergia(object):
         Aergia.last_signal_time = Aergia.gettime()
 
     def stop():
+        os.environ["PYTHONASYNCIODEBUG"] = "0"
         Aergia.disable_signals()
         Aergia.print_samples()
 
@@ -240,10 +265,6 @@ class Aergia(object):
         '''Given an asyncio event loop, returns the list of idle task frames.
         A task is considered 'idle' if it is not currently executing.'''
         idle = []
-        if not loop.get_debug():
-            # ensure the loop is in debug mode. Otherwise,
-            # _source_traceback is NULL.
-            loop.set_debug(True)
         for th in loop._scheduled:
             st = th._source_traceback
             if st: