104 lines
3.9 KiB
Python
104 lines
3.9 KiB
Python
import time
|
|
import os, sys
|
|
import signal
|
|
import subprocess
|
|
import threading
|
|
import queue
|
|
import select
|
|
from functools import wraps
|
|
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
|
|
import core.logger as logger
|
|
log = logger.get_logger()
|
|
|
|
def timeit(func):
|
|
@wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
start_time = time.time()
|
|
result = func(*args, **kwargs)
|
|
end_time = time.time()
|
|
execution_time = end_time - start_time
|
|
cmd = args[0]
|
|
log.trace(f"Spent [{execution_time:.4f}s] Run: [ {cmd} ]")
|
|
return result
|
|
return wrapper
|
|
|
|
@timeit
|
|
def shell(cmd, *kwargs, timeout:float=10, checkret:bool=True, async_handle=None, handle_param=None):
|
|
executer = ShellExecute(cmd, timeout, async_handle, handle_param)
|
|
ret, output, error = executer.execute()
|
|
if checkret and ret!= 0:
|
|
log.error(f"Execute [\n{cmd}\n] failed Ret:{ret}")
|
|
log.debug(f"Stdout:{output}")
|
|
log.error(f"Stderr:{error}")
|
|
raise Exception("error execute shell command")
|
|
return ret, output, error
|
|
|
|
class ShellExecute(threading.Thread):
|
|
def __init__(self, cmd_line, timeout:int, async_handle, handle_param):
|
|
super().__init__(target=self.monitor)
|
|
self.cmd_line = cmd_line
|
|
self.timeout = timeout
|
|
self.async_handle = async_handle
|
|
self.handle_param = handle_param
|
|
|
|
def monitor(self):
|
|
if self.timeout is None:
|
|
log.debug("Block to execute [{}] ".format(self.cmd_line))
|
|
return
|
|
cnt = self.timeout / 0.1
|
|
while self.process.poll() is None and cnt > 0:
|
|
time.sleep(0.1)
|
|
cnt -= 1
|
|
if self.process.poll() is not None: # 进程已经结束
|
|
return
|
|
os.kill(self.process.pid, signal.SIGSEGV)
|
|
log.error("execute [{}] timeout, Force kill it(pid:{})...".format(self.cmd_line, self.process.pid))
|
|
raise TimeoutError("Execute cmd timeout")
|
|
|
|
def execute(self):
|
|
if self.async_handle is not None:
|
|
return self.async_execute()
|
|
else:
|
|
return self.sync_execute()
|
|
def sync_execute(self):
|
|
process = subprocess.Popen(self.cmd_line, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
self.process = process
|
|
self.start() #启动监控线程
|
|
stdout, stderr = process.communicate()
|
|
output = stdout.decode('utf-8')
|
|
error = stderr.decode('utf-8')
|
|
return_code = process.returncode
|
|
return return_code, output, error
|
|
|
|
def async_execute(self):
|
|
output_queue = queue.Queue()
|
|
# read_fd, write_fd = os.pipe()
|
|
thread = threading.Thread(target=self.async_handle, args=(output_queue,self.handle_param), daemon=True)
|
|
thread.start()
|
|
process = subprocess.Popen(self.cmd_line, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, shell=True)
|
|
self.process = process
|
|
self.start() #启动监控线程
|
|
try:
|
|
epoll = select.epoll()
|
|
epoll.register(process.stdout.fileno(), select.EPOLLIN)
|
|
while process.poll() is None:
|
|
events = epoll.poll()
|
|
for fileno, _ in events:
|
|
if fileno != process.stdout.fileno():
|
|
continue
|
|
while True:
|
|
std_out = process.stdout.readline()
|
|
if not std_out:
|
|
break
|
|
output_queue.put(std_out) # 将输出放入队列
|
|
except KeyboardInterrupt:
|
|
log.warning("Interrupted Process")
|
|
os.kill(self.process.pid, signal.SIGSEGV)
|
|
log.trace("Process finished.")
|
|
output_queue.put(None) # 结束队列
|
|
thread.join() # 等待线程结束
|
|
stdout, stderr = process.communicate()
|
|
output = stdout#.decode('utf-8')
|
|
error = stderr#.decode('utf-8')
|
|
return_code = process.returncode
|
|
return return_code, output, error |