From 8cd6a2f342431d70afa9433b75d026ec7bb78e0d Mon Sep 17 00:00:00 2001 From: Jakub Trllo <43494761+iLLiCiTiT@users.noreply.github.com> Date: Wed, 31 Jul 2024 17:26:34 +0200 Subject: [PATCH 1/4] wrap tray info into object --- client/ayon_core/tools/tray/lib.py | 311 +++++++++++++++++++++-------- 1 file changed, 233 insertions(+), 78 deletions(-) diff --git a/client/ayon_core/tools/tray/lib.py b/client/ayon_core/tools/tray/lib.py index ad190482a8..e6de884cb5 100644 --- a/client/ayon_core/tools/tray/lib.py +++ b/client/ayon_core/tools/tray/lib.py @@ -144,17 +144,6 @@ def get_tray_storage_dir() -> str: return get_ayon_appdirs("tray") -def _get_tray_information(tray_url: str) -> Optional[Dict[str, Any]]: - if not tray_url: - return None - try: - response = requests.get(f"{tray_url}/tray") - response.raise_for_status() - return response.json() - except (requests.HTTPError, requests.ConnectionError): - return None - - def _get_tray_info_filepath( server_url: Optional[str] = None, variant: Optional[str] = None @@ -165,6 +154,51 @@ def _get_tray_info_filepath( return os.path.join(hash_dir, filename) +def _get_tray_file_info( + server_url: Optional[str] = None, + variant: Optional[str] = None +) -> Tuple[Optional[Dict[str, Any]], Optional[float]]: + filepath = _get_tray_info_filepath(server_url, variant) + if not os.path.exists(filepath): + return None, None + file_modified = os.path.getmtime(filepath) + try: + with open(filepath, "r") as stream: + data = json.load(stream) + except Exception: + return None, file_modified + + return data, file_modified + + +def _remove_tray_server_url( + server_url: Optional[str], + variant: Optional[str], + file_modified: Optional[float], +): + """Remove tray information file. + + Called from tray logic, do not use on your own. + + Args: + server_url (Optional[str]): AYON server url. + variant (Optional[str]): Settings variant. + file_modified (Optional[float]): File modified timestamp. Is validated + against current state of file. + + """ + filepath = _get_tray_info_filepath(server_url, variant) + if not os.path.exists(filepath): + return + + if ( + file_modified is not None + and os.path.getmtime(filepath) != file_modified + ): + return + os.remove(filepath) + + def get_tray_file_info( server_url: Optional[str] = None, variant: Optional[str] = None @@ -182,15 +216,156 @@ def get_tray_file_info( Optional[Dict[str, Any]]: Tray information. """ - filepath = _get_tray_info_filepath(server_url, variant) - if not os.path.exists(filepath): + file_info, _ = _get_tray_file_info(server_url, variant) + return file_info + + +def _get_tray_rest_information(tray_url: str) -> Optional[Dict[str, Any]]: + if not tray_url: return None try: - with open(filepath, "r") as stream: - data = json.load(stream) - except Exception: + response = requests.get(f"{tray_url}/tray") + response.raise_for_status() + return response.json() + except (requests.HTTPError, requests.ConnectionError): return None - return data + + +class TrayInfo: + def __init__( + self, + server_url: str, + variant: str, + timeout: Optional[int] = None + ): + self.server_url = server_url + self.variant = variant + + if timeout is None: + timeout = 10 + + self._timeout = timeout + + self._file_modified = None + self._file_info = None + self._file_info_cached = False + self._tray_info = None + self._tray_info_cached = False + self._file_state = None + self._state = None + + @classmethod + def new( + cls, + server_url: Optional[str] = None, + variant: Optional[str] = None, + timeout: Optional[int] = None, + wait_to_start: Optional[bool] = True + ) -> "TrayInfo": + server_url, variant = _get_server_and_variant(server_url, variant) + obj = cls(server_url, variant, timeout=timeout) + if wait_to_start: + obj.wait_to_start() + return obj + + def get_pid(self) -> Optional[int]: + file_info = self.get_file_info() + if file_info: + return file_info.get("pid") + return None + + def reset(self): + self._file_modified = None + self._file_info = None + self._file_info_cached = False + self._tray_info = None + self._tray_info_cached = False + self._state = None + self._file_state = None + + def get_file_info(self) -> Optional[Dict[str, Any]]: + if not self._file_info_cached: + file_info, file_modified = _get_tray_file_info( + self.server_url, self.variant + ) + self._file_info = file_info + self._file_modified = file_modified + self._file_info_cached = True + return self._file_info + + def get_file_url(self) -> Optional[str]: + file_info = self.get_file_info() + if file_info: + return file_info.get("url") + return None + + def get_tray_url(self) -> Optional[str]: + info = self.get_tray_info() + if info: + return self.get_file_url() + return None + + def get_tray_info(self) -> Optional[Dict[str, Any]]: + if self._tray_info_cached: + return self._tray_info + + tray_url = self.get_file_url() + tray_info = None + if tray_url: + tray_info = _get_tray_rest_information(tray_url) + + self._tray_info = tray_info + self._tray_info_cached = True + return self._tray_info + + def get_file_state(self) -> int: + if self._file_state is not None: + return self._file_state + + state = TrayState.NOT_RUNNING + file_info = self.get_file_info() + if file_info: + state = TrayState.STARTING + if file_info.get("started") is True: + state = TrayState.RUNNING + self._file_state = state + return self._file_state + + def get_state(self) -> int: + if self._state is not None: + return self._state + + state = self.get_file_state() + if state == TrayState.RUNNING and not self.get_tray_info(): + state = TrayState.NOT_RUNNING + pid = self.pid + if pid: + _kill_tray_process(pid) + # Remove the file as tray is not running anymore and update + # the state of this object. + _remove_tray_server_url( + self.server_url, self.variant, self._file_modified + ) + self.reset() + + self._state = state + return self._state + + def get_ayon_username(self) -> Optional[str]: + tray_info = self.get_tray_info() + if tray_info: + return tray_info.get("username") + return None + + def wait_to_start(self) -> bool: + _wait_for_starting_tray( + self.server_url, self.variant, self._timeout + ) + self.reset() + return self.get_file_state() == TrayState.RUNNING + + pid = property(get_pid) + state = property(get_state) def get_tray_server_url( @@ -214,25 +389,12 @@ def get_tray_server_url( Optional[str]: Tray server url. """ - data = get_tray_file_info(server_url, variant) - if data is None: - return None - - if data.get("started") is False: - data = _wait_for_starting_tray(server_url, variant, timeout) - if data is None: - return None - - url = data.get("url") - if not url: - return None - - if not validate: - return url - - if _get_tray_information(url): - return url - return None + tray_info = TrayInfo.new( + server_url, variant, timeout, wait_to_start=True + ) + if validate: + return tray_info.get_tray_url() + return tray_info.get_file_url() def set_tray_server_url(tray_url: Optional[str], started: bool): @@ -246,10 +408,13 @@ def set_tray_server_url(tray_url: Optional[str], started: bool): that tray is starting up. """ - file_info = get_tray_file_info() - if file_info and file_info["pid"] != os.getpid(): - if not file_info["started"] or _get_tray_information(file_info["url"]): - raise TrayIsRunningError("Tray is already running.") + info = TrayInfo.new(wait_to_start=False) + if ( + info.pid + and info.pid != os.getpid() + and info.state in (TrayState.RUNNING, TrayState.STARTING) + ): + raise TrayIsRunningError("Tray is already running.") filepath = _get_tray_info_filepath() os.makedirs(os.path.dirname(filepath), exist_ok=True) @@ -292,20 +457,21 @@ def remove_tray_server_url(force: Optional[bool] = False): def get_tray_information( server_url: Optional[str] = None, - variant: Optional[str] = None -) -> Optional[Dict[str, Any]]: + variant: Optional[str] = None, + timeout: Optional[int] = None, +) -> TrayInfo: """Get information about tray. Args: server_url (Optional[str]): AYON server url. variant (Optional[str]): Settings variant. + timeout (Optional[int]): Timeout for tray start-up. Returns: - Optional[Dict[str, Any]]: Tray information. + TrayInfo: Tray information. """ - tray_url = get_tray_server_url(server_url, variant) - return _get_tray_information(tray_url) + return TrayInfo.new(server_url, variant, timeout) def get_tray_state( @@ -322,20 +488,8 @@ def get_tray_state( int: Tray state. """ - file_info = get_tray_file_info(server_url, variant) - if file_info is None: - return TrayState.NOT_RUNNING - - if file_info.get("started") is False: - return TrayState.STARTING - - tray_url = file_info.get("url") - info = _get_tray_information(tray_url) - if not info: - # Remove the information as the tray is not running - remove_tray_server_url(force=True) - return TrayState.NOT_RUNNING - return TrayState.RUNNING + tray_info = get_tray_information(server_url, variant) + return tray_info.state def is_tray_running( @@ -435,38 +589,39 @@ def main(force=False): Logger.set_process_name("Tray") - state = get_tray_state() - if force and state in (TrayState.RUNNING, TrayState.STARTING): - file_info = get_tray_file_info() or {} - pid = file_info.get("pid") + tray_info = TrayInfo.new(wait_to_start=False) + file_state = tray_info.get_file_state() + if force and file_state in (TrayState.RUNNING, TrayState.STARTING): + pid = tray_info.pid if pid is not None: _kill_tray_process(pid) remove_tray_server_url(force=True) - state = TrayState.NOT_RUNNING + file_state = TrayState.NOT_RUNNING - if state == TrayState.RUNNING: - show_message_in_tray( - "Tray is already running", - "Your AYON tray application is already running." - ) - print("Tray is already running.") - return + if file_state == TrayState.RUNNING: + if tray_info.get_state() == TrayState.RUNNING: + show_message_in_tray( + "Tray is already running", + "Your AYON tray application is already running." + ) + print("Tray is already running.") + return + file_state = tray_info.get_file_state() - if state == TrayState.STARTING: + if file_state == TrayState.STARTING: print("Tray is starting. Waiting for it to start.") - _wait_for_starting_tray() - state = get_tray_state() - if state == TrayState.RUNNING: + tray_info.wait_to_start() + file_state = tray_info.get_file_state() + if file_state == TrayState.RUNNING: print("Tray started. Exiting.") return - if state == TrayState.STARTING: + if file_state == TrayState.STARTING: print( "Tray did not start in expected time." " Killing the process and starting new." ) - file_info = get_tray_file_info() or {} - pid = file_info.get("pid") + pid = tray_info.pid if pid is not None: _kill_tray_process(pid) remove_tray_server_url(force=True) From e10a78d761210ae8f9484c653b049d02a16444bc Mon Sep 17 00:00:00 2001 From: Jakub Trllo <43494761+iLLiCiTiT@users.noreply.github.com> Date: Tue, 6 Aug 2024 11:22:49 +0200 Subject: [PATCH 2/4] shutdown existing tray if is running under different user --- client/ayon_core/tools/tray/lib.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/client/ayon_core/tools/tray/lib.py b/client/ayon_core/tools/tray/lib.py index e6de884cb5..f09f400dfa 100644 --- a/client/ayon_core/tools/tray/lib.py +++ b/client/ayon_core/tools/tray/lib.py @@ -13,7 +13,12 @@ from typing import Optional, Dict, Tuple, Any import ayon_api import requests -from ayon_core.lib import Logger, get_ayon_launcher_args, run_detached_process +from ayon_core.lib import ( + Logger, + get_ayon_launcher_args, + run_detached_process, + get_ayon_username, +) from ayon_core.lib.local_settings import get_ayon_appdirs @@ -590,6 +595,7 @@ def main(force=False): Logger.set_process_name("Tray") tray_info = TrayInfo.new(wait_to_start=False) + file_state = tray_info.get_file_state() if force and file_state in (TrayState.RUNNING, TrayState.STARTING): pid = tray_info.pid @@ -598,6 +604,17 @@ def main(force=False): remove_tray_server_url(force=True) file_state = TrayState.NOT_RUNNING + if file_state in (TrayState.RUNNING, TrayState.STARTING): + expected_username = get_ayon_username() + username = tray_info.get_ayon_username() + # TODO probably show some message to the user??? + if expected_username != username: + pid = tray_info.pid + if pid is not None: + _kill_tray_process(pid) + remove_tray_server_url(force=True) + file_state = TrayState.NOT_RUNNING + if file_state == TrayState.RUNNING: if tray_info.get_state() == TrayState.RUNNING: show_message_in_tray( From d7ab1328cfeab813e55ca0c0c5dd5f590e03e6cc Mon Sep 17 00:00:00 2001 From: Jakub Trllo <43494761+iLLiCiTiT@users.noreply.github.com> Date: Tue, 6 Aug 2024 16:42:22 +0200 Subject: [PATCH 3/4] 'make_sure_tray_is_running' is also handling username --- client/ayon_core/tools/tray/lib.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/client/ayon_core/tools/tray/lib.py b/client/ayon_core/tools/tray/lib.py index f09f400dfa..0d0ead85d4 100644 --- a/client/ayon_core/tools/tray/lib.py +++ b/client/ayon_core/tools/tray/lib.py @@ -551,6 +551,7 @@ def show_message_in_tray( def make_sure_tray_is_running( ayon_url: Optional[str] = None, variant: Optional[str] = None, + username: Optional[str] = None, env: Optional[Dict[str, str]] = None ): """Make sure that tray for AYON url and variant is running. @@ -558,19 +559,26 @@ def make_sure_tray_is_running( Args: ayon_url (Optional[str]): AYON server url. variant (Optional[str]): Settings variant. + username (Optional[str]): Username under which should be tray running. env (Optional[Dict[str, str]]): Environment variables for the process. """ - state = get_tray_state(ayon_url, variant) - if state == TrayState.RUNNING: - return + tray_info = TrayInfo.new( + ayon_url, variant, wait_to_start=False + ) + if tray_info.state == TrayState.STARTING: + tray_info.wait_to_start() - if state == TrayState.STARTING: - _wait_for_starting_tray(ayon_url, variant) - state = get_tray_state(ayon_url, variant) - if state == TrayState.RUNNING: + if tray_info.state == TrayState.RUNNING: + if not username: + username = get_ayon_username() + if tray_info.get_ayon_username() == username: return + pid = tray_info.pid + if pid is not None: + _kill_tray_process(pid) + args = get_ayon_launcher_args("tray", "--force") if env is None: env = os.environ.copy() From 0a05731413b411a0fa554b727288e74b13f08aeb Mon Sep 17 00:00:00 2001 From: Jakub Trllo <43494761+iLLiCiTiT@users.noreply.github.com> Date: Tue, 6 Aug 2024 16:43:15 +0200 Subject: [PATCH 4/4] don't kill tray in make sure function --- client/ayon_core/tools/tray/lib.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/client/ayon_core/tools/tray/lib.py b/client/ayon_core/tools/tray/lib.py index 0d0ead85d4..7fec47ee56 100644 --- a/client/ayon_core/tools/tray/lib.py +++ b/client/ayon_core/tools/tray/lib.py @@ -575,10 +575,6 @@ def make_sure_tray_is_running( if tray_info.get_ayon_username() == username: return - pid = tray_info.pid - if pid is not None: - _kill_tray_process(pid) - args = get_ayon_launcher_args("tray", "--force") if env is None: env = os.environ.copy()