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 [], "{} 不是一个目录".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("参数异常") if not self.task: raise RuntimeError("任务不存在") 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("初始化文件列表失败", 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("初始化文件列表失败", 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("节点索引超出范围") 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, 队列已空" % worker_id) break except Exception as e: print("worker_id: %s, 获取任务失败" % 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": "数据源节点执行传输异常，请检查节点是否正常"}) 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 "任务启动超时" 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 "任务不存在" # 远程节点任务 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 "日志ID不能为空" log = fdb.TransferLog.get_byid(log_id) if not log: return "日志不存在" if log.status != 3: return "任务状态不是异常, 无需重试" if log.transfer_task_id != task_id: return "日志ID与任务ID不匹配" file_task = FiletransferTask(task, print, exclude_nodes=[], the_log_id=log_id) file_task.start() return file_task.status_dict