fixed shutdown trigger for case where it calls task.unique(); see #117

Author: craigbarratt

custom_components/pyscript/__init__.py

@@ -166,12 +166,11 @@ async def jupyter_kernel_start(call):
             global_ctx_name, global_sym_table={"__name__": global_ctx_name}, manager=GlobalContextMgr
         )
         global_ctx.set_auto_start(True)
-
         GlobalContextMgr.set(global_ctx_name, global_ctx)
 
         ast_ctx = AstEval(global_ctx_name, global_ctx)
         Function.install_ast_funcs(ast_ctx)
-        kernel = Kernel(call.data, ast_ctx, global_ctx_name)
+        kernel = Kernel(call.data, ast_ctx, global_ctx, global_ctx_name)
         await kernel.session_start()
         hass.states.async_set(call.data["state_var"], json.dumps(kernel.get_ports()))
 
@@ -215,7 +214,9 @@ async def hass_started(event):
     async def hass_stop(event):
         _LOGGER.debug("stopping global contexts")
         await unload_scripts(unload_all=True)
-        # tell reaper task to exit (after other tasks are cancelled)
+        # sync with waiter, and then tell waiter and reaper tasks to exit
+        await Function.waiter_sync()
+        await Function.waiter_stop()
         await Function.reaper_stop()
 
     # Store callbacks to event listeners so we can unsubscribe on unload
@@ -255,7 +256,7 @@ async def unload_scripts(global_ctx_only=None, unload_all=False):
         ctx_delete[global_ctx_name] = global_ctx
     for global_ctx_name, global_ctx in ctx_delete.items():
         GlobalContextMgr.delete(global_ctx_name)
-    await Function.reaper_sync()
+    await Function.waiter_sync()
 
 
 @bind_hass
@@ -501,7 +502,7 @@ def import_recurse(ctx_name, visited, ctx2imports):
             if global_ctx_name not in ctx2files or not ctx2files[global_ctx_name].autoload:
                 _LOGGER.info("Unloaded %s", global_ctx.get_file_path())
             GlobalContextMgr.delete(global_ctx_name)
-    await Function.reaper_sync()
+    await Function.waiter_sync()
 
     #
     # now load the requested files, and files that depend on loaded files

custom_components/pyscript/function.py

@@ -49,9 +49,15 @@ class Function:
 
     #
     # task id of the task that cancels and waits for other tasks,
-    # and also awaits on coros
     #
     task_reaper = None
+    task_reaper_q = None
+
+    #
+    # task id of the task that awaits for coros (used by shutdown triggers)
+    #
+    task_waiter = None
+    task_waiter_q = None
 
     def __init__(self):
         """Warn on Function instantiation."""
@@ -97,10 +103,6 @@ async def task_reaper(reaper_q):
                             await cmd[1]
                         except asyncio.CancelledError:
                             pass
-                    elif cmd[0] == "await":
-                        await cmd[1]
-                    elif cmd[0] == "sync":
-                        await cmd[1].put(0)
                     else:
                         _LOGGER.error("task_reaper: unknown command %s", cmd[0])
                 except asyncio.CancelledError:
@@ -112,36 +114,71 @@ async def task_reaper(reaper_q):
             cls.task_reaper_q = asyncio.Queue(0)
             cls.task_reaper = Function.create_task(task_reaper(cls.task_reaper_q))
 
+        #
+        # start a task which creates tasks to run coros, and then syncs on their completion;
+        # this is used by the shutdown trigger
+        #
+        async def task_waiter(waiter_q):
+            aws = []
+            while True:
+                try:
+                    cmd = await waiter_q.get()
+                    if cmd[0] == "exit":
+                        return
+                    if cmd[0] == "await":
+                        aws.append(cls.create_task(cmd[1]))
+                    elif cmd[0] == "sync":
+                        if len(aws) > 0:
+                            await asyncio.gather(*aws)
+                            aws = []
+                        await cmd[1].put(0)
+                    else:
+                        _LOGGER.error("task_waiter: unknown command %s", cmd[0])
+                except asyncio.CancelledError:
+                    raise
+                except Exception:
+                    _LOGGER.error("task_waiter: got exception %s", traceback.format_exc(-1))
+
+        if not cls.task_waiter:
+            cls.task_waiter_q = asyncio.Queue(0)
+            cls.task_waiter = Function.create_task(task_waiter(cls.task_waiter_q))
+            _LOGGER.debug("task_waiter: started")
+
+    @classmethod
+    def reaper_cancel(cls, task):
+        """Send a task to be canceled by the reaper."""
+        cls.task_reaper_q.put_nowait(["cancel", task])
+
     @classmethod
     async def reaper_stop(cls):
-        """Tell the reaper task to exit by sending a special task None."""
+        """Tell the reaper task to exit."""
         if cls.task_reaper:
             cls.task_reaper_q.put_nowait(["exit"])
             await cls.task_reaper
             cls.task_reaper = None
             cls.task_reaper_q = None
 
     @classmethod
-    def reaper_cancel(cls, task):
-        """Send a task to be canceled by the reaper."""
-        cls.task_reaper_q.put_nowait(["cancel", task])
-
-    @classmethod
-    def reaper_await(cls, coro):
-        """Send a coro to be awaited by the reaper."""
-        cls.task_reaper_q.put_nowait(["await", coro])
+    def waiter_await(cls, coro):
+        """Send a coro to be awaited by the waiter task."""
+        cls.task_waiter_q.put_nowait(["await", coro])
 
     @classmethod
-    async def reaper_sync(cls):
-        """Wait until the reaper queue is empty."""
-        sync_q = asyncio.Queue(0)
-        sync_q.put_nowait(["sync", sync_q])
-        await sync_q.get()
+    async def waiter_sync(cls):
+        """Wait until the waiter queue is empty."""
+        if cls.task_waiter:
+            sync_q = asyncio.Queue(0)
+            cls.task_waiter_q.put_nowait(["sync", sync_q])
+            await sync_q.get()
 
     @classmethod
-    def reaper_exit(cls):
-        """Send an exit request to the reaper."""
-        cls.task_reaper_q.put_nowait(["exit"])
+    async def waiter_stop(cls):
+        """Tell the waiter task to exit."""
+        if cls.task_waiter:
+            cls.task_waiter_q.put_nowait(["exit"])
+            await cls.task_waiter
+            cls.task_waiter = None
+            cls.task_waiter_q = None
 
     @classmethod
     async def async_sleep(cls, duration):

custom_components/pyscript/jupyter_kernel.py

@@ -189,9 +189,10 @@ def close(self):
 class Kernel:
     """Define a Jupyter Kernel class."""
 
-    def __init__(self, config, ast_ctx, global_ctx_name):
+    def __init__(self, config, ast_ctx, global_ctx, global_ctx_name):
         """Initialize a Kernel object, one instance per session."""
         self.config = config.copy()
+        self.global_ctx = global_ctx
         self.global_ctx_name = global_ctx_name
         self.ast_ctx = ast_ctx
 
@@ -346,11 +347,15 @@ async def shell_handler(self, shell_socket, wire_msg):
             if code.startswith("_rwho_ls = %who_ls"):
                 code = "print([])"
 
+            self.global_ctx.set_auto_start(False)
             self.ast_ctx.parse(code)
             exc = self.ast_ctx.get_exception_obj()
             if exc is None:
                 result = await self.ast_ctx.eval()
                 exc = self.ast_ctx.get_exception_obj()
+            await Function.waiter_sync()
+            self.global_ctx.set_auto_start(True)
+            self.global_ctx.start()
             if exc:
                 traceback_mesg = self.ast_ctx.get_exception_long().split("\n")
 

custom_components/pyscript/trigger.py

@@ -832,11 +832,12 @@ def stop(self):
                 Mqtt.notify_del(self.mqtt_trigger[0], self.notify_q)
             if self.task:
                 Function.reaper_cancel(self.task)
+                self.task = None
         if self.run_on_shutdown:
             notify_type = "shutdown"
             notify_info = {"trigger_type": "time", "trigger_time": "shutdown"}
             action_future = self.call_action(notify_type, notify_info, run_task=False)
-            Function.reaper_await(action_future)
+            Function.waiter_await(action_future)
 
     def start(self):
         """Start this trigger task."""

docs/new_features.rst

@@ -36,3 +36,6 @@ The new features since 1.1.0 in master include:
 Breaking changes since 1.1.0 include:
 
 Bug fixes since 1.1.0 include:
+
+- Fixed shutdown trigger for case where it calls ``task.unique()`` (#117).
+- Added error message for invalid ``@time_active`` argument, by @dlashua (#118).

tests/test_function.py

@@ -1134,10 +1134,11 @@ async def test_service_call_params(hass):
         assert call.call_args[1] == {"context": Context(id="test"), "blocking": False}
 
     # Stop all tasks to avoid conflicts with other tests
+    await Function.waiter_stop()
     await Function.reaper_stop()
 
 
-async def test_serive_call_blocking(hass, caplog):
+async def test_service_call_blocking(hass, caplog):
     """Test that service calls with blocking=True actually block."""
     notify_q = asyncio.Queue(0)
 
@@ -1171,6 +1172,9 @@ def func_startup():
     service.call("pyscript", "long_sleep", blocking=True, limit=1e-6)
     pyscript.done = [seq_num, pyscript.var1]
 
+    seq_num += 1
+    pyscript.done = [seq_num]
+
 @service
 def long_sleep():
     task.sleep(10000)
@@ -1187,3 +1191,4 @@ def service1():
     assert literal_eval(await wait_until_done(notify_q)) == [2, "3"]
     assert literal_eval(await wait_until_done(notify_q)) == [3, "4"]
     assert literal_eval(await wait_until_done(notify_q)) == [4, "5"]
+    assert literal_eval(await wait_until_done(notify_q)) == [5]