meerschaum.jobs.systemd

  1#! /usr/bin/env python3
  2# vim:fenc=utf-8
  3
  4"""
  5Manage `meerschaum.jobs.Job` via `systemd`.
  6"""
  7
  8import os
  9import pathlib
 10import shlex
 11import sys
 12import asyncio
 13import json
 14import time
 15import traceback
 16import shutil
 17from datetime import datetime, timezone
 18from functools import partial
 19
 20import meerschaum as mrsm
 21from meerschaum.jobs import Job, Executor, make_executor
 22from meerschaum.utils.typing import Dict, Any, List, SuccessTuple, Union, Optional, Callable
 23from meerschaum.config import get_config
 24from meerschaum.config.static import STATIC_CONFIG
 25from meerschaum.utils.warnings import warn, dprint
 26from meerschaum._internal.arguments._parse_arguments import parse_arguments
 27
 28JOB_METADATA_CACHE_SECONDS: int = STATIC_CONFIG['api']['jobs']['metadata_cache_seconds']
 29
 30
 31@make_executor
 32class SystemdExecutor(Executor):
 33    """
 34    Execute Meerschaum jobs via `systemd`.
 35    """
 36
 37    def get_job_names(self, debug: bool = False) -> List[str]:
 38        """
 39        Return a list of existing jobs, including hidden ones.
 40        """
 41        from meerschaum.config.paths import SYSTEMD_USER_RESOURCES_PATH
 42        return [
 43            service_name[len('mrsm-'):(-1 * len('.service'))]
 44            for service_name in os.listdir(SYSTEMD_USER_RESOURCES_PATH)
 45            if (
 46                service_name.startswith('mrsm-')
 47                and service_name.endswith('.service')
 48                ### Check for broken symlinks.
 49                and (SYSTEMD_USER_RESOURCES_PATH / service_name).exists()
 50            )
 51        ]
 52
 53    def get_job_exists(self, name: str, debug: bool = False) -> bool:
 54        """
 55        Return whether a job exists.
 56        """
 57        user_services = self.get_job_names(debug=debug)
 58        if debug:
 59            dprint(f'Existing services: {user_services}')
 60        return name in user_services
 61
 62    def get_jobs(self, debug: bool = False) -> Dict[str, Job]:
 63        """
 64        Return a dictionary of `systemd` Jobs (including hidden jobs).
 65        """
 66        user_services = self.get_job_names(debug=debug)
 67        jobs = {
 68            name: Job(name, executor_keys=str(self))
 69            for name in user_services
 70        }
 71        return {
 72            name: job
 73            for name, job in jobs.items()
 74        }
 75
 76    def get_service_name(self, name: str, debug: bool = False) -> str:
 77        """
 78        Return a job's service name.
 79        """
 80        return f"mrsm-{name.replace(' ', '-')}.service"
 81
 82    def get_service_job_path(self, name: str, debug: bool = False) -> pathlib.Path:
 83        """
 84        Return the path for the job's files under the root directory.
 85        """
 86        from meerschaum.config.paths import SYSTEMD_JOBS_RESOURCES_PATH
 87        return SYSTEMD_JOBS_RESOURCES_PATH / name
 88
 89    def get_service_symlink_file_path(self, name: str, debug: bool = False) -> pathlib.Path:
 90        """
 91        Return the path to where to create the service symlink.
 92        """
 93        from meerschaum.config.paths import SYSTEMD_USER_RESOURCES_PATH
 94        return SYSTEMD_USER_RESOURCES_PATH / self.get_service_name(name, debug=debug)
 95
 96    def get_service_file_path(self, name: str, debug: bool = False) -> pathlib.Path:
 97        """
 98        Return the path to a Job's service file.
 99        """
100        return (
101            self.get_service_job_path(name, debug=debug)
102            / self.get_service_name(name, debug=debug)
103        )
104
105    def get_service_logs_path(self, name: str, debug: bool = False) -> pathlib.Path:
106        """
107        Return the path to direct service logs to.
108        """
109        from meerschaum.config.paths import SYSTEMD_LOGS_RESOURCES_PATH
110        return SYSTEMD_LOGS_RESOURCES_PATH / (self.get_service_name(name, debug=debug) + '.log')
111
112    def get_socket_path(self, name: str, debug: bool = False) -> pathlib.Path:
113        """
114        Return the path to the FIFO file.
115        """
116        return (
117            self.get_service_job_path(name, debug=debug)
118            / (self.get_service_name(name, debug=debug) + '.stdin')
119        )
120
121    def get_result_path(self, name: str, debug: bool = False) -> pathlib.Path:
122        """
123        Return the path to the result file.
124        """
125        return (
126            self.get_service_job_path(name, debug=debug)
127            / (self.get_service_name(name, debug=debug) + '.result.json')
128        )
129
130    def get_service_file_text(self, name: str, sysargs: List[str], debug: bool = False) -> str:
131        """
132        Return the contents of the unit file.
133        """
134        service_logs_path = self.get_service_logs_path(name, debug=debug)
135        socket_path = self.get_socket_path(name, debug=debug)
136        result_path = self.get_result_path(name, debug=debug)
137        job = self.get_hidden_job(name, debug=debug)
138
139        sysargs_str = shlex.join(sysargs)
140        exec_str = f'{sys.executable} -m meerschaum {sysargs_str}'
141        mrsm_env_var_names = set([var for var in STATIC_CONFIG['environment'].values()])
142        mrsm_env_vars = {
143            key: val
144            for key, val in os.environ.items()
145            if key in mrsm_env_var_names
146        }
147
148        ### Add new environment variables for the service process.
149        mrsm_env_vars.update({
150            STATIC_CONFIG['environment']['daemon_id']: name,
151            STATIC_CONFIG['environment']['systemd_log_path']: service_logs_path.as_posix(),
152            STATIC_CONFIG['environment']['systemd_result_path']: result_path.as_posix(),
153            STATIC_CONFIG['environment']['systemd_stdin_path']: socket_path.as_posix(),
154            STATIC_CONFIG['environment']['systemd_delete_job']: (
155                '1'
156                if job.delete_after_completion
157                else '0',
158            ),
159        })
160
161        ### Allow for user-defined environment variables.
162        mrsm_env_vars.update(job.env)
163
164        environment_lines = [
165            f"Environment={key}={shlex.quote(str(val))}"
166            for key, val in mrsm_env_vars.items()
167        ]
168        environment_str = '\n'.join(environment_lines)
169        service_name = self.get_service_name(name, debug=debug)
170
171        service_text = (
172            "[Unit]\n"
173            f"Description=Run the job '{name}'\n"
174            "\n"
175            "[Service]\n"
176            f"ExecStart={exec_str}\n"
177            "KillSignal=SIGTERM\n"
178            "Restart=always\n"
179            "RestartPreventExitStatus=0\n"
180            f"SyslogIdentifier={service_name}\n"
181            f"{environment_str}\n"
182            "\n"
183            "[Install]\n"
184            "WantedBy=default.target\n"
185        )
186        return service_text
187
188    def get_socket_file_text(self, name: str, debug: bool = False) -> str:
189        """
190        Return the contents of the socket file.
191        """
192        service_name = self.get_service_name(name, debug=debug)
193        socket_path = self.get_socket_path(name, debug=debug)
194        socket_text = (
195            "[Unit]\n"
196            f"BindsTo={service_name}\n"
197            "\n"
198            "[Socket]\n"
199            f"ListenFIFO={socket_path.as_posix()}\n"
200            "FileDescriptorName=stdin\n"
201            "RemoveOnStop=true\n"
202            "SocketMode=0660\n"
203        )
204        return socket_text
205
206    def get_hidden_job(
207        self,
208        name: str,
209        sysargs: Optional[List[str]] = None,
210        properties: Optional[Dict[str, Any]] = None,
211        debug: bool = False,
212    ):
213        """
214        Return the hidden "sister" job to store a job's parameters.
215        """
216        job = Job(
217            name,
218            sysargs,
219            executor_keys='local',
220            _properties=properties,
221            _rotating_log=self.get_job_rotating_file(name, debug=debug),
222            _stdin_file=self.get_job_stdin_file(name, debug=debug),
223            _status_hook=partial(self.get_job_status, name),
224            _result_hook=partial(self.get_job_result, name),
225            _externally_managed=True,
226        )
227        return job
228
229
230    def get_job_metadata(self, name: str, debug: bool = False) -> Dict[str, Any]:
231        """
232        Return metadata about a job.
233        """
234        now = time.perf_counter()
235
236        if '_jobs_metadata' not in self.__dict__:
237            self._jobs_metadata: Dict[str, Any] = {}
238
239        if name in self._jobs_metadata:
240            ts = self._jobs_metadata[name].get('timestamp', None)
241
242            if ts is not None and (now - ts) <= JOB_METADATA_CACHE_SECONDS:
243                if debug:
244                    dprint(f"Retuning cached metadata for job '{name}'.")
245                return self._jobs_metadata[name]['metadata']
246
247        metadata = {
248            'sysargs': self.get_job_sysargs(name, debug=debug),
249            'result': self.get_job_result(name, debug=debug),
250            'restart': self.get_job_restart(name, debug=debug),
251            'daemon': {
252                'status': self.get_job_status(name, debug=debug),
253                'pid': self.get_job_pid(name, debug=debug),
254                'properties': self.get_job_properties(name, debug=debug),
255            },
256        }
257        self._jobs_metadata[name] = {
258            'timestamp': now,
259            'metadata': metadata,
260        }
261        return metadata
262
263    def get_job_restart(self, name: str, debug: bool = False) -> bool:
264        """
265        Return whether a job restarts.
266        """
267        from meerschaum.jobs._Job import RESTART_FLAGS
268        sysargs = self.get_job_sysargs(name, debug=debug)
269        if not sysargs:
270            return False
271
272        for flag in RESTART_FLAGS:
273            if flag in sysargs:
274                return True
275
276        return False
277
278    def get_job_properties(self, name: str, debug: bool = False) -> Dict[str, Any]:
279        """
280        Return the properties for a job.
281        """
282        job = self.get_hidden_job(name, debug=debug)
283        return {
284            k: v for k, v in job.daemon.properties.items()
285            if k != 'externally_managed'
286        }
287
288    def get_job_process(self, name: str, debug: bool = False):
289        """
290        Return a `psutil.Process` for the job's PID.
291        """
292        pid = self.get_job_pid(name, debug=debug)
293        if pid is None:
294            return None
295
296        psutil = mrsm.attempt_import('psutil')
297        try:
298            return psutil.Process(pid)
299        except Exception:
300            return None
301
302    def get_job_status(self, name: str, debug: bool = False) -> str:
303        """
304        Return the job's service status.
305        """
306        output = self.run_command(
307            ['is-active', self.get_service_name(name, debug=debug)],
308            as_output=True,
309            debug=debug,
310        )
311
312        if output == 'activating':
313            return 'running'
314
315        if output == 'active':
316            process = self.get_job_process(name, debug=debug)
317            if process is None:
318                return 'stopped'
319
320            try:
321                if process.status() == 'stopped':
322                    return 'paused'
323            except Exception:
324                return 'stopped'
325
326            return 'running'
327
328        return 'stopped'
329
330    def get_job_pid(self, name: str, debug: bool = False) -> Union[int, None]:
331        """
332        Return the job's service PID.
333        """
334        from meerschaum.utils.misc import is_int
335
336        output = self.run_command(
337            [
338                'show',
339                self.get_service_name(name, debug=debug),
340                '--property=MainPID',
341            ],
342            as_output=True,
343            debug=debug,
344        )
345        if not output.startswith('MainPID='):
346            return None
347
348        pid_str = output[len('MainPID='):]
349        if pid_str == '0':
350            return None
351
352        if is_int(pid_str):
353            return int(pid_str)
354        
355        return None
356
357    def get_job_began(self, name: str, debug: bool = False) -> Union[str, None]:
358        """
359        Return when a job began running.
360        """
361        output = self.run_command(
362            [
363                'show',
364                self.get_service_name(name, debug=debug),
365                '--property=ActiveEnterTimestamp'
366            ],
367            as_output=True,
368            debug=debug,
369        )
370        if not output.startswith('ActiveEnterTimestamp'):
371            return None
372
373        dt_str = output.split('=')[-1]
374        if not dt_str:
375            return None
376
377        dateutil_parser = mrsm.attempt_import('dateutil.parser')
378        try:
379            dt = dateutil_parser.parse(dt_str)
380        except Exception as e:
381            warn(f"Cannot parse '{output}' as a datetime:\n{e}")
382            return None
383
384        return dt.astimezone(timezone.utc).isoformat()
385
386    def get_job_ended(self, name: str, debug: bool = False) -> Union[str, None]:
387        """
388        Return when a job began running.
389        """
390        output = self.run_command(
391            [
392                'show',
393                self.get_service_name(name, debug=debug),
394                '--property=InactiveEnterTimestamp'
395            ],
396            as_output=True,
397            debug=debug,
398        )
399        if not output.startswith('InactiveEnterTimestamp'):
400            return None
401
402        dt_str = output.split('=')[-1]
403        if not dt_str:
404            return None
405
406        dateutil_parser = mrsm.attempt_import('dateutil.parser')
407
408        try:
409            dt = dateutil_parser.parse(dt_str)
410        except Exception as e:
411            warn(f"Cannot parse '{output}' as a datetime:\n{e}")
412            return None
413        return dt.astimezone(timezone.utc).isoformat()
414
415    def get_job_paused(self, name: str, debug: bool = False) -> Union[str, None]:
416        """
417        Return when a job was paused.
418        """
419        job = self.get_hidden_job(name, debug=debug)
420        if self.get_job_status(name, debug=debug) != 'paused':
421            return None
422
423        stop_time = job.stop_time
424        if stop_time is None:
425            return None
426
427        return stop_time.isoformat()
428
429    def get_job_result(self, name: str, debug: bool = False) -> SuccessTuple:
430        """
431        Return the job's result SuccessTuple.
432        """
433        result_path = self.get_result_path(name, debug=debug)
434        if not result_path.exists():
435            return False, "No result available."
436
437        try:
438            with open(result_path, 'r', encoding='utf-8') as f:
439                result = json.load(f)
440        except Exception:
441            return False, f"Could not read result for Job '{name}'."
442
443        return tuple(result)
444
445    def get_job_sysargs(self, name: str, debug: bool = False) -> Union[List[str], None]:
446        """
447        Return the sysargs from the service file.
448        """
449        service_file_path = self.get_service_file_path(name, debug=debug)
450        if not service_file_path.exists():
451            return []
452
453        with open(service_file_path, 'r', encoding='utf-8') as f:
454            service_lines = f.readlines()
455
456        for line in service_lines:
457            if line.startswith('ExecStart='):
458                sysargs_str = line.split(' -m meerschaum ')[-1].split('<')[0]
459                return shlex.split(sysargs_str)
460
461        return []
462
463    def run_command(
464        self,
465        command_args: List[str],
466        as_output: bool = False,
467        debug: bool = False,
468    ) -> Union[SuccessTuple, str]:
469        """
470        Run a `systemd` command and return success.
471
472        Parameters
473        ----------
474        command_args: List[str]
475            The command to pass to `systemctl --user`.
476
477        as_output: bool, default False
478            If `True`, return the process stdout output.
479            Defaults to a `SuccessTuple`.
480
481        Returns
482        -------
483        A `SuccessTuple` indicating success or a str for the process output.
484        """
485        from meerschaum.utils.process import run_process
486
487        if command_args[:2] != ['systemctl', '--user']:
488            command_args = ['systemctl', '--user'] + command_args
489
490        if debug:
491            dprint(shlex.join(command_args))
492
493        proc = run_process(
494            command_args,
495            foreground=False,
496            as_proc=True,
497            capture_output=True,
498            text=True,
499        )
500        stdout, stderr = proc.communicate()
501        if debug:
502            dprint(f"{stdout}")
503
504        if as_output:
505            return stdout.strip()
506            
507        command_success = proc.wait() == 0
508        command_msg = (
509            "Success"
510            if command_success
511            else f"Failed to execute command `{shlex.join(command_args)}`."
512        )
513        return command_success, command_msg
514
515    def get_job_stdin_file(self, name: str, debug: bool = False):
516        """
517        Return a `StdinFile` for the job.
518        """
519        from meerschaum.utils.daemon import StdinFile
520        if '_stdin_files' not in self.__dict__:
521            self._stdin_files: Dict[str, StdinFile] = {}
522
523        if name not in self._stdin_files:
524            socket_path = self.get_socket_path(name, debug=debug)
525            socket_path.parent.mkdir(parents=True, exist_ok=True)
526            self._stdin_files[name] = StdinFile(socket_path)
527
528        return self._stdin_files[name]
529
530    def create_job(
531        self,
532        name: str,
533        sysargs: List[str],
534        properties: Optional[Dict[str, Any]] = None,
535        debug: bool = False,
536    ) -> SuccessTuple:
537        """
538        Create a job as a service to be run by `systemd`.
539        """
540        from meerschaum.utils.misc import make_symlink
541        service_name = self.get_service_name(name, debug=debug)
542        service_file_path = self.get_service_file_path(name, debug=debug)
543        service_symlink_file_path = self.get_service_symlink_file_path(name, debug=debug)
544        socket_stdin = self.get_job_stdin_file(name, debug=debug)
545        _ = socket_stdin.file_handler
546
547        ### Init the externally_managed file.
548        ### NOTE: We must write the pickle file in addition to the properties file.
549        job = self.get_hidden_job(name, sysargs=sysargs, properties=properties, debug=debug)
550        job._set_externally_managed()
551        pickle_success, pickle_msg = job.daemon.write_pickle()
552        if not pickle_success:
553            return pickle_success, pickle_msg
554        properties_success, properties_msg = job.daemon.write_properties()
555        if not properties_success:
556            return properties_success, properties_msg
557
558        service_file_path.parent.mkdir(parents=True, exist_ok=True)
559
560        with open(service_file_path, 'w+', encoding='utf-8') as f:
561            f.write(self.get_service_file_text(name, sysargs, debug=debug))
562
563        symlink_success, symlink_msg = make_symlink(service_file_path, service_symlink_file_path)
564        if not symlink_success:
565            return symlink_success, symlink_msg
566
567        commands = [
568            ['daemon-reload'],
569            ['enable', service_name],
570            ['start', service_name],
571        ]
572
573        fails = 0
574        for command_list in commands:
575            command_success, command_msg = self.run_command(command_list, debug=debug)
576            if not command_success:
577                fails += 1
578
579        if fails > 1:
580            return False, "Failed to reload systemd."
581
582        return True, f"Started job '{name}' via systemd."
583
584    def start_job(self, name: str, debug: bool = False) -> SuccessTuple:
585        """
586        Stop a job's service.
587        """
588        job = self.get_hidden_job(name, debug=debug)
589        job.daemon._remove_stop_file()
590
591        status = self.get_job_status(name, debug=debug)
592        if status == 'paused':
593            return self.run_command(
594                ['kill', '-s', 'SIGCONT', self.get_service_name(name, debug=debug)],
595                debug=debug,
596            )
597
598        return self.run_command(
599            ['start', self.get_service_name(name, debug=debug)],
600            debug=debug,
601        )
602
603    def stop_job(self, name: str, debug: bool = False) -> SuccessTuple:
604        """
605        Stop a job's service.
606        """
607        job = self.get_hidden_job(name, debug=debug)
608        job.daemon._write_stop_file('quit')
609        sigint_success, sigint_msg = self.run_command(
610            ['kill', '-s', 'SIGINT', self.get_service_name(name, debug=debug)],
611            debug=debug,
612        )
613
614        check_timeout_interval = get_config('jobs', 'check_timeout_interval_seconds')
615        loop_start = time.perf_counter()
616        timeout_seconds = get_config('jobs', 'timeout_seconds')
617        while (time.perf_counter() - loop_start) < timeout_seconds:
618            if self.get_job_status(name, debug=debug) == 'stopped':
619                return True, 'Success'
620
621            time.sleep(check_timeout_interval)
622
623        return self.run_command(
624            ['stop', self.get_service_name(name, debug=debug)],
625            debug=debug,
626        )
627
628    def pause_job(self, name: str, debug: bool = False) -> SuccessTuple:
629        """
630        Pause a job's service.
631        """
632        job = self.get_hidden_job(name, debug=debug)
633        job.daemon._write_stop_file('pause')
634        return self.run_command(
635            ['kill', '-s', 'SIGSTOP', self.get_service_name(name, debug=debug)],
636            debug=debug,
637        )
638
639    def delete_job(self, name: str, debug: bool = False) -> SuccessTuple:
640        """
641        Delete a job's service.
642        """
643        from meerschaum.config.paths import SYSTEMD_LOGS_RESOURCES_PATH
644        job = self.get_hidden_job(name, debug=debug)
645
646        if not job.delete_after_completion:
647            _ = self.stop_job(name, debug=debug)
648            _ = self.run_command(
649                ['disable', self.get_service_name(name, debug=debug)],
650                debug=debug,
651            )
652
653        service_job_path = self.get_service_job_path(name, debug=debug)
654        try:
655            if service_job_path.exists():
656                shutil.rmtree(service_job_path)
657        except Exception as e:
658            warn(e)
659            return False, str(e)
660
661        service_logs_path = self.get_service_logs_path(name, debug=debug)
662        logs_paths = [
663            (SYSTEMD_LOGS_RESOURCES_PATH / name)
664            for name in os.listdir(SYSTEMD_LOGS_RESOURCES_PATH)
665            if name.startswith(service_logs_path.name + '.')
666        ]
667        paths = [
668            self.get_service_file_path(name, debug=debug),
669            self.get_service_symlink_file_path(name, debug=debug),
670            self.get_socket_path(name, debug=debug),
671            self.get_result_path(name, debug=debug),
672        ] + logs_paths
673
674        for path in paths:
675            if path.exists():
676                try:
677                    path.unlink()
678                except Exception as e:
679                    warn(e)
680                    return False, str(e)
681
682        _ = job.delete()
683
684        return self.run_command(['daemon-reload'], debug=debug)
685
686    def get_logs(self, name: str, debug: bool = False) -> str:
687        """
688        Return a job's journal logs.
689        """
690        rotating_file = self.get_job_rotating_file(name, debug=debug)
691        return rotating_file.read()
692
693    def get_job_stop_time(self, name: str, debug: bool = False) -> Union[datetime, None]:
694        """
695        Return a job's stop time.
696        """
697        job = self.get_hidden_job(name, debug=debug)
698        return job.stop_time
699
700    def get_job_is_blocking_on_stdin(self, name: str, debug: bool = False) -> bool:
701        """
702        Return whether a job is blocking on stdin.
703        """
704        socket_path = self.get_socket_path(name, debug=debug)
705        blocking_path = socket_path.parent / (socket_path.name + '.block')
706        return blocking_path.exists()
707
708    def get_job_rotating_file(self, name: str, debug: bool = False):
709        """
710        Return a `RotatingFile` for the job's log output.
711        """
712        from meerschaum.utils.daemon import RotatingFile
713        service_logs_path = self.get_service_logs_path(name, debug=debug)
714        return RotatingFile(service_logs_path)
715
716    async def monitor_logs_async(
717        self,
718        name: str,
719        *args,
720        debug: bool = False,
721        **kwargs
722    ):
723        """
724        Monitor a job's output.
725        """
726        from meerschaum.config.paths import SYSTEMD_LOGS_RESOURCES_PATH
727        job = self.get_hidden_job(name, debug=debug)
728        kwargs.update({
729            '_logs_path': SYSTEMD_LOGS_RESOURCES_PATH,
730            '_log': self.get_job_rotating_file(name, debug=debug),
731            '_stdin_file': self.get_job_stdin_file(name, debug=debug),
732            'debug': debug,
733        })
734        await job.monitor_logs_async(*args, **kwargs)
735
736    def monitor_logs(self, *args, **kwargs):
737        """
738        Monitor a job's output.
739        """
740        asyncio.run(self.monitor_logs_async(*args, **kwargs))