mirror of
https://github.com/aaPanel/aaPanel.git
synced 2026-08-20 06:35:47 +02:00
485 lines
18 KiB
Python
485 lines
18 KiB
Python
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
|