diff --git a/client/ayon_core/tools/tray/lib.py b/client/ayon_core/tools/tray/lib.py index ad190482a8..7fec47ee56 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 @@ -144,17 +149,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 +159,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 +221,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 +394,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 +413,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 +462,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 +493,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( @@ -392,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. @@ -399,17 +559,20 @@ 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 args = get_ayon_launcher_args("tray", "--force") @@ -435,38 +598,51 @@ 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 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 state == TrayState.STARTING: + 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 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)