hardware-dashboard/task_exec.py

105 lines
5.7 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#!/usr/bin/python3
"""Run one approved, root-owned task definition; record bounded output and result."""
import fcntl
import hashlib
import json
import os
import selectors
import shlex
import signal
import subprocess
import sys
import time
import uuid
import task_core as core
LIMIT=512*1024
def main(task_id):
task=core.get_task(task_id)
definition=task['definition']
lock=(core.BASE/(task_id+'.lock')).open('a')
try:fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB)
except BlockingIOError:return 0
run_id=uuid.uuid4().hex;started=time.time();start_mono=time.monotonic()
boot=core.read_boot() if hasattr(core,'read_boot') else open('/proc/sys/kernel/random/boot_id').read().strip()
with core.db() as con:
requested=con.execute('SELECT * FROM requests WHERE task_id=?',(task_id,)).fetchone()
trigger='manual' if requested and requested['expires']>=started else 'scheduled'
con.execute('DELETE FROM requests WHERE task_id=?',(task_id,))
body=definition['body']
fingerprint=hashlib.sha256((body if definition['mode']!='file' else definition['path']+' '+definition['arguments']).encode()).hexdigest()
con.execute('INSERT INTO runs(id,task_id,task_name,started,state,trigger,boot,script_hash,revision) VALUES(?,?,?,?,?,?,?,?,?)',
(run_id,task_id,task['name'],started,'running',trigger,boot,fingerprint,task['revision']))
cancelled=[False]
def stop(*_):cancelled[0]=True
signal.signal(signal.SIGTERM,stop);signal.signal(signal.SIGINT,stop)
state='failed';exit_code=None;proc=None;written=0;truncated=False
log_path=core.LOGS/(run_id+'.log')
with log_path.open('wb',buffering=0) as output:
os.chmod(log_path,0o600)
def write(data):
nonlocal written,truncated
if isinstance(data,str):data=data.encode('utf-8')
available=max(0,LIMIT-written)
if available:
piece=data[:available];output.write(piece);written+=len(piece)
if len(data)>available:truncated=True
write(f'开始:{core.iso(started)}\n任务:{task["name"]}\n触发:{"手动" if trigger=="manual" else "定时"}\n执行身份:root\n修订:{task["revision"]}\n'+'─'*32+'\n')
try:
if definition['mode']=='command':
argv=['/bin/bash','-e','-o','pipefail','-c',body]
elif definition['mode']=='script':
script=core.BASE/'scripts'/(run_id+('.py' if definition['interpreter']=='python3' else '.sh'))
core.atom_write(script,body+'\n')
argv=['/usr/bin/python3',str(script)] if definition['interpreter']=='python3' else ['/bin/bash','-e','-o','pipefail',str(script)]
else:
argv=[definition['path'],*shlex.split(definition['arguments'])]
if definition['interpreter']=='bash':argv=['/bin/bash','-e','-o','pipefail',*argv]
elif definition['interpreter']=='python3':argv=['/usr/bin/python3',*argv]
env={'HOME':'/root','USER':'root','LOGNAME':'root','LANG':'C.UTF-8','TZ':'Asia/Shanghai',
'PATH':'/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/snap/bin',
'HARDWARE_TASK_ID':task_id,'HARDWARE_TASK_RUN_ID':run_id}
proc=subprocess.Popen(argv,cwd=definition['cwd'],env=env,stdout=subprocess.PIPE,stderr=subprocess.STDOUT,start_new_session=True)
selector=selectors.DefaultSelector();selector.register(proc.stdout,selectors.EVENT_READ)
deadline=start_mono+definition['timeout'];exited_at=None
while selector.get_map() or proc.poll() is None:
if cancelled[0] or time.monotonic()>=deadline:
state='cancelled' if cancelled[0] else 'timeout'
try:os.killpg(proc.pid,signal.SIGTERM)
except ProcessLookupError:pass
try:proc.wait(timeout=3)
except subprocess.TimeoutExpired:
try:os.killpg(proc.pid,signal.SIGKILL)
except ProcessLookupError:pass
proc.wait(timeout=3)
break
for key,_ in selector.select(.25):
chunk=os.read(key.fileobj.fileno(),32768)
if chunk:write(chunk)
else:selector.unregister(key.fileobj)
if proc.poll() is not None:
exited_at=exited_at or time.monotonic()
if time.monotonic()-exited_at>1:break
selector.close();proc.stdout.close()
exit_code=proc.wait(timeout=3)
if state not in ('cancelled','timeout'):state='success' if exit_code==0 else 'failed'
except Exception as exc:
write('\n执行错误:'+str(exc)+'\n');state='failed'
if proc and proc.poll() is None:
try:os.killpg(proc.pid,signal.SIGKILL);proc.wait(timeout=3)
except (ProcessLookupError,subprocess.TimeoutExpired):pass
elapsed=time.monotonic()-start_mono
tail=f'\n'+('─'*32)+f'\n结束:{core.iso()}\n结果:{state}\n退出码:{exit_code}\n耗时:{elapsed:.2f} 秒\n'
if truncated:output.write('\n[输出超过 512 KiB,后续内容已省略。]\n'.encode())
output.write(tail.encode())
os.fsync(output.fileno())
with core.db() as con:
con.execute('UPDATE runs SET finished=?,state=?,exit_code=?,log_bytes=?,truncated=?,duration=? WHERE id=?',
(time.time(),state,exit_code,log_path.stat().st_size,int(truncated),elapsed,run_id))
lock.close()
return 0 if state=='success' else 1
if __name__=='__main__':
raise SystemExit(main(core.validate_id(sys.argv[1])))