From 2019da1e10bb693b4cacc71734b71c10bbb558b5 Mon Sep 17 00:00:00 2001 From: Inex Code Date: Mon, 12 Feb 2024 18:17:18 +0300 Subject: [PATCH] feat: Track the status of the nixos rebuild systemd unit --- selfprivacy_api/actions/system.py | 26 +++- .../graphql/mutations/system_mutations.py | 26 ++-- selfprivacy_api/jobs/upgrade_system.py | 120 ++++++++++++++++++ selfprivacy_api/migrations/__init__.py | 4 + .../check_for_system_rebuild_jobs.py | 47 +++++++ 5 files changed, 206 insertions(+), 17 deletions(-) create mode 100644 selfprivacy_api/jobs/upgrade_system.py create mode 100644 selfprivacy_api/migrations/check_for_system_rebuild_jobs.py diff --git a/selfprivacy_api/actions/system.py b/selfprivacy_api/actions/system.py index 13c3708..9b52497 100644 --- a/selfprivacy_api/actions/system.py +++ b/selfprivacy_api/actions/system.py @@ -4,6 +4,8 @@ import subprocess import pytz from typing import Optional, List from pydantic import BaseModel +from selfprivacy_api.jobs import Job, JobStatus, Jobs +from selfprivacy_api.jobs.upgrade_system import rebuild_system_task from selfprivacy_api.utils import WriteUserData, ReadUserData @@ -87,10 +89,16 @@ def run_blocking(cmd: List[str], new_session: bool = False) -> str: return stdout -def rebuild_system() -> int: +def rebuild_system() -> Job: """Rebuild the system""" - run_blocking(["systemctl", "start", "sp-nixos-rebuild.service"], new_session=True) - return 0 + job = Jobs.add( + type_id="system.nixos.rebuild", + name="Rebuild system", + description="Applying the new system configuration by building the new NixOS generation.", + status=JobStatus.CREATED, + ) + rebuild_system_task(job) + return job def rollback_system() -> int: @@ -99,10 +107,16 @@ def rollback_system() -> int: return 0 -def upgrade_system() -> int: +def upgrade_system() -> Job: """Upgrade the system""" - run_blocking(["systemctl", "start", "sp-nixos-upgrade.service"], new_session=True) - return 0 + job = Jobs.add( + type_id="system.nixos.upgrade", + name="Upgrade system", + description="Upgrading the system to the latest version.", + status=JobStatus.CREATED, + ) + rebuild_system_task(job, upgrade=True) + return job def reboot_system() -> None: diff --git a/selfprivacy_api/graphql/mutations/system_mutations.py b/selfprivacy_api/graphql/mutations/system_mutations.py index 13ac16b..5740a0d 100644 --- a/selfprivacy_api/graphql/mutations/system_mutations.py +++ b/selfprivacy_api/graphql/mutations/system_mutations.py @@ -3,7 +3,9 @@ import typing import strawberry from selfprivacy_api.graphql import IsAuthenticated +from selfprivacy_api.graphql.common_types.jobs import job_to_api_job from selfprivacy_api.graphql.mutations.mutation_interface import ( + GenericJobMutationReturn, GenericMutationReturn, MutationReturnInterface, ) @@ -114,16 +116,17 @@ class SystemMutations: ) @strawberry.mutation(permission_classes=[IsAuthenticated]) - def run_system_rebuild(self) -> GenericMutationReturn: + def run_system_rebuild(self) -> GenericJobMutationReturn: try: - system_actions.rebuild_system() - return GenericMutationReturn( + job = system_actions.rebuild_system() + return GenericJobMutationReturn( success=True, - message="Starting rebuild system", + message="Starting system rebuild", code=200, + job=job_to_api_job(job), ) except system_actions.ShellException as e: - return GenericMutationReturn( + return GenericJobMutationReturn( success=False, message=str(e), code=500, @@ -135,7 +138,7 @@ class SystemMutations: try: return GenericMutationReturn( success=True, - message="Starting rebuild system", + message="Starting system rollback", code=200, ) except system_actions.ShellException as e: @@ -146,16 +149,17 @@ class SystemMutations: ) @strawberry.mutation(permission_classes=[IsAuthenticated]) - def run_system_upgrade(self) -> GenericMutationReturn: - system_actions.upgrade_system() + def run_system_upgrade(self) -> GenericJobMutationReturn: try: - return GenericMutationReturn( + job = system_actions.upgrade_system() + return GenericJobMutationReturn( success=True, - message="Starting rebuild system", + message="Starting system upgrade", code=200, + job=job_to_api_job(job), ) except system_actions.ShellException as e: - return GenericMutationReturn( + return GenericJobMutationReturn( success=False, message=str(e), code=500, diff --git a/selfprivacy_api/jobs/upgrade_system.py b/selfprivacy_api/jobs/upgrade_system.py new file mode 100644 index 0000000..2d645b8 --- /dev/null +++ b/selfprivacy_api/jobs/upgrade_system.py @@ -0,0 +1,120 @@ +""" +A task to start the system upgrade or rebuild by starting a systemd unit. +After starting, track the status of the systemd unit and update the Job +status accordingly. +""" +import subprocess +from selfprivacy_api.utils.huey import huey +from selfprivacy_api.jobs import JobStatus, Jobs, Job +import time + + +@huey.task() +def rebuild_system_task(job: Job, upgrade: bool = False): + """Rebuild the system""" + try: + if upgrade: + command = ["systemctl", "start", "sp-nixos-upgrade.service"] + else: + command = ["systemctl", "start", "sp-nixos-rebuild.service"] + subprocess.run( + command, + check=True, + start_new_session=True, + shell=False, + ) + Jobs.update( + job=job, + status=JobStatus.RUNNING, + status_text="Rebuilding the system...", + ) + # Get current time to handle timeout + start_time = time.time() + # Wait for the systemd unit to start + while True: + try: + status = subprocess.run( + ["systemctl", "is-active", "selfprivacy-upgrade"], + check=True, + capture_output=True, + text=True, + ) + if status.stdout.strip() == "active": + log_line = subprocess.run( + [ + "journalctl", + "-u", + "selfprivacy-upgrade", + "-n", + "1", + "-o", + "cat", + ], + check=True, + capture_output=True, + text=True, + ).stdout.strip() + Jobs.update( + job=job, + status=JobStatus.RUNNING, + status_text=f"Rebuilding the system... Latest log line: {log_line}", + ) + break + # Timeount after 5 minutes + if time.time() - start_time > 300: + Jobs.update( + job=job, + status=JobStatus.ERROR, + error="System rebuild timed out.", + ) + return + time.sleep(1) + except subprocess.CalledProcessError: + pass + Jobs.update( + job=job, + status=JobStatus.RUNNING, + status_text="Rebuilding the system...", + ) + # Wait for the systemd unit to finish + while True: + try: + status = subprocess.run( + ["systemctl", "is-active", "selfprivacy-upgrade"], + check=True, + capture_output=True, + text=True, + ) + if status.stdout.strip() == "inactive": + Jobs.update( + job=job, + status=JobStatus.FINISHED, + result="System rebuilt.", + progress=100, + ) + elif status.stdout.strip() == "failed": + Jobs.update( + job=job, + status=JobStatus.ERROR, + error="System rebuild failed.", + ) + break + # Timeout of 60 minutes + if time.time() - start_time > 3600: + Jobs.update( + job=job, + status=JobStatus.ERROR, + error="System rebuild timed out.", + ) + break + except subprocess.CalledProcessError: + pass + + time.sleep(5) + + except subprocess.CalledProcessError as e: + Jobs.update( + job=job, + status=JobStatus.ERROR, + status_text=str(e), + ) diff --git a/selfprivacy_api/migrations/__init__.py b/selfprivacy_api/migrations/__init__.py index 5e05b2d..2a2cbaa 100644 --- a/selfprivacy_api/migrations/__init__.py +++ b/selfprivacy_api/migrations/__init__.py @@ -11,9 +11,13 @@ Adding DISABLE_ALL to that array disables the migrations module entirely. from selfprivacy_api.utils import ReadUserData, UserDataFiles from selfprivacy_api.migrations.write_token_to_redis import WriteTokenToRedis +from selfprivacy_api.migrations.check_for_system_rebuild_jobs import ( + CheckForSystemRebuildJobs, +) migrations = [ WriteTokenToRedis(), + CheckForSystemRebuildJobs(), ] diff --git a/selfprivacy_api/migrations/check_for_system_rebuild_jobs.py b/selfprivacy_api/migrations/check_for_system_rebuild_jobs.py new file mode 100644 index 0000000..9bbac8a --- /dev/null +++ b/selfprivacy_api/migrations/check_for_system_rebuild_jobs.py @@ -0,0 +1,47 @@ +from selfprivacy_api.migrations.migration import Migration +from selfprivacy_api.jobs import JobStatus, Jobs + + +class CheckForSystemRebuildJobs(Migration): + """Check if there are unfinished system rebuild jobs and finish them""" + + def get_migration_name(self): + return "check_for_system_rebuild_jobs" + + def get_migration_description(self): + return "Check if there are unfinished system rebuild jobs and finish them" + + def is_migration_needed(self): + # Check if there are any unfinished system rebuild jobs + for job in Jobs.get_jobs(): + if ( + job.type_id + in [ + "system.nixos.rebuild", + "system.nixos.upgrade", + ] + ) and job.status in [ + JobStatus.CREATED, + JobStatus.RUNNING, + ]: + return True + + def migrate(self): + # As the API is restarted, we assume that the jobs are finished + for job in Jobs.get_jobs(): + if ( + job.type_id + in [ + "system.nixos.rebuild", + "system.nixos.upgrade", + ] + ) and job.status in [ + JobStatus.CREATED, + JobStatus.RUNNING, + ]: + Jobs.update( + job=job, + status=JobStatus.FINISHED, + result="System rebuilt.", + progress=100, + )