diff --git a/pywebio/session/base.py b/pywebio/session/base.py index 05cb231a..1687d8db 100644 --- a/pywebio/session/base.py +++ b/pywebio/session/base.py @@ -29,6 +29,7 @@ class Session: next_client_event on_task_exception register_callback + need_keep_alive defer_call @@ -159,6 +160,9 @@ def defer_call(self, func): """设置会话结束时调用的函数。可以用于资源清理。""" self.deferred_functions.append(func) + def need_keep_alive(self) -> bool: + raise NotImplementedError + def get_session_info_from_headers(headers): """从Http请求头中获取会话信息 diff --git a/pywebio/session/coroutinebased.py b/pywebio/session/coroutinebased.py index e544f06c..750f2c10 100644 --- a/pywebio/session/coroutinebased.py +++ b/pywebio/session/coroutinebased.py @@ -87,14 +87,22 @@ def __init__(self, target, session_info, on_task_command=None, on_session_close= self._closed = False + self._need_keep_alive = False + # 当前会话未结束运行(已创建和正在运行的)的协程数量。当 _alive_coro_cnt 变为 0 时,会话结束。 self._alive_coro_cnt = 1 - main_task = Task(target(), session=self, on_coro_stop=self._on_task_finish) + main_task = Task(self._start_main_task(target), session=self, on_coro_stop=self._on_task_finish) self.coros[main_task.coro_id] = main_task self._step_task(main_task) + async def _start_main_task(self, target): + await target() + if self.need_keep_alive(): + from ..session import hold + await hold() + def _step_task(self, task, result=None): asyncio.get_event_loop().call_soon_threadsafe(partial(task.step, result)) @@ -203,6 +211,8 @@ async def callback_coro(): callback_task.coro.send(None) cls.get_current_session().coros[callback_task.coro_id] = callback_task + self._need_keep_alive = True + return callback_task.coro_id def run_async(self, coro_obj): @@ -227,6 +237,9 @@ async def run_asyncio_coroutine(self, coro_obj): res = await WebIOFuture(coro=coro_obj) return res + def need_keep_alive(self) -> bool: + return self._need_keep_alive + class TaskHandler: """The handler of coroutine task diff --git a/pywebio/session/threadbased.py b/pywebio/session/threadbased.py index c0bf628c..50c32d3c 100644 --- a/pywebio/session/threadbased.py +++ b/pywebio/session/threadbased.py @@ -91,10 +91,15 @@ def main_task(target): for t in self.threads: if t.is_alive() and t is not threading.current_thread(): t.join() - try: - self.send_task_command(dict(command='close_session')) - except SessionClosedException: - pass + + if self.need_keep_alive(): + from ..session import hold + hold() + else: + try: + self.send_task_command(dict(command='close_session')) + except SessionClosedException: + pass self._trigger_close_event() self.close() @@ -281,6 +286,10 @@ def register_thread(self, t: threading.Thread): event_mq = queue.Queue(maxsize=self.event_mq_maxsize) # 线程内的用户事件队列 self.task_mqs[self._get_task_id(t)] = event_mq + def need_keep_alive(self) -> bool: + # if callback thread is activated, then the session need to keep alive + return self.callback_thread is not None + class ScriptModeSession(ThreadBasedSession): """Script mode的会话实现""" diff --git a/test/1.basic.py b/test/1.basic.py index 9510751a..875541cb 100644 --- a/test/1.basic.py +++ b/test/1.basic.py @@ -32,6 +32,8 @@ def test(server_proc: subprocess.Popen, browser: Chrome): time.sleep(1) template.save_output(browser, '1.basic.html') + browser.get('http://localhost:8080/') # to close current session + template.test_defer_call()