mirror of
https://github.com/aaPanel/aaPanel.git
synced 2026-09-06 01:27:38 +02:00
Update to v7.65.0
This commit is contained in:
aapanel.com
@@ -0,0 +1,484 @@
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
import queue
|
||||
import time
|
||||
import traceback
|
||||
import itertools
|
||||
from datetime import datetime
|
||||
from typing import List, Dict, Callable, Any, Tuple, Union, Optional
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
|
||||
from mod.base.ssh_executor import SSHExecutor
|
||||
from mod.project.node.dbutil import ServerNodeDB, TaskFlowsDB, TransferTask, TransferFile, TransferLog
|
||||
from mod.project.node.nodeutil import ServerNode, LPanelNode, SSHApi
|
||||
from mod.project.node.filetransfer.socket_server import StatusServer, StatusClient, register_cleanup
|
||||
|
||||
|
||||
_SOCKET_FILE_DIR = "/tmp/flow_task"
|
||||
if not os.path.exists(_SOCKET_FILE_DIR):
|
||||
os.mkdir(_SOCKET_FILE_DIR)
|
||||
|
||||
def _dir_walk(path: str) -> Tuple[List[dict], str]:
|
||||
if not os.path.isdir(path):
|
||||
return [], "{} Not a directory".format(path)
|
||||
res_file = []
|
||||
count = 0
|
||||
empty_dir = []
|
||||
for root, dirs, files in os.walk(path):
|
||||
if not files:
|
||||
empty_dir.append(root)
|
||||
for f in files:
|
||||
count += 1
|
||||
try:
|
||||
res_file.append({
|
||||
"path": os.path.join(root, f),
|
||||
"size": os.path.getsize(os.path.join(root, f)),
|
||||
"is_dir": 0
|
||||
})
|
||||
except:
|
||||
pass
|
||||
return [{"path": d, "size": 0, "is_dir": 1} for d in empty_dir] + res_file, ""
|
||||
|
||||
|
||||
class FiletransferTask(object):
|
||||
|
||||
def __init__(self, task: Union[int, TransferTask],
|
||||
call_update: Callable[[Any], None],
|
||||
exclude_nodes: List[int] = None,
|
||||
the_log_id: int = None,
|
||||
):
|
||||
self._fdb = TaskFlowsDB()
|
||||
if isinstance(task, int):
|
||||
self.task = self._fdb.TransferTask.get_byid(task)
|
||||
elif isinstance(task, TransferTask):
|
||||
self.task = task
|
||||
else:
|
||||
raise ValueError("Parameter exception")
|
||||
|
||||
if not self.task:
|
||||
raise RuntimeError("Task does not exist")
|
||||
|
||||
self.exclude_nodes = exclude_nodes or []
|
||||
self.the_log_id = max(the_log_id, 0) if isinstance(the_log_id, int) else 0
|
||||
self.event_queue = queue.Queue()
|
||||
self.trans_queue = queue.Queue()
|
||||
self.mut = threading.Lock()
|
||||
self._srv_cache: Dict[int, Union[SSHApi, LPanelNode, ServerNode]] = {}
|
||||
self.status_dict: Dict[str, Any] = {
|
||||
"task_id": self.task.id,
|
||||
"task_type": "file",
|
||||
"flow_idx": self.task.step_index -1,
|
||||
"count": 0,
|
||||
"complete": 0,
|
||||
"error": 0,
|
||||
"error_nodes": [],
|
||||
"exclude_nodes": self.exclude_nodes,
|
||||
"data": None,
|
||||
}
|
||||
self.is_trans_end = False
|
||||
self.call_update = call_update
|
||||
|
||||
def _init_files(self): # 初始化文件列表
|
||||
has_file = self._fdb.TransferFile.find("flow_id = ? AND transfer_task_id = ?", (self.task.flow_id, self.task.id))
|
||||
# 判断文件列表是否已经初始化
|
||||
if has_file:
|
||||
return
|
||||
|
||||
file_list = []
|
||||
for src_item in self.task.path_list:
|
||||
dst_path = src_item["dst_path"].rstrip("/")
|
||||
src_item["path"] = src_item["path"].rstrip("/")
|
||||
if not os.path.exists(src_item["path"]):
|
||||
continue
|
||||
src_item["is_dir"] = os.path.isdir(src_item["path"])
|
||||
if src_item["is_dir"]:
|
||||
f_list, err = _dir_walk(src_item["path"])
|
||||
if not f_list:
|
||||
src_item["dst_file"] = os.path.join(dst_path, os.path.basename(src_item["path"]))
|
||||
file_list.append(src_item)
|
||||
else:
|
||||
for f_item in f_list:
|
||||
f_item["dst_file"] = f_item["path"].replace(os.path.dirname(src_item["path"]), dst_path)
|
||||
file_list.extend(f_list)
|
||||
else:
|
||||
if not os.path.isfile(src_item["path"]):
|
||||
continue
|
||||
src_item["dst_file"] = os.path.join(dst_path, os.path.basename(src_item["path"]))
|
||||
src_item["size"] = os.path.getsize(src_item["path"])
|
||||
file_list.append(src_item)
|
||||
|
||||
t_list = []
|
||||
for f_item in file_list:
|
||||
fl = TransferFile(
|
||||
flow_id=self.task.flow_id,
|
||||
transfer_task_id=self.task.id,
|
||||
src_file=f_item["path"],
|
||||
dst_file=f_item["dst_file"],
|
||||
file_size=f_item["size"],
|
||||
is_dir=f_item["is_dir"],
|
||||
)
|
||||
t_list.append(fl)
|
||||
try:
|
||||
self._fdb.TransferFile.create(t_list)
|
||||
except:
|
||||
print("Failed to initialize file list", traceback.format_exc())
|
||||
|
||||
def _init_files_log(self):
|
||||
tf_list = self._fdb.TransferFile.query("flow_id = ? AND transfer_task_id = ?", (self.task.flow_id, self.task.id))
|
||||
if not tf_list:
|
||||
return []
|
||||
has_fl = self._fdb.TransferLog.query("transfer_task_id = ? AND transfer_file_id = ?", (self.task.id, tf_list[0].id))
|
||||
if has_fl:
|
||||
return self._fdb.TransferLog.query("transfer_task_id = ?", (self.task.id,))
|
||||
|
||||
fl_list = []
|
||||
for (tf, dst_node_id) in itertools.product(tf_list, self.task.dst_nodes.keys()):
|
||||
fl = TransferLog(
|
||||
flow_id=self.task.flow_id,
|
||||
transfer_task_id=self.task.id,
|
||||
transfer_file_id=tf.id,
|
||||
dst_node_idx=int(dst_node_id),
|
||||
status=0,
|
||||
progress=0,
|
||||
message=""
|
||||
)
|
||||
fl_list.append(fl)
|
||||
|
||||
try:
|
||||
self._fdb.TransferLog.create(fl_list)
|
||||
except:
|
||||
print("Failed to initialize file list", traceback.format_exc())
|
||||
|
||||
|
||||
def _get_srv(self, idx: int) -> Union[SSHApi, LPanelNode, ServerNode]:
|
||||
idx = int(idx)
|
||||
if idx in self._srv_cache:
|
||||
return self._srv_cache[idx]
|
||||
with self.mut:
|
||||
if idx in self._srv_cache:
|
||||
return self._srv_cache[idx]
|
||||
if idx not in self.task.dst_nodes:
|
||||
raise RuntimeError("Node index is out of range")
|
||||
srv_data: dict = self.task.dst_nodes[idx]
|
||||
if srv_data.get("lpver", None):
|
||||
srv = LPanelNode(srv_data["address"], srv_data["api_key"], srv_data["lpver"])
|
||||
elif srv_data["api_key"] or srv_data["app_key"]:
|
||||
srv = ServerNode(srv_data["address"], srv_data["api_key"], srv_data["app_key"])
|
||||
else:
|
||||
srv_data["ssh_conf"]["threading_mod"] = True # 线程模式, 在不同线程中使用同一个ssh链接的不同会话
|
||||
srv = SSHApi(**srv_data["ssh_conf"])
|
||||
self._srv_cache[idx] = srv
|
||||
return srv
|
||||
|
||||
def start(self):
|
||||
self.task.status = 1
|
||||
self._fdb.TransferTask.update(self.task)
|
||||
self._init_files()
|
||||
self._init_files_log()
|
||||
if self.the_log_id > 0: # 重试某个固定任务
|
||||
query_where = "transfer_task_id = ? and id = ?"
|
||||
files_logs = self._fdb.TransferLog.query(query_where, (self.task.id, self.the_log_id))
|
||||
else:
|
||||
if self.exclude_nodes:
|
||||
# 获取未完成文件列表
|
||||
query_where = "transfer_task_id = ? and status not in (2, 4) and dst_node_idx not in ({})".format(
|
||||
",".join(["?"] * len(self.exclude_nodes))
|
||||
)
|
||||
else:
|
||||
query_where = "transfer_task_id = ? and status not in (2, 4)"
|
||||
files_logs = self._fdb.TransferLog.query(query_where, (self.task.id, *self.exclude_nodes))
|
||||
files_list = self._fdb.TransferFile.query("transfer_task_id = ?", (self.task.id,))
|
||||
if not files_logs:
|
||||
return
|
||||
files_map = {fl.id: fl for fl in files_list}
|
||||
for (idx, fl) in enumerate(files_logs):
|
||||
fl.log_idx = idx
|
||||
fl.tf = files_map[fl.transfer_file_id]
|
||||
self.trans_queue.put(fl)
|
||||
|
||||
self.status_dict["count"] = len(files_logs)
|
||||
|
||||
th_event = threading.Thread(target=self.event_func,)
|
||||
th_event.start()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=8) as executor:
|
||||
futures = [executor.submit(self.once_trans, worker_id) for worker_id in range(8)]
|
||||
for i in range(8):
|
||||
executor.submit(self.once_trans)
|
||||
for future in as_completed(futures):
|
||||
print("Completed result:", future.result())
|
||||
|
||||
self.is_trans_end = True
|
||||
th_event.join()
|
||||
if self.the_log_id > 0:
|
||||
# 如果有未完成或错误的文件， 则任务完成
|
||||
if self._fdb.TransferLog.count("transfer_task_id = ? and status in (0,1,3)", (self.task.id, )) == 0:
|
||||
self.task.status = 3
|
||||
else:
|
||||
self.task.status = 2
|
||||
else:
|
||||
self.task.status = 2 if self.status_dict["error"] == 0 else 3
|
||||
self._fdb.TransferTask.update(self.task)
|
||||
self._fdb.close()
|
||||
|
||||
def once_trans(self, worker_id: int):
|
||||
while True:
|
||||
try:
|
||||
tl = self.trans_queue.get(block=False)
|
||||
except queue.Empty:
|
||||
print("worker_id: %s, The queue is empty" % worker_id)
|
||||
break
|
||||
except Exception as e:
|
||||
print("worker_id: %s, Failed to obtain task" % worker_id)
|
||||
print(traceback.format_exc())
|
||||
break
|
||||
|
||||
# 执行一次文件传输
|
||||
try:
|
||||
if tl.status == 2: # 跳过已完成的文件
|
||||
self.event_queue.put(tl)
|
||||
continue
|
||||
srv = self._get_srv(tl.dst_node_idx)
|
||||
if tl.tf.is_dir: # 处理空目录
|
||||
exits, err = srv.target_file_exits(tl.tf.dst_file)
|
||||
if err: # 获取文件状态错误
|
||||
tl.message = err
|
||||
tl.status = 3
|
||||
self.event_queue.put(tl)
|
||||
elif exits: # 目标文件已存在
|
||||
tl.status = 4
|
||||
tl.progress = 100
|
||||
self.event_queue.put(tl)
|
||||
else: # 目标文件不存在， 创建目录
|
||||
res, err = srv.create_dir(tl.tf.dst_file)
|
||||
if err:
|
||||
tl.message = err
|
||||
tl.status = 3
|
||||
elif isinstance(res, dict):
|
||||
if res["status"]:
|
||||
tl.status = 2
|
||||
tl.message = ""
|
||||
tl.progress = 100
|
||||
else:
|
||||
tl.message = res["msg"]
|
||||
tl.status = 3
|
||||
else:
|
||||
tl.status = 2
|
||||
tl.message = ""
|
||||
tl.progress = 100
|
||||
|
||||
self.event_queue.put(tl)
|
||||
else: # 处理文件上传
|
||||
tl.status = 1
|
||||
self.event_queue.put(tl)
|
||||
def _call_log(progress, log):
|
||||
tl.progress = progress
|
||||
self.event_queue.put(tl)
|
||||
|
||||
err = srv.upload_file(
|
||||
filename=tl.tf.src_file,
|
||||
target_path=os.path.dirname(tl.tf.dst_file),
|
||||
mode="cover",
|
||||
call_log=_call_log)
|
||||
|
||||
if err:
|
||||
tl.status = 3
|
||||
tl.message = err
|
||||
else:
|
||||
tl.status = 2
|
||||
tl.message = ""
|
||||
tl.progress = 100
|
||||
|
||||
self.event_queue.put(tl)
|
||||
except Exception as e:
|
||||
err = traceback.format_exc()
|
||||
tl.status = 3
|
||||
tl.message = str(e) + "\n" + err
|
||||
self.event_queue.put(tl)
|
||||
|
||||
def event_func(self):
|
||||
fdb = TaskFlowsDB()
|
||||
last_time = time.time()
|
||||
tmp_dict = {}
|
||||
update_fields = ("status", "message", "progress", "completed_at", "started_at")
|
||||
complete_set, error_set = set(), set()
|
||||
error_node_set = set()
|
||||
while True:
|
||||
try:
|
||||
tl: TransferLog = self.event_queue.get(timeout=0.1)
|
||||
except queue.Empty:
|
||||
if self.is_trans_end:
|
||||
break
|
||||
else:
|
||||
continue
|
||||
except Exception as e:
|
||||
print(e)
|
||||
break
|
||||
if tl.status in (2, 4):
|
||||
complete_set.add(tl.id)
|
||||
self.status_dict["complete"] = len(complete_set)
|
||||
if not tl.started_at:
|
||||
tl.started_at = tl.started_at or datetime.now()
|
||||
tl.completed_at = tl.completed_at or datetime.now()
|
||||
elif tl.status == 3:
|
||||
error_set.add(tl.id)
|
||||
self.status_dict["error"] = len(error_set)
|
||||
tl.completed_at = datetime.now()
|
||||
error_node_set.add(tl.dst_node_idx)
|
||||
elif tl.status == 1:
|
||||
tl.started_at = datetime.now()
|
||||
|
||||
tmp_dict[tl.id] = tl
|
||||
if time.time() - last_time > 0.5:
|
||||
fdb.TransferLog.bath_update(tmp_dict.values(), update_fields=update_fields)
|
||||
last_time = time.time()
|
||||
|
||||
self.status_dict["data"] = [i.to_show_data() for i in tmp_dict.values()]
|
||||
self.status_dict["error_nodes"] = list(error_node_set)
|
||||
self.call_update(self.status_dict)
|
||||
tmp_dict.clear()
|
||||
|
||||
|
||||
if tmp_dict:
|
||||
fdb.TransferLog.bath_update(tmp_dict.values(), update_fields=update_fields)
|
||||
self.status_dict["data"] = [i.to_show_data() for i in tmp_dict.values()]
|
||||
self.status_dict["error_nodes"] = list(error_node_set)
|
||||
self.call_update(self.status_dict)
|
||||
|
||||
fdb.close()
|
||||
|
||||
|
||||
# 在远程节点上执行文件传输
|
||||
class NodeFiletransferTask(object):
|
||||
|
||||
def __init__(self, task: TransferTask,
|
||||
call_update: Callable[[Any], None],
|
||||
exclude_nodes: List[int] = None,
|
||||
the_log_id: int = None,
|
||||
):
|
||||
self.task = task
|
||||
src_node = task.src_node
|
||||
self.exclude_nodes = exclude_nodes or []
|
||||
self.srv = ServerNode(src_node["address"],src_node["api_key"], src_node["app_key"], src_node["name"])
|
||||
self.the_log_id = max(the_log_id, 0) if isinstance(the_log_id, int) else 0
|
||||
self.call_update = call_update
|
||||
self.default_status_data = {
|
||||
"task_id": self.task.id,
|
||||
"task_type": "file",
|
||||
}
|
||||
self.status_dict = dict() # 状态数据
|
||||
|
||||
def start(self):
|
||||
fdb = TaskFlowsDB()
|
||||
self.task.status = 1
|
||||
fdb.TransferTask.update(self.task)
|
||||
err = self.srv.proxy_transferfile_status(self.task.src_node_task_id, self.exclude_nodes, self.the_log_id, self.handle_proxy_data)
|
||||
if err:
|
||||
self.task.status = 3
|
||||
self.task.message += ";" + err
|
||||
else:
|
||||
if self.status_dict and self.status_dict.get("error", 0):
|
||||
self.task.status = 3
|
||||
else:
|
||||
self.task.status = 2
|
||||
if self.task.message:
|
||||
self.task.status = 3
|
||||
|
||||
fdb.TransferTask.update(self.task)
|
||||
|
||||
def handle_proxy_data(self, data):
|
||||
ret = {"count": 0,"complete": 0,"error": 0, "error_nodes":[], "data": []}
|
||||
try:
|
||||
data_dict = json.loads(data)
|
||||
if "type" not in data_dict:
|
||||
return
|
||||
|
||||
if data_dict["type"] == "status":
|
||||
if "init" in data_dict["data"]: # 初始化状态跳过
|
||||
return
|
||||
ret.update(data_dict["data"])
|
||||
ret.update(self.default_status_data)
|
||||
else: # end / error 状态 获取历史数据或错误信息
|
||||
if "data" in data_dict:
|
||||
ret.update(data_dict["data"])
|
||||
ret.update(self.default_status_data)
|
||||
elif "msg" in data_dict:
|
||||
self.task.message = data_dict["msg"]
|
||||
return
|
||||
except:
|
||||
print(traceback.format_exc())
|
||||
ret["data"].append({"message": "Data source node execution transmission exception, please check if the node is functioning properly"})
|
||||
ret.update(self.default_status_data)
|
||||
|
||||
self.status_dict = ret
|
||||
self.call_update(ret)
|
||||
|
||||
|
||||
# 本机执行文件传输，返回信息到远程节点
|
||||
class SelfFiletransferTask(object):
|
||||
|
||||
def __init__(self, task_id: int, exclude_nodes: List[int] = None, the_log_id: int = None):
|
||||
self.status_server = StatusServer(self.get_status, (_SOCKET_FILE_DIR + "/file_task_" + str(task_id)))
|
||||
self.f_task = FiletransferTask(task_id, self.update_status, exclude_nodes, the_log_id)
|
||||
|
||||
@staticmethod
|
||||
def get_status( init: bool = False) -> Dict:
|
||||
return {"init": True }
|
||||
|
||||
def start_status_server(self):
|
||||
t = threading.Thread(target=self.status_server.start_server, args=(), daemon=True)
|
||||
t.start()
|
||||
register_cleanup(self.status_server)
|
||||
|
||||
def update_status(self, update_data: Dict):
|
||||
self.status_server.update_status(update_data)
|
||||
|
||||
def start(self):
|
||||
self.start_status_server()
|
||||
self.f_task.start()
|
||||
return
|
||||
|
||||
|
||||
def self_file_running_log(task_id: int, call_log: Callable[[Union[str,dict]], None], timeout:float = 3.0) -> str:
|
||||
socket_file = _SOCKET_FILE_DIR + "/file_task_" + str(task_id)
|
||||
while not os.path.exists(socket_file):
|
||||
if timeout <= 0:
|
||||
return "Task startup timeout"
|
||||
timeout -= 0.05
|
||||
time.sleep(0.05)
|
||||
|
||||
s_client = StatusClient(socket_file, callback=call_log)
|
||||
s_client.connect()
|
||||
s_client.wait_receive()
|
||||
return ""
|
||||
|
||||
|
||||
# 同步执行文件相关任务的重试
|
||||
def file_task_run_sync(task_id: int, log_id: int) -> Union[str, Dict[str, Any]]:
|
||||
fdb = TaskFlowsDB()
|
||||
task = fdb.TransferTask.get_byid(task_id)
|
||||
if not task:
|
||||
return "Task does not exist"
|
||||
|
||||
# 远程节点任务
|
||||
if task.src_node_task_id > 0:
|
||||
node_file_task = NodeFiletransferTask(task, print, exclude_nodes=[], the_log_id=log_id)
|
||||
node_file_task.start()
|
||||
return node_file_task.status_dict
|
||||
|
||||
if not log_id:
|
||||
return "The log ID cannot be empty"
|
||||
log = fdb.TransferLog.get_byid(log_id)
|
||||
if not log:
|
||||
return "log does not exist"
|
||||
|
||||
if log.status != 3:
|
||||
return "The task status is not abnormal, no need to retry"
|
||||
|
||||
if log.transfer_task_id != task_id:
|
||||
return "The log ID does not match the task ID"
|
||||
|
||||
file_task = FiletransferTask(task, print, exclude_nodes=[], the_log_id=log_id)
|
||||
file_task.start()
|
||||
return file_task.status_dict
|
||||
Reference in New Issue
Block a user