Browse Source

feat: process management module

master
dnomd343 2 years ago
parent
commit
e6cf772c4c
  1. 201
      Basis/Process.py

201
Basis/Process.py

@ -0,0 +1,201 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os
import copy
import time
import ctypes
import signal
from Basis.Logger import logging
from Basis.Functions import genFlag
from subprocess import Popen, STDOUT, DEVNULL
libcPaths = [
'/usr/lib/libc.so.6', # CentOS
'/usr/lib64/libc.so.6',
'/lib/libc.musl-i386.so.1', # Alpine
'/lib/libc.musl-x86_64.so.1',
'/lib/libc.musl-aarch64.so.1',
'/lib/i386-linux-gnu/libc.so.6', # Debian / Ubuntu
'/lib/x86_64-linux-gnu/libc.so.6',
'/lib/aarch64-linux-gnu/libc.so.6',
]
libcPath = None
for libc in libcPaths:
if os.path.exists(libc): # try to locate libc.so
libcPath = libc
break
if libcPath is None: # lost libc.so -> unable to utilize prctl
logging.warning('libc.so not found')
else:
logging.info('libc.so -> ' + str(libcPath))
class Process(object):
""" Manage a sub process and it's file.
Arguments:
cmd: Command list, which use to start the program.
env: Environment variables for sub process.
file: dict or list or None, include path and content option.
workDir: A directory for storing log files and configuration files.
taskId: Task id, defaults to 12 random characters length.
Attributes:
id, workDir, output
"""
output = None # sub process output if capture is True
__capture = None # capture the sub process output or not
__logfile = None # the log file which sub process output into STDOUT and STDERR
__process = None # CompletedProcess object of subprocess module
@staticmethod
def __preExec() -> None:
# TODO: remember to remove test code
print('WARNING: pre exec at static method')
print('libc.so: ' + str(libcPath))
ctypes.CDLL(libcPath).prctl(1, signal.SIGTERM) # sub process killed when father process exit
os.setpgrp() # new process group
def __checkWorkDir(self) -> None: # check if the working directory is normal
if os.path.isdir(self.workDir):
return
logging.warning('[%s] Work directory %s not exist' % (self.id, self.workDir))
try:
os.makedirs(self.workDir) # just like `mkdir -p ...`
logging.info('[%s] New directory -> %s' % (self.id, self.workDir))
except:
if os.path.exists(self.workDir):
logging.error('[%s] %s already exist but not folder' % (self.id, self.workDir))
else:
logging.error('[%s] Unable to create new folder -> %s' % (self.id, self.workDir))
raise RuntimeError('working directory error') # fatal error
def __killProcess(self, killSignal: int) -> None:
try:
pgid = os.getpgid(self.__process.pid) # progress group id
os.killpg(pgid, killSignal) # kill sub process group
logging.debug('[%s] Send kill signal to PGID %i' % (self.id, pgid))
except:
logging.warning('[%s] Failed to get PGID of sub process (PID = %i)' % (self.id, self.__process.pid))
def __deleteFile(self, filePath: str) -> None:
if not os.path.isfile(filePath): # file not found (or not a file)
logging.warning('[%s] File %s not found' % (self.id, filePath))
return
try:
os.remove(filePath) # remove config file
logging.debug('[%s] File %s deleted successfully' % (self.id, filePath))
except:
logging.error('[%s] Unable to delete file %s' % (self.id, filePath))
def __init__(self, workDir: str, taskId: str = '', isStart: bool = True,
cmd: str or list or None = None, env: dict or None = None, file: dict or list or None = None) -> None:
self.id = genFlag(length = 12) if taskId == '' else taskId
self.workDir = workDir
self.__env = copy.copy(env) # depth = 1
self.__cmd = copy.copy([cmd] if type(cmd) == str else cmd) # depth = 1
self.__file = copy.deepcopy([file] if type(file) == dict else file) # depth = 2
self.__checkWorkDir() # ensure the working direction is normal
logging.debug('[%s] Process command -> %s (%s)' % (self.id, self.__cmd, self.__env))
if self.__file is not None:
if len(self.__file) > 1:
logging.debug('[%s] Manage %i files' % (self.id, len(self.__file)))
for file in self.__file: # traverse all files
if not isStart: # don't print log twice
logging.debug('[%s] File %s -> %s' % (self.id, file['path'], file['content']))
if isStart:
self.start()
def setCmd(self, cmd: str or list) -> None:
self.__cmd = copy.copy([cmd] if type(cmd) == str else cmd)
logging.info('[%s] Process setting command -> %s' % (self.id, self.__cmd))
def setEnv(self, env: dict or None) -> None:
self.__env = copy.copy(env)
logging.info('[%s] Process setting environ -> %s' % (self.id, self.__env))
def setFile(self, file: dict or list or None) -> None:
self.__file = copy.deepcopy([file] if type(file) == dict else file)
if self.__file is None:
logging.info('[%s] Process setting file -> None' % self.id)
return
for file in self.__file: # traverse all files
logging.info('[%s] Process setting file %s -> %s' % (self.id, file['path'], file['content']))
def start(self, isCapture: bool = True) -> None:
self.__capture = isCapture
logging.debug('[%s] Process ready to start (%s)' % (self.id, (
'with output capture' if self.__capture else 'without output capture'
)))
if self.__cmd is None: # ERROR CASE
logging.error('[%s] Process miss start command' % self.id)
raise RuntimeError('miss start command')
if self.__process is not None and self.__process.poll() is None: # ERROR CASE
logging.error('[%s] Sub process is still running' % self.id)
raise RuntimeError('sub process is still running')
if self.__env is not None and 'PATH' not in self.__env and '/' not in self.__cmd[0]: # WARNING CASE
logging.warning('[%s] Executable file in relative path but miss PATH in environ' % self.id)
if self.__file is not None: # create and write file contents
for file in self.__file:
with open(file['path'], 'w', encoding = 'utf-8') as fileObject: # save file content
fileObject.write(file['content'])
logging.debug('[%s] File %s -> %s' % (self.id, file['path'], file['content']))
if self.__capture: # with output capture
self.__logfile = os.path.join(self.workDir, self.id + '.log')
logging.debug('[%s] Output capture file -> %s' % (self.id, self.__logfile))
stdout = open(self.__logfile, 'w', encoding = 'utf-8')
stderr = STDOUT # combine the stderr with stdout
else: # discard all the output of sub process
stdout = DEVNULL
stderr = DEVNULL
self.__process = Popen(self.__cmd, env = self.__env, stdout = stdout,
stderr = stderr, preexec_fn = None if libcPath is None else Process.__preExec)
logging.info('[%s] Process running -> PID = %i' % (self.id, self.__process.pid))
def signal(self, signalNum: int) -> None: # send specified signal to sub process
try:
signalName = signal.Signals(signalNum).name
except:
signalName = 'unknown'
logging.info('[%s] Send signal -> %i (%s)' % (self.id, signalNum, signalName))
self.__process.send_signal(signalNum)
def status(self) -> bool: # check if the sub process is still running
status = self.__process.poll() is None
logging.debug('[%s] Check status -> %s' % (self.id, 'running' if status else 'exit'))
return status
def wait(self, timeout: int or None = None) -> None: # blocking wait sub process
logging.info('[%s] Process wait -> timeout = %s' % (self.id, str(timeout)))
try:
self.__process.wait(timeout = timeout)
logging.info('[%s] Process wait timeout -> sub process exit' % self.id)
except:
logging.info('[%s] Process wait timeout -> sub process still running' % self.id)
def quit(self, isForce: bool = False, waitTime: int = 50) -> None: # wait 50ms in default
killSignal = signal.SIGKILL if isForce else signal.SIGTERM # 9 -> force kill / 15 -> terminate signal
logging.debug('[%s] Kill signal = %i (%s)' % (self.id, killSignal, signal.Signals(killSignal).name))
self.__killProcess(killSignal)
time.sleep(waitTime / 1000) # sleep (ms -> s)
while self.__process.poll() is None: # confirm sub process exit
self.__killProcess(killSignal)
time.sleep(waitTime / 1000)
logging.info('[%s] Process terminated -> PID = %i' % (self.id, self.__process.pid))
if self.__capture:
try:
with open(self.__logfile, 'r', encoding = 'utf-8') as fileObject: # read sub process output
self.output = fileObject.read()
logging.debug('[%s] Process output -> length = %s' % (self.id, len(self.output)))
self.__deleteFile(self.__logfile)
except:
logging.error('[%s] Failed to read capture log file -> %s' % (self.id, self.__logfile))
if self.__file is not None: # with config file
for file in self.__file:
self.__deleteFile(file['path'])
Loading…
Cancel
Save