commit fea70def640471732bd022b7962bf894ef058acd Author: yufengzhen Date: Wed Jun 17 19:50:39 2026 +0800 init diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..7832012 --- /dev/null +++ b/.gitignore @@ -0,0 +1,26 @@ +*.csv +*.log +**/__pycache__/ +*.pyc +.vscode/ +.DS_Store +logs/ +.python-version +pyproject.toml +uv.lock +runtime/* +worker_inventory.json.bak +worker_inventory.json +config.yaml +config/config.json +config/current_env.txt +config/local.json +config/prod.json +config/test.json +**/DroidBot/guiagent_core/credentials.json +**/DroidBot/guiagent_core/token.json +**/guiagent_core/credentials.json +**/guiagent_core/token.json +AGENTS.md +.claude/ +tmp/ diff --git a/analytics.py b/analytics.py new file mode 100644 index 0000000..621593f --- /dev/null +++ b/analytics.py @@ -0,0 +1,3655 @@ +import csv +import glob +import json +import os +import sqlite3 +import threading +import time +from contextlib import contextmanager +from datetime import datetime +from typing import Any, Dict, Iterable, List, Optional, Tuple +from zoneinfo import ZoneInfo + +from config import ( + ANALYTICS_ARTIFACT_WAIT_SECONDS, + ANALYTICS_ARTIFACT_FILE_WAIT_SECONDS, + ANALYTICS_JOB_POLL_SECONDS, + ANALYTICS_TPDPI_APP_LIST, + ANALYTICS_TPDPI_URL_LIB, + ANALYTICS_TRAFFIC_ROOT, + ANALYTICS_TRAVERSAL_ROOT, + MODEL_TRAFFIC_THRESHOLD, + MONITORING_DB_PATH, + MONITORING_TIMEZONE, + normalize_worker_id, +) +from log_manager import logger +from result_codes import ( + APP_ERROR_DESC, + BUSINESS_ERROR_DESC, + DOWNLOAD_ERROR_DESC, + INFRA_ERROR_DESC, + AppError, + BusinessError, + DownloadError, + ErrorCategory, + ErrorInfo, + InfraError, +) + + +SHANGHAI_TZ = ZoneInfo(MONITORING_TIMEZONE) +NON_RETRYABLE_DOWNLOAD_BUCKETS = ( + ("lt_1k", "<1K", 0, 1_000), + ("1k_to_10k", "1K-10K", 1_000, 10_000), + ("10k_to_100k", "10K-100K", 10_000, 100_000), + ("100k_to_1m", "100K-1M", 100_000, 1_000_000), + ("1m_to_10m", "1M-10M", 1_000_000, 10_000_000), + ("10m_to_100m", "10M-100M", 10_000_000, 100_000_000), + ("100m_plus", "100M+", 100_000_000, None), +) +NON_RETRYABLE_DOWNLOAD_SPLIT_THRESHOLD = 10_000 +LIGHT_RESTRICTED_MIN_NUM_NODES = 5 +AUTO_PENDING_PHYSICAL_REASON = "auto_pending_physical_reroute" +ACTIVE_CATALOG_WHERE = "COALESCE(is_active, 1) = 1" +DEFAULT_INCREMENTAL_TOP_N = 3000 +AUTO_PENDING_PHYSICAL_FAILURE_TYPES = { + f"APP_ERROR/{int(AppError.CRASH)}", + f"APP_ERROR/{int(AppError.ROOT_MODE_UNSUPPORTED)}", +} + + +def _safe_json_dumps(payload: Any) -> str: + return json.dumps(payload, ensure_ascii=False, sort_keys=True) + + +def _safe_json_loads(payload: Optional[str], fallback: Optional[Any] = None) -> Any: + if not payload: + return fallback + try: + return json.loads(payload) + except (TypeError, ValueError): + return fallback + + +def _format_ts(value: Optional[float]) -> str: + if not value: + return "" + return datetime.fromtimestamp(float(value), SHANGHAI_TZ).strftime("%Y-%m-%d %H:%M:%S") + + +def _normalize_incremental_batch_tag(value: Any) -> str: + """返回单个标签字符串（向后兼容）。""" + return str(value or "").strip() + + +def _normalize_incremental_batch_tags(value: Any) -> List[str]: + """将存储的标签值解析为去重后的有序列表。支持 JSON 数组或旧版单标签字符串。""" + if value is None: + return [] + if isinstance(value, (list, tuple)): + result: List[str] = [] + seen = set() + for item in value: + tag = str(item or "").strip() + if tag and tag not in seen: + seen.add(tag) + result.append(tag) + return result + text = str(value).strip() + if not text: + return [] + if text.startswith("["): + try: + parsed = json.loads(text) + if isinstance(parsed, list): + return _normalize_incremental_batch_tags(parsed) + except (TypeError, ValueError): + pass + return [text] + + +def _incremental_batch_tags_to_json(tags: List[str]) -> str: + """将标签列表序列化为 JSON 数组字符串。""" + normalized = _normalize_incremental_batch_tags(tags) + return json.dumps(normalized, ensure_ascii=False) + + +def _incremental_batch_tag_in_tags_sql(tag: str) -> str: + """返回在 SQL 中检查 incremental_batch_tag JSON 数组是否包含指定 tag 的表达式。""" + escaped = str(tag or "").replace("'", "''") + return f"incremental_batch_tag LIKE '%\"{escaped}\"%'" + + +def _add_incremental_batch_tag(existing_tags: Any, new_tag: str) -> List[str]: + """向现有标签列表中添加一个新标签（去重）。""" + tags = _normalize_incremental_batch_tags(existing_tags) + normalized_new = str(new_tag or "").strip() + if normalized_new and normalized_new not in tags: + tags.append(normalized_new) + return tags + + +def _remove_incremental_batch_tag(existing_tags: Any, tag_to_remove: str) -> List[str]: + """从现有标签列表中移除指定标签。""" + tags = _normalize_incremental_batch_tags(existing_tags) + normalized_remove = str(tag_to_remove or "").strip() + if not normalized_remove: + return tags + return [t for t in tags if t != normalized_remove] + + +def _normalize_app_magic_label(value: Any) -> str: + return str(value or "").strip() + + +def _normalize_last_updated(value: Any) -> str: + text = str(value or "").strip() + if not text: + return "" + normalized = text.replace("-", "/").replace(".", "/") + for fmt in ("%Y/%m/%d", "%Y/%-m/%-d"): + try: + parsed = datetime.strptime(normalized, fmt) + return parsed.strftime("%Y/%m/%d") + except ValueError: + continue + parts = normalized.split("/") + if len(parts) == 3: + try: + year, month, day = [int(part) for part in parts] + return datetime(year, month, day).strftime("%Y/%m/%d") + except (TypeError, ValueError): + return "" + return "" + + +def _last_updated_ts(value: Any) -> float: + normalized = _normalize_last_updated(value) + if not normalized: + return 0.0 + try: + return datetime.strptime(normalized, "%Y/%m/%d").timestamp() + except ValueError: + return 0.0 + + +def _last_updated_interval_days(new_value: Any, old_value: Any) -> int: + new_ts = _last_updated_ts(new_value) + old_ts = _last_updated_ts(old_value) + if not new_ts or not old_ts or new_ts <= old_ts: + return 0 + return max(1, int((new_ts - old_ts) // 86400)) + + +def _normalize_top_n(value: Any, default: int = DEFAULT_INCREMENTAL_TOP_N) -> int: + try: + normalized = int(value) + except (TypeError, ValueError): + normalized = default + return normalized if normalized > 0 else default + + +def _top_n_label(top_n: int) -> str: + return f"榜单 Top{int(top_n)}" + + +def _decode_mixed_csv_line(raw_line: bytes) -> str: + for encoding in ("utf-8-sig", "utf-8", "gb18030"): + try: + return raw_line.decode(encoding) + except UnicodeDecodeError: + continue + return raw_line.decode("gb18030", errors="replace") + + +def _read_csv_lines(csv_path: str) -> List[str]: + with open(csv_path, "rb") as handle: + return [_decode_mixed_csv_line(line) for line in handle] + + +def _read_dict_rows(csv_path: str) -> Iterable[Dict[str, str]]: + reader = csv.DictReader(_read_csv_lines(csv_path)) + for row in reader: + normalized = {} + for key, value in (row or {}).items(): + normalized[str(key or "").strip()] = str(value or "").strip() + yield normalized + + +def _parse_int_value(value: Any) -> int: + try: + return int(float(value or 0)) + except (TypeError, ValueError): + return 0 + + +def _parse_float_value(value: Any) -> float: + try: + return float(value or 0.0) + except (TypeError, ValueError): + return 0.0 + + +def _parse_downloads_value(value: Any) -> Optional[int]: + text = str(value or "").strip().upper() + if not text: + return None + normalized = text.replace(",", "").replace(" ", "") + if normalized.endswith("+"): + normalized = normalized[:-1] + multiplier = 1 + if normalized.endswith("K"): + multiplier = 1_000 + normalized = normalized[:-1] + elif normalized.endswith("M"): + multiplier = 1_000_000 + normalized = normalized[:-1] + elif normalized.endswith("B"): + multiplier = 1_000_000_000 + normalized = normalized[:-1] + elif normalized.endswith("W") or normalized.endswith("万"): + multiplier = 10_000 + normalized = normalized[:-1] + elif normalized.endswith("亿"): + multiplier = 100_000_000 + normalized = normalized[:-1] + try: + return int(float(normalized) * multiplier) + except (TypeError, ValueError): + return None + + +def _extract_downloads_from_task_payload(task_payload: Any) -> Optional[int]: + if isinstance(task_payload, dict): + payload = task_payload + else: + payload = _safe_json_loads(task_payload, {}) or {} + if not isinstance(payload, dict): + return None + original_row = payload.get("original_row") + if isinstance(original_row, dict): + for key in ("downloads", "download", "下载量"): + if key in original_row: + downloads = _parse_downloads_value(original_row.get(key)) + if downloads is not None: + return downloads + for key in ("downloads", "download", "下载量"): + if key in payload: + downloads = _parse_downloads_value(payload.get(key)) + if downloads is not None: + return downloads + return None + + +def _extract_downloads_from_row(row: Dict[str, Any]) -> Optional[int]: + if not isinstance(row, dict): + return None + return _parse_downloads_value( + row.get("downloads") + or row.get("download") + or row.get("下载量") + ) + + +def _normalize_download_errors(download_errors: Any) -> Dict[str, Dict[str, Any]]: + return download_errors if isinstance(download_errors, dict) else {} + + +def _resolve_downloads_for_package( + package_name: str, + task_payload: Any = None, + *, + downloads: Any = None, +) -> Optional[int]: + del package_name + normalized_downloads = _parse_downloads_value(downloads) + if normalized_downloads is not None: + return normalized_downloads + return _extract_downloads_from_task_payload(task_payload) + + +def _find_download_bucket_spec(downloads: int) -> Tuple[str, str, int, Optional[int]]: + normalized_downloads = max(0, int(downloads or 0)) + for bucket_key, label, lower_bound, upper_bound in NON_RETRYABLE_DOWNLOAD_BUCKETS: + if upper_bound is None and normalized_downloads >= lower_bound: + return bucket_key, label, lower_bound, upper_bound + if upper_bound is not None and lower_bound <= normalized_downloads < upper_bound: + return bucket_key, label, lower_bound, upper_bound + return NON_RETRYABLE_DOWNLOAD_BUCKETS[0] + + +def _build_non_retryable_download_distribution( + rows: Iterable[sqlite3.Row], +) -> Dict[str, Any]: + rows = list(rows) + if not rows: + return { + "distribution": [], + "matched_downloads_apps": 0, + "missing_downloads_apps": 0, + } + + distribution = [ + { + "bucket_key": bucket_key, + "label": label, + "min_downloads": lower_bound, + "max_downloads": upper_bound, + "app_count": 0, + } + for bucket_key, label, lower_bound, upper_bound in NON_RETRYABLE_DOWNLOAD_BUCKETS + ] + bucket_index = {item["bucket_key"]: item for item in distribution} + unknown_bucket = { + "bucket_key": "unknown", + "label": "未知", + "min_downloads": None, + "max_downloads": None, + "app_count": 0, + } + + matched_downloads_apps = 0 + for row in rows: + package_name = str(row["package_name"] or "").strip() + downloads = _resolve_downloads_for_package( + package_name, + row["task_payload_json"], + downloads=row["downloads"], + ) + if downloads is None: + unknown_bucket["app_count"] += 1 + continue + matched_downloads_apps += 1 + bucket_key, _, _, _ = _find_download_bucket_spec(downloads) + bucket_index[bucket_key]["app_count"] += 1 + + total_count = len(rows) + distribution.append(unknown_bucket) + return { + "distribution": [ + { + **item, + "app_count": int(item["app_count"] or 0), + "share_percent": round((int(item["app_count"] or 0) / total_count) * 100, 2) + if total_count > 0 + else 0.0, + } + for item in distribution + ], + "matched_downloads_apps": matched_downloads_apps, + "missing_downloads_apps": int(unknown_bucket["app_count"] or 0), + } + + +def _merge_detail_text(*parts: Any) -> str: + normalized_parts = [] + seen = set() + for part in parts: + text = str(part or "").strip() + if not text or text in seen: + continue + seen.add(text) + normalized_parts.append(text) + return " | ".join(normalized_parts) + + +def _parse_traffic_size(size_str: str) -> int: + size_str = str(size_str or "").strip() + if not size_str: + return 0 + number = [] + unit = [] + for char in size_str: + if char.isdigit() or char == ".": + number.append(char) + elif not char.isspace(): + unit.append(char) + if not number: + return 0 + try: + value = float("".join(number)) + except ValueError: + return 0 + normalized_unit = "".join(unit).upper() + multipliers = { + "": 1, + "B": 1, + "K": 1024, + "KB": 1024, + "M": 1024 ** 2, + "MB": 1024 ** 2, + "G": 1024 ** 3, + "GB": 1024 ** 3, + "T": 1024 ** 4, + "TB": 1024 ** 4, + } + return int(value * multipliers.get(normalized_unit, 1)) + + +def _is_model_traffic_line(fields: List[str]) -> bool: + if len(fields) < 7 or str(fields[2] or "").strip() != "model_data:": + return False + protocol = str(fields[5] or "").strip().upper() + if protocol not in ("TCP", "UDP"): + return False + flow_info = str(fields[3] or "").strip() + if not flow_info: + return False + parts = [part.strip() for part in flow_info.split("-") if part.strip()] + if len(parts) < 2: + return False + dst_part = parts[1] + dst_ip = dst_part.rsplit(":", 1)[0].strip() if ":" in dst_part else "" + if not dst_ip: + return False + for prefix in ("10.", "192.", "198.18", "198.19"): + if dst_ip.startswith(prefix): + return False + return True + + +def _determine_domain_or_model(fields: List[str]) -> str: + domain = fields[2].strip() if len(fields) > 2 else "" + flow_info = fields[3].strip() if len(fields) > 3 else "" + if domain != "model_data:": + return domain + return f"model_data:{flow_info}" + + +def _extract_second_level_domain(domain: str) -> str: + cleaned = str(domain or "").strip().lower() + if not cleaned: + return "" + cleaned = cleaned.replace("http://", "").replace("https://", "").strip().strip("\"'") + parts = [part for part in cleaned.split(".") if part] + if len(parts) <= 2: + return cleaned + return "*." + ".".join(parts[-2:]) + + +def _pick_match_state(states: Iterable[str]) -> str: + normalized_states = {str(item or "").strip() for item in states if str(item or "").strip()} + if "self" in normalized_states: + return "self" + if "server" in normalized_states: + return "server" + if "unmatched" in normalized_states: + return "unmatched" + if "ip" in normalized_states: + return "ip" + return "" + + +def _resolve_data_root(path: str) -> str: + raw_path = str(path or "").strip() + if not raw_path or os.path.isdir(raw_path): + return raw_path + normalized = raw_path.replace("\\", "/").lstrip("/") + parts = [part for part in normalized.split("/") if part] + if len(parts) >= 2 and "." in parts[0]: + # macOS: /Volumes//... + candidate = os.path.join("/Volumes", parts[1], *parts[2:]) + if os.path.isdir(candidate): + return candidate + # Linux: try common SMB mount root + candidate = os.path.join("/srv/samba", parts[1], *parts[2:]) + if os.path.isdir(candidate): + return candidate + if len(parts) >= 3 and parts[0].endswith(":") and parts[1].lower() == "share": + candidate = os.path.join("/Volumes", "dpi-sync", *parts[2:]) + if os.path.isdir(candidate): + return candidate + return raw_path + + +def _aggregate_second_level_domain_rows(domain_rows: Iterable[sqlite3.Row]) -> List[Dict[str, Any]]: + aggregated: Dict[str, Dict[str, Any]] = {} + for row in domain_rows: + domain = str(row["domain"] or "") + if str(row["domain_type"] or "") == "model_data" or domain.startswith("model_data:"): + continue + second_level_domain = _extract_second_level_domain(domain) + if not second_level_domain: + continue + bucket = aggregated.setdefault( + second_level_domain, + { + "domain": second_level_domain, + "traffic_bytes": 0, + "flow_count": 0, + "organization": "", + "matched_pattern": "", + "states": set(), + }, + ) + bucket["traffic_bytes"] += int(row["traffic_bytes"] or 0) + bucket["flow_count"] += int(row["flow_count"] or 0) + if row["organization"] and not bucket["organization"]: + bucket["organization"] = row["organization"] + if row["matched_pattern"] and not bucket["matched_pattern"]: + bucket["matched_pattern"] = row["matched_pattern"] + if row["match_state"]: + bucket["states"].add(row["match_state"]) + aggregated_rows = [] + for bucket in aggregated.values(): + aggregated_rows.append( + { + "domain": bucket["domain"], + "traffic_bytes": int(bucket["traffic_bytes"]), + "flow_count": int(bucket["flow_count"]), + "organization": bucket["organization"], + "matched_pattern": bucket["matched_pattern"], + "match_state": _pick_match_state(bucket["states"]), + } + ) + return sorted(aggregated_rows, key=lambda item: (-int(item["traffic_bytes"]), item["domain"])) + + +def _classify_restriction_status( + latest_status: str, + self_ratio: float, + latest_failure_type: str, + *, + num_nodes: int = 0, + download_errors: Any = None, +) -> Tuple[str, str]: + normalized_status = str(latest_status or "").strip().lower() + if normalized_status == "success": + return "success", "not_applicable" + try: + normalized_self_ratio = float(self_ratio or 0.0) + except (TypeError, ValueError): + normalized_self_ratio = 0.0 + normalized_num_nodes = _parse_int_value(num_nodes) + if normalized_self_ratio > 0 or normalized_num_nodes >= LIGHT_RESTRICTED_MIN_NUM_NODES: + return "light_restricted", "not_applicable" + + retryability = "retryable" + error_type = str(latest_failure_type or "").strip() + if error_type: + category_name, _, code_text = error_type.partition("/") + try: + error_info = ErrorInfo(ErrorCategory[category_name], int(code_text), error_type) + retryability = "non_retryable" if error_info.is_no_retry(download_errors=download_errors) else "retryable" + except (KeyError, TypeError, ValueError): + retryability = "retryable" + return "severe_restricted", retryability + + +def _make_task_key(app_name: str, package_name: str) -> str: + return f"{str(app_name or '').strip()}_{str(package_name or '').strip()}" + + +def _normalize_country_codes(raw_country_code: Any) -> List[str]: + if isinstance(raw_country_code, (list, tuple)): + raw_items = raw_country_code + else: + raw_items = str(raw_country_code or "").split(",") + normalized: List[str] = [] + seen = set() + for item in raw_items: + candidate = str(item or "").strip().upper() + if not candidate or candidate in seen: + continue + seen.add(candidate) + normalized.append(candidate) + return normalized + + +def _normalize_device_type(raw_device_type: Any) -> str: + candidate = str(raw_device_type or "").strip().lower() + if candidate in {"emulator", "physical"}: + return candidate + return "" + + +def _normalize_task_device_type(raw_device_type: Any) -> str: + candidate = str(raw_device_type or "").strip().lower() + if candidate in {"emulator", "physical", "any"}: + return candidate + return "emulator" + + +def _build_task_payload_with_device_type( + task_payload: Any, + *, + app_name: str, + package_name: str, + country_code: str, + device_type: str, +) -> Dict[str, Any]: + normalized_device_type = _normalize_task_device_type(device_type) + if isinstance(task_payload, dict) and task_payload: + payload = dict(task_payload) + else: + payload = build_task_payload_from_row( + { + "app_name": app_name, + "package_name": package_name, + "country_code": country_code, + "device_type": normalized_device_type, + } + ) + payload["task_key"] = str(payload.get("task_key") or _make_task_key(app_name, package_name)).strip() + payload["app_name"] = app_name + payload["package_name"] = package_name + payload["country_code"] = country_code + payload["country_codes"] = _normalize_country_codes(payload.get("country_codes") or country_code) + payload["device_type"] = normalized_device_type + payload["available_sources"] = payload.get("available_sources") or ["google_play", "local"] + original_row = payload.get("original_row") + normalized_original_row = dict(original_row) if isinstance(original_row, dict) else {} + normalized_original_row["app_name"] = app_name + normalized_original_row["package_name"] = package_name + normalized_original_row["country_code"] = country_code + normalized_original_row["device_type"] = normalized_device_type + payload["original_row"] = normalized_original_row + return payload + + +def build_task_payload_from_row(row: Dict[str, Any]) -> Dict[str, Any]: + app_name = str(row.get("app_name") or row.get("应用名称") or "").strip() + package_name = str(row.get("package_name") or row.get("包名") or "").strip() + country_code = str(row.get("country_code") or row.get("国家") or "").strip() + app_magic_label = _normalize_app_magic_label(row.get("app_magic_label") or row.get("应用魔法标签") or "") + last_updated = _normalize_last_updated(row.get("last_updated") or row.get("最后更新") or row.get("更新时间") or "") + device_type = _normalize_task_device_type(row.get("device_type") or row.get("设备类型") or "") + downloads = _extract_downloads_from_row(row) + available_sources = row.get("available_sources") + if isinstance(available_sources, str): + normalized_sources = [item.strip() for item in available_sources.split(",") if item.strip()] + elif isinstance(available_sources, (list, tuple)): + normalized_sources = [str(item).strip() for item in available_sources if str(item).strip()] + else: + normalized_sources = [] + if not normalized_sources: + normalized_sources = ["google_play", "local"] + normalized_row = {str(key or "").strip(): str(value or "").strip() for key, value in (row or {}).items()} + normalized_row["device_type"] = device_type + payload = { + "task_key": _make_task_key(app_name, package_name), + "app_name": app_name, + "package_name": package_name, + "country_code": country_code, + "country_codes": _normalize_country_codes(country_code), + "app_magic_label": app_magic_label, + "last_updated": last_updated, + "device_type": device_type, + "available_sources": normalized_sources, + "original_row": normalized_row, + } + if downloads is not None: + payload["downloads"] = downloads + return payload + + +def _derive_collection_status( + restriction_status: str, + retryability: str, + latest_failure_type: str, + *, + allow_legacy_download_error_1_pending: bool = False, +) -> Tuple[str, str]: + normalized_restriction = str(restriction_status or "").strip() + normalized_retryability = str(retryability or "").strip() + normalized_failure_type = str(latest_failure_type or "").strip().upper() + if allow_legacy_download_error_1_pending and normalized_failure_type == "DOWNLOAD_ERROR/1": + return "pending", "legacy_download_error_1" + if normalized_restriction in {"success", "light_restricted"}: + return "qualified", normalized_restriction or "qualified" + if normalized_restriction == "severe_restricted" and normalized_retryability == "non_retryable": + return "failed_terminal", normalized_failure_type or "severe_non_retryable" + return "pending", normalized_restriction or "pending" + + +def _should_preserve_non_retryable_reason( + existing_row: Dict[str, Any], + new_restriction_status: str, + new_retryability: str, +) -> bool: + previous_restriction_status = str(existing_row.get("restriction_status") or "").strip() + previous_retryability = str(existing_row.get("retryability") or "").strip() + previous_failure_type = str(existing_row.get("latest_failure_type") or "").strip() + normalized_new_restriction = str(new_restriction_status or "").strip() + normalized_new_retryability = str(new_retryability or "").strip() + return ( + previous_restriction_status == "severe_restricted" + and previous_retryability == "non_retryable" + and normalized_new_restriction == "severe_restricted" + and normalized_new_retryability == "retryable" + and bool(previous_failure_type) + ) + + +def _should_auto_reroute_to_physical( + current_device_type: str, + latest_failure_type: str, + restriction_status: str, + retryability: str, +) -> bool: + normalized_device_type = _normalize_task_device_type(current_device_type) + normalized_failure_type = str(latest_failure_type or "").strip().upper() + normalized_restriction = str(restriction_status or "").strip() + normalized_retryability = str(retryability or "").strip() + return ( + normalized_device_type != "physical" + and normalized_failure_type in AUTO_PENDING_PHYSICAL_FAILURE_TYPES + and normalized_restriction == "severe_restricted" + and normalized_retryability == "non_retryable" + ) + + +def _humanize_error_type(error_type: str) -> str: + normalized = str(error_type or "").strip() + if not normalized: + return "未标记错误" + category_name, _, code_text = normalized.partition("/") + try: + code = int(code_text) + except (TypeError, ValueError): + return normalized + try: + if category_name == "INFRA_ERROR": + return INFRA_ERROR_DESC.get(InfraError(code), normalized) + if category_name == "APP_ERROR": + return APP_ERROR_DESC.get(AppError(code), normalized) + if category_name == "BUSINESS_ERROR": + return BUSINESS_ERROR_DESC.get(BusinessError(code), normalized) + if category_name == "DOWNLOAD_ERROR": + return DOWNLOAD_ERROR_DESC.get(DownloadError(code), normalized) + except ValueError: + return normalized + return normalized + + +class DpiRuleMatcher: + def __init__(self, app_list_path: str, url_lib_path: str): + self.app_list_path = app_list_path + self.url_lib_path = url_lib_path + self._loaded = False + self._lock = threading.Lock() + self.tp_packages: Dict[str, set] = {} + self.tp_pkg_string: Dict[str, str] = {} + self.pkg_to_tp_mark: Dict[str, str] = {} + self._trie: Dict[str, Any] = {} + + def ensure_loaded(self): + if self._loaded: + return + with self._lock: + if self._loaded: + return + self._load_app_list() + self._load_url_rules() + self._loaded = True + + def _load_app_list(self): + if not os.path.exists(self.app_list_path): + logger.warning("Analytics app list is missing: %s", self.app_list_path) + return + for row in _read_dict_rows(self.app_list_path): + tp_mark = row.get("app_name", "").strip() + package_string = row.get("package_name", "").strip() + if not tp_mark: + continue + self.tp_pkg_string[tp_mark] = package_string + packages = {item.strip() for item in package_string.split(",") if item.strip()} + self.tp_packages.setdefault(tp_mark, set()).update(packages) + for package_name in packages: + self.pkg_to_tp_mark[package_name] = tp_mark + + def _load_url_rules(self): + if not os.path.exists(self.url_lib_path): + logger.warning("Analytics url lib is missing: %s", self.url_lib_path) + return + rule_count = 0 + for row in _read_dict_rows(self.url_lib_path): + suffix = row.get("url", "").strip().lower() + tp_mark = row.get("app", "").strip() + if not suffix or not tp_mark: + continue + self._insert_suffix(suffix, tp_mark) + rule_count += 1 + logger.info("[Analytics] loaded %s dpi url rules", rule_count) + + def _insert_suffix(self, suffix: str, tp_mark: str): + node = self._trie + for part in reversed([item for item in suffix.split(".") if item]): + node = node.setdefault(part, {}) + node["$"] = (suffix, tp_mark) + + def match_domain(self, domain: str) -> Tuple[Optional[str], Optional[str]]: + self.ensure_loaded() + cleaned = str(domain or "").strip().lower() + if not cleaned or cleaned.startswith("model_data:"): + return None, None + cleaned = cleaned.replace("http://", "").replace("https://", "").strip().strip("\"'") + node = self._trie + best = None + for part in reversed([item for item in cleaned.split(".") if item]): + node = node.get(part) + if node is None: + break + if "$" in node: + best = node["$"] + if not best: + return None, None + return best[1], best[0] + + +class AnalyticsRepository: + def __init__(self, db_path: str = MONITORING_DB_PATH): + self.db_path = db_path + self._write_lock = threading.Lock() + self._initialize_schema() + + @contextmanager + def _connect(self): + connection = sqlite3.connect(self.db_path, check_same_thread=False) + connection.row_factory = sqlite3.Row + try: + yield connection + connection.commit() + except Exception: + connection.rollback() + raise + finally: + connection.close() + + def _table_exists(self, connection: sqlite3.Connection, table_name: str) -> bool: + row = connection.execute( + """ + SELECT 1 + FROM sqlite_master + WHERE type = 'table' AND name = ? + """, + (table_name,), + ).fetchone() + return bool(row) + + def _table_columns(self, connection: sqlite3.Connection, table_name: str) -> set: + if not self._table_exists(connection, table_name): + return set() + return {row["name"] for row in connection.execute(f"PRAGMA table_info({table_name})").fetchall()} + + def _initialize_schema(self): + """ + 初始化数据库 schema。 + + analytics 直接使用新表 app_catalog / collection_task，不创建兼容视图。 + """ + os.makedirs(os.path.dirname(self.db_path) or ".", exist_ok=True) + with self._connect() as connection: + connection.execute("PRAGMA journal_mode=WAL") + connection.executescript( + """ + CREATE TABLE IF NOT EXISTS collection_task ( + package_name TEXT NOT NULL, + batch_tag TEXT NOT NULL, + run_kind TEXT NOT NULL DEFAULT 'ranking' + CHECK(run_kind IN ('ranking', 'block', 'model', 'manual')), + attempt INTEGER NOT NULL DEFAULT 1, + task_key TEXT, + app_name TEXT, + app_magic_label TEXT, + is_new_app INTEGER DEFAULT 0, + task_status TEXT DEFAULT 'pending', + execution_status TEXT, + worker_id TEXT, + created_at TEXT DEFAULT (datetime('now', 'localtime')), + started_at TEXT, + completed_at TEXT, + duration_seconds REAL, + download_duration_seconds REAL, + execution_duration_seconds REAL, + analysis_duration_seconds REAL, + error_category TEXT, + error_code INTEGER, + error_reason TEXT, + error_details TEXT, + crashed_source TEXT, + droidbot_steps INTEGER DEFAULT 0, + gui_agent_steps INTEGER DEFAULT 0, + total_steps INTEGER DEFAULT 0, + num_nodes INTEGER DEFAULT 0, + num_reached_activities INTEGER DEFAULT 0, + app_num_total_activities INTEGER DEFAULT 0, + total_traffic_bytes INTEGER DEFAULT 0, + self_traffic_bytes INTEGER DEFAULT 0, + server_traffic_bytes INTEGER DEFAULT 0, + unrecognized_traffic_bytes INTEGER DEFAULT 0, + model_flow_count INTEGER DEFAULT 0, + model_traffic_bytes INTEGER DEFAULT 0, + login_count INTEGER DEFAULT 0, + register_count INTEGER DEFAULT 0, + stuck_reason_code INTEGER, + guiagent_message TEXT, + scenario_triggered INTEGER DEFAULT 0, + download_source TEXT, + is_retry INTEGER DEFAULT 0, + exit_code INTEGER, + trace_json TEXT, + traffic_file_paths TEXT, + PRIMARY KEY (package_name, batch_tag, run_kind, attempt) + ); + CREATE INDEX IF NOT EXISTS idx_collection_task_status + ON collection_task(task_status, run_kind, batch_tag); + CREATE INDEX IF NOT EXISTS idx_collection_task_execution + ON collection_task(execution_status, run_kind); + CREATE INDEX IF NOT EXISTS idx_collection_task_pkg + ON collection_task(package_name, run_kind); + CREATE INDEX IF NOT EXISTS idx_collection_task_time + ON collection_task(created_at); + CREATE INDEX IF NOT EXISTS idx_collection_task_worker + ON collection_task(worker_id, completed_at); + CREATE INDEX IF NOT EXISTS idx_collection_task_magic + ON collection_task(app_magic_label); + CREATE INDEX IF NOT EXISTS idx_collection_task_batch + ON collection_task(batch_tag, run_kind); + + CREATE TABLE IF NOT EXISTS app_catalog ( + package_name TEXT PRIMARY KEY, + app_name TEXT, + batch_tags TEXT, + app_magic_label TEXT, + last_updated TEXT, + country_code TEXT, + device_type TEXT, + task_payload_json TEXT, + last_update_interval_days INTEGER DEFAULT 0, + task_queue TEXT DEFAULT 'default', + task_priority INTEGER DEFAULT 50, + source_order INTEGER, + is_active INTEGER DEFAULT 1, + is_blocked INTEGER DEFAULT 0, + category TEXT, + downloads INTEGER, + created_at TEXT DEFAULT (datetime('now', 'localtime')), + updated_at TEXT DEFAULT (datetime('now', 'localtime')) + ); + CREATE INDEX IF NOT EXISTS idx_app_catalog_priority + ON app_catalog(task_priority DESC); + CREATE INDEX IF NOT EXISTS idx_app_catalog_magic + ON app_catalog(app_magic_label); + CREATE INDEX IF NOT EXISTS idx_app_catalog_active + ON app_catalog(is_active, source_order); + + CREATE TABLE IF NOT EXISTS apk_registry ( + package_name TEXT PRIMARY KEY, + download_date TEXT NOT NULL, + download_time REAL NOT NULL, + local_dir TEXT, + smb_dir TEXT, + version_name TEXT, + source TEXT, + apk_files_json TEXT, + updated_at REAL NOT NULL + ); + + CREATE TABLE IF NOT EXISTS worker_activity_summary ( + worker_id TEXT NOT NULL, + stat_date TEXT NOT NULL, + stat_hour INTEGER NOT NULL, + idle_duration_seconds INTEGER DEFAULT 0, + busy_duration_seconds INTEGER DEFAULT 0, + offline_duration_seconds INTEGER DEFAULT 0, + task_count INTEGER DEFAULT 0, + success_count INTEGER DEFAULT 0, + failed_count INTEGER DEFAULT 0, + PRIMARY KEY (worker_id, stat_date, stat_hour) + ); + CREATE INDEX IF NOT EXISTS idx_worker_activity_date + ON worker_activity_summary(stat_date DESC); + + CREATE TABLE IF NOT EXISTS analytics_job ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + job_type TEXT NOT NULL, + status TEXT NOT NULL, + package_name TEXT, + task_key TEXT, + worker_id TEXT, + trigger_source TEXT, + report_time REAL, + attempt_count INTEGER DEFAULT 0, + error_message TEXT, + payload_json TEXT, + scheduled_after REAL DEFAULT 0, + created_at REAL NOT NULL, + updated_at REAL NOT NULL, + finished_at REAL + ); + CREATE INDEX IF NOT EXISTS idx_analytics_job_status_created_at + ON analytics_job(status, created_at); + CREATE INDEX IF NOT EXISTS idx_analytics_job_package_name + ON analytics_job(package_name, created_at); + + CREATE TABLE IF NOT EXISTS analytics_source_file ( + file_path TEXT PRIMARY KEY, + file_type TEXT NOT NULL, + package_name TEXT, + worker_id TEXT, + size_bytes INTEGER, + mtime REAL, + checksum TEXT, + last_seen_at REAL, + last_processed_at REAL + ); + CREATE INDEX IF NOT EXISTS idx_analytics_source_file_package_name + ON analytics_source_file(package_name, file_type); + + CREATE TABLE IF NOT EXISTS app_domain_traffic ( + package_name TEXT NOT NULL, + domain TEXT NOT NULL, + domain_type TEXT, + traffic_bytes INTEGER DEFAULT 0, + flow_count INTEGER DEFAULT 0, + traffic_ratio REAL DEFAULT 0, + domain_traffic_ratio REAL DEFAULT 0, + organization TEXT, + matched_tp_mark TEXT, + matched_pattern TEXT, + match_state TEXT, + updated_at REAL NOT NULL, + PRIMARY KEY(package_name, domain) + ); + CREATE INDEX IF NOT EXISTS idx_app_domain_traffic_package_name + ON app_domain_traffic(package_name, traffic_bytes DESC); + + CREATE TABLE IF NOT EXISTS app_traffic_component ( + package_name TEXT NOT NULL, + component_name TEXT NOT NULL, + component_package_names TEXT, + is_self INTEGER DEFAULT 0, + traffic_bytes INTEGER DEFAULT 0, + share_percent REAL DEFAULT 0, + updated_at REAL NOT NULL, + PRIMARY KEY(package_name, component_name) + ); + CREATE INDEX IF NOT EXISTS idx_app_traffic_component_package_name + ON app_traffic_component(package_name, traffic_bytes DESC); + """ + ) + self._ensure_columns( + connection, + "app_catalog", + { + "last_updated": "TEXT", + "country_code": "TEXT", + "device_type": "TEXT", + "task_payload_json": "TEXT", + "last_update_interval_days": "INTEGER DEFAULT 0", + }, + ) + + def _ensure_columns(self, connection: sqlite3.Connection, table_name: str, columns: Dict[str, str]) -> None: + existing_columns = self._table_columns(connection, table_name) + for column_name, column_type in columns.items(): + if column_name in existing_columns: + continue + connection.execute(f"ALTER TABLE {table_name} ADD COLUMN {column_name} {column_type}") + def create_job( + self, + *, + job_type: str, + status: str, + package_name: Optional[str] = None, + task_key: Optional[str] = None, + worker_id: Optional[str] = None, + trigger_source: Optional[str] = None, + report_time: Optional[float] = None, + payload: Optional[Dict[str, Any]] = None, + scheduled_after: Optional[float] = None, + ) -> Dict[str, Any]: + now = time.time() + payload_json = _safe_json_dumps(payload or {}) + with self._write_lock, self._connect() as connection: + cursor = connection.execute( + """ + INSERT INTO analytics_job ( + job_type, status, package_name, task_key, worker_id, trigger_source, + report_time, attempt_count, error_message, payload_json, scheduled_after, + created_at, updated_at, finished_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, 0, '', ?, ?, ?, ?, NULL) + """, + ( + job_type, + status, + package_name, + task_key, + worker_id, + trigger_source or "", + report_time, + payload_json, + float(scheduled_after or 0), + now, + now, + ), + ) + job_id = int(cursor.lastrowid) + return self.get_job(job_id) + + def get_job(self, job_id: int) -> Dict[str, Any]: + with self._connect() as connection: + row = connection.execute( + "SELECT * FROM analytics_job WHERE id = ?", + (job_id,), + ).fetchone() + return self._serialize_job(row) if row else {} + + def _serialize_job(self, row: sqlite3.Row) -> Dict[str, Any]: + return { + "id": int(row["id"]), + "job_type": row["job_type"], + "status": row["status"], + "package_name": row["package_name"] or "", + "task_key": row["task_key"] or "", + "worker_id": row["worker_id"] or "", + "trigger_source": row["trigger_source"] or "", + "report_time": _format_ts(row["report_time"]), + "attempt_count": int(row["attempt_count"] or 0), + "error_message": row["error_message"] or "", + "payload": _safe_json_loads(row["payload_json"], {}) or {}, + "created_at": _format_ts(row["created_at"]), + "updated_at": _format_ts(row["updated_at"]), + "finished_at": _format_ts(row["finished_at"]), + } + + def list_jobs( + self, + *, + job_type: Optional[str] = None, + status: Optional[str] = None, + package_name: Optional[str] = None, + limit: int = 50, + ) -> List[Dict[str, Any]]: + clauses = [] + params: List[Any] = [] + if job_type: + clauses.append("job_type = ?") + params.append(job_type) + if status: + clauses.append("status = ?") + params.append(status) + if package_name: + clauses.append("package_name = ?") + params.append(package_name) + where_clause = f"WHERE {' AND '.join(clauses)}" if clauses else "" + with self._connect() as connection: + rows = connection.execute( + f""" + SELECT * + FROM analytics_job + {where_clause} + ORDER BY created_at DESC, id DESC + LIMIT ? + """, + [*params, max(1, min(limit, 500))], + ).fetchall() + return [self._serialize_job(row) for row in rows] + + def claim_next_job(self, artifact_wait_seconds: int) -> Optional[Dict[str, Any]]: + now = time.time() + with self._write_lock, self._connect() as connection: + row = connection.execute( + """ + SELECT * + FROM analytics_job + WHERE (status = 'queued' + AND COALESCE(scheduled_after, 0) <= ?) + OR ( + status = 'waiting_artifacts' + AND COALESCE(report_time, created_at) <= ? + ) + ORDER BY CASE status WHEN 'queued' THEN 0 ELSE 1 END, created_at ASC, id ASC + LIMIT 1 + """, + (now, now - max(0, artifact_wait_seconds)), + ).fetchone() + if not row: + return None + connection.execute( + """ + UPDATE analytics_job + SET status = 'running', + attempt_count = COALESCE(attempt_count, 0) + 1, + updated_at = ? + WHERE id = ? + """, + (now, row["id"]), + ) + return self.get_job(int(row["id"])) + + def mark_job_waiting(self, job_id: int, error_message: str = ""): + now = time.time() + with self._write_lock, self._connect() as connection: + connection.execute( + """ + UPDATE analytics_job + SET status = 'waiting_artifacts', + error_message = ?, + updated_at = ?, + finished_at = NULL + WHERE id = ? + """, + (error_message, now, job_id), + ) + + def mark_job_ready(self, *, task_key: Optional[str] = None, package_name: Optional[str] = None, worker_id: Optional[str] = None) -> int: + clauses = ["status = 'waiting_artifacts'"] + params: List[Any] = [] + if task_key: + clauses.append("task_key = ?") + params.append(task_key) + if package_name: + clauses.append("package_name = ?") + params.append(package_name) + if worker_id: + clauses.append("worker_id = ?") + params.append(worker_id) + if len(clauses) == 1: + return 0 + now = time.time() + with self._write_lock, self._connect() as connection: + cursor = connection.execute( + f""" + UPDATE analytics_job + SET status = 'queued', + scheduled_after = ?, + updated_at = ?, + error_message = '' + WHERE {' AND '.join(clauses)} + """, + [now + ANALYTICS_ARTIFACT_FILE_WAIT_SECONDS, now, *params], + ) + return int(cursor.rowcount or 0) + + def finish_job(self, job_id: int, *, status: str, error_message: str = ""): + now = time.time() + with self._write_lock, self._connect() as connection: + connection.execute( + """ + UPDATE analytics_job + SET status = ?, + error_message = ?, + updated_at = ?, + finished_at = ? + WHERE id = ? + """, + (status, error_message, now, now, job_id), + ) + + def get_latest_task_execution(self, package_name: str) -> Optional[Dict[str, Any]]: + with self._connect() as connection: + if not self._table_exists(connection, "task_execution"): + return None + row = connection.execute( + """ + SELECT + execution_id, + task_key, + app_name, + package_name, + last_updated, + collection_task_type, + worker_id, + status, + error_type, + error_message, + result_detail, + download_duration_seconds, + collect_duration_seconds, + total_duration_seconds, + failed_total_duration_seconds, + droidbot_steps, + guiagent_steps, + total_steps, + num_nodes, + num_reached_activities, + app_num_total_activities, + trace_json, + COALESCE(task_ended_at, task_started_at, last_updated_at) AS latest_time + FROM task_execution + WHERE package_name = ? + ORDER BY COALESCE(task_ended_at, task_started_at, last_updated_at) DESC, execution_id DESC + LIMIT 1 + """, + (package_name,), + ).fetchone() + if not row: + return None + trace = _safe_json_loads(row["trace_json"], {}) + download_errors = trace.get("download_errors") if isinstance(trace, dict) else {} + return { + "execution_id": row["execution_id"], + "task_key": row["task_key"] or "", + "app_name": row["app_name"] or "", + "package_name": row["package_name"] or package_name, + "last_updated": row["last_updated"] or "", + "collection_task_type": row["collection_task_type"] or "", + "worker_id": row["worker_id"] or "", + "status": row["status"] or "", + "error_type": row["error_type"] or "", + "error_message": row["error_message"] or "", + "result_detail": row["result_detail"] or row["error_message"] or "", + "download_duration_seconds": float(row["download_duration_seconds"] or 0.0), + "collect_duration_seconds": float(row["collect_duration_seconds"] or 0.0), + "total_duration_seconds": float(row["total_duration_seconds"] or row["failed_total_duration_seconds"] or 0.0), + "droidbot_steps": int(row["droidbot_steps"] or 0), + "guiagent_steps": int(row["guiagent_steps"] or 0), + "total_steps": int(row["total_steps"] or 0), + "num_nodes": int(row["num_nodes"] or 0), + "num_reached_activities": int(row["num_reached_activities"] or 0), + "app_num_total_activities": int(row["app_num_total_activities"] or 0), + "latest_time": float(row["latest_time"] or 0.0), + "download_errors": download_errors if isinstance(download_errors, dict) else {}, + "trace": trace if isinstance(trace, dict) else {}, + } + + def _load_latest_download_errors_by_packages( + self, + connection: sqlite3.Connection, + package_names: List[str], + ) -> Dict[str, Dict[str, Any]]: + normalized = [str(item or "").strip() for item in package_names if str(item or "").strip()] + if not normalized or not self._table_exists(connection, "task_execution"): + return {} + placeholders = ",".join("?" for _ in normalized) + rows = connection.execute( + f""" + SELECT package_name, trace_json + FROM ( + SELECT + package_name, + trace_json, + ROW_NUMBER() OVER ( + PARTITION BY package_name + ORDER BY COALESCE(task_ended_at, task_started_at, last_updated_at) DESC, execution_id DESC + ) AS row_num + FROM task_execution + WHERE package_name IN ({placeholders}) + ) + WHERE row_num = 1 + """, + normalized, + ).fetchall() + latest_download_errors: Dict[str, Dict[str, Any]] = {} + for row in rows: + package_name = str(row["package_name"] or "").strip() + if not package_name: + continue + trace = _safe_json_loads(row["trace_json"], {}) or {} + if not isinstance(trace, dict): + continue + latest_download_errors[package_name] = _normalize_download_errors(trace.get("download_errors")) + return latest_download_errors + + def list_all_task_packages(self) -> List[str]: + with self._connect() as connection: + if not self._table_exists(connection, "task_execution"): + return [] + rows = connection.execute( + """ + SELECT DISTINCT package_name + FROM task_execution + WHERE package_name IS NOT NULL AND package_name != '' + """ + ).fetchall() + return sorted({row["package_name"] for row in rows if row["package_name"]}) + + def _parse_db_datetime(self, value: Any) -> float: + text = str(value or "").strip() + if not text: + return 0.0 + for fmt in ("%Y-%m-%d %H:%M:%S", "%Y-%m-%dT%H:%M:%S"): + try: + return datetime.strptime(text[:19], fmt).replace(tzinfo=SHANGHAI_TZ).timestamp() + except ValueError: + continue + try: + return float(text) + except (TypeError, ValueError): + return 0.0 + + def _latest_task_row(self, connection: sqlite3.Connection, package_name: str) -> Optional[sqlite3.Row]: + return connection.execute( + """ + SELECT * + FROM collection_task + WHERE package_name = ? + ORDER BY + datetime(COALESCE(NULLIF(completed_at, ''), NULLIF(started_at, ''), created_at)) DESC, + attempt DESC, + batch_tag DESC, + run_kind DESC + LIMIT 1 + """, + (package_name,), + ).fetchone() + + def _failure_type_from_task(self, task_row: Optional[sqlite3.Row]) -> str: + if not task_row: + return "" + category = str(task_row["error_category"] or "").strip() + code = task_row["error_code"] + if not category: + return "" + if code is None or str(code).strip() == "": + return category + return f"{category}/{int(code)}" + + def _download_errors_from_task(self, task_row: Optional[sqlite3.Row]) -> Dict[str, Any]: + if not task_row: + return {} + trace = _safe_json_loads(task_row["trace_json"], {}) or {} + if not isinstance(trace, dict): + return {} + return _normalize_download_errors(trace.get("download_errors")) + + def _ratio_from_task(self, task_row: Optional[sqlite3.Row], numerator_key: str) -> float: + if not task_row: + return 0.0 + total = int(task_row["total_traffic_bytes"] or 0) + if total <= 0: + return 0.0 + numerator = int(task_row[numerator_key] or 0) + return round((numerator / total) * 100, 2) + + def _recognition_ratio_from_task(self, task_row: Optional[sqlite3.Row]) -> float: + if not task_row: + return 0.0 + total = int(task_row["total_traffic_bytes"] or 0) + if total <= 0: + return 0.0 + recognized = int(task_row["self_traffic_bytes"] or 0) + int(task_row["server_traffic_bytes"] or 0) + return round((recognized / total) * 100, 2) + + def _state_from_task(self, task_row: Optional[sqlite3.Row]) -> Dict[str, str]: + if not task_row: + return { + "collection_status": "pending", + "collection_status_reason": "no_collection_task", + "restriction_status": "", + "retryability": "not_applicable", + } + task_status = str(task_row["task_status"] or "").strip().lower() + execution_status = str(task_row["execution_status"] or "").strip().lower() + if task_status in {"pending", "running"} or not execution_status: + return { + "collection_status": "pending", + "collection_status_reason": str(task_row["error_reason"] or "pending"), + "restriction_status": "", + "retryability": "not_applicable", + } + + failure_type = self._failure_type_from_task(task_row) + restriction_status, retryability = _classify_restriction_status( + execution_status, + self._ratio_from_task(task_row, "self_traffic_bytes"), + failure_type, + num_nodes=int(task_row["num_nodes"] or 0), + download_errors=self._download_errors_from_task(task_row), + ) + if ( + execution_status == "success" + and self._ratio_from_task(task_row, "self_traffic_bytes") <= 0 + and int(task_row["num_nodes"] or 0) < LIGHT_RESTRICTED_MIN_NUM_NODES + ): + restriction_status = "light_restricted" + retryability = "not_applicable" + collection_status, derived_reason = _derive_collection_status( + restriction_status, + retryability, + failure_type, + ) + return { + "collection_status": collection_status, + "collection_status_reason": str(task_row["error_reason"] or derived_reason or ""), + "restriction_status": restriction_status, + "retryability": retryability, + } + + def _catalog_payload(self, catalog_row: sqlite3.Row, collection_task_type: str = "new_app") -> Dict[str, Any]: + payload = _safe_json_loads(catalog_row["task_payload_json"], {}) or {} + if not isinstance(payload, dict) or not payload: + payload = build_task_payload_from_row( + { + "app_name": catalog_row["app_name"] or catalog_row["package_name"], + "package_name": catalog_row["package_name"], + "app_magic_label": catalog_row["app_magic_label"] or "", + "last_updated": catalog_row["last_updated"] or "", + "country_code": catalog_row["country_code"] or "", + "device_type": catalog_row["device_type"] or "", + "downloads": catalog_row["downloads"], + } + ) + else: + payload = dict(payload) + payload["app_name"] = payload.get("app_name") or catalog_row["app_name"] or catalog_row["package_name"] + payload["package_name"] = payload.get("package_name") or catalog_row["package_name"] + payload["app_magic_label"] = payload.get("app_magic_label") or catalog_row["app_magic_label"] or "" + payload["last_updated"] = payload.get("last_updated") or catalog_row["last_updated"] or "" + payload["country_code"] = payload.get("country_code") or catalog_row["country_code"] or "" + payload["device_type"] = _normalize_task_device_type(payload.get("device_type") or catalog_row["device_type"] or "") + original_row = payload.get("original_row") + if isinstance(original_row, dict): + original_row_copy = dict(original_row) + original_row_copy["app_magic_label"] = original_row_copy.get("app_magic_label") or payload["app_magic_label"] + original_row_copy["last_updated"] = original_row_copy.get("last_updated") or payload["last_updated"] + original_row_copy["country_code"] = original_row_copy.get("country_code") or payload["country_code"] + original_row_copy["device_type"] = original_row_copy.get("device_type") or payload["device_type"] + payload["original_row"] = original_row_copy + payload["collection_task_type"] = collection_task_type or "new_app" + payload["last_update_interval_days"] = int(catalog_row["last_update_interval_days"] or 0) + return payload + + def _task_recency_epoch(self, task_row: Optional[sqlite3.Row]) -> float: + if not task_row: + return 0.0 + return max( + self._parse_db_datetime(task_row["completed_at"]), + self._parse_db_datetime(task_row["started_at"]), + self._parse_db_datetime(task_row["created_at"]), + ) + + def _source_file_count(self, connection: sqlite3.Connection, package_name: str) -> int: + row = connection.execute( + "SELECT COUNT(*) AS total FROM analytics_source_file WHERE package_name = ?", + (package_name,), + ).fetchone() + return int((row["total"] or 0) if row else 0) + + def _artifact_status(self, connection: sqlite3.Connection, package_name: str, task_row: Optional[sqlite3.Row]) -> str: + has_task = bool(task_row and str(task_row["task_status"] or "").strip()) + traffic_file_paths = _safe_json_loads(task_row["traffic_file_paths"], []) if task_row else [] + has_files = bool(traffic_file_paths) or self._source_file_count(connection, package_name) > 0 + if has_task and has_files: + return "complete" + if has_task or has_files: + return "partial" + return "missing" + + def _catalog_summary_from_rows( + self, + connection: sqlite3.Connection, + catalog_row: sqlite3.Row, + task_row: Optional[sqlite3.Row], + ) -> Dict[str, Any]: + package_name = str(catalog_row["package_name"] or "").strip() + state = self._state_from_task(task_row) + collection_task_type = "new_app" + if task_row: + collection_task_type = "new_app" if int(task_row["is_new_app"] or 0) else "app_update" + payload = self._catalog_payload(catalog_row, collection_task_type) + latest_test_time = self._parse_db_datetime(task_row["completed_at"]) if task_row else 0.0 + latest_status = str(task_row["execution_status"] or "").strip() if task_row else "" + latest_failure_type = self._failure_type_from_task(task_row) + domain_rows = connection.execute( + """ + SELECT domain, domain_type + FROM app_domain_traffic + WHERE package_name = ? + """, + (package_name,), + ).fetchall() + unique_domains = sorted( + { + str(row["domain"] or "").strip() + for row in domain_rows + if str(row["domain"] or "").strip() + and str(row["domain_type"] or "") != "model_data" + and not str(row["domain"] or "").startswith("model_data:") + } + ) + unique_second_level_domains = sorted( + { + _extract_second_level_domain(domain) + for domain in unique_domains + if _extract_second_level_domain(domain) + } + ) + return { + "package_name": package_name, + "app_name": catalog_row["app_name"] or "", + "app_magic_label": _normalize_app_magic_label(catalog_row["app_magic_label"] or ""), + "last_updated": _normalize_last_updated(catalog_row["last_updated"] or ""), + "country_code": catalog_row["country_code"] or "", + "device_type": _normalize_task_device_type(catalog_row["device_type"] or ""), + "source_order": catalog_row["source_order"], + "task_payload": payload, + "task_payload_json": _safe_json_dumps(payload), + "downloads": _parse_downloads_value(catalog_row["downloads"]), + "collection_task_type": collection_task_type, + "last_update_interval_days": int(catalog_row["last_update_interval_days"] or 0), + "incremental_batch_tag": catalog_row["batch_tags"] or "[]", + "incremental_batch_tags": _normalize_incremental_batch_tags(catalog_row["batch_tags"]), + "incremental_batch_marked_at": 0.0, + "catalog_active": bool(catalog_row["is_active"]), + "is_active": int(catalog_row["is_active"] if catalog_row["is_active"] is not None else 1), + "task_queue": str(catalog_row["task_queue"] or "default").strip() or "default", + "collection_status": state["collection_status"], + "collection_status_reason": state["collection_status_reason"], + "restriction_status": state["restriction_status"], + "retryability": state["retryability"], + "latest_task_key": task_row["task_key"] or "" if task_row else "", + "latest_status": latest_status, + "latest_test_time": latest_test_time, + "latest_worker_id": task_row["worker_id"] or "" if task_row else "", + "latest_task_detail": task_row["error_details"] or "" if task_row else "", + "latest_failure_type": latest_failure_type, + "unique_domain_count": len(unique_domains), + "unique_domain_names": unique_domains, + "unique_second_level_domain_count": len(unique_second_level_domains), + "unique_second_level_domains": unique_second_level_domains, + "droidbot_steps": int(task_row["droidbot_steps"] or 0) if task_row else 0, + "gui_agent_steps": int(task_row["gui_agent_steps"] or 0) if task_row else 0, + "duration_seconds": float(task_row["duration_seconds"] or 0.0) if task_row else 0.0, + "num_nodes": int(task_row["num_nodes"] or 0) if task_row else 0, + "num_reached_activities": int(task_row["num_reached_activities"] or 0) if task_row else 0, + "app_num_total_activities": int(task_row["app_num_total_activities"] or 0) if task_row else 0, + "total_traffic_bytes": int(task_row["total_traffic_bytes"] or 0) if task_row else 0, + "self_traffic_bytes": int(task_row["self_traffic_bytes"] or 0) if task_row else 0, + "server_traffic_bytes": int(task_row["server_traffic_bytes"] or 0) if task_row else 0, + "unrecognized_traffic_bytes": int(task_row["unrecognized_traffic_bytes"] or 0) if task_row else 0, + "self_ratio": self._ratio_from_task(task_row, "self_traffic_bytes"), + "recognition_ratio": self._recognition_ratio_from_task(task_row), + "model_flow_count": int(task_row["model_flow_count"] or 0) if task_row else 0, + "model_traffic_bytes": int(task_row["model_traffic_bytes"] or 0) if task_row else 0, + "model_eligible": 1 + if task_row and 0 < int(task_row["model_flow_count"] or 0) <= MODEL_TRAFFIC_THRESHOLD + else 0, + "artifact_status": self._artifact_status(connection, package_name, task_row), + "updated_at": self._parse_db_datetime(catalog_row["updated_at"]), + "latest_task_row": task_row, + } + + def _list_catalog_summaries(self, connection: sqlite3.Connection) -> List[Dict[str, Any]]: + rows = connection.execute( + """ + SELECT * + FROM app_catalog + ORDER BY + CASE WHEN source_order IS NULL THEN 1 ELSE 0 END ASC, + source_order ASC, + package_name ASC + """ + ).fetchall() + return [ + self._catalog_summary_from_rows(connection, row, self._latest_task_row(connection, row["package_name"])) + for row in rows + ] + + def _latest_catalog_status(self, connection: sqlite3.Connection, package_name: str) -> Dict[str, Any]: + catalog_row = connection.execute( + "SELECT * FROM app_catalog WHERE package_name = ?", + (package_name,), + ).fetchone() + if not catalog_row: + return {} + return self._catalog_summary_from_rows(connection, catalog_row, self._latest_task_row(connection, package_name)) + + def _next_attempt(self, connection: sqlite3.Connection, package_name: str, batch_tag: str, run_kind: str) -> int: + row = connection.execute( + """ + SELECT COALESCE(MAX(attempt), 0) AS max_attempt + FROM collection_task + WHERE package_name = ? + AND batch_tag = ? + AND run_kind = ? + """, + (package_name, batch_tag, run_kind), + ).fetchone() + return int((row["max_attempt"] or 0) if row else 0) + 1 + + def _catalog_batch_tag(self, catalog_row: sqlite3.Row, fallback: str = "manual") -> str: + tags = _normalize_incremental_batch_tags(catalog_row["batch_tags"]) + if tags: + return tags[-1] + return fallback + + def _insert_pending_collection_task( + self, + connection: sqlite3.Connection, + catalog_row: sqlite3.Row, + *, + reason: str, + run_kind: str = "ranking", + batch_tag: str = "", + collection_task_type: str = "", + ) -> None: + package_name = str(catalog_row["package_name"] or "").strip() + if not package_name: + return + batch_tag = batch_tag or self._catalog_batch_tag(catalog_row, "manual") + existing = self._latest_catalog_status(connection, package_name) + existing_task_type = collection_task_type or str(existing.get("collection_task_type") or "new_app") + is_new_app = 1 if existing_task_type == "new_app" else 0 + attempt = self._next_attempt(connection, package_name, batch_tag, run_kind) + app_name = str(catalog_row["app_name"] or package_name).strip() or package_name + task_key = _make_task_key(app_name, package_name) + now_iso = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + connection.execute( + """ + INSERT INTO collection_task ( + package_name, batch_tag, run_kind, attempt, + task_key, app_name, app_magic_label, is_new_app, + task_status, execution_status, error_reason, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'pending', NULL, ?, ?) + """, + ( + package_name, + batch_tag, + run_kind, + attempt, + task_key, + app_name, + _normalize_app_magic_label(catalog_row["app_magic_label"] or ""), + is_new_app, + reason, + now_iso, + ), + ) + + def get_collection_row(self, package_name: str) -> Optional[Dict[str, Any]]: + """查询应用完整信息（新架构：直接读取 app_catalog + collection_task）""" + with self._connect() as connection: + summary = self._latest_catalog_status(connection, package_name) + if not summary: + return None + summary.pop("latest_task_row", None) + return summary + + def upsert_catalog_entries(self, entries: List[Dict[str, Any]]) -> int: + if not entries: + return 0 + now_iso = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + count = 0 + with self._write_lock, self._connect() as connection: + for entry in entries: + package_name = str(entry.get("package_name") or "").strip() + if not package_name: + continue + payload = entry.get("task_payload") or {} + downloads = entry.get("downloads") + if downloads is None: + downloads = _extract_downloads_from_task_payload(payload) + last_updated = _normalize_last_updated(entry.get("last_updated") or payload.get("last_updated")) + country_code = str(entry.get("country_code") or payload.get("country_code") or "").strip() + device_type = _normalize_task_device_type(entry.get("device_type") or payload.get("device_type") or "") + last_update_interval_days = int(entry.get("last_update_interval_days") or payload.get("last_update_interval_days") or 0) + connection.execute( + """ + INSERT INTO app_catalog ( + package_name, app_name, app_magic_label, batch_tags, + last_updated, country_code, device_type, task_payload_json, + last_update_interval_days, task_queue, task_priority, source_order, is_active, + downloads, created_at, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(package_name) DO UPDATE SET + app_name = COALESCE(NULLIF(excluded.app_name, ''), app_catalog.app_name), + app_magic_label = COALESCE(NULLIF(excluded.app_magic_label, ''), app_catalog.app_magic_label), + batch_tags = COALESCE(NULLIF(excluded.batch_tags, '[]'), app_catalog.batch_tags), + last_updated = COALESCE(NULLIF(excluded.last_updated, ''), app_catalog.last_updated), + country_code = excluded.country_code, + device_type = excluded.device_type, + task_payload_json = excluded.task_payload_json, + last_update_interval_days = COALESCE(excluded.last_update_interval_days, app_catalog.last_update_interval_days), + task_queue = COALESCE(NULLIF(excluded.task_queue, ''), app_catalog.task_queue), + task_priority = COALESCE(excluded.task_priority, app_catalog.task_priority), + source_order = excluded.source_order, + is_active = COALESCE(excluded.is_active, app_catalog.is_active), + downloads = COALESCE(excluded.downloads, app_catalog.downloads), + updated_at = excluded.updated_at + """, + ( + package_name, + entry.get("app_name", "") or package_name, + _normalize_app_magic_label(entry.get("app_magic_label", "")), + _incremental_batch_tags_to_json(entry.get("batch_tags") or entry.get("incremental_batch_tags") or []), + last_updated, + country_code, + device_type, + _safe_json_dumps(payload), + last_update_interval_days, + str(entry.get("task_queue") or "default").strip() or "default", + int(entry.get("task_priority") or 50), + entry.get("source_order"), + entry.get("catalog_active", entry.get("is_active", 1)), + downloads, + now_iso, + now_iso, + ), + ) + catalog_row = connection.execute( + "SELECT * FROM app_catalog WHERE package_name = ?", + (package_name,), + ).fetchone() + if catalog_row and str(entry.get("collection_status") or "pending").strip() == "pending": + self._insert_pending_collection_task( + connection, + catalog_row, + reason=str(entry.get("collection_status_reason") or "catalog_sync"), + collection_task_type=str(entry.get("collection_task_type") or "new_app"), + ) + count += 1 + return count + + def replace_catalog_entries(self, entries: List[Dict[str, Any]]) -> Dict[str, int]: + normalized_entries: List[Dict[str, Any]] = [] + seen_packages = set() + for entry in entries: + package_name = str(entry.get("package_name") or "").strip() + if not package_name or package_name in seen_packages: + continue + seen_packages.add(package_name) + normalized_entries.append(dict(entry)) + + if not normalized_entries: + return {"listed": 0, "upserted": 0, "deactivated": 0} + + now_iso = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + upserted = 0 + listed_packages = [str(entry["package_name"]).strip() for entry in normalized_entries] + placeholders = ",".join("?" for _ in listed_packages) + with self._write_lock, self._connect() as connection: + cursor = connection.execute( + f""" + UPDATE app_catalog + SET is_active = 0, + updated_at = ? + WHERE {ACTIVE_CATALOG_WHERE} + AND package_name NOT IN ({placeholders}) + """, + [now_iso, *listed_packages], + ) + deactivated = int(cursor.rowcount or 0) + + for entry in normalized_entries: + package_name = str(entry.get("package_name") or "").strip() + payload = entry.get("task_payload") or {} + downloads = entry.get("downloads") + if downloads is None: + downloads = _extract_downloads_from_task_payload(payload) + existing_summary = self._latest_catalog_status(connection, package_name) + existing_status = str(existing_summary.get("collection_status") or "").strip() + collection_task_type = str(entry.get("collection_task_type") or existing_summary.get("collection_task_type") or "new_app") + should_mark_pending = not existing_summary or existing_status in {"pending", "failed_terminal"} + pending_reason = str(entry.get("collection_status_reason") or "catalog_sync") + incoming_last_updated = _normalize_last_updated(entry.get("last_updated") or (payload or {}).get("last_updated")) + previous_last_updated = _normalize_last_updated(existing_summary.get("last_updated", "")) if existing_summary else "" + effective_last_updated = incoming_last_updated or previous_last_updated + last_update_interval_days = int(entry.get("last_update_interval_days") or 0) + if incoming_last_updated and previous_last_updated: + incoming_ts = _last_updated_ts(incoming_last_updated) + previous_ts = _last_updated_ts(previous_last_updated) + if incoming_ts > previous_ts: + last_update_interval_days = _last_updated_interval_days(incoming_last_updated, previous_last_updated) + previous_success = ( + existing_summary.get("collection_status") == "qualified" + or existing_summary.get("latest_status") == "success" + or existing_summary.get("restriction_status") in {"success", "light_restricted"} + ) + collection_task_type = "app_update" if previous_success else "new_app" + should_mark_pending = True + pending_reason = f"catalog_version_update:{previous_last_updated}->{incoming_last_updated}" + elif incoming_last_updated and not previous_last_updated and existing_summary: + should_mark_pending = True + pending_reason = f"catalog_last_updated_fill:{incoming_last_updated}" + if isinstance(payload, dict): + payload = dict(payload) + payload["last_updated"] = effective_last_updated + payload["collection_task_type"] = collection_task_type + payload["last_update_interval_days"] = last_update_interval_days + original_row = payload.get("original_row") + if isinstance(original_row, dict): + original_row_copy = dict(original_row) + original_row_copy["last_updated"] = effective_last_updated + payload["original_row"] = original_row_copy + country_code = str(entry.get("country_code") or (payload or {}).get("country_code") or "").strip() + device_type = _normalize_task_device_type(entry.get("device_type") or (payload or {}).get("device_type") or "") + connection.execute( + """ + INSERT INTO app_catalog ( + package_name, app_name, app_magic_label, batch_tags, + last_updated, country_code, device_type, task_payload_json, + last_update_interval_days, task_queue, task_priority, source_order, is_active, + downloads, created_at, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 1, ?, ?, ?) + ON CONFLICT(package_name) DO UPDATE SET + app_name = COALESCE(NULLIF(excluded.app_name, ''), app_catalog.app_name), + app_magic_label = COALESCE(NULLIF(excluded.app_magic_label, ''), app_catalog.app_magic_label), + batch_tags = COALESCE(NULLIF(excluded.batch_tags, '[]'), app_catalog.batch_tags), + last_updated = COALESCE(NULLIF(excluded.last_updated, ''), app_catalog.last_updated), + country_code = excluded.country_code, + device_type = excluded.device_type, + task_payload_json = excluded.task_payload_json, + last_update_interval_days = COALESCE(excluded.last_update_interval_days, app_catalog.last_update_interval_days), + task_queue = COALESCE(NULLIF(excluded.task_queue, ''), app_catalog.task_queue), + task_priority = COALESCE(excluded.task_priority, app_catalog.task_priority), + source_order = excluded.source_order, + is_active = 1, + downloads = COALESCE(excluded.downloads, app_catalog.downloads), + updated_at = excluded.updated_at + """, + ( + package_name, + entry.get("app_name", "") or package_name, + _normalize_app_magic_label(entry.get("app_magic_label", "")), + _incremental_batch_tags_to_json(entry.get("batch_tags") or entry.get("incremental_batch_tags") or []), + effective_last_updated, + country_code, + device_type, + _safe_json_dumps(payload), + last_update_interval_days, + str(entry.get("task_queue") or "default").strip() or "default", + int(entry.get("task_priority") or 50), + entry.get("source_order"), + downloads, + now_iso, + now_iso, + ), + ) + catalog_row = connection.execute( + "SELECT * FROM app_catalog WHERE package_name = ?", + (package_name,), + ).fetchone() + if catalog_row and should_mark_pending: + self._insert_pending_collection_task( + connection, + catalog_row, + reason=pending_reason, + collection_task_type=collection_task_type, + ) + upserted += 1 + + return { + "listed": len(normalized_entries), + "upserted": upserted, + "deactivated": deactivated, + } + + def list_active_catalog_packages(self) -> List[str]: + with self._connect() as connection: + rows = connection.execute( + """ + SELECT package_name + FROM app_catalog + WHERE COALESCE(is_active, 1) = 1 + ORDER BY + CASE WHEN source_order IS NULL THEN 1 ELSE 0 END ASC, + source_order ASC, + package_name ASC + """ + ).fetchall() + return [str(row["package_name"]).strip() for row in rows if row["package_name"]] + + def clear_incremental_batch_tags(self, batch_tag: str = "") -> int: + normalized_tag = _normalize_incremental_batch_tag(batch_tag) + now_iso = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + with self._write_lock, self._connect() as connection: + if not normalized_tag: + cursor = connection.execute( + """ + UPDATE app_catalog + SET batch_tags = '[]', + updated_at = ? + WHERE COALESCE(batch_tags, '') NOT IN ('', '[]') + """, + (now_iso,), + ) + return int(cursor.rowcount or 0) + rows = connection.execute( + """ + SELECT package_name, batch_tags + FROM app_catalog + WHERE COALESCE(batch_tags, '') != '' + """ + ).fetchall() + count = 0 + for row in rows: + next_tags = _remove_incremental_batch_tag(row["batch_tags"], normalized_tag) + if next_tags == _normalize_incremental_batch_tags(row["batch_tags"]): + continue + connection.execute( + """ + UPDATE app_catalog + SET batch_tags = ?, + updated_at = ? + WHERE package_name = ? + """, + (_incremental_batch_tags_to_json(next_tags), now_iso, row["package_name"]), + ) + count += 1 + return count + + def set_incremental_batch_tag_for_packages(self, package_names: List[str], batch_tag: str) -> int: + normalized_tag = _normalize_incremental_batch_tag(batch_tag) + normalized_packages = [str(item or "").strip() for item in package_names if str(item or "").strip()] + if not normalized_tag or not normalized_packages: + return 0 + now_iso = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + count = 0 + with self._write_lock, self._connect() as connection: + for package_name in normalized_packages: + row = connection.execute( + "SELECT batch_tags FROM app_catalog WHERE package_name = ?", + (package_name,), + ).fetchone() + if not row: + continue + next_tags = _add_incremental_batch_tag(row["batch_tags"], normalized_tag) + connection.execute( + """ + UPDATE app_catalog + SET batch_tags = ?, + updated_at = ? + WHERE package_name = ? + """, + (_incremental_batch_tags_to_json(next_tags), now_iso, package_name), + ) + count += 1 + return count + + def update_collection_statuses(self, updates: List[Dict[str, Any]]) -> int: + if not updates: + return 0 + count = 0 + with self._write_lock, self._connect() as connection: + for item in updates: + package_name = str(item.get("package_name") or "").strip() + if not package_name: + continue + if str(item.get("collection_status") or "").strip() != "pending": + continue + catalog_row = connection.execute( + "SELECT * FROM app_catalog WHERE package_name = ?", + (package_name,), + ).fetchone() + if not catalog_row: + continue + self._insert_pending_collection_task( + connection, + catalog_row, + reason=str(item.get("collection_status_reason") or "manual_pending"), + ) + count += 1 + return count + + def upsert_high_priority_app(self, package_name: str, app_name: str, payload: Dict[str, Any]) -> None: + now_iso = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + app_magic_label = _normalize_app_magic_label(payload.get("app_magic_label", "")) + downloads = payload.get("downloads") + if downloads is None: + downloads = _extract_downloads_from_task_payload(payload) + last_updated = _normalize_last_updated(payload.get("last_updated", "")) + country_code = str(payload.get("country_code") or "").strip() + device_type = _normalize_task_device_type(payload.get("device_type") or "") + collection_task_type = "new_app" + if isinstance(payload, dict): + payload = dict(payload) + payload["collection_task_type"] = collection_task_type + payload["last_update_interval_days"] = int(payload.get("last_update_interval_days") or 0) + with self._write_lock, self._connect() as connection: + connection.execute( + """ + INSERT INTO app_catalog ( + package_name, app_name, app_magic_label, batch_tags, + last_updated, country_code, device_type, task_payload_json, + last_update_interval_days, + task_queue, task_priority, source_order, is_active, + downloads, created_at, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, 'high', 100, 0, 1, ?, ?, ?) + ON CONFLICT(package_name) DO UPDATE SET + app_name = COALESCE(NULLIF(excluded.app_name, ''), app_catalog.app_name), + app_magic_label = COALESCE(NULLIF(excluded.app_magic_label, ''), app_catalog.app_magic_label), + batch_tags = COALESCE(app_catalog.batch_tags, excluded.batch_tags), + last_updated = COALESCE(NULLIF(excluded.last_updated, ''), app_catalog.last_updated), + country_code = excluded.country_code, + device_type = excluded.device_type, + task_payload_json = excluded.task_payload_json, + last_update_interval_days = COALESCE(excluded.last_update_interval_days, app_catalog.last_update_interval_days), + source_order = COALESCE(excluded.source_order, app_catalog.source_order), + downloads = COALESCE(excluded.downloads, app_catalog.downloads), + task_queue = 'high', + task_priority = 100, + is_active = 1, + updated_at = excluded.updated_at + """, + ( + package_name, + app_name or package_name, + app_magic_label, + "[]", + last_updated, + country_code, + device_type, + _safe_json_dumps(payload), + int(payload.get("last_update_interval_days") or 0), + downloads, + now_iso, + now_iso, + ), + ) + catalog_row = connection.execute( + "SELECT * FROM app_catalog WHERE package_name = ?", + (package_name,), + ).fetchone() + if catalog_row: + self._insert_pending_collection_task( + connection, + catalog_row, + reason="high_priority_enqueue", + collection_task_type=collection_task_type, + ) + + def set_collection_status_for_packages(self, package_names: List[str], collection_status: str, reason: str = "") -> int: + normalized = [str(item or "").strip() for item in package_names if str(item or "").strip()] + if not normalized: + return 0 + if str(collection_status or "").strip() != "pending": + return 0 + count = 0 + with self._write_lock, self._connect() as connection: + for package_name in normalized: + catalog_row = connection.execute( + "SELECT * FROM app_catalog WHERE package_name = ?", + (package_name,), + ).fetchone() + if not catalog_row: + continue + self._insert_pending_collection_task( + connection, + catalog_row, + reason=reason or "manual_pending", + ) + count += 1 + return count + + def has_successful_related_magic_label(self, app_magic_label: str, package_name: str) -> bool: + normalized_label = _normalize_app_magic_label(app_magic_label) + normalized_package = str(package_name or "").strip() + if not normalized_label or not normalized_package: + return False + with self._connect() as connection: + row = connection.execute( + """ + SELECT 1 + FROM app_catalog ac + JOIN collection_task ct + ON ct.package_name = ac.package_name + WHERE COALESCE(ac.app_magic_label, '') = ? + AND ac.package_name != ? + AND ct.execution_status = 'success' + LIMIT 1 + """, + (normalized_label, normalized_package), + ).fetchone() + return bool(row) + + def list_pending_collection_tasks(self) -> List[Dict[str, Any]]: + with self._connect() as connection: + summaries = self._list_catalog_summaries(connection) + pending = [ + item for item in summaries + if item["catalog_active"] and item["collection_status"] == "pending" + ] + pending.sort( + key=lambda item: ( + 0 if item["collection_task_type"] == "new_app" else 1, + -int(item["last_update_interval_days"] or 0), + -int(item["downloads"] or 0), + 1 if item["source_order"] is None else 0, + int(item["source_order"] or 0), + item["package_name"], + ) + ) + return [ + { + "package_name": item["package_name"], + "app_name": item["app_name"], + "source_order": item["source_order"], + "collection_task_type": item["collection_task_type"], + "last_update_interval_days": int(item["last_update_interval_days"] or 0), + "task_payload": item["task_payload"], + "task_queue": item["task_queue"], + } + for item in pending + ] + + def list_qualified_apps(self) -> List[Dict[str, Any]]: + with self._connect() as connection: + summaries = self._list_catalog_summaries(connection) + return [ + { + "package_name": item["package_name"], + "app_name": item["app_name"], + "task_payload": item["task_payload"], + } + for item in sorted(summaries, key=lambda value: value["package_name"]) + if item["collection_status"] == "qualified" + ] + + def list_all_catalog_apps(self) -> List[Dict[str, Any]]: + with self._connect() as connection: + summaries = self._list_catalog_summaries(connection) + return [ + { + "package_name": item["package_name"], + "app_name": item["app_name"], + "task_payload": item["task_payload"], + } + for item in sorted(summaries, key=lambda value: value["package_name"]) + ] + + def list_model_eligible_apps(self) -> List[Dict[str, Any]]: + with self._connect() as connection: + summaries = self._list_catalog_summaries(connection) + eligible = [item for item in summaries if item["model_eligible"]] + eligible.sort(key=lambda item: (-int(item["downloads"] or 0), int(item["source_order"] or 0), item["package_name"])) + return [ + { + "package_name": item["package_name"], + "app_name": item["app_name"], + "downloads": int(item["downloads"] or 0), + "task_payload": item["task_payload"], + } + for item in eligible + ] + + def replace_package_snapshot( + self, + package_name: str, + summary: Dict[str, Any], + domain_rows: List[Dict[str, Any]], + component_rows: List[Dict[str, Any]], + source_rows: List[Dict[str, Any]], + ): + """ + 替换应用快照（新架构）。 + + 写入 collection_task（执行记录） + app_catalog（元数据） + """ + now = time.time() + now_iso = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(now)) + latest_test_time = _parse_float_value(summary.get("latest_test_time")) + completed_at = ( + time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(latest_test_time)) + if latest_test_time > 0 + else now_iso + ) + + # 解析主键字段从 latest_task_key + task_key = summary.get("latest_task_key", "") + batch_tag, run_kind, attempt = self._parse_task_key( + task_key, + batch_tag=summary.get("batch_tag"), + run_kind=summary.get("run_kind"), + attempt=summary.get("attempt"), + ) + + # 解析错误信息 + error_category = None + error_code = None + failure_type = summary.get("latest_failure_type", "") + if "/" in failure_type: + parts = failure_type.split("/", 1) + error_category = parts[0].strip() or None + try: + error_code = int(parts[1].strip()) + except (ValueError, IndexError): + pass + elif str(failure_type or "").strip(): + error_category = str(failure_type).strip() + traffic_file_paths = summary.get("traffic_file_paths") + if not traffic_file_paths: + traffic_file_paths = [ + str(row.get("file_path") or "").strip() + for row in source_rows + if str(row.get("file_path") or "").strip() + ] + traffic_file_paths_json = _safe_json_dumps(traffic_file_paths or []) + + with self._write_lock, self._connect() as connection: + # 1. 删除旧的详细数据 + connection.execute("DELETE FROM app_domain_traffic WHERE package_name = ?", (package_name,)) + connection.execute("DELETE FROM app_traffic_component WHERE package_name = ?", (package_name,)) + connection.execute("DELETE FROM analytics_source_file WHERE package_name = ?", (package_name,)) + + # 2. 插入域名流量数据 + for row in domain_rows: + connection.execute(""" + INSERT INTO app_domain_traffic ( + package_name, domain, domain_type, traffic_bytes, flow_count, + traffic_ratio, domain_traffic_ratio, organization, + matched_tp_mark, matched_pattern, match_state, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, ( + package_name, row["domain"], row.get("domain_type", ""), + int(row.get("traffic_bytes", 0)), int(row.get("flow_count", 0)), + float(row.get("traffic_ratio", 0)), float(row.get("domain_traffic_ratio", 0)), + row.get("organization", ""), row.get("matched_tp_mark", ""), + row.get("matched_pattern", ""), row.get("match_state", ""), + now, + )) + + # 3. 插入组件流量数据 + for row in component_rows: + connection.execute(""" + INSERT INTO app_traffic_component ( + package_name, component_name, component_package_names, + is_self, traffic_bytes, share_percent, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?) + """, ( + package_name, row["component_name"], + row.get("component_package_names", ""), + 1 if row.get("is_self") else 0, + int(row.get("traffic_bytes", 0)), float(row.get("share_percent", 0)), + now, + )) + + # 4. 插入 source 文件索引 + for row in source_rows: + file_path = str(row.get("file_path") or "").strip() + if not file_path: + continue + connection.execute( + """ + INSERT INTO analytics_source_file ( + file_path, file_type, package_name, worker_id, + size_bytes, mtime, checksum, last_seen_at, last_processed_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(file_path) DO UPDATE SET + file_type = excluded.file_type, + package_name = excluded.package_name, + worker_id = excluded.worker_id, + size_bytes = excluded.size_bytes, + mtime = excluded.mtime, + checksum = excluded.checksum, + last_seen_at = excluded.last_seen_at, + last_processed_at = excluded.last_processed_at + """, + ( + file_path, + row.get("file_type", "traffic"), + package_name, + normalize_worker_id(row.get("worker_id", "")), + int(row.get("size_bytes", 0) or 0), + float(row.get("mtime", 0.0) or 0.0), + row.get("checksum", ""), + now, + now, + ), + ) + + # 5. 插入执行记录到 collection_task + connection.execute(""" + INSERT INTO collection_task ( + package_name, batch_tag, run_kind, attempt, + task_key, app_name, app_magic_label, is_new_app, + task_status, execution_status, worker_id, created_at, completed_at, + error_category, error_code, error_reason, error_details, + duration_seconds, download_duration_seconds, execution_duration_seconds, + analysis_duration_seconds, droidbot_steps, gui_agent_steps, total_steps, + num_nodes, num_reached_activities, app_num_total_activities, + total_traffic_bytes, self_traffic_bytes, server_traffic_bytes, + unrecognized_traffic_bytes, model_flow_count, model_traffic_bytes, + login_count, register_count, stuck_reason_code, guiagent_message, + scenario_triggered, download_source, is_retry, exit_code, trace_json, + traffic_file_paths + ) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?) + ON CONFLICT(package_name, batch_tag, run_kind, attempt) DO UPDATE SET + task_key=excluded.task_key, + app_name=excluded.app_name, + app_magic_label=excluded.app_magic_label, + is_new_app=excluded.is_new_app, + task_status=excluded.task_status, + execution_status=excluded.execution_status, + worker_id=excluded.worker_id, + completed_at=excluded.completed_at, + error_category=excluded.error_category, + error_code=excluded.error_code, + error_reason=excluded.error_reason, + error_details=excluded.error_details, + duration_seconds=excluded.duration_seconds, + download_duration_seconds=excluded.download_duration_seconds, + execution_duration_seconds=excluded.execution_duration_seconds, + analysis_duration_seconds=excluded.analysis_duration_seconds, + droidbot_steps=excluded.droidbot_steps, + gui_agent_steps=excluded.gui_agent_steps, + total_steps=excluded.total_steps, + total_traffic_bytes=excluded.total_traffic_bytes, + self_traffic_bytes=excluded.self_traffic_bytes, + server_traffic_bytes=excluded.server_traffic_bytes, + unrecognized_traffic_bytes=excluded.unrecognized_traffic_bytes, + model_flow_count=excluded.model_flow_count, + model_traffic_bytes=excluded.model_traffic_bytes, + num_nodes=excluded.num_nodes, + num_reached_activities=excluded.num_reached_activities, + app_num_total_activities=excluded.app_num_total_activities, + traffic_file_paths=excluded.traffic_file_paths + """, ( + package_name, batch_tag, run_kind, attempt, + task_key, summary.get("app_name"), summary.get("app_magic_label"), + 1 if summary.get("collection_task_type") == "new_app" else 0, + "completed", summary.get("latest_status"), normalize_worker_id(summary.get("latest_worker_id")), now_iso, completed_at, + error_category, error_code, summary.get("collection_status_reason"), summary.get("latest_task_detail"), + summary.get("duration_seconds", 0), + summary.get("download_duration_seconds", 0), + summary.get("execution_duration_seconds", summary.get("collect_duration_seconds", 0)), + summary.get("analysis_duration_seconds", 0), + summary.get("droidbot_steps", 0), summary.get("gui_agent_steps", 0), + int(summary.get("droidbot_steps", 0) or 0) + int(summary.get("gui_agent_steps", 0) or 0), + summary.get("num_nodes", 0), summary.get("num_reached_activities", 0), summary.get("app_num_total_activities", 0), + summary.get("total_traffic_bytes", 0), summary.get("self_traffic_bytes", 0), summary.get("server_traffic_bytes", 0), + summary.get("unrecognized_traffic_bytes", 0), summary.get("model_flow_count", 0), summary.get("model_traffic_bytes", 0), + summary.get("login_count", 0), summary.get("register_count", 0), summary.get("stuck_reason_code"), + summary.get("guiagent_message", ""), 1 if summary.get("scenario_triggered") else 0, + summary.get("download_source", ""), 1 if summary.get("is_retry") else 0, + summary.get("exit_code"), _safe_json_dumps(summary.get("trace") or {}), + traffic_file_paths_json, + )) + + # 6. 更新 app_catalog（元数据） + existing_catalog = connection.execute( + "SELECT * FROM app_catalog WHERE package_name = ?", + (package_name,), + ).fetchone() + batch_tags = _normalize_incremental_batch_tags(existing_catalog["batch_tags"] if existing_catalog else []) + if batch_tag and batch_tag not in batch_tags: + batch_tags.append(batch_tag) + task_payload = summary.get("task_payload") + if not isinstance(task_payload, dict): + task_payload = _safe_json_loads(existing_catalog["task_payload_json"], {}) if existing_catalog else {} + last_updated = _normalize_last_updated( + summary.get("last_updated") + or (task_payload or {}).get("last_updated") + or (existing_catalog["last_updated"] if existing_catalog else "") + ) + country_code = str( + summary.get("country_code") + or (task_payload or {}).get("country_code") + or (existing_catalog["country_code"] if existing_catalog else "") + or "" + ).strip() + device_type = _normalize_task_device_type( + summary.get("device_type") + or (task_payload or {}).get("device_type") + or (existing_catalog["device_type"] if existing_catalog else "") + or "" + ) + last_update_interval_days = int( + summary.get("last_update_interval_days") + or ((task_payload or {}).get("last_update_interval_days") if isinstance(task_payload, dict) else 0) + or (existing_catalog["last_update_interval_days"] if existing_catalog else 0) + or 0 + ) + if isinstance(task_payload, dict) and task_payload: + task_payload = dict(task_payload) + task_payload["last_updated"] = task_payload.get("last_updated") or last_updated + task_payload["country_code"] = task_payload.get("country_code") or country_code + task_payload["device_type"] = _normalize_task_device_type(task_payload.get("device_type") or device_type) + task_payload["collection_task_type"] = summary.get("collection_task_type") or task_payload.get("collection_task_type") or "new_app" + task_payload["last_update_interval_days"] = last_update_interval_days + connection.execute(""" + INSERT INTO app_catalog ( + package_name, app_name, app_magic_label, batch_tags, + last_updated, country_code, device_type, task_payload_json, + last_update_interval_days, downloads, source_order, is_active, + created_at, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(package_name) DO UPDATE SET + app_name=COALESCE(excluded.app_name, app_catalog.app_name), + app_magic_label=COALESCE(NULLIF(excluded.app_magic_label, ''), app_catalog.app_magic_label), + batch_tags=excluded.batch_tags, + last_updated=COALESCE(NULLIF(excluded.last_updated, ''), app_catalog.last_updated), + country_code=COALESCE(NULLIF(excluded.country_code, ''), app_catalog.country_code), + device_type=COALESCE(NULLIF(excluded.device_type, ''), app_catalog.device_type), + task_payload_json=COALESCE(NULLIF(excluded.task_payload_json, '{}'), app_catalog.task_payload_json), + last_update_interval_days=COALESCE(excluded.last_update_interval_days, app_catalog.last_update_interval_days), + downloads=COALESCE(excluded.downloads, app_catalog.downloads), + source_order=COALESCE(excluded.source_order, app_catalog.source_order), + is_active=1, + updated_at=excluded.updated_at + """, ( + package_name, summary.get('app_name'), summary.get('app_magic_label'), + _incremental_batch_tags_to_json(batch_tags), + last_updated, country_code, device_type, _safe_json_dumps(task_payload or {}), + last_update_interval_days, summary.get('downloads'), summary.get('source_order'), 1, now_iso, now_iso, + )) + + def _parse_task_key( + self, + task_key: str, + *, + batch_tag: Any = None, + run_kind: Any = None, + attempt: Any = None, + ): + """从 task_key 解析 batch_tag, run_kind, attempt""" + explicit_batch_tag = str(batch_tag or "").strip() + explicit_run_kind = str(run_kind or "").strip() + explicit_attempt = _parse_int_value(attempt) + if explicit_batch_tag and explicit_run_kind in {"ranking", "block", "model", "manual"} and explicit_attempt > 0: + return explicit_batch_tag, explicit_run_kind, explicit_attempt + + parts = str(task_key or "").split("_") + if len(parts) >= 3: + try: + attempt = int(parts[-1]) + except (ValueError, IndexError): + attempt = 1 + + run_kind_candidate = parts[-2] + if run_kind_candidate in ('ranking', 'block', 'model', 'manual'): + run_kind = run_kind_candidate + batch_tag = "_".join(parts[:-2]) or "unknown" + return batch_tag, run_kind, attempt + else: + run_kind = 'ranking' + + return 'unknown', run_kind, attempt + + if str(task_key or "").endswith("_block"): + return 'unknown', 'block', 1 + if str(task_key or "").endswith("_model"): + return 'unknown', 'model', 1 + return 'unknown', 'ranking', 1 + + def delete_snapshots_not_in(self, package_names: List[str]) -> int: + """删除不在列表中的应用快照（新架构：从 app_catalog 删除）""" + normalized = [str(item or "").strip() for item in package_names if str(item or "").strip()] + if not normalized: + return 0 + placeholders = ",".join("?" for _ in normalized) + with self._write_lock, self._connect() as connection: + stale_rows = connection.execute( + f""" + SELECT package_name + FROM app_catalog + WHERE package_name NOT IN ({placeholders}) + """, + normalized, + ).fetchall() + stale_packages = [row["package_name"] for row in stale_rows if row["package_name"]] + if not stale_packages: + return 0 + stale_placeholders = ",".join("?" for _ in stale_packages) + for table_name in ( + "app_domain_traffic", + "app_traffic_component", + "analytics_source_file", + "app_catalog", + ): + connection.execute( + f"DELETE FROM {table_name} WHERE package_name IN ({stale_placeholders})", + stale_packages, + ) + return len(stale_packages) + + def clear_all_snapshots(self) -> int: + """清空所有快照数据（新架构：只清理辅助表，激活 app_catalog）""" + with self._write_lock, self._connect() as connection: + row = connection.execute("SELECT COUNT(*) AS total FROM app_catalog").fetchone() + cleared = int((row["total"] or 0) if row else 0) + + # 删除快照辅助表数据 + for table_name in ( + "app_domain_traffic", + "app_traffic_component", + "analytics_source_file", + ): + connection.execute(f"DELETE FROM {table_name}") + + # 激活所有应用以便重新采集 + # 注意：不修改 batch_tags（保持原有批次标签） + connection.execute( + """ + UPDATE app_catalog + SET is_active = 1, + updated_at = datetime('now','localtime') + """, + ) + return cleared + + def get_overview( + self, + *, + incremental_batch_tag: str = "", + top_n: Any = DEFAULT_INCREMENTAL_TOP_N, + ) -> Dict[str, Any]: + normalized_batch_tag = _normalize_incremental_batch_tag(incremental_batch_tag) + normalized_top_n = _normalize_top_n(top_n) + top_n_label = _top_n_label(normalized_top_n) + pending_preview_limit = 12 + with self._connect() as connection: + summaries = self._list_catalog_summaries(connection) + + if normalized_batch_tag: + base_scope = [item for item in summaries if normalized_batch_tag in item["incremental_batch_tags"]] + else: + base_scope = [item for item in summaries if item["catalog_active"]] + filtered = [ + item for item in base_scope + if item["source_order"] is not None and int(item["source_order"]) < normalized_top_n + ] + + def _count(predicate) -> int: + return sum(1 for item in filtered if predicate(item)) + + severe_non_retryable_count = _count( + lambda item: item["restriction_status"] == "severe_restricted" + and item["retryability"] == "non_retryable" + ) + breakdown_source_rows = [ + { + "package_name": item["package_name"], + "error_type": item["latest_failure_type"], + "downloads": item["downloads"], + "task_payload_json": item["task_payload_json"], + } + for item in filtered + if item["restriction_status"] == "severe_restricted" + and item["retryability"] == "non_retryable" + ] + non_retryable_download_rows = sorted( + breakdown_source_rows, + key=lambda item: (item.get("package_name") or ""), + ) + pending_preview_rows = sorted( + [item for item in filtered if item["collection_status"] == "pending"], + key=lambda item: ( + 1 if item["source_order"] is None else 0, + int(item["source_order"] or 0), + item["package_name"], + ), + )[:pending_preview_limit] + + batch_buckets: Dict[str, Dict[str, Any]] = {} + for item in summaries: + for tag in item["incremental_batch_tags"]: + bucket = batch_buckets.setdefault(tag, {"tag": tag, "app_count": 0, "marked_at": 0.0}) + bucket["app_count"] += 1 + bucket["marked_at"] = max(float(bucket["marked_at"] or 0.0), float(item["updated_at"] or 0.0)) + available_batch_rows = sorted( + batch_buckets.values(), + key=lambda item: (-float(item["marked_at"] or 0.0), item["tag"]), + ) + + non_retryable_breakdown_buckets: Dict[str, Dict[str, Any]] = {} + for breakdown_row in breakdown_source_rows: + normalized_error_type = str(breakdown_row["error_type"] or "").strip() or "UNMARKED" + bucket = non_retryable_breakdown_buckets.setdefault( + normalized_error_type, + { + "error_type": normalized_error_type, + "label": _humanize_error_type(normalized_error_type), + "app_count": 0, + "downloads_below_10k_count": 0, + "downloads_10k_or_above_count": 0, + "downloads_unknown_count": 0, + }, + ) + bucket["app_count"] += 1 + downloads = _resolve_downloads_for_package( + breakdown_row["package_name"], + breakdown_row["task_payload_json"], + downloads=breakdown_row["downloads"], + ) + if downloads is None: + bucket["downloads_unknown_count"] += 1 + elif downloads < NON_RETRYABLE_DOWNLOAD_SPLIT_THRESHOLD: + bucket["downloads_below_10k_count"] += 1 + else: + bucket["downloads_10k_or_above_count"] += 1 + non_retryable_breakdown = sorted( + ( + { + "error_type": item["error_type"], + "label": item["label"], + "app_count": int(item["app_count"] or 0), + "downloads_below_10k_count": int(item["downloads_below_10k_count"] or 0), + "downloads_10k_or_above_count": int(item["downloads_10k_or_above_count"] or 0), + "downloads_unknown_count": int(item["downloads_unknown_count"] or 0), + "share_percent": round((int(item["app_count"] or 0) / severe_non_retryable_count) * 100, 2) + if severe_non_retryable_count > 0 + else 0.0, + } + for item in non_retryable_breakdown_buckets.values() + ), + key=lambda item: (-int(item["app_count"]), str(item["error_type"])), + ) + download_distribution = _build_non_retryable_download_distribution(non_retryable_download_rows) + self_ratios = [float(item["self_ratio"] or 0.0) for item in filtered] + last_updated_at = max([float(item["updated_at"] or 0.0) for item in filtered] or [0.0]) + return { + "scope": "incremental_topn" if normalized_batch_tag else "topn_catalog", + "incremental_batch_tag": normalized_batch_tag, + "top_n": normalized_top_n, + "top_n_label": top_n_label, + "base_scope_app_count": len(base_scope), + "base_scope_missing_source_order_count": sum(1 for item in base_scope if item["source_order"] is None), + "app_count": len(filtered), + "pending_count": _count(lambda item: item["collection_status"] == "pending"), + "qualified_count": _count(lambda item: item["collection_status"] == "qualified"), + "failed_terminal_count": _count(lambda item: item["collection_status"] == "failed_terminal"), + "analyzed_app_count": _count(lambda item: bool(item["restriction_status"])), + "success_app_count": _count(lambda item: item["latest_status"] == "success"), + "non_success_app_count": _count(lambda item: item["latest_status"] not in {"", "success"}), + "restriction_success_count": _count(lambda item: item["restriction_status"] == "success"), + "light_restricted_count": _count(lambda item: item["restriction_status"] == "light_restricted"), + "severe_restricted_count": _count(lambda item: item["restriction_status"] == "severe_restricted"), + "severe_retryable_count": _count( + lambda item: item["restriction_status"] == "severe_restricted" + and item["retryability"] == "retryable" + ), + "severe_non_retryable_count": severe_non_retryable_count, + "pending_non_retryable_count": _count( + lambda item: item["collection_status"] == "pending" + and item["restriction_status"] == "severe_restricted" + and item["retryability"] == "non_retryable" + ), + "partial_app_count": _count(lambda item: item["artifact_status"] == "partial"), + "domain_count": sum(int(item["unique_domain_count"] or 0) for item in filtered), + "total_traffic_bytes": sum(int(item["total_traffic_bytes"] or 0) for item in filtered), + "avg_self_ratio": round((sum(self_ratios) / len(self_ratios)) if self_ratios else 0.0, 2), + "last_updated_at": _format_ts(last_updated_at), + "pending_preview_limit": pending_preview_limit, + "pending_preview": [ + { + "package_name": pending_row["package_name"], + "app_name": pending_row["app_name"] or "", + "source_order": int(pending_row["source_order"]) + if pending_row["source_order"] is not None + else None, + } + for pending_row in pending_preview_rows + ], + "non_retryable_breakdown": non_retryable_breakdown, + "non_retryable_download_split_threshold": NON_RETRYABLE_DOWNLOAD_SPLIT_THRESHOLD, + "non_retryable_download_distribution": download_distribution["distribution"], + "non_retryable_downloads_matched_count": download_distribution["matched_downloads_apps"], + "non_retryable_downloads_missing_count": download_distribution["missing_downloads_apps"], + "available_incremental_batches": [ + { + "tag": batch_row["tag"] or "", + "app_count": int(batch_row["app_count"] or 0), + "marked_at": _format_ts(batch_row["marked_at"]), + } + for batch_row in available_batch_rows + if str(batch_row["tag"] or "").strip() + ], + } + + def list_apps( + self, + *, + q: str = "", + latest_status: str = "", + artifact_status: str = "", + collection_status: str = "", + restriction_status: str = "", + retryability: str = "", + incremental_batch_tag: str = "", + top_n: Any = DEFAULT_INCREMENTAL_TOP_N, + sort: str = "updated_at", + order: str = "desc", + page: int = 1, + page_size: int = 50, + ) -> Dict[str, Any]: + sort_map = { + "updated_at": "updated_at", + "app_name": "app_name", + "package_name": "package_name", + "collection_status": "collection_status", + "latest_status": "latest_status", + "restriction_status": "restriction_status", + "retryability": "retryability", + "latest_test_time": "latest_test_time", + "self_ratio": "self_ratio", + "recognition_ratio": "recognition_ratio", + "total_traffic_bytes": "total_traffic_bytes", + "unique_domain_count": "unique_domain_count", + "source_order": "source_order", + } + sort_key = sort_map.get(sort, "updated_at") + reverse_order = str(order).lower() != "asc" + normalized_batch_tag = _normalize_incremental_batch_tag(incremental_batch_tag) + normalized_top_n = _normalize_top_n(top_n) + top_n_label = _top_n_label(normalized_top_n) + page = max(1, int(page or 1)) + page_size = max(1, min(int(page_size or 50), 200)) + offset = (page - 1) * page_size + with self._connect() as connection: + summaries = self._list_catalog_summaries(connection) + if normalized_batch_tag: + filtered = [item for item in summaries if normalized_batch_tag in item["incremental_batch_tags"]] + else: + filtered = [item for item in summaries if item["catalog_active"]] + filtered = [ + item for item in filtered + if item["source_order"] is not None and int(item["source_order"]) < normalized_top_n + ] + if q: + needle = str(q).lower() + filtered = [ + item for item in filtered + if needle in str(item["app_name"] or "").lower() + or needle in str(item["package_name"] or "").lower() + ] + if latest_status: + filtered = [item for item in filtered if item["latest_status"] == latest_status] + if artifact_status: + filtered = [item for item in filtered if item["artifact_status"] == artifact_status] + if collection_status: + filtered = [item for item in filtered if item["collection_status"] == collection_status] + if restriction_status: + filtered = [item for item in filtered if item["restriction_status"] == restriction_status] + if retryability: + filtered = [item for item in filtered if item["retryability"] == retryability] + + def _sortable_value(item: Dict[str, Any]) -> Any: + value = item.get(sort_key) + if sort_key in { + "updated_at", + "latest_test_time", + "self_ratio", + "recognition_ratio", + "total_traffic_bytes", + "unique_domain_count", + "source_order", + }: + return float(value or 0) + return str(value or "") + + if sort_key == "source_order": + filtered.sort( + key=lambda item: ( + 1 if item["source_order"] is None else 0, + float(item["source_order"] or 0), + item["package_name"], + ), + reverse=False, + ) + if reverse_order: + filtered = list(reversed(filtered)) + else: + filtered.sort(key=lambda item: (_sortable_value(item), item["package_name"]), reverse=reverse_order) + total = len(filtered) + rows = filtered[offset:offset + page_size] + items = [] + for row in rows: + items.append( + { + "app_name": row["app_name"] or "", + "package_name": row["package_name"], + "app_magic_label": row["app_magic_label"] or "", + "last_updated": row["last_updated"] or "", + "collection_task_type": row["collection_task_type"] or "new_app", + "last_update_interval_days": int(row["last_update_interval_days"] or 0), + "device_type": _normalize_task_device_type(row["device_type"] or ""), + "source_order": int(row["source_order"]) if row["source_order"] is not None else None, + "collection_status": row["collection_status"] or "pending", + "latest_status": row["latest_status"] or "", + "restriction_status": row["restriction_status"] or "", + "retryability": row["retryability"] or "not_applicable", + "latest_test_time": _format_ts(row["latest_test_time"]), + "latest_worker_id": row["latest_worker_id"] or "", + "downloads": _parse_downloads_value(row["downloads"]), + "incremental_batch_tag": ", ".join(_normalize_incremental_batch_tags(row["incremental_batch_tag"])) or "", + "incremental_batch_tags": _normalize_incremental_batch_tags(row["incremental_batch_tag"]), + "incremental_batch_marked_at": _format_ts(row["incremental_batch_marked_at"]), + "self_traffic_bytes": int(row["self_traffic_bytes"] or 0), + "server_traffic_bytes": int(row["server_traffic_bytes"] or 0), + "self_ratio": round(float(row["self_ratio"] or 0.0), 2), + "recognition_ratio": round(float(row["recognition_ratio"] or 0.0), 2), + "total_traffic_bytes": int(row["total_traffic_bytes"] or 0), + "unique_domain_count": int(row["unique_domain_count"] or 0), + "unique_second_level_domain_count": int(row["unique_second_level_domain_count"] or 0), + "duration_seconds": round(float(row["duration_seconds"] or 0.0), 2), + "num_nodes": int(row["num_nodes"] or 0), + "num_reached_activities": int(row["num_reached_activities"] or 0), + "app_num_total_activities": int(row["app_num_total_activities"] or 0), + "artifact_status": row["artifact_status"] or "missing", + "updated_at": _format_ts(row["updated_at"]), + } + ) + return { + "items": items, + "page": page, + "page_size": page_size, + "total": total, + "incremental_batch_tag": normalized_batch_tag, + "top_n": normalized_top_n, + "top_n_label": top_n_label, + } + + def get_app_detail(self, package_name: str) -> Optional[Dict[str, Any]]: + with self._connect() as connection: + summary_row = self._latest_catalog_status(connection, package_name) + if not summary_row: + return None + domain_rows = connection.execute( + """ + SELECT domain, domain_type, traffic_bytes, flow_count, + organization, matched_tp_mark, matched_pattern, match_state + FROM app_domain_traffic + WHERE package_name = ? + ORDER BY traffic_bytes DESC, domain ASC + """, + (package_name,), + ).fetchall() + component_rows = connection.execute( + """ + SELECT * + FROM app_traffic_component + WHERE package_name = ? + ORDER BY traffic_bytes DESC, component_name ASC + """, + (package_name,), + ).fetchall() + raw_domains = [ + { + "domain": row["domain"], + "domain_type": row["domain_type"] or "", + "traffic_bytes": int(row["traffic_bytes"] or 0), + "flow_count": int(row["flow_count"] or 0), + "organization": row["organization"] or "", + "matched_pattern": row["matched_pattern"] or "", + "match_state": row["match_state"] or "", + } + for row in domain_rows + ] + top_raw_domains = raw_domains[:12] + top_second_level_domains = _aggregate_second_level_domain_rows(domain_rows)[:12] + components = [ + { + "component_name": row["component_name"], + "component_package_names": row["component_package_names"] or "", + "is_self": bool(row["is_self"]), + "traffic_bytes": int(row["traffic_bytes"] or 0), + "share_percent": round(float(row["share_percent"] or 0.0), 2), + "updated_at": _format_ts(row["updated_at"]), + } + for row in component_rows + ] + summary = { + "app_name": summary_row["app_name"] or "", + "package_name": summary_row["package_name"], + "app_magic_label": summary_row["app_magic_label"] or "", + "last_updated": summary_row["last_updated"] or "", + "collection_task_type": summary_row["collection_task_type"] or "new_app", + "last_update_interval_days": int(summary_row["last_update_interval_days"] or 0), + "country_code": summary_row["country_code"] or "", + "device_type": _normalize_task_device_type(summary_row["device_type"] or ""), + "downloads": _parse_downloads_value(summary_row["downloads"]), + "source_order": int(summary_row["source_order"]) if summary_row["source_order"] is not None else None, + "incremental_batch_tag": ", ".join(_normalize_incremental_batch_tags(summary_row["incremental_batch_tag"])) or "", + "incremental_batch_tags": _normalize_incremental_batch_tags(summary_row["incremental_batch_tag"]), + "incremental_batch_marked_at": _format_ts(summary_row["incremental_batch_marked_at"]), + "collection_status": summary_row["collection_status"] or "pending", + "collection_status_reason": summary_row["collection_status_reason"] or "", + "latest_task_key": summary_row["latest_task_key"] or "", + "latest_status": summary_row["latest_status"] or "", + "restriction_status": summary_row["restriction_status"] or "", + "retryability": summary_row["retryability"] or "not_applicable", + "latest_test_time": _format_ts(summary_row["latest_test_time"]), + "latest_worker_id": summary_row["latest_worker_id"] or "", + "latest_task_detail": summary_row["latest_task_detail"] or "", + "latest_failure_type": summary_row["latest_failure_type"] or "", + "unique_domain_count": int(summary_row["unique_domain_count"] or 0), + "unique_domain_names": summary_row["unique_domain_names"] or [], + "unique_second_level_domain_count": int(summary_row["unique_second_level_domain_count"] or 0), + "unique_second_level_domains": summary_row["unique_second_level_domains"] or [], + "droidbot_steps": int(summary_row["droidbot_steps"] or 0), + "gui_agent_steps": int(summary_row["gui_agent_steps"] or 0), + "duration_seconds": round(float(summary_row["duration_seconds"] or 0.0), 2), + "num_nodes": int(summary_row["num_nodes"] or 0), + "num_reached_activities": int(summary_row["num_reached_activities"] or 0), + "app_num_total_activities": int(summary_row["app_num_total_activities"] or 0), + "total_traffic_bytes": int(summary_row["total_traffic_bytes"] or 0), + "self_traffic_bytes": int(summary_row["self_traffic_bytes"] or 0), + "server_traffic_bytes": int(summary_row["server_traffic_bytes"] or 0), + "unrecognized_traffic_bytes": int(summary_row["unrecognized_traffic_bytes"] or 0), + "self_ratio": round(float(summary_row["self_ratio"] or 0.0), 2), + "recognition_ratio": round(float(summary_row["recognition_ratio"] or 0.0), 2), + "artifact_status": summary_row["artifact_status"] or "missing", + "updated_at": _format_ts(summary_row["updated_at"]), + } + return { + "summary": summary, + "top_raw_domains": top_raw_domains, + "top_second_level_domains": top_second_level_domains, + "components": components, + "top_unmatched_domains": [item for item in top_second_level_domains if item["match_state"] == "unmatched"][:20], + } + + +class AnalyticsService: + def __init__( + self, + *, + db_path: str = MONITORING_DB_PATH, + traffic_root: str = ANALYTICS_TRAFFIC_ROOT, + traffic_root_block: str = "", + traversal_root: str = ANALYTICS_TRAVERSAL_ROOT, + app_list_path: str = ANALYTICS_TPDPI_APP_LIST, + url_lib_path: str = ANALYTICS_TPDPI_URL_LIB, + artifact_wait_seconds: int = ANALYTICS_ARTIFACT_WAIT_SECONDS, + job_poll_seconds: int = ANALYTICS_JOB_POLL_SECONDS, + on_snapshot_callback=None, + start_worker: bool = True, + ): + self.db_path = db_path + self.traffic_root = _resolve_data_root(traffic_root) + self.traffic_root_block = _resolve_data_root(traffic_root_block) if traffic_root_block else "" + self.traversal_root = _resolve_data_root(traversal_root) + self.app_list_path = app_list_path + self.url_lib_path = url_lib_path + self.artifact_wait_seconds = max(0, int(artifact_wait_seconds or 0)) + self.job_poll_seconds = max(0.2, float(job_poll_seconds or 1)) + self.repo = AnalyticsRepository(db_path=db_path) + self.on_snapshot_callback = on_snapshot_callback + self._rules: Optional[DpiRuleMatcher] = None + self._rules_lock = threading.Lock() + self._stop_event = threading.Event() + self._wake_event = threading.Event() + self._thread: Optional[threading.Thread] = None + if start_worker: + self.start() + + def start(self): + if self._thread and self._thread.is_alive(): + return + self._stop_event.clear() + self._thread = threading.Thread(target=self._run_loop, name="analytics-worker", daemon=True) + self._thread.start() + + def close(self): + self._stop_event.set() + self._wake_event.set() + if self._thread: + self._thread.join(timeout=3) + + def _ensure_rules(self) -> DpiRuleMatcher: + if self._rules is not None: + return self._rules + with self._rules_lock: + if self._rules is None: + self._rules = DpiRuleMatcher(self.app_list_path, self.url_lib_path) + return self._rules + + def enqueue_incremental( + self, + *, + package_name: str, + task_key: Optional[str] = None, + worker_id: Optional[str] = None, + report_time: Optional[float] = None, + trigger_source: str = "worker_report", + waiting_for_artifacts: bool = True, + scheduled_after: Optional[float] = None, + ) -> Dict[str, Any]: + if not package_name: + return {} + job = self.repo.create_job( + job_type="incremental", + status="waiting_artifacts" if waiting_for_artifacts else "queued", + package_name=package_name, + task_key=task_key, + worker_id=worker_id, + trigger_source=trigger_source, + report_time=report_time, + scheduled_after=scheduled_after, + ) + self._wake_event.set() + return job + + def enqueue_backfill(self, packages: Optional[List[str]] = None, trigger_source: str = "api") -> Dict[str, Any]: + payload = {"scope": "packages" if packages else "all", "packages": packages or []} + job = self.repo.create_job( + job_type="backfill", + status="queued", + trigger_source=trigger_source, + payload=payload, + ) + self._wake_event.set() + return job + + def enqueue_manual_rebuild(self, package_name: str, trigger_source: str = "api") -> Dict[str, Any]: + job = self.repo.create_job( + job_type="manual_rebuild", + status="queued", + package_name=package_name, + trigger_source=trigger_source, + ) + self._wake_event.set() + return job + + def handle_worker_event(self, payload: Dict[str, Any]) -> int: + if str(payload.get("event_type") or "").strip() != "artifacts_synced": + return 0 + package_name = str(payload.get("package_name") or "").strip() + task_key = str(payload.get("task_key") or "").strip() + worker_id = normalize_worker_id(payload.get("worker_id")) + promoted = self.repo.mark_job_ready(task_key=task_key or None, package_name=package_name or None, worker_id=worker_id or None) + if promoted == 0 and package_name: + self.enqueue_incremental( + package_name=package_name, + task_key=task_key or None, + worker_id=worker_id or None, + report_time=float(payload.get("event_time") or time.time()), + trigger_source="artifacts_synced", + waiting_for_artifacts=False, + scheduled_after=time.time() + ANALYTICS_ARTIFACT_FILE_WAIT_SECONDS, + ) + self._wake_event.set() + return promoted + + def get_overview(self, **kwargs) -> Dict[str, Any]: + return self.repo.get_overview(**kwargs) + + def list_apps(self, **kwargs) -> Dict[str, Any]: + return self.repo.list_apps(**kwargs) + + def get_app_detail(self, package_name: str) -> Optional[Dict[str, Any]]: + return self.repo.get_app_detail(package_name) + + def clear_incremental_batch_tags(self, batch_tag: str = "") -> int: + return self.repo.clear_incremental_batch_tags(batch_tag) + + def set_incremental_batch_tag_for_packages(self, package_names: List[str], batch_tag: str) -> int: + return self.repo.set_incremental_batch_tag_for_packages(package_names, batch_tag) + + def list_jobs(self, **kwargs) -> List[Dict[str, Any]]: + return self.repo.list_jobs(**kwargs) + + def list_pending_collection_tasks(self) -> List[Dict[str, Any]]: + return self.repo.list_pending_collection_tasks() + + def list_model_eligible_apps(self) -> List[Dict[str, Any]]: + return self.repo.list_model_eligible_apps() + + def run_backfill_now(self, packages: Optional[List[str]] = None) -> Tuple[str, str]: + return self._process_backfill_job({"payload": {"scope": "packages" if packages else "all", "packages": packages or []}}) + + def sync_catalog_from_rows(self, rows: List[Dict[str, Any]]) -> int: + result = self.replace_catalog_from_rows(rows) + return int(result.get("upserted") or 0) + + def replace_catalog_from_rows(self, rows: List[Dict[str, Any]]) -> Dict[str, int]: + entries = [] + seen = set() + for index, row in enumerate(rows): + payload = build_task_payload_from_row(row) + package_name = payload["package_name"] + if not package_name or package_name in seen: + continue + seen.add(package_name) + entries.append( + { + "package_name": package_name, + "app_name": payload["app_name"] or package_name, + "app_magic_label": payload.get("app_magic_label", ""), + "last_updated": payload.get("last_updated", ""), + "country_code": payload["country_code"], + "device_type": payload["device_type"], + "source_order": index, + "task_payload": payload, + "downloads": payload.get("downloads"), + "collection_status": "pending", + "collection_status_reason": "catalog_sync", + } + ) + return self.repo.replace_catalog_entries(entries) + + def refresh_collection_statuses( + self, + *, + package_names: Optional[List[str]] = None, + allow_legacy_download_error_1_pending: bool = False, + ) -> int: + del package_names, allow_legacy_download_error_1_pending + # 新 schema 不落库保存派生状态；读取时直接从 collection_task 计算。 + return 0 + + def add_catalog_app( + self, + *, + app_name: str, + package_name: str, + country_code: str = "", + device_type: str = "", + app_magic_label: str = "", + last_updated: str = "", + ) -> int: + normalized_device_type = _normalize_task_device_type(device_type) + payload = build_task_payload_from_row( + { + "app_name": app_name, + "package_name": package_name, + "app_magic_label": app_magic_label, + "last_updated": last_updated, + "country_code": country_code, + "device_type": normalized_device_type, + } + ) + with self.repo._connect() as connection: + row = connection.execute("SELECT COALESCE(MAX(source_order), -1) AS max_order FROM app_catalog").fetchone() + max_order = int((row["max_order"] or -1) if row else -1) + return self.repo.upsert_catalog_entries( + [ + { + "package_name": package_name, + "app_name": app_name or package_name, + "app_magic_label": _normalize_app_magic_label(app_magic_label), + "last_updated": _normalize_last_updated(last_updated), + "country_code": country_code, + "device_type": normalized_device_type, + "source_order": max_order + 1, + "task_payload": payload, + "downloads": payload.get("downloads"), + "collection_status": "pending", + "collection_status_reason": "manual_add", + } + ] + ) + + def set_packages_pending(self, package_names: List[str], reason: str = "manual_pending") -> int: + return self.repo.set_collection_status_for_packages(package_names, "pending", reason) + + def _run_loop(self): + while not self._stop_event.is_set(): + try: + job = self.repo.claim_next_job(self.artifact_wait_seconds) + if not job: + self._wake_event.wait(timeout=self.job_poll_seconds) + self._wake_event.clear() + continue + self._process_job(job) + except Exception: + logger.exception("Analytics background loop failed") + self._wake_event.wait(timeout=self.job_poll_seconds) + self._wake_event.clear() + + def _process_job(self, job: Dict[str, Any]): + job_id = int(job["id"]) + try: + if job["job_type"] == "backfill": + status, message = self._process_backfill_job(job) + else: + package_name = str(job.get("package_name") or "").strip() + if not package_name: + raise ValueError("package_name is required for package rebuild jobs") + summary = self.rebuild_package_now(package_name) + if callable(self.on_snapshot_callback): + try: + self.on_snapshot_callback(package_name, summary) + except Exception: + logger.exception("Analytics snapshot callback failed for %s", package_name) + status = "succeeded" if summary.get("artifact_status") == "complete" else "partial" + message = f"rebuilt {package_name} ({summary.get('artifact_status', 'missing')})" + self.repo.finish_job(job_id, status=status, error_message=message) + except Exception as exc: + logger.exception("Analytics job %s failed", job_id) + self.repo.finish_job(job_id, status="failed", error_message=str(exc)) + + def _process_backfill_job(self, job: Dict[str, Any]) -> Tuple[str, str]: + payload = job.get("payload") or {} + packages = payload.get("packages") or [] + cleared = 0 + if not packages: + packages = self.repo.list_active_catalog_packages() + if not packages: + return "failed", "no packages found in active catalog" + cleared = self.repo.clear_all_snapshots() + rebuilt = 0 + partial = 0 + failed_packages: List[str] = [] + for package_name in packages: + try: + summary = self.rebuild_package_now(package_name) + if callable(self.on_snapshot_callback): + try: + self.on_snapshot_callback(package_name, summary) + except Exception: + logger.exception("Analytics snapshot callback failed for %s", package_name) + rebuilt += 1 + if summary.get("artifact_status") != "complete": + partial += 1 + except Exception: + logger.exception("Analytics backfill failed for package %s", package_name) + failed_packages.append(package_name) + if failed_packages and rebuilt == 0: + return "failed", f"failed packages: {', '.join(failed_packages[:10])}" + if failed_packages or partial: + message_parts = [f"rebuilt={rebuilt}", f"partial={partial}"] + if cleared: + message_parts.append(f"cleared={cleared}") + if failed_packages: + message_parts.append(f"failed={len(failed_packages)}") + return "partial", ", ".join(message_parts) + message_parts = [f"rebuilt={rebuilt}"] + if cleared: + message_parts.append(f"cleared={cleared}") + return "succeeded", ", ".join(message_parts) + + def _discover_packages(self) -> set: + packages = set(self.repo.list_all_task_packages()) + for root in (self.traffic_root, self.traffic_root_block): + if not root or not os.path.isdir(root): + continue + for worker_dir in os.scandir(root): + if not worker_dir.is_dir(): + continue + try: + for package_dir in os.scandir(worker_dir.path): + if package_dir.is_dir(): + packages.add(package_dir.name) + except OSError: + continue + return {item for item in packages if item} + + def _find_traffic_files(self, package_name: str) -> List[str]: + files: List[str] = [] + for root in (self.traffic_root, self.traffic_root_block): + if not root or not os.path.isdir(root): + continue + pattern = os.path.join(root, "*", package_name, "traffic_count*.txt") + files += [path for path in glob.glob(pattern) if os.path.isfile(path)] + return sorted(files) + + def _build_source_rows(self, package_name: str, traffic_files: List[str]) -> List[Dict[str, Any]]: + rows = [] + for file_path in traffic_files: + if not file_path or not os.path.exists(file_path): + continue + stat = os.stat(file_path) + worker_id = "" + try: + if self.traffic_root_block and os.path.abspath(file_path).startswith(os.path.abspath(self.traffic_root_block)): + relative_path = os.path.relpath(file_path, self.traffic_root_block) + else: + relative_path = os.path.relpath(file_path, self.traffic_root) + worker_id = normalize_worker_id(relative_path.split(os.sep, 1)[0]) + except ValueError: + worker_id = "" + rows.append( + { + "file_path": os.path.abspath(file_path), + "file_type": "traffic", + "package_name": package_name, + "worker_id": worker_id, + "size_bytes": int(stat.st_size), + "mtime": float(stat.st_mtime), + "checksum": f"{int(stat.st_size)}:{int(stat.st_mtime)}", + } + ) + return rows + + def _parse_package_traffic(self, package_name: str, traffic_files: List[str]) -> Dict[str, Any]: + aggregated: Dict[str, Dict[str, Any]] = {} + unique_domains = set() + unique_second_level_domains = set() + app_name = "" + model_flow_count = 0 + model_traffic_bytes = 0 + for file_path in traffic_files: + try: + with open(file_path, "r", encoding="utf-8", errors="ignore") as handle: + for raw_line in handle: + line = raw_line.strip() + if not line: + continue + fields = [item.strip() for item in line.split(",")] + if len(fields) < 7: + continue + app_id = fields[0].strip() + if package_name and app_id and app_id != package_name: + continue + if fields[1].strip() and not app_name: + app_name = fields[1].strip() + domain = _determine_domain_or_model(fields) + traffic_bytes = _parse_traffic_size(fields[6]) + if _is_model_traffic_line(fields): + model_flow_count += 1 + model_traffic_bytes += traffic_bytes + organization = fields[7].strip() if len(fields) > 7 else "" + bucket = aggregated.setdefault( + domain, + {"traffic_bytes": 0, "flow_count": 0, "organization": ""}, + ) + bucket["traffic_bytes"] += traffic_bytes + bucket["flow_count"] += 1 + if organization and not bucket["organization"]: + bucket["organization"] = organization + if not domain.startswith("model_data:"): + unique_domains.add(domain) + second_level = _extract_second_level_domain(domain) + if second_level: + unique_second_level_domains.add(second_level) + except OSError as exc: + logger.warning("Failed to read traffic file %s: %s", file_path, exc) + total_traffic_bytes = sum(item["traffic_bytes"] for item in aggregated.values()) + total_domain_bytes = sum(item["traffic_bytes"] for key, item in aggregated.items() if not key.startswith("model_data:")) + return { + "app_name": app_name, + "aggregated": aggregated, + "unique_domains": sorted(unique_domains), + "unique_second_level_domains": sorted(unique_second_level_domains), + "total_traffic_bytes": total_traffic_bytes, + "total_domain_bytes": total_domain_bytes, + "model_flow_count": model_flow_count, + "model_traffic_bytes": model_traffic_bytes, + } + + def rebuild_package_now( + self, + package_name: str, + *, + latest_task_override: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + rules = self._ensure_rules() + rules.ensure_loaded() + existing_row = self.repo.get_collection_row(package_name) or {} + latest_task = self.repo.get_latest_task_execution(package_name) or {} + if latest_task_override: + latest_task = {**latest_task, **latest_task_override} + traffic_files = self._find_traffic_files(package_name) + traffic = self._parse_package_traffic(package_name, traffic_files) + + domain_rows: List[Dict[str, Any]] = [] + component_buckets: Dict[str, Dict[str, Any]] = {} + total_traffic_bytes = int(traffic["total_traffic_bytes"] or 0) + total_domain_bytes = int(traffic["total_domain_bytes"] or 0) + self_traffic_bytes = 0 + server_traffic_bytes = 0 + unrecognized_traffic_bytes = 0 + self_component_name = ( + rules.pkg_to_tp_mark.get(package_name) + or latest_task.get("app_name") + or traffic.get("app_name") + or package_name + ) + self_component_packages = rules.tp_pkg_string.get(self_component_name, package_name) + + sorted_domains = sorted( + traffic["aggregated"].items(), + key=lambda item: (-int(item[1]["traffic_bytes"]), item[0]), + ) + for domain, bucket in sorted_domains: + domain_type = "model_data" if domain.startswith("model_data:") else "domain" + traffic_bytes = int(bucket["traffic_bytes"] or 0) + flow_count = int(bucket["flow_count"] or 0) + matched_tp_mark = "" + matched_pattern = "" + match_state = "" + if domain.startswith("model_data:"): + match_state = "ip" + unrecognized_traffic_bytes += traffic_bytes + else: + matched_tp_mark, matched_pattern = rules.match_domain(domain) + matched_tp_mark = matched_tp_mark or "" + matched_pattern = matched_pattern or "" + if matched_tp_mark: + owner_packages = rules.tp_packages.get(matched_tp_mark, set()) + is_self = package_name in owner_packages + match_state = "self" if is_self else "server" + if is_self: + self_traffic_bytes += traffic_bytes + else: + server_traffic_bytes += traffic_bytes + component = component_buckets.setdefault( + matched_tp_mark, + { + "component_name": matched_tp_mark, + "component_package_names": rules.tp_pkg_string.get(matched_tp_mark, ""), + "is_self": is_self, + "traffic_bytes": 0, + }, + ) + component["traffic_bytes"] += traffic_bytes + component["is_self"] = bool(component["is_self"]) or is_self + else: + # 未匹配三方规则的域名默认按应用自有域名处理。 + match_state = "self" + self_traffic_bytes += traffic_bytes + component = component_buckets.setdefault( + self_component_name, + { + "component_name": self_component_name, + "component_package_names": self_component_packages, + "is_self": True, + "traffic_bytes": 0, + }, + ) + component["traffic_bytes"] += traffic_bytes + component["is_self"] = True + domain_rows.append( + { + "domain": domain, + "domain_type": domain_type, + "traffic_bytes": traffic_bytes, + "flow_count": flow_count, + "traffic_ratio": round((traffic_bytes / total_traffic_bytes) * 100, 2) if total_traffic_bytes > 0 else 0.0, + "domain_traffic_ratio": round((traffic_bytes / total_domain_bytes) * 100, 2) + if (not domain.startswith("model_data:") and total_domain_bytes > 0) + else 0.0, + "organization": bucket.get("organization", ""), + "matched_tp_mark": matched_tp_mark, + "matched_pattern": matched_pattern, + "match_state": match_state, + } + ) + + component_rows = [] + for component in sorted(component_buckets.values(), key=lambda item: (-int(item["traffic_bytes"]), item["component_name"])): + component_rows.append( + { + "component_name": component["component_name"], + "component_package_names": component["component_package_names"], + "is_self": bool(component["is_self"]), + "traffic_bytes": int(component["traffic_bytes"]), + "share_percent": round((int(component["traffic_bytes"]) / total_traffic_bytes) * 100, 2) + if total_traffic_bytes > 0 + else 0.0, + } + ) + + artifact_status = "missing" + if traffic_files and latest_task: + artifact_status = "complete" + elif traffic_files or latest_task: + artifact_status = "partial" + + summary = { + "app_name": latest_task.get("app_name") or traffic.get("app_name") or existing_row.get("app_name") or package_name, + "app_magic_label": _normalize_app_magic_label(existing_row.get("app_magic_label", "")), + "last_updated": _normalize_last_updated(existing_row.get("last_updated", "")), + "country_code": existing_row.get("country_code"), + "device_type": _normalize_task_device_type(existing_row.get("device_type")), + "source_order": existing_row.get("source_order"), + "task_payload": existing_row.get("task_payload"), + "downloads": existing_row.get("downloads"), + "collection_task_type": existing_row.get("collection_task_type", "new_app"), + "last_update_interval_days": int(existing_row.get("last_update_interval_days") or 0), + "latest_task_key": latest_task.get("task_key", ""), + "latest_status": latest_task.get("status", ""), + "latest_test_time": latest_task.get("latest_time", 0.0), + "latest_worker_id": normalize_worker_id(latest_task.get("worker_id", "")), + "latest_task_detail": latest_task.get("result_detail", ""), + "latest_failure_type": latest_task.get("error_type", ""), + "unique_domain_count": len(traffic["unique_domains"]), + "unique_domain_names": traffic["unique_domains"], + "unique_second_level_domain_count": len(traffic["unique_second_level_domains"]), + "unique_second_level_domains": traffic["unique_second_level_domains"], + "droidbot_steps": int(latest_task.get("droidbot_steps") or 0), + "gui_agent_steps": int(latest_task.get("guiagent_steps") or 0), + "duration_seconds": round(float(latest_task.get("total_duration_seconds") or 0.0), 2), + "num_nodes": int(latest_task.get("num_nodes") or 0), + "num_reached_activities": int(latest_task.get("num_reached_activities") or 0), + "app_num_total_activities": int(latest_task.get("app_num_total_activities") or 0), + "total_traffic_bytes": total_traffic_bytes, + "self_traffic_bytes": self_traffic_bytes, + "server_traffic_bytes": server_traffic_bytes, + "unrecognized_traffic_bytes": unrecognized_traffic_bytes, + "self_ratio": round((self_traffic_bytes / total_traffic_bytes) * 100, 2) if total_traffic_bytes > 0 else 0.0, + "recognition_ratio": round(((self_traffic_bytes + server_traffic_bytes) / total_traffic_bytes) * 100, 2) + if total_traffic_bytes > 0 + else 0.0, + "artifact_status": artifact_status, + "model_flow_count": int(traffic.get("model_flow_count", 0)), + "model_traffic_bytes": int(traffic.get("model_traffic_bytes", 0)), + } + summary["model_eligible"] = 1 if 0 < summary["model_flow_count"] <= MODEL_TRAFFIC_THRESHOLD else 0 + restriction_status, retryability = _classify_restriction_status( + summary["latest_status"], + summary["self_ratio"], + summary["latest_failure_type"], + num_nodes=summary["num_nodes"], + download_errors=latest_task.get("download_errors"), + ) + if _should_preserve_non_retryable_reason(existing_row, restriction_status, retryability): + summary["latest_failure_type"] = str(existing_row.get("latest_failure_type") or "").strip() + if str(existing_row.get("latest_task_detail") or "").strip(): + summary["latest_task_detail"] = str(existing_row.get("latest_task_detail") or "").strip() + restriction_status = "severe_restricted" + retryability = "non_retryable" + related_magic_label_collected = ( + str(summary["latest_status"] or "").strip().lower() != "success" + and bool(summary["app_magic_label"]) + and self.repo.has_successful_related_magic_label(summary["app_magic_label"], package_name) + ) + if related_magic_label_collected: + restriction_status = "light_restricted" + retryability = "not_applicable" + summary["restriction_status"] = restriction_status + summary["retryability"] = retryability + if not related_magic_label_collected and _should_auto_reroute_to_physical( + summary["device_type"], + summary["latest_failure_type"], + restriction_status, + retryability, + ): + summary["device_type"] = "physical" + summary["task_payload"] = _build_task_payload_with_device_type( + summary.get("task_payload"), + app_name=summary["app_name"], + package_name=package_name, + country_code=summary.get("country_code") or "", + device_type="physical", + ) + collection_status = "pending" + collection_reason = f"{AUTO_PENDING_PHYSICAL_REASON}:{summary['latest_failure_type']}" + else: + collection_status, collection_reason = _derive_collection_status( + restriction_status, + retryability, + summary["latest_failure_type"], + ) + if related_magic_label_collected: + collection_reason = "关联应用已采集" + summary["collection_status"] = collection_status + summary["collection_status_reason"] = collection_reason + source_rows = self._build_source_rows(package_name, traffic_files) + self.repo.replace_package_snapshot(package_name, summary, domain_rows, component_rows, source_rows) + return summary diff --git a/analytics_replace_snapshot_v2.py b/analytics_replace_snapshot_v2.py new file mode 100644 index 0000000..a00f899 --- /dev/null +++ b/analytics_replace_snapshot_v2.py @@ -0,0 +1,206 @@ +""" +replace_package_snapshot 新实现 - 适配新架构 + +核心改动： +1. 拆分写入：collection_task（执行记录） + app_catalog（元数据） +2. 从 latest_task_key 解析 batch_tag, run_kind, attempt +3. 事务保护确保原子性 +""" + +def _parse_task_key(task_key: str): + """ + 从 task_key 解析 batch_tag, run_kind, attempt + + 格式: {batch_tag}_{package_name}_{run_kind}_{attempt} + 示例: 2026-6-15_com.test.app_ranking_1 + """ + parts = str(task_key or "").split("_") + if len(parts) >= 4: + # 最后两个是 run_kind 和 attempt + attempt = 1 + try: + attempt = int(parts[-1]) + except (ValueError, IndexError): + pass + + run_kind = parts[-2] if parts[-2] in ('ranking', 'block', 'model', 'manual') else 'ranking' + + # batch_tag 是除了最后3部分之外的（去掉package_name, run_kind, attempt） + # 但 package_name 可能包含多个_，所以简化：取第一个部分作为batch + batch_tag = parts[0] if parts else 'unknown' + + return batch_tag, run_kind, attempt + + # 降级处理 + if '_block' in task_key: + return 'unknown', 'block', 1 + return 'unknown', 'ranking', 1 + + +def replace_package_snapshot_v2( + self, + package_name: str, + summary: dict, + domain_rows: list, + component_rows: list, + source_rows: list, +): + """ + 替换应用快照（新架构版本） + + 核心逻辑： + 1. 写入 collection_task（执行记录） + 2. 更新 app_catalog（元数据） + 3. 写入 app_domain_traffic, app_traffic_component（详细数据） + """ + import time + now = time.time() + now_iso = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(now)) + + # 解析主键字段 + task_key = summary.get('latest_task_key', '') + batch_tag, run_kind, attempt = _parse_task_key(task_key) + + # 准备 collection_task 数据 + task_data = { + 'package_name': package_name, + 'batch_tag': batch_tag, + 'run_kind': run_kind, + 'attempt': attempt, + 'task_key': task_key, + 'app_name': summary.get('app_name'), + 'app_magic_label': summary.get('app_magic_label'), + 'is_new_app': 1 if summary.get('collection_task_type') == 'new_app' else 0, + 'task_status': 'completed', + 'execution_status': summary.get('latest_status'), # 'success' / 'failed' + 'worker_id': summary.get('latest_worker_id'), + 'completed_at': now_iso, + # 错误信息（从 latest_failure_type 解析） + 'error_category': None, + 'error_code': None, + 'error_reason': summary.get('collection_status_reason'), + 'error_details': summary.get('latest_task_detail'), + # 执行统计 + 'duration_seconds': summary.get('duration_seconds', 0), + 'droidbot_steps': summary.get('droidbot_steps', 0), + 'gui_agent_steps': summary.get('gui_agent_steps', 0), + 'num_nodes': summary.get('num_nodes', 0), + 'num_reached_activities': summary.get('num_reached_activities', 0), + 'app_num_total_activities': summary.get('app_num_total_activities', 0), + # 流量统计 + 'total_traffic_bytes': summary.get('total_traffic_bytes', 0), + 'self_traffic_bytes': summary.get('self_traffic_bytes', 0), + 'server_traffic_bytes': summary.get('server_traffic_bytes', 0), + 'unrecognized_traffic_bytes': summary.get('unrecognized_traffic_bytes', 0), + 'model_flow_count': summary.get('model_flow_count', 0), + 'model_traffic_bytes': summary.get('model_traffic_bytes', 0), + } + + # 解析 error_category 和 error_code + failure_type = summary.get('latest_failure_type', '') + if '/' in failure_type: + cat, code_str = failure_type.split('/', 1) + task_data['error_category'] = cat.strip() + try: + task_data['error_code'] = int(code_str.strip()) + except ValueError: + pass + + # 准备 app_catalog 数据（元数据） + catalog_data = { + 'package_name': package_name, + 'app_name': summary.get('app_name'), + 'app_magic_label': summary.get('app_magic_label'), + 'downloads': summary.get('downloads'), + 'source_order': summary.get('source_order'), + 'is_active': 1, + } + + with self._write_lock, self._connect() as connection: + # 1. 删除旧的详细数据 + connection.execute("DELETE FROM app_domain_traffic WHERE package_name = ?", (package_name,)) + connection.execute("DELETE FROM app_traffic_component WHERE package_name = ?", (package_name,)) + connection.execute("DELETE FROM analytics_source_file WHERE package_name = ?", (package_name,)) + + # 2. 插入域名流量数据 + for row in domain_rows: + connection.execute(""" + INSERT INTO app_domain_traffic ( + package_name, domain, domain_type, traffic_bytes, flow_count, + traffic_ratio, domain_traffic_ratio, organization, + matched_tp_mark, matched_pattern, match_state, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, ( + package_name, row["domain"], row["domain_type"], + int(row["traffic_bytes"]), int(row["flow_count"]), + float(row["traffic_ratio"]), float(row["domain_traffic_ratio"]), + row.get("organization", ""), row.get("matched_tp_mark", ""), + row.get("matched_pattern", ""), row.get("match_state", ""), + now, + )) + + # 3. 插入组件流量数据 + for row in component_rows: + connection.execute(""" + INSERT INTO app_traffic_component ( + package_name, component_name, component_package_names, + is_self, traffic_bytes, share_percent, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?) + """, ( + package_name, row["component_name"], + row.get("component_package_names", ""), + 1 if row.get("is_self") else 0, + int(row["traffic_bytes"]), float(row["share_percent"]), + now, + )) + + # 4. 插入执行记录到 collection_task + connection.execute(""" + INSERT INTO collection_task ( + package_name, batch_tag, run_kind, attempt, + task_key, app_name, app_magic_label, is_new_app, + task_status, execution_status, worker_id, completed_at, + error_category, error_code, error_reason, error_details, + duration_seconds, droidbot_steps, gui_agent_steps, + num_nodes, num_reached_activities, app_num_total_activities, + total_traffic_bytes, self_traffic_bytes, server_traffic_bytes, + unrecognized_traffic_bytes, model_flow_count, model_traffic_bytes, + created_at + ) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?) + ON CONFLICT(package_name, batch_tag, run_kind, attempt) DO UPDATE SET + execution_status=excluded.execution_status, + error_category=excluded.error_category, + error_code=excluded.error_code, + total_traffic_bytes=excluded.total_traffic_bytes, + self_traffic_bytes=excluded.self_traffic_bytes, + num_nodes=excluded.num_nodes, + updated_at=datetime('now','localtime') + """, ( + task_data['package_name'], task_data['batch_tag'], task_data['run_kind'], task_data['attempt'], + task_data['task_key'], task_data['app_name'], task_data['app_magic_label'], task_data['is_new_app'], + task_data['task_status'], task_data['execution_status'], task_data['worker_id'], task_data['completed_at'], + task_data['error_category'], task_data['error_code'], task_data['error_reason'], task_data['error_details'], + task_data['duration_seconds'], task_data['droidbot_steps'], task_data['gui_agent_steps'], + task_data['num_nodes'], task_data['num_reached_activities'], task_data['app_num_total_activities'], + task_data['total_traffic_bytes'], task_data['self_traffic_bytes'], task_data['server_traffic_bytes'], + task_data['unrecognized_traffic_bytes'], task_data['model_flow_count'], task_data['model_traffic_bytes'], + now_iso, + )) + + # 5. 更新 app_catalog（元数据） + connection.execute(""" + INSERT INTO app_catalog ( + package_name, app_name, app_magic_label, downloads, source_order, is_active, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, datetime('now','localtime')) + ON CONFLICT(package_name) DO UPDATE SET + app_name=COALESCE(excluded.app_name, app_catalog.app_name), + app_magic_label=COALESCE(NULLIF(excluded.app_magic_label, ''), app_catalog.app_magic_label), + downloads=COALESCE(excluded.downloads, app_catalog.downloads), + source_order=COALESCE(excluded.source_order, app_catalog.source_order), + updated_at=datetime('now','localtime') + """, ( + catalog_data['package_name'], catalog_data['app_name'], catalog_data['app_magic_label'], + catalog_data['downloads'], catalog_data['source_order'], catalog_data['is_active'], + )) + + connection.commit() diff --git a/apk_cloud/MINIO_USAGE.md b/apk_cloud/MINIO_USAGE.md new file mode 100644 index 0000000..789818f --- /dev/null +++ b/apk_cloud/MINIO_USAGE.md @@ -0,0 +1,167 @@ +# MinioStorage 使用指南 + +## 概述 + +`MinioStorage` 提供 APK 云存储功能，支持通过 `MINIO_ENABLED` 配置开关控制是否启用。 + +## 配置开关 + +在配置文件（如 `config/local.yaml` 或环境变量）中设置： + +```yaml +# 启用 Minio（默认） +MINIO_ENABLED: true +MINIO_ENDPOINT: "minio.example.com:9000" +MINIO_ACCESS_KEY: "your-access-key" +MINIO_SECRET_KEY: "your-secret-key" +MINIO_BUCKET: "apk-storage" +MINIO_SECURE: false + +# 禁用 Minio（跨团队部署或无 Minio 服务时） +MINIO_ENABLED: false +``` + +## 使用方式 + +### 方式 1：直接实例化（不推荐） + +```python +from apk_cloud.storage import MinioStorage + +try: + storage = MinioStorage() + # 使用 storage 进行操作 +except RuntimeError as e: + print(f"Minio 功能已禁用: {e}") + # 降级到直接下载模式 +``` + +**缺点**：当 `MINIO_ENABLED=False` 时会抛出异常，需要额外的异常处理。 + +### 方式 2：工厂方法（推荐） + +```python +from apk_cloud.storage import MinioStorage + +storage = MinioStorage.create() +if storage is None: + print("Minio 已禁用，使用直接下载模式") + # 执行降级逻辑，例如直接从 APK 源下载 +else: + # 使用 storage 进行 Minio 操作 + storage.push_download_queue(tasks) +``` + +**优点**：无需异常处理，代码更简洁清晰。 + +### 方式 3：预检查 + +```python +from apk_cloud.storage import MinioStorage + +if MinioStorage.is_enabled(): + storage = MinioStorage() + # 使用 Minio 功能 +else: + # 使用降级方案 + print("Minio 已禁用，跳过云存储同步") +``` + +## 降级方案 + +当 `MINIO_ENABLED=False` 时，推荐的降级方案： + +### 1. APK 直接下载模式 + +配置 `APK_DOWNLOAD_MODE='direct'`，直接从 APK 源（如 Google Play、APKPure）下载到本地，不经过 Minio 中转。 + +```python +from config import APK_DOWNLOAD_MODE + +if APK_DOWNLOAD_MODE == 'direct': + # 直接下载到 worker 本地 + apk_path = download_apk_from_source(package_name) +else: + # 从 Minio 下载 + storage = MinioStorage.create() + if storage: + apk_path = storage.download_apk(package_name, manifest, dest_dir) +``` + +### 2. 队列同步跳过 + +当 Minio 禁用时，跳过下载队列的云端同步，改为本地队列管理： + +```python +storage = MinioStorage.create() +if storage: + # 同步云端队列 + storage.push_download_queue(tasks) +else: + # 使用本地队列文件 + with open("local_queue.json", "w") as f: + json.dump(tasks, f) +``` + +### 3. 结果存储本地化 + +下载结果存储到本地文件系统，而非 Minio： + +```python +storage = MinioStorage.create() +if storage: + storage.write_download_result(package_name, result) +else: + # 本地存储 + result_path = f"results/{package_name}.json" + os.makedirs("results", exist_ok=True) + with open(result_path, "w") as f: + json.dump(result, f) +``` + +## 向后兼容性 + +- **默认行为不变**：`MINIO_ENABLED` 默认为 `True`，现有代码无需修改即可继续使用。 +- **旧代码兼容**：直接调用 `MinioStorage()` 的旧代码仍然有效（启用时），禁用时会抛出友好的错误提示。 +- **新代码建议**：新代码推荐使用 `MinioStorage.create()` 工厂方法，以便更好地处理禁用场景。 + +## 部署场景示例 + +### 场景 1：单团队部署（有 Minio 服务） + +```yaml +# config/prod.yaml +MINIO_ENABLED: true +MINIO_ENDPOINT: "minio.internal:9000" +APK_DOWNLOAD_MODE: "minio" # 通过 Minio 分发 APK +``` + +### 场景 2：跨团队部署（无 Minio 服务） + +```yaml +# config/prod.yaml +MINIO_ENABLED: false +APK_DOWNLOAD_MODE: "direct" # 每个 worker 直接下载 +``` + +### 场景 3：开发环境（可选 Minio） + +```yaml +# config/local.yaml +MINIO_ENABLED: false # 本地开发不依赖 Minio +APK_DOWNLOAD_MODE: "direct" +``` + +## 常见问题 + +### Q: 为什么不用环境变量 `MINIO_ENDPOINT` 空值判断？ + +A: 配置项可能因网络分区、配置错误等原因为空，不能作为功能开关的可靠信号。显式的 `MINIO_ENABLED` 开关语义更清晰，避免歧义。 + +### Q: 禁用 Minio 后，已有的云端数据会丢失吗？ + +A: 不会。禁用仅影响新操作，Minio 服务端的数据保持不变。重新启用后可继续访问。 + +### Q: 可以在运行时动态切换吗？ + +A: 不建议。`MINIO_ENABLED` 应在启动时通过配置文件或环境变量设置，运行时修改可能导致状态不一致。 diff --git a/apk_cloud/__init__.py b/apk_cloud/__init__.py new file mode 100644 index 0000000..4fbc458 --- /dev/null +++ b/apk_cloud/__init__.py @@ -0,0 +1,2 @@ +from apk_cloud.storage import MinioStorage +from apk_cloud.registry import ApkRegistry \ No newline at end of file diff --git a/apk_cloud/adb_client.py b/apk_cloud/adb_client.py new file mode 100644 index 0000000..27c8951 --- /dev/null +++ b/apk_cloud/adb_client.py @@ -0,0 +1,113 @@ +# -*- coding: utf-8 -*- +"""通用 ADB 命令封装，不依赖特定框架层级。""" + +from __future__ import annotations + +import logging +import subprocess +import time +from typing import List, Optional + +logger = logging.getLogger(__name__) + + +class ADBClient: + """通过 subprocess 调用 adb 的通用封装。 + + 与具体框架解耦，仅依赖系统 PATH 中的 adb。 + """ + + def __init__(self, serial: Optional[str] = None, adb_path: str = "adb"): + """ + Args: + serial: 目标设备序列号。为 ``None`` 时使用 adb 默认设备。 + adb_path: adb 可执行文件路径。 + """ + self.serial = serial + self.adb_path = adb_path + self.logger = logging.getLogger(self.__class__.__name__) + + def _build_cmd(self, args: List[str]) -> List[str]: + cmd = [self.adb_path] + if self.serial: + cmd.extend(["-s", self.serial]) + cmd.extend(args) + return cmd + + def run( + self, + args: List[str], + *, + check: bool = False, + timeout: int = 30, + retries: int = 1, + retry_delay: float = 1.0, + ) -> subprocess.CompletedProcess[str]: + """执行 adb 命令并返回 :class:`subprocess.CompletedProcess`。 + + Args: + args: adb 子命令及参数列表，如 ``["shell", "ls", "/sdcard"]`` + check: 若为 ``True``，返回码非零时抛出 :exc:`RuntimeError`。 + timeout: 单条命令超时秒数。 + retries: 最大重试次数（仅对"空输出 + 非零退出码"的瞬态错误重试）。 + retry_delay: 重试间隔秒数。 + + Returns: + 命令执行结果。 + + Raises: + RuntimeError: adb 未找到，或 ``check=True`` 且命令最终失败。 + """ + cmd = self._build_cmd(args) + last_result: Optional[subprocess.CompletedProcess[str]] = None + + for attempt in range(1, retries + 1): + try: + result = subprocess.run( + cmd, + capture_output=True, + text=True, + encoding="utf-8", + errors="replace", + check=False, + timeout=max(timeout, 1), + ) + except FileNotFoundError as exc: + raise RuntimeError("adb not found. Please ensure adb is installed and in PATH.") from exc + + last_result = result + if result.returncode == 0: + return result + + # If there is meaningful output, no need to retry — the command + # executed but reported an error (e.g. package not found). + stdout = (result.stdout or "").strip() + stderr = (result.stderr or "").strip() + if stdout or stderr: + break + + # Empty output with non-zero exit code is likely a transient ADB + # connection issue; retry after a short delay. + if attempt < retries: + self.logger.warning( + "ADB command returned %d with no output, retrying (%d/%d): %s", + result.returncode, attempt, retries, " ".join(cmd), + ) + time.sleep(retry_delay) + + assert last_result is not None + if check and last_result.returncode != 0: + error = (last_result.stderr or last_result.stdout or "adb command failed").strip() + raise RuntimeError(f"ADB command failed ({' '.join(cmd)}): {error}") + return last_result + + def shell(self, command: str | List[str], *, check: bool = False, timeout: int = 30) -> str: + """执行 ``adb shell `` 并返回 stdout 文本。""" + if isinstance(command, list): + # Pass arguments directly; adb will forward them to the device shell. + return self.run(["shell"] + command, check=check, timeout=timeout).stdout.strip() + return self.run(["shell", command], check=check, timeout=timeout).stdout.strip() + + def pull(self, remote: str, local: str, *, check: bool = True, timeout: int = 30) -> str: + """执行 ``adb pull`` 并返回 stdout 文本。""" + return self.run(["pull", "-a", remote, local], check=check, timeout=timeout).stdout.strip() diff --git a/apk_cloud/android_google_play_downloader.py b/apk_cloud/android_google_play_downloader.py new file mode 100644 index 0000000..387b077 --- /dev/null +++ b/apk_cloud/android_google_play_downloader.py @@ -0,0 +1,664 @@ +from __future__ import annotations + +import logging +import os +import re +import time +import xml.etree.ElementTree as ET +from dataclasses import dataclass +from typing import Optional + +from adb_client import ADBClient +from android_utils import get_current_package + +logger = logging.getLogger(__name__) + +_BOUNDS_RE = re.compile(r"\[(\d+),(\d+)\]\[(\d+),(\d+)\]") + +STORE_GOOGLE_PLAY = "google_play" +STORE_AURORA = "aurora" + + +@dataclass(frozen=True) +class UiNode: + text: str + desc: str + package: str + enabled: bool + clickable: bool + bounds: tuple[int, int, int, int] + + @property + def center(self) -> tuple[int, int]: + left, top, right, bottom = self.bounds + return ((left + right) // 2, (top + bottom) // 2) + + def labels(self) -> tuple[str, ...]: + values = [] + if self.text.strip(): + values.append(self.text.strip()) + if self.desc.strip(): + values.append(self.desc.strip()) + return tuple(values) + + +class GooglePlayDownloader: + GOOGLE_PLAY_PACKAGE = "com.android.vending" + + def __init__(self, serial: Optional[str] = None, store_type: str = STORE_GOOGLE_PLAY): + self._adb = ADBClient(serial=serial) + self._serial = serial + self._store_type = store_type + if store_type == STORE_AURORA: + self._store_package = "com.aurora.store" + self._store_name = "Aurora Store" + self._build_url = self._build_aurora_url + else: + self._store_package = "com.android.vending" + self._store_name = "Google Play" + self._build_url = self._build_google_play_url + + OPEN_TIMEOUT_SEC = 30 + INSTALL_TIMEOUT_SEC = 200 + POLL_INTERVAL_SEC = 2 + # 页面无有效内容的最大等待时间（Aurora 找不到包名会一直卡在加载界面） + LOADING_TIMEOUT_SEC = 30 + # Aurora Store 点击 Install 后页面无变化的超时时间 + # 正常情况点击 Install 后数秒内应出现系统安装弹窗或下载进度， + # 若 Install 按钮始终存在且无任何变化，说明 Aurora 未能触发下载 + AURORA_INSTALL_STUCK_SEC = 8 + + INSTALL_BUTTONS = ("Install", "Update") + SUCCESS_BUTTONS = ("Open", "Play") + INSTALLED_STATE_BUTTONS = ("Open", "Play", "Uninstall") + DISMISS_BUTTONS = ("Continue", "Accept", "Allow", "Got it", "Skip", "No thanks", "Not now", "OK", "Done") + + # Aurora Store 系统安装对话框的包名 + _PACKAGE_INSTALLER = "com.android.packageinstaller" + + ACCOUNT_BANNED_TEXTS = ( + "Authentication is required. You need to sign in to your Google Account.", + ) + REGION_BLOCKED_TEXTS = ( + "This item isn't available in your country.", + ) + NOT_FOUND_TEXTS = ( + "Item not found", + ) + INCOMPATIBLE_TEXTS = ( + "This app is available only for your other devices", + "Your device isn't compatible with this version.", + ) + PAGE_LOAD_FAILED_TEXTS = ( + "Try again", + ) + + AURORA_FAILURE_TEXTS = ( + "Session expired", + "Token expired", + "Too many requests", + "Rate limit", + "Not available", + "App not found", + ) + + # 系统安装失败弹窗中常见的错误文本 + # 这类弹窗如果不关闭，会遮挡后续所有应用的安装界面 + SYSTEM_INSTALL_ERROR_TEXTS = ( + "App not installed", + "Installation failed", + "Can't install", + "Install failed", + "not installed", + "Couldn't install", + ) + + def start(self, package_name, max_retry=1, target_account=None, **kwargs): + del max_retry, target_account, kwargs + + try: + # Aurora Store 使用 Compose UI，需要关闭动画才能正常执行 uiautomator dump + if self._store_type == STORE_AURORA: + self._disable_animations() + + # 清理上一次残留的弹窗/状态，避免遮挡当前任务的 UI + self._clear_screen() + + if self.is_installed(package_name): + logger.info("%s download skipped, app already installed: %s", + self._store_name, package_name) + return True, "app already installed" + + opened, message = self._open_store_page(package_name) + if not opened: + return False, message + + return self._install_from_store(package_name) + except Exception as exc: + logger.error("%s download failed for %s: %s", self._store_name, package_name, exc) + return False, str(exc) + + def stop(self): + return None + + def is_installed(self, package_name: str) -> bool: + result = self._adb.run(["shell", "pm", "path", package_name], check=False, timeout=10) + if result.returncode != 0: + return False + return any(line.strip().startswith("package:") for line in (result.stdout or "").splitlines()) + + def _install_from_store(self, package_name: str) -> tuple[bool, str]: + install_clicked = False + system_install_clicked = False + deadline = time.time() + self.INSTALL_TIMEOUT_SEC + # 页面无进展计时：如果长时间没有出现 Install 按钮或安装状态， + # 说明页面可能卡在加载界面（Aurora 找不到包名时的典型表现） + no_progress_since: float | None = None + # Aurora 点击 Install 后无反应计时： + # 点击 Install 后按钮始终存在且无系统安装弹窗，说明 Aurora 未能触发下载 + install_stuck_since: float | None = None + + while time.time() < deadline: + if self.is_installed(package_name): + return True, "install success" if install_clicked else "app already installed" + + nodes = self._dump_ui_nodes() + if nodes: + # ---- 失败检测 ---- + failure = self._detect_failure(nodes) + if failure is not None: + # 尝试关闭错误弹窗，避免遗留弹窗遮挡后续任务的 UI + self._dismiss_dialog(nodes) + return False, failure + + # ---- 系统安装失败弹窗检测 ---- + # 某些应用安装失败时，系统弹出 "App not installed" 等弹窗 + # 如果不关闭，会遮挡后续所有应用的安装界面 + sys_error = self._detect_system_install_error(nodes) + if sys_error is not None: + logger.warning("检测到系统安装失败弹窗: %s", sys_error) + self._dismiss_dialog(nodes) + return False, sys_error + + # ---- Aurora Store: 系统安装对话框确认 ---- + # 仅在点击 Aurora 的 Install 按钮后才检测系统安装弹窗， + # 避免识别到上一次下载遗留的弹窗导致 system_install_clicked 误置 + if self._store_type == STORE_AURORA and install_clicked and not system_install_clicked: + pkg_installer_node = self._find_node_by_package( + nodes, self._PACKAGE_INSTALLER, ("Install",) + ) + if pkg_installer_node is not None: + logger.info("Aurora: 检测到系统安装确认对话框，点击 INSTALL") + if self._tap_node(pkg_installer_node): + system_install_clicked = True + no_progress_since = None + time.sleep(5) + continue + + # ---- 弹窗自动关闭 ---- + if self._click_first(nodes, self.DISMISS_BUTTONS): + no_progress_since = None + time.sleep(2) + continue + + # ---- 安装完成检测（双重校验） ---- + if self._has_installed_state(nodes): + # 双重校验：UI 显示已安装，再用 pm path 确认 + # 防止 Aurora 安装过程中短暂出现 Open 按钮导致误判 + if self.is_installed(package_name): + return True, "install success" if install_clicked else "app already installed" + logger.debug("UI 显示已安装但 pm path 未确认，继续等待: %s", package_name) + + # ---- 点击 Install/Update 按钮 ---- + # 仅首次点击，避免重复点击时 continue 跳过下方的卡住检测 + install_button = self._find_first(nodes, self.INSTALL_BUTTONS) + if install_button is not None and not install_clicked: + if self._tap_node(install_button): + install_clicked = True + no_progress_since = None + logger.info("Clicked '%s' button in %s", + install_button.text or install_button.desc, + self._store_name) + time.sleep(3) + continue + + # ---- Aurora 点击 Install 后无反应检测 ---- + # 点击 Install 后按钮始终存在、无系统安装弹窗 → Aurora 未能触发下载 + if (install_clicked and install_button is not None + and not system_install_clicked): + if install_stuck_since is None: + install_stuck_since = time.time() + elif time.time() - install_stuck_since > self.AURORA_INSTALL_STUCK_SEC: + logger.warning( + "Aurora Install 点击后 %ds 无变化（按钮仍为'%s'），下载未触发: %s", + self.AURORA_INSTALL_STUCK_SEC, + install_button.text or install_button.desc, + package_name, + ) + return False, "aurora install not triggered" + else: + install_stuck_since = None + + # ---- 加载超时检测 ---- + # 判断是否有"有效进展"：有 Install/Update 按钮、安装状态、或失败信息 + has_progress = ( + install_clicked + or install_button is not None + or system_install_clicked + ) + if has_progress: + no_progress_since = None + else: + if no_progress_since is None: + no_progress_since = time.time() + elif time.time() - no_progress_since > self.LOADING_TIMEOUT_SEC: + logger.warning("%s 页面加载超时 (%ds)，包名可能不存在: %s", + self._store_name, self.LOADING_TIMEOUT_SEC, + package_name) + return False, "page loading timeout - app may not exist" + + time.sleep(self.POLL_INTERVAL_SEC) + + return False, "install timeout" + + def _open_store_page(self, package_name: str) -> tuple[bool, str]: + url = self._build_url(package_name) + # 构建 intent 命令，通过 -p 参数显式指定目标商店包名 + # 避免 market:// 被其他应用拦截（如设备上同时存在 Play Store 和 Aurora Store） + am_cmd = [ + "shell", + "am", + "start", + "-a", + "android.intent.action.VIEW", + "-d", + url, + ] + if self._store_package: + am_cmd.extend(["-p", self._store_package]) + + result = self._adb.run(am_cmd, check=False, timeout=15) + if result.returncode != 0: + error = (result.stderr or result.stdout or "jump failed").strip() + return False, error + + deadline = time.time() + self.OPEN_TIMEOUT_SEC + while time.time() < deadline: + if self._is_store_in_foreground(): + return True, "opened" + + nodes = self._dump_ui_nodes() + if nodes and ( + self._detect_failure(nodes) is not None + or self._find_first(nodes, self.INSTALL_BUTTONS) is not None + or self._has_installed_state(nodes) + ): + return True, "opened" + + time.sleep(2) + + return False, "jump timeout" + + def _is_store_in_foreground(self) -> bool: + current = get_current_package(lambda cmd: self._adb.shell(cmd, timeout=10)) + return current == self._store_package + + def _ensure_single_device(self) -> None: + if self._serial: + return + result = self._adb.run(["devices"], check=False, timeout=10) + lines = [line.strip() for line in (result.stdout or "").splitlines()] + online = [line for line in lines[1:] if line.endswith("\tdevice")] + if len(online) != 1: + raise RuntimeError(f"Expected exactly 1 online adb device, found {len(online)}.") + + def _is_google_play_in_foreground(self) -> bool: + return self._is_store_in_foreground() + + def _dump_ui_nodes(self, max_retries: int = 3) -> list[UiNode]: + """执行 uiautomator dump 并解析 UI 节点树。 + + Aurora Store 使用 Compose UI，偶尔会导致 dump 失败（ERROR: could not + get idle state），因此增加重试逻辑。 + """ + for attempt in range(max_retries): + dump_result = self._adb.run( + ["shell", "uiautomator", "dump", "/sdcard/uidump.xml"], + check=False, + timeout=15, + ) + if dump_result.returncode != 0: + dump_err = (dump_result.stdout or "").strip() + if attempt < max_retries - 1: + logger.debug("uiautomator dump 失败 (%d/%d): %s", + attempt + 1, max_retries, dump_err) + time.sleep(1) + continue + return [] + + read_result = self._adb.run( + ["exec-out", "cat", "/sdcard/uidump.xml"], + check=False, + timeout=15, + ) + xml_text = (read_result.stdout or "").strip() + if read_result.returncode != 0 or not xml_text: + if attempt < max_retries - 1: + time.sleep(1) + continue + return [] + + try: + root = ET.fromstring(xml_text) + except ET.ParseError: + if attempt < max_retries - 1: + time.sleep(1) + continue + return [] + + nodes: list[UiNode] = [] + for raw in root.iter("node"): + bounds = self._parse_bounds(raw.attrib.get("bounds", "")) + if bounds is None: + continue + nodes.append( + UiNode( + text=str(raw.attrib.get("text", "") or ""), + desc=str(raw.attrib.get("content-desc", "") or ""), + package=str(raw.attrib.get("package", "") or ""), + enabled=str(raw.attrib.get("enabled", "true")).lower() == "true", + clickable=str(raw.attrib.get("clickable", "false")).lower() == "true", + bounds=bounds, + ) + ) + return nodes + + return [] + + def _detect_failure(self, nodes: list[UiNode]) -> str | None: + labels = [label.lower() for node in nodes for label in node.labels()] + + if self._contains_text(labels, self.ACCOUNT_BANNED_TEXTS): + return "account banned" + if self._contains_text(labels, self.REGION_BLOCKED_TEXTS): + return "app not available" + if self._contains_text(labels, self.NOT_FOUND_TEXTS): + return "app not found" + if self._contains_text(labels, self.INCOMPATIBLE_TEXTS): + return "app incompatible" + if self._contains_text(labels, self.PAGE_LOAD_FAILED_TEXTS): + return "app page load failed" + if self._store_type == STORE_AURORA: + if self._contains_text(labels, self.AURORA_FAILURE_TEXTS): + return "aurora session/server error" + return None + + def _detect_system_install_error(self, nodes: list[UiNode]) -> str | None: + """检测系统安装失败弹窗。 + + 某些应用安装失败时，Android 系统会弹出 "App not installed" 等错误对话框。 + 这类弹窗如果不关闭，会遮挡后续所有应用的安装界面，导致脚本无法获取 + Install 按钮而误判为无法下载。 + """ + labels = [label.lower() for node in nodes for label in node.labels()] + if self._contains_text(labels, self.SYSTEM_INSTALL_ERROR_TEXTS): + return "system install error dialog" + return None + + def _dismiss_dialog(self, nodes: list[UiNode] | None = None) -> None: + """尝试关闭当前屏幕上的弹窗。 + + 先尝试通过 UI 节点点击常见的关闭按钮（OK/Close/Got it 等）， + 再按 Back 键作为兜底。 + """ + if nodes: + dismiss_labels = ("OK", "Close", "Got it", "Done", "Cancel") + self._click_first(nodes, dismiss_labels) + time.sleep(1) + # Back 键兜底关闭弹窗 + self._adb.run(["shell", "input", "keyevent", "KEYCODE_BACK"], + check=False, timeout=5) + time.sleep(0.5) + + def _has_installed_state(self, nodes: list[UiNode]) -> bool: + """判断应用是否已安装完成。 + + 注意：点击 Install/Update 后，Google Play 页面可能会短暂显示灰色的 + Open 按钮，此时如果同时存在 Cancel 按钮，说明安装仍在进行中。 + 必须以 Uninstall 按钮出现，或 Open/Play 且没有 Cancel/Install/Update + 作为真正完成的标志。 + """ + # Uninstall 是最可靠的安装完成标志 + if self._find_first(nodes, ("Uninstall",)) is not None: + return True + + # 如果还有 Cancel 或 Install/Update 按钮，说明安装还在进行中 + if self._find_first(nodes, ("Cancel", "Install", "Update")) is not None: + return False + + # 只有 Open/Play 且没有上述安装相关按钮，算完成 + if self._find_first(nodes, ("Open", "Play")) is not None: + return True + + return False + + def _click_first(self, nodes: list[UiNode], labels: tuple[str, ...]) -> bool: + node = self._find_first(nodes, labels) + if node is None: + return False + return self._tap_node(node) + + def _find_first(self, nodes: list[UiNode], labels: tuple[str, ...]) -> UiNode | None: + wanted = {label.strip().lower() for label in labels} + matches: list[UiNode] = [] + for node in nodes: + if not node.enabled: + continue + for label in node.labels(): + if label.strip().lower() in wanted: + matches.append(node) + break + + if not matches: + return None + + return min(matches, key=lambda node: (node.center[1], node.center[0], not node.clickable)) + + def _find_node_by_package( + self, + nodes: list[UiNode], + package: str, + labels: tuple[str, ...], + ) -> UiNode | None: + """在指定包名的节点中查找匹配 label 的控件。 + + 用于在 Aurora Store 场景中精确匹配系统安装对话框 + (com.android.packageinstaller) 的按钮，避免与商店内同名按钮混淆。 + """ + wanted = {label.strip().lower() for label in labels} + for node in nodes: + if not node.enabled: + continue + if node.package != package: + continue + for label in node.labels(): + if label.strip().lower() in wanted: + return node + return None + + def _disable_animations(self) -> None: + """关闭设备动画。 + + Aurora Store 使用 Jetpack Compose 构建 UI，Compose 的动画/过渡效果 + 会导致 uiautomator dump 报 'ERROR: could not get idle state' 而失败。 + 关闭全局动画可以有效解决此问题。 + """ + for setting in ( + "window_animation_scale", + "transition_animation_scale", + "animator_duration_scale", + ): + self._adb.run( + ["shell", "settings", "put", "global", setting, "0"], + check=False, + timeout=10, + ) + logger.debug("已关闭设备动画 (serial=%s)", self._serial) + + def _clear_screen(self) -> None: + """清理屏幕上可能残留的弹窗/遮挡。 + + 在每次下载开始前调用，确保商店界面不被之前遗留的系统弹窗 + （如安装失败提示）遮挡。 + """ + # 先回到主屏幕 + self._adb.run(["shell", "input", "keyevent", "KEYCODE_HOME"], + check=False, timeout=5) + # 关闭商店应用（清除残留状态） + self._adb.run(["shell", "am", "force-stop", self._store_package], + check=False, timeout=5) + # 关闭系统安装器（可能有残留的安装失败弹窗） + self._adb.run(["shell", "am", "force-stop", self._PACKAGE_INSTALLER], + check=False, timeout=5) + time.sleep(0.5) + + def reset_state(self) -> None: + """完成一个应用的下载/导出/卸载后重置设备状态。 + + 在 us_download_worker 中每次处理完一个包后调用，确保： + 1. 关闭残留的商店页面和弹窗 + 2. 回到主屏幕 + 3. 不影响下一个包的下载 + """ + self._adb.run(["shell", "input", "keyevent", "KEYCODE_HOME"], + check=False, timeout=5) + self._adb.run(["shell", "am", "force-stop", self._store_package], + check=False, timeout=5) + self._adb.run(["shell", "am", "force-stop", self._PACKAGE_INSTALLER], + check=False, timeout=5) + # 也关闭可能弹出的权限管理器等系统组件 + self._adb.run(["shell", "am", "force-stop", "com.google.android.permissioncontroller"], + check=False, timeout=5) + time.sleep(0.5) + logger.debug("设备状态已重置 (serial=%s)", self._serial) + + def _tap_node(self, node: UiNode) -> bool: + x, y = node.center + if x <= 0 and y <= 0: + return False + result = self._adb.run( + ["shell", "input", "tap", str(x), str(y)], + check=False, + timeout=10, + ) + return result.returncode == 0 + + @staticmethod + def _contains_text(labels: list[str], expected_values: tuple[str, ...]) -> bool: + for expected in expected_values: + expected_lower = expected.lower() + if any(expected_lower in label for label in labels): + return True + return False + + @staticmethod + def _parse_bounds(value: str) -> tuple[int, int, int, int] | None: + match = _BOUNDS_RE.fullmatch(str(value or "").strip()) + if match is None: + return None + left, top, right, bottom = (int(part) for part in match.groups()) + return left, top, right, bottom + + def export_apk(self, package_name: str, output_dir: str) -> list: + os.makedirs(output_dir, exist_ok=True) + + result = self._adb.run( + ["shell", "pm", "path", package_name], + check=False, + timeout=10, + ) + if result.returncode != 0: + logger.warning("pm path failed for %s: %s", package_name, (result.stderr or "").strip()) + return [] + + apk_paths = [] + for line in (result.stdout or "").splitlines(): + line = line.strip() + if line.startswith("package:"): + apk_paths.append(line.split(":", 1)[1].strip()) + + if not apk_paths: + logger.warning("No APK paths found for %s", package_name) + return [] + + exported = [] + for apk_path in apk_paths: + filename = os.path.basename(apk_path) + dest = os.path.join(output_dir, filename) + pull_result = self._adb.run( + ["pull", apk_path, dest], + check=False, + timeout=60, + ) + if pull_result.returncode != 0 or not os.path.isfile(dest): + logger.warning("adb pull failed for %s: %s", package_name, apk_path) + continue + exported.append(dest) + + return exported + + def get_apk_version(self, package_name: str) -> str: + result = self._adb.run( + ["shell", "dumpsys", "package", package_name], + check=False, + timeout=15, + ) + if result.returncode != 0: + return "" + for line in (result.stdout or "").splitlines(): + if "versionName=" in line: + val = line.split("versionName=", 1)[1].strip() + if val: + return val + return "" + + @staticmethod + def _build_google_play_url(package_name: str) -> str: + return f"https://play.google.com/store/apps/details?id={package_name}" + + @staticmethod + def _build_aurora_url(package_name: str) -> str: + return f"market://details?id={package_name}" + + +def main() -> int: + import argparse + + parser = argparse.ArgumentParser(description="Google Play / Aurora Store Downloader") + parser.add_argument("--package", default="com.google.android.youtube", + help="Package name to download") + parser.add_argument("--serial", default="", + help="Target device serial (optional)") + parser.add_argument("--aurora", action="store_true", + help="Use Aurora Store instead of Google Play") + args = parser.parse_args() + + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(message)s", + datefmt="%H:%M:%S", + ) + + store_type = STORE_AURORA if args.aurora else STORE_GOOGLE_PLAY + downloader = GooglePlayDownloader( + serial=args.serial or None, + store_type=store_type, + ) + success, message = downloader.start(args.package) + print(f"success={success}") + print(f"message={message}") + return 0 if success else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apk_cloud/android_utils.py b/apk_cloud/android_utils.py new file mode 100644 index 0000000..2309cea --- /dev/null +++ b/apk_cloud/android_utils.py @@ -0,0 +1,71 @@ +# -*- coding: utf-8 -*- +"""共享的 Android 工具函数，可被框架内外复用。""" + +from __future__ import annotations + +import logging +import re +import time +from typing import Callable, Optional + +logger = logging.getLogger(__name__) + + +def get_current_package( + adb_shell_func: Callable[[str], str], + max_attempts: int = 3, + retry_delay: float = 0.5, +) -> Optional[str]: + """Read the foreground package via ``dumpsys``, retrying transient failures. + + Args: + adb_shell_func: A callable that accepts a shell command string and + returns the stdout text. Example:: + + lambda cmd: adb_client.shell(cmd) + max_attempts: Maximum retry attempts for transient failures. + retry_delay: Seconds to wait between retries. + + Returns: + The foreground package name, or ``None`` if it could not be determined. + """ + for attempt in range(max_attempts): + try: + focused_app = adb_shell_func("dumpsys window | grep mFocusedApp") + if focused_app: + matches = re.findall( + r"mFocusedApp=ActivityRecord\{[a-f0-9]+\s+\S+\s+([^/]+)/", + focused_app, + ) + if matches: + return matches[-1] + + current_focus = adb_shell_func("dumpsys window | grep mCurrentFocus") + if current_focus: + matches = re.findall( + r"Window\{[a-f0-9]+\s+\S+\s+([^/]+)/", + current_focus, + ) + if matches: + return matches[-1] + + resumed_activity = adb_shell_func("dumpsys activity activities | grep mResumedActivity") + if resumed_activity: + matches = re.findall( + r"\{[a-f0-9]+\s+\S+\s+([^/]+)/", + resumed_activity, + ) + if matches: + return matches[-1] + + if attempt < max_attempts - 1: + time.sleep(retry_delay) + except Exception as exc: + if attempt < max_attempts - 1: + time.sleep(retry_delay) + continue + logger.error("Failed to get current package: %s", exc) + return None + + logger.warning("Failed to get package name after %d attempts", max_attempts) + return None diff --git a/apk_cloud/apkpure_downloader.py b/apk_cloud/apkpure_downloader.py new file mode 100644 index 0000000..6073cc3 --- /dev/null +++ b/apk_cloud/apkpure_downloader.py @@ -0,0 +1,212 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +APKPure 下载器 + +通过 APKPure 的直连 URL 下载应用安装包，作为 US Download Worker 的最终回退源。 +当 Google Play 和 Aurora Store 都失败时使用此模块。 + +参考 scripts/download.py 中的下载逻辑，适配 us_download_worker 的返回格式。 +""" + +from __future__ import annotations + +import logging +import os +import time +from datetime import datetime +from typing import Any, Dict, List, Optional, Tuple + +import requests +from requests.adapters import HTTPAdapter +from urllib3.util.retry import Retry + +logger = logging.getLogger("apkpure_downloader") + +# APKPure 直连下载 URL 模板 +_APKPURE_DOWNLOAD_URL = "https://d.apkpure.net/b/XAPK/{package_name}?version=latest" + +# 请求头 +_HEADERS = { + "User-Agent": ( + "Mozilla/5.0 (Windows NT 10.0; Win64; x64) " + "AppleWebKit/537.36 (KHTML, like Gecko) " + "Chrome/108.0.0.0 Safari/537.36" + ), +} + +# 下载超时（秒）- 对 stream 下载的每个 chunk 读取生效 +_DOWNLOAD_CONNECT_TIMEOUT = 15 +_DOWNLOAD_READ_TIMEOUT = 30 +_DOWNLOAD_PROGRESS_TIMEOUT = 200 +# 单次下载最大重试次数 +_MAX_RETRIES = 2 +# 每次下载间的延迟（秒），避免被限流 +DOWNLOAD_DELAY_SECONDS = 3 + + +def _create_session(retries: int = 3, backoff_factor: float = 0.3, + pool_connections: int = 20, pool_maxsize: int = 20) -> requests.Session: + """创建带自动重试和连接池的 requests Session。""" + session = requests.Session() + retry = Retry( + total=retries, + read=retries, + connect=retries, + backoff_factor=backoff_factor, + status_forcelist=(500, 502, 504), + ) + adapter = HTTPAdapter(max_retries=retry, pool_connections=pool_connections, + pool_maxsize=pool_maxsize) + session.mount("http://", adapter) + session.mount("https://", adapter) + return session + + +# 模块级共享 session，避免每次下载重复创建连接池 +_shared_session: Optional[requests.Session] = None +_session_lock = __import__('threading').Lock() + + +def _get_shared_session() -> requests.Session: + global _shared_session + if _shared_session is None: + with _session_lock: + if _shared_session is None: + _shared_session = _create_session() + return _shared_session + + +def download_from_apkpure( + package_name: str, + export_dir: str, +) -> Tuple[Optional[Dict[str, Any]], Optional[str]]: + """从 APKPure 下载应用安装包。 + + Args: + package_name: 应用包名 + export_dir: 导出目录（APK 文件保存到 {export_dir}/{package_name}/） + + Returns: + (result_dict, failure_reason) + - 成功时 result_dict 包含 package_name, download_date, version_name, files, local_dir + - 失败时 result_dict 为 None，failure_reason 为失败原因 + """ + download_url = _APKPURE_DOWNLOAD_URL.format(package_name=package_name) + pkg_dir = os.path.join(export_dir, package_name) + + final_filename = os.path.join(pkg_dir, f"{package_name}.xapk") + if os.path.exists(final_filename) and os.path.getsize(final_filename) > 0: + logger.info("[APKPure] 文件已存在，跳过下载: %s", package_name) + return _build_result(package_name, final_filename, export_dir), None + + for old_file in os.listdir(pkg_dir) if os.path.isdir(pkg_dir) else []: + if old_file.endswith(('.apk', '.xapk', '.apkm')): + try: + os.remove(os.path.join(pkg_dir, old_file)) + except OSError: + pass + os.makedirs(pkg_dir, exist_ok=True) + + partial_filename = final_filename + ".part" + + headers = dict(_HEADERS) + headers["Referer"] = f"https://apkpure.net/1/{package_name}" + + # 支持断点续传 + resumable_size = 0 + if os.path.exists(partial_filename): + resumable_size = os.path.getsize(partial_filename) + headers["Range"] = f"bytes={resumable_size}-" + + try: + logger.info("[APKPure] 开始下载: %s (URL: %s)", package_name, download_url) + start_time = time.time() + with _get_shared_session().get( + download_url, headers=headers, stream=True, + timeout=(_DOWNLOAD_CONNECT_TIMEOUT, _DOWNLOAD_READ_TIMEOUT), + verify=False, + ) as response: + # 检查 HTTP 状态 + if response.status_code == 404: + return None, f"APKPure 404: 应用不存在或不可用" + if response.status_code == 403: + return None, f"APKPure 403: 访问被拒绝" + response.raise_for_status() + + # 获取文件总大小 + content_length = response.headers.get("Content-Length") + total_size = int(content_length) + resumable_size if content_length else 0 + + mode = "ab" if resumable_size > 0 else "wb" + with open(partial_filename, mode) as f: + downloaded = resumable_size + for chunk in response.iter_content(chunk_size=1024 * 1024): + if not chunk: + break + f.write(chunk) + downloaded += len(chunk) + if time.time() - start_time > _DOWNLOAD_PROGRESS_TIMEOUT: + logger.warning("[APKPure] 下载总时长超限: %s (%.0fs)", package_name, + time.time() - start_time) + return None, f"下载总时长超限: {time.time() - start_time:.0f}s" + if total_size and downloaded % (10 * 1024 * 1024) == 0: + logger.debug("[APKPure] 下载进度: %s %.0f%%", package_name, + downloaded / total_size * 100) + + # 验证文件完整性 + actual_size = os.path.getsize(partial_filename) + if total_size > 0 and actual_size != total_size: + return None, f"文件下载不完整: 期望 {total_size} 字节, 实际 {actual_size} 字节" + + if actual_size == 0: + # 空文件，清理并报告失败 + try: + os.remove(partial_filename) + except OSError: + pass + return None, "下载文件为空" + + # 重命名为最终文件 + os.rename(partial_filename, final_filename) + logger.info("[APKPure] 下载完成: %s (%.1f MB)", + package_name, actual_size / (1024 * 1024)) + + return _build_result(package_name, final_filename, export_dir), None + + except requests.exceptions.ConnectionError as e: + logger.warning("[APKPure] 连接错误: %s - %s", package_name, e) + return None, f"连接错误: {e}" + except requests.exceptions.ReadTimeout as e: + logger.warning("[APKPure] 读取超时(卡住): %s - %s", package_name, e) + return None, f"下载卡住超时: {e}" + except requests.exceptions.Timeout as e: + logger.warning("[APKPure] 超时: %s - %s", package_name, e) + return None, f"下载超时: {e}" + except requests.exceptions.HTTPError as e: + logger.warning("[APKPure] HTTP错误: %s - %s", package_name, e) + return None, f"HTTP错误: {e}" + except Exception as e: + logger.warning("[APKPure] 下载异常: %s - %s", package_name, e) + return None, f"下载异常: {e}" + + +def _build_result( + package_name: str, file_path: str, export_dir: str, +) -> Dict[str, Any]: + """构建与 download_and_export 兼容的返回结果。""" + file_size = os.path.getsize(file_path) if os.path.isfile(file_path) else 0 + download_date = datetime.now().strftime("%Y-%m-%d %H:%M:%S") + + return { + "package_name": package_name, + "download_date": download_date, + "version_name": "", # APKPure 直连无法获取版本号 + "files": [ + { + "filename": os.path.basename(file_path), + "size": file_size, + } + ], + "local_dir": export_dir, + } diff --git a/apk_cloud/batch_download.py b/apk_cloud/batch_download.py new file mode 100644 index 0000000..fcc0c95 --- /dev/null +++ b/apk_cloud/batch_download.py @@ -0,0 +1,352 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +批量 Google Play 下载脚本 + +按名单逐个下载安装应用，下完一个自动换下一个。 +从 CSV 文件读取包名列表，默认读取 config/streaming.csv， +解析 target 列作为包名，并根据 enabled 列过滤。 + +用法: + python batch_download.py + python batch_download.py --export-apk /path/to/apk/export +""" + +from __future__ import annotations + +import csv +import hashlib +import json +import logging +import os +import sys +import time +from dataclasses import asdict, dataclass, field +from datetime import datetime +from pathlib import Path +from typing import List, Optional + +# 复用项目现有的下载器 +from android_google_play_downloader import GooglePlayDownloader + +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if PROJECT_ROOT not in sys.path: + sys.path.insert(0, PROJECT_ROOT) + +from config import APK_LOCAL_STORAGE_DIR + +logger = logging.getLogger("batch_download") + + +@dataclass +class Config: + """运行参数配置，直接在这里修改即可。""" + + # CSV 文件路径（相对于项目根目录或绝对路径） + csv_path: str = "config/streaming.csv" + + # 结果输出 JSON 文件路径（留空则不保存） + output_path: str = "" + + # 是否跳过已安装的应用 + skip_installed: bool = True + + # 是否只下载 enabled=1 的行 + only_enabled: bool = True + + # APK 导出目录（留空则跳过导出） + export_apk_dir: str = "" + + # 下载成功后是否删除模拟器中的APK（释放空间） + uninstall_after_export: bool = False + + +@dataclass +class DownloadResult: + package: str + success: bool + message: str + start_time: str = "" + end_time: str = "" + duration_sec: float = 0.0 + exported_apks: List[str] = field(default_factory=list) + version_name: str = "" + + +@dataclass +class ApkFileInfo: + filename: str + size: int + md5: str + + +@dataclass +class DownloadManifest: + package_name: str + download_date: str + version_name: str + version_code: str + files: List[ApkFileInfo] = field(default_factory=list) + + +@dataclass +class BatchReport: + total: int = 0 + success: int = 0 + failed: int = 0 + skipped: int = 0 + results: List[DownloadResult] = field(default_factory=list) + + +def load_packages_from_csv(path: str, only_enabled: bool = True) -> List[str]: + """从 CSV 文件加载包名列表，读取 target 列。 + + CSV 表头预期: + enabled,task_id,category,platform,target + """ + packages: List[str] = [] + csv_file = Path(path) + if not csv_file.exists(): + print(f"错误: CSV 文件不存在: {path}", file=sys.stderr) + return packages + + with csv_file.open("r", encoding="utf-8", newline="") as f: + reader = csv.DictReader(f) + for row in reader: + # 根据 enabled 过滤 + if only_enabled: + enabled = row.get("enabled", "").strip() + if enabled not in ("1", "true", "True", "TRUE", "yes", "Yes"): + continue + + target = row.get("target", "").strip() + if target: + packages.append(target) + + return packages + + +def _file_md5(filepath: str) -> str: + h = hashlib.md5() + with open(filepath, "rb") as fh: + for chunk in iter(lambda: fh.read(65536), b""): + h.update(chunk) + return h.hexdigest() + + +def build_download_manifest( + package_name: str, apk_files: List[str], version_name: str, download_date: str +) -> DownloadManifest: + file_infos = [] + for fp in apk_files: + size = os.path.getsize(fp) if os.path.isfile(fp) else 0 + md5 = _file_md5(fp) if os.path.isfile(fp) else "" + file_infos.append(ApkFileInfo( + filename=os.path.basename(fp), size=size, md5=md5, + )) + return DownloadManifest( + package_name=package_name, + download_date=download_date, + version_name=version_name, + version_code="", + files=file_infos, + ) + + +def save_manifest(manifest: DownloadManifest, output_dir: str) -> str: + os.makedirs(output_dir, exist_ok=True) + path = os.path.join(output_dir, f"{manifest.package_name}.json") + data = { + "package_name": manifest.package_name, + "download_date": manifest.download_date, + "version_name": manifest.version_name, + "version_code": manifest.version_code, + "files": [asdict(f) for f in manifest.files], + } + with open(path, "w", encoding="utf-8") as fh: + json.dump(data, fh, ensure_ascii=False, indent=2) + return path + + +def run_batch( + packages: List[str], + output_path: Optional[str] = None, + skip_installed: bool = True, + export_apk_dir: str = "", + uninstall_after_export: bool = False, +) -> BatchReport: + """执行批量下载。""" + report = BatchReport(total=len(packages)) + downloader = GooglePlayDownloader() + + print(f"\n{'=' * 60}") + print(f"批量下载开始，共 {len(packages)} 个应用") + if export_apk_dir: + print(f"APK导出目录: {export_apk_dir}") + print(f"{'=' * 60}\n") + + for idx, package in enumerate(packages, start=1): + result = DownloadResult(package=package, success=False, message="") + result.start_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S") + t0 = time.time() + + print(f"[{idx}/{len(packages)}] 正在处理: {package}") + + try: + # 如果设置了跳过已安装，先检查一下 + if skip_installed and downloader.is_installed(package): + result.success = True + result.message = "already installed, skipped" + report.skipped += 1 + print(f" -> 已安装，跳过") + else: + success, message = downloader.start(package) + result.success = success + result.message = message + if success: + report.success += 1 + print(f" -> 成功: {message}") + else: + report.failed += 1 + print(f" -> 失败: {message}") + + # 导出APK + if result.success and export_apk_dir: + version = downloader.get_apk_version(package) + result.version_name = version + # 导出到 export_apk_dir/package/ 子目录，与 registry.refresh_from_local_disk 期望的结构一致 + pkg_export_dir = os.path.join(export_apk_dir, package) + apk_files = downloader.export_apk(package, pkg_export_dir) + if apk_files: + result.exported_apks = apk_files + download_date = datetime.now().strftime("%Y-%m-%d") + manifest = build_download_manifest( + package, apk_files, version, download_date, + ) + # manifest 也保存到 package 子目录下: export_apk_dir/package/package.json + manifest_path = save_manifest(manifest, pkg_export_dir) + print(f" -> 已导出APK: {len(apk_files)}个文件") + print(f" -> 版本: {version or 'N/A'}") + print(f" -> manifest: {manifest_path}") + + if uninstall_after_export: + downloader._adb.run( + ["shell", "pm", "uninstall", package], + check=False, timeout=15, + ) + print(f" -> 已卸载") + else: + print(f" -> 警告: APK导出为空") + + except Exception as exc: + result.success = False + result.message = str(exc) + report.failed += 1 + print(f" -> 异常: {exc}") + + result.end_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S") + result.duration_sec = round(time.time() - t0, 2) + report.results.append(result) + + # 每个之间稍微停顿一下，避免操作过快 + if idx < len(packages): + print(" (等待 2 秒切换下一个...)\n") + time.sleep(2) + else: + print() + + print(f"{'=' * 60}") + print(f"批量下载结束") + print(f" 总计: {report.total}") + print(f" 成功: {report.success}") + print(f" 失败: {report.failed}") + print(f" 跳过(已安装): {report.skipped}") + if export_apk_dir: + exported_count = sum(1 for r in report.results if r.exported_apks) + print(f" 已导出APK: {exported_count}") + print(f"{'=' * 60}") + + # 保存结果到文件 + if output_path: + save_report(report, output_path) + print(f"\n详细结果已保存到: {output_path}") + + return report + + +def save_report(report: BatchReport, path: str) -> None: + """将结果保存为 JSON。""" + data = { + "summary": { + "total": report.total, + "success": report.success, + "failed": report.failed, + "skipped": report.skipped, + }, + "results": [asdict(r) for r in report.results], + } + with open(path, "w", encoding="utf-8") as f: + json.dump(data, f, ensure_ascii=False, indent=2) + + +def main() -> int: + import argparse + + parser = argparse.ArgumentParser(description="批量 Google Play 下载") + parser.add_argument("--csv", default="", help="CSV文件路径") + parser.add_argument("--output", default="", help="结果输出JSON路径") + parser.add_argument("--skip-installed", action="store_true", default=True, + help="跳过已安装的应用") + parser.add_argument("--no-skip-installed", dest="skip_installed", + action="store_false", help="不跳过已安装的应用") + parser.add_argument("--export-apk", default="", help="APK导出目录") + parser.add_argument("--uninstall-after-export", action="store_true", + help="导出后卸载模拟器中的应用") + parser.set_defaults(skip_installed=True) + args = parser.parse_args() + + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(message)s", + datefmt="%H:%M:%S", + ) + + config = Config( + csv_path=args.csv or "config/streaming.csv", + output_path=args.output, + skip_installed=args.skip_installed, + export_apk_dir=args.export_apk, + uninstall_after_export=args.uninstall_after_export, + ) + + # 从 CSV 加载包名列表 + packages = load_packages_from_csv( + config.csv_path, + only_enabled=config.only_enabled, + ) + + if not packages: + print("错误: 没有加载到任何包名，请检查 CSV 文件路径和 enabled 配置。", file=sys.stderr) + return 1 + + # 去重，保持顺序 + seen = set() + unique_packages: List[str] = [] + for p in packages: + if p not in seen: + seen.add(p) + unique_packages.append(p) + + run_batch( + packages=unique_packages, + output_path=config.output_path or None, + skip_installed=config.skip_installed, + export_apk_dir=config.export_apk_dir, + uninstall_after_export=config.uninstall_after_export, + ) + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apk_cloud/registry.py b/apk_cloud/registry.py new file mode 100644 index 0000000..a0f32b2 --- /dev/null +++ b/apk_cloud/registry.py @@ -0,0 +1,1067 @@ +#!/usr/bin/env python3 +""" +本地 APK 注册表 + +跟踪本地已缓存的 APK 信息，提供检查 APK 是否足够新的方法。 +比较逻辑：download_date >= last_updated 表示此 APK 可用。 + +包含以下子模块： +- apk_registry: 本地 APK 文件注册 (package_name -> download_date / local_dir / files) +- apk_pending_download: 等待下载的包队列 +- app_download_record: 下载结果跟踪 (per-package, per-country, per-source) + +数据存储在 monitoring.sqlite3 中。 +表 schema 由 monitoring.py 的 MonitoringRepository._initialize_schema() 管理。 +""" + +from __future__ import annotations + +import json +import os +import sqlite3 +import sys +import threading +import time +from contextlib import contextmanager +from datetime import datetime +from typing import Any, Dict, List, Optional + +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if PROJECT_ROOT not in sys.path: + sys.path.insert(0, PROJECT_ROOT) + +from config import MONITORING_DB_PATH, APK_LOCAL_STORAGE_DIR, APK_SMB_DIR + +_DEFAULT_SMB_DIR = "" + + +def _parse_date(value: Any) -> Optional[datetime]: + text = str(value or "").strip() + if not text: + return None + for fmt in ("%Y-%m-%d %H:%M:%S", "%Y-%m-%dT%H:%M:%S", "%Y-%m-%d", "%Y/%m/%d", "%Y-%m"): + try: + return datetime.strptime(text, fmt) + except (ValueError, IndexError): + continue + return None + + +def _last_updated_ts(value: Any) -> float: + dt = _parse_date(value) + return dt.timestamp() if dt else 0.0 + + +def _is_apk_stale(download_date: Any, last_updated: Any) -> bool: + return _last_updated_ts(download_date) < _last_updated_ts(last_updated) + + +def _smb_path(package_name: str, smb_root: str = "") -> str: + root = str(smb_root or APK_SMB_DIR or _DEFAULT_SMB_DIR).rstrip("\\") + if not root: + return "" + return f"{root}\\{package_name}" + + +class ApkRegistry: + def __init__(self, db_path: str = "", storage_dir: str = "", + smb_dir: str = ""): + self.db_path = db_path or MONITORING_DB_PATH + self.storage_dir = storage_dir or APK_LOCAL_STORAGE_DIR + self.smb_dir = smb_dir or APK_SMB_DIR or _DEFAULT_SMB_DIR + self._write_lock = threading.Lock() + self._ensure_tables() + + @contextmanager + def _connect(self): + """每次操作使用独立连接，避免跨线程共享连接导致游标状态串扰。""" + connection = sqlite3.connect(self.db_path, timeout=30) + connection.row_factory = sqlite3.Row + connection.execute("PRAGMA busy_timeout=30000") + connection.execute("PRAGMA synchronous=NORMAL") + connection.execute("PRAGMA journal_mode=WAL") + try: + yield connection + connection.commit() + except Exception: + try: + connection.rollback() + except Exception: + pass + raise + finally: + connection.close() + + def _ensure_tables(self) -> None: + """确保 apk_registry、apk_pending_download、app_download_record 表存在。 + + 通常由 MonitoringRepository._initialize_schema() 创建， + 此处做兜底确认，避免单独使用 ApkRegistry 时表不存在。 + """ + os.makedirs(os.path.dirname(self.db_path) or ".", exist_ok=True) + with self._connect() as conn: + conn.executescript( + """ + CREATE TABLE IF NOT EXISTS apk_registry ( + package_name TEXT PRIMARY KEY, + download_date TEXT NOT NULL, + download_time REAL NOT NULL, + local_dir TEXT, + smb_dir TEXT, + version_name TEXT, + source TEXT, + apk_files_json TEXT, + updated_at REAL NOT NULL + ); + CREATE TABLE IF NOT EXISTS apk_pending_download ( + package_name TEXT PRIMARY KEY, + created_at REAL NOT NULL + ); + CREATE TABLE IF NOT EXISTS app_download_record ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + package_name TEXT NOT NULL, + country_code TEXT NOT NULL DEFAULT '', + app_name TEXT DEFAULT '', + last_updated TEXT NOT NULL DEFAULT '', + downloads INTEGER DEFAULT 0, + + worker_play_status TEXT NOT NULL DEFAULT 'pending', + worker_play_error TEXT, + worker_play_date TEXT, + + us_play_status TEXT NOT NULL DEFAULT 'pending', + us_play_error TEXT, + us_play_date TEXT, + + us_aurora_status TEXT NOT NULL DEFAULT 'pending', + us_aurora_error TEXT, + us_aurora_date TEXT, + + apkpure_status TEXT NOT NULL DEFAULT 'pending', + apkpure_error TEXT, + apkpure_date TEXT, + + overall_status TEXT NOT NULL DEFAULT 'pending', + overall_error TEXT, + apk_version_name TEXT DEFAULT '', + + local_apk_dir TEXT DEFAULT '', + local_apk_file_count INTEGER DEFAULT 0, + local_apk_is_stale INTEGER DEFAULT 1, + + created_at REAL NOT NULL, + updated_at REAL NOT NULL, + + UNIQUE(package_name, country_code) + ); + CREATE INDEX IF NOT EXISTS idx_download_record_pkg + ON app_download_record(package_name); + CREATE INDEX IF NOT EXISTS idx_download_record_country + ON app_download_record(country_code); + CREATE INDEX IF NOT EXISTS idx_download_record_overall + ON app_download_record(overall_status); + CREATE INDEX IF NOT EXISTS idx_download_record_stale + ON app_download_record(local_apk_is_stale); + """ + ) + # 为已存在的表安全添加 apkpure 列（列已存在时忽略错误） + for col_def in ( + "apkpure_status TEXT NOT NULL DEFAULT 'pending'", + "apkpure_error TEXT", + "apkpure_date TEXT", + ): + try: + conn.execute( + f"ALTER TABLE app_download_record ADD COLUMN {col_def}" + ) + except sqlite3.OperationalError: + pass # 列已存在，忽略 + + # ── query ────────────────────────────────────────────────────── + + def get(self, package_name: str) -> Optional[Dict[str, Any]]: + with self._connect() as conn: + row = conn.execute( + "SELECT * FROM apk_registry WHERE package_name = ?", + (package_name,), + ).fetchone() + if not row: + return None + return self._row_to_entry(row) + + def get_local_apk_dir(self, package_name: str) -> Optional[str]: + entry = self.get(package_name) + if not entry: + return None + local_dir = entry.get("local_dir", "") + if local_dir and os.path.isdir(local_dir): + return local_dir + pkg_dir = os.path.join(self.storage_dir, package_name) + if os.path.isdir(pkg_dir): + return pkg_dir + return None + + def get_smb_apk_dir(self, package_name: str) -> Optional[str]: + """返回 Windows Worker 可访问的 SMB 路径。""" + entry = self.get(package_name) + if not entry: + return None + smb = entry.get("smb_dir", "") + if smb: + return smb + return _smb_path(package_name, self.smb_dir) or None + + def is_fresh_enough(self, package_name: str, last_updated: Any) -> bool: + with self._connect() as conn: + row = conn.execute( + "SELECT download_date FROM apk_registry WHERE package_name = ?", + (package_name,), + ).fetchone() + if not row: + return False + download_date = row["download_date"] or "" + if not download_date: + return False + download_ts = _last_updated_ts(download_date) + last_ts = _last_updated_ts(last_updated) + return download_ts >= last_ts + + def list_all_packages(self) -> List[str]: + with self._connect() as conn: + rows = conn.execute( + "SELECT package_name FROM apk_registry ORDER BY package_name" + ).fetchall() + return [row["package_name"] for row in rows] + + def get_download_time(self, package_name: str) -> float: + with self._connect() as conn: + row = conn.execute( + "SELECT download_time FROM apk_registry WHERE package_name = ?", + (package_name,), + ).fetchone() + if not row: + return 0.0 + return float(row["download_time"] or 0.0) + + # ── pending download tracking ────────────────────────────────── + + def mark_pending_download(self, package_name: str) -> None: + with self._write_lock, self._connect() as conn: + conn.execute( + "INSERT OR IGNORE INTO apk_pending_download (package_name, created_at) VALUES (?, ?)", + (package_name, time.time()), + ) + + def clear_pending_download(self, package_name: str) -> None: + with self._write_lock, self._connect() as conn: + conn.execute( + "DELETE FROM apk_pending_download WHERE package_name = ?", + (package_name,), + ) + + def is_pending_download(self, package_name: str) -> bool: + with self._connect() as conn: + row = conn.execute( + "SELECT 1 FROM apk_pending_download WHERE package_name = ?", + (package_name,), + ).fetchone() + return bool(row) + + def get_pending_download_packages(self) -> List[str]: + with self._connect() as conn: + rows = conn.execute( + "SELECT package_name FROM apk_pending_download" + ).fetchall() + return [row["package_name"] for row in rows] + + def get_pending_that_are_available(self) -> List[str]: + """返回已标记为 pending 且在 apk_registry 中有记录（有 local_dir）的包名。""" + with self._connect() as conn: + rows = conn.execute( + """ + SELECT p.package_name + FROM apk_pending_download p + JOIN apk_registry r ON p.package_name = r.package_name + WHERE r.local_dir IS NOT NULL AND r.local_dir != '' + """, + ).fetchall() + return [row["package_name"] for row in rows] + + def get_all_pending_downloads(self) -> List[str]: + with self._connect() as conn: + rows = conn.execute( + "SELECT DISTINCT package_name FROM apk_pending_download" + ).fetchall() + return [row["package_name"] for row in rows] + + def has_failed_with_same_version(self, package_name: str, last_updated: str) -> bool: + with self._connect() as conn: + row = conn.execute( + """ + SELECT 1 FROM app_download_record + WHERE package_name = ? + AND last_updated = ? + AND (worker_play_status = 'failed' + OR us_play_status = 'failed' + OR us_aurora_status = 'failed' + OR apkpure_status = 'failed') + LIMIT 1 + """, + (package_name, last_updated), + ).fetchone() + return bool(row) + + def verify_apk_versions(self, package_name: str) -> bool: + pkg_dir = os.path.join(self.storage_dir, package_name) + base = os.path.join(pkg_dir, "base.apk") + if not os.path.isfile(base): + return True + try: + import subprocess, re + from log_manager import logger + out = subprocess.check_output( + ["aapt", "dump", "badging", base], + stderr=subprocess.DEVNULL, timeout=15, + ).decode() + m = re.search(r"versionCode='(\d+)'", out) + base_vc = int(m.group(1)) if m else None + if base_vc is None: + return True + except Exception: + return True + + for f in os.listdir(pkg_dir): + if not f.startswith("split_") or not f.endswith(".apk"): + continue + fp = os.path.join(pkg_dir, f) + try: + out = subprocess.check_output( + ["aapt", "dump", "badging", fp], + stderr=subprocess.DEVNULL, timeout=15, + ).decode() + m = re.search(r"versionCode='(\d+)'", out) + split_vc = int(m.group(1)) if m else None + if split_vc is not None and split_vc != base_vc: + logger.warning("[APK] 版本不匹配 %s: base=%d %s=%d", + package_name, base_vc, f, split_vc) + return False + except Exception: + pass + return True + + # ── update ───────────────────────────────────────────────────── + + def mark_available(self, package_name: str, download_date: str, + local_dir: str = "", version_name: str = "", + apk_files: Optional[List[Dict[str, Any]]] = None, + source: str = "") -> None: + apk_files_json = json.dumps(apk_files or [], ensure_ascii=False) + smb_dir = _smb_path(package_name, self.smb_dir) + now = time.time() + with self._write_lock, self._connect() as conn: + conn.execute( + """ + INSERT INTO apk_registry ( + package_name, download_date, download_time, local_dir, + smb_dir, version_name, source, apk_files_json, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(package_name) DO UPDATE SET + download_date = excluded.download_date, + download_time = excluded.download_time, + local_dir = excluded.local_dir, + smb_dir = excluded.smb_dir, + version_name = excluded.version_name, + source = excluded.source, + apk_files_json = excluded.apk_files_json, + updated_at = excluded.updated_at + """, + (package_name, download_date, now, local_dir, + smb_dir, version_name, source or "", apk_files_json, now), + ) + # 一旦本地有 APK，清除 pending 状态 + conn.execute( + "DELETE FROM apk_pending_download WHERE package_name = ?", + (package_name,), + ) + # 同步刷新 app_download_record 的 stale 标志 + self._sync_stale_for_package(conn, package_name, download_date) + + def mark_unavailable(self, package_name: str) -> None: + with self._write_lock, self._connect() as conn: + conn.execute( + "DELETE FROM apk_registry WHERE package_name = ?", + (package_name,), + ) + conn.execute( + "UPDATE app_download_record SET local_apk_is_stale = 1, " + "local_apk_file_count = 0, updated_at = ? " + "WHERE package_name = ?", + (time.time(), package_name), + ) + + def _sync_stale_for_package(self, conn: sqlite3.Connection, + package_name: str, + download_date: str) -> None: + rows = conn.execute( + "SELECT country_code, last_updated FROM app_download_record " + "WHERE package_name = ?", + (package_name,), + ).fetchall() + now = time.time() + for row in rows: + is_stale = 1 if _is_apk_stale(download_date, row["last_updated"]) else 0 + conn.execute( + """ + UPDATE app_download_record + SET local_apk_is_stale = ?, + local_apk_file_count = ( + SELECT COUNT(*) + FROM json_each( + COALESCE( + (SELECT apk_files_json FROM apk_registry + WHERE package_name = app_download_record.package_name), + '[]' + ) + ) + ), + local_apk_dir = ?, + apk_version_name = COALESCE( + (SELECT version_name FROM apk_registry + WHERE package_name = app_download_record.package_name), + apk_version_name + ), + updated_at = ? + WHERE package_name = ? AND country_code = ? + """, + (is_stale, self.storage_dir, now, + package_name, row["country_code"]), + ) + + def import_from_download_result(self, package_name: str, + result: Dict[str, Any], + local_dir: str = "", + source: str = "minio") -> None: + self.mark_available( + package_name=package_name, + download_date=result.get("download_date", ""), + local_dir=local_dir or self.storage_dir, + version_name=result.get("version_name", ""), + apk_files=result.get("files", result.get("uploaded_files")), + source=source, + ) + + # ── stale detection ──────────────────────────────────────────── + + def get_stale_packages(self, catalog_last_updated: Dict[str, str]) -> List[str]: + stale = [] + for package_name, last_updated in catalog_last_updated.items(): + if not self.is_fresh_enough(package_name, last_updated): + stale.append(package_name) + return stale + + # ── download record management ───────────────────────────────── + + def ensure_download_record(self, package_name: str, country_code: str, + app_name: str = "", last_updated: str = "", + downloads: int = 0) -> None: + """懒初始化：从新目录任务信息同步一条下载记录。 + + 仅在 Redis 任务队列中需要处理此包时由中控调用。 + 如果记录已存在则更新 catalog 字段（last_updated / downloads）。 + """ + package_name = str(package_name or "").strip() + country_code = str(country_code or "").strip() + if not package_name: + return + now = time.time() + with self._write_lock, self._connect() as conn: + conn.execute( + """ + INSERT INTO app_download_record ( + package_name, country_code, app_name, last_updated, + downloads, created_at, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(package_name, country_code) DO UPDATE SET + app_name = COALESCE(excluded.app_name, app_download_record.app_name), + last_updated = excluded.last_updated, + downloads = excluded.downloads, + updated_at = excluded.updated_at + """, + (package_name, country_code, app_name or "", + last_updated or "", int(downloads or 0), now, now), + ) + + def update_source_result(self, package_name: str, country_code: str, + source: str, status: str, error: str = "", + download_date: str = "") -> None: + """更新单个下载源 (worker_play / us_play / us_aurora / apkpure) 的结果。""" + package_name = str(package_name or "").strip() + country_code = str(country_code or "").strip() + source = str(source or "").strip().lower() + status = str(status or "pending").strip() + now = time.time() + + source_col_map = { + "worker_play": ("worker_play_status", "worker_play_error", "worker_play_date"), + "us_play": ("us_play_status", "us_play_error", "us_play_date"), + "us_aurora": ("us_aurora_status", "us_aurora_error", "us_aurora_date"), + "apkpure": ("apkpure_status", "apkpure_error", "apkpure_date"), + } + if source not in source_col_map: + return + + status_col, error_col, date_col = source_col_map[source] + with self._write_lock, self._connect() as conn: + conn.execute( + f""" + UPDATE app_download_record + SET {status_col} = ?, + {error_col} = CASE WHEN ? = 'failed' THEN ? ELSE NULL END, + {date_col} = CASE WHEN ? = 'success' THEN ? ELSE NULL END, + updated_at = ? + WHERE package_name = ? AND country_code = ? + """, + (status, status, error or "", + status, download_date or "", now, + package_name, country_code), + ) + self._recompute_overall(conn, package_name, country_code) + + def _recompute_overall(self, conn: sqlite3.Connection, + package_name: str, country_code: str) -> None: + """根据四个来源的状态计算 overall_status 和 overall_error。""" + row = conn.execute( + """ + SELECT worker_play_status, us_play_status, us_aurora_status, + apkpure_status, + worker_play_error, us_play_error, us_aurora_error, + apkpure_error + FROM app_download_record + WHERE package_name = ? AND country_code = ? + """, + (package_name, country_code), + ).fetchone() + if not row: + return + + wp_s, up_s, ua_s, ap_s = ( + row["worker_play_status"] or "pending", + row["us_play_status"] or "pending", + row["us_aurora_status"] or "pending", + row["apkpure_status"] or "pending", + ) + successes = sum(1 for s in (wp_s, up_s, ua_s, ap_s) if s == "success") + failures = [(col, e) for s, e, col in ( + (wp_s, row["worker_play_error"], "worker_play"), + (up_s, row["us_play_error"], "us_play"), + (ua_s, row["us_aurora_error"], "us_aurora"), + (ap_s, row["apkpure_error"], "apkpure"), + ) if s == "failed"] + + if successes >= 1: + overall_status = "success" + overall_error = "" + elif failures: + overall_status = "all_failed" + error_parts = [] + for src, err in failures: + label = { + "worker_play": "Worker Play", + "us_play": "US Play", + "us_aurora": "US Aurora", + "apkpure": "APKPure", + }.get(src, src) + error_parts.append(f"{label}: {err or 'unknown'}") + overall_error = "; ".join(error_parts) + else: + overall_status = "pending" + overall_error = "" + + conn.execute( + """ + UPDATE app_download_record + SET overall_status = ?, + overall_error = ?, + updated_at = ? + WHERE package_name = ? AND country_code = ? + """, + (overall_status, overall_error, time.time(), + package_name, country_code), + ) + + def record_download_from_result(self, package_name: str, + country_code: str, + result: Dict[str, Any]) -> None: + """从 MinIO download result 解析 source_details 写入下载记录。""" + package_name = str(package_name or "").strip() + country_code = str(country_code or "").strip() + if not package_name: + return + + existing = self.get_download_record(package_name, country_code) + if not existing: + self.ensure_download_record( + package_name=package_name, + country_code=country_code, + app_name=result.get("app_name", ""), + last_updated=result.get("last_updated", ""), + downloads=int(result.get("downloads") or 0), + ) + + source_details = result.get("source_details") + if isinstance(source_details, dict): + for source, detail in source_details.items(): + if not isinstance(detail, dict): + continue + self.update_source_result( + package_name=package_name, + country_code=country_code, + source=source, + status=detail.get("status", "not_attempted"), + error=detail.get("error", ""), + download_date=detail.get("date", ""), + ) + + status = result.get("status", "") + if status == "ok": + self._update_local_apk_info(package_name, country_code, + result.get("files", []), + result.get("version_name", ""), + result.get("download_date", "")) + + self._refresh_stale_flag(package_name, country_code) + + def _update_local_apk_info(self, package_name: str, country_code: str, + files: List[Dict[str, Any]], + version_name: str = "", + download_date: str = "") -> None: + """更新本地 APK 文件信息和 stale 标志。""" + now = time.time() + with self._write_lock, self._connect() as conn: + conn.execute( + """ + UPDATE app_download_record + SET apk_version_name = ?, + local_apk_file_count = ?, + local_apk_dir = ?, + local_apk_is_stale = ?, + updated_at = ? + WHERE package_name = ? AND country_code = ? + """, + (version_name, len(files), self.storage_dir, + 1 if _is_apk_stale(download_date, self._get_last_updated(conn, package_name, country_code)) else 0, + now, package_name, country_code), + ) + + def _get_last_updated(self, conn: sqlite3.Connection, + package_name: str, country_code: str) -> str: + row = conn.execute( + """ + SELECT last_updated FROM app_download_record + WHERE package_name = ? AND country_code = ? + """, + (package_name, country_code), + ).fetchone() + return row["last_updated"] if row else "" + + def _refresh_stale_flag(self, package_name: str, country_code: str) -> None: + """重新从 apk_registry 检查并刷新 local_apk_is_stale。""" + entry = self.get(package_name) + download_date = entry.get("download_date", "") if entry else "" + last_updated = "" + with self._connect() as conn: + last_updated = self._get_last_updated(conn, package_name, country_code) + + is_stale = 1 if _is_apk_stale(download_date, last_updated) else 0 + now = time.time() + with self._write_lock, self._connect() as conn: + conn.execute( + """ + UPDATE app_download_record + SET local_apk_is_stale = ?, + updated_at = ? + WHERE package_name = ? AND country_code = ? + """, + (is_stale, now, package_name, country_code), + ) + + def get_stale_records(self) -> List[Dict[str, Any]]: + """返回 local_apk_is_stale = 1 的下载记录列表。""" + with self._connect() as conn: + rows = conn.execute( + """ + SELECT package_name, country_code, app_name, last_updated, downloads, + worker_play_status, us_play_status, us_aurora_status, + overall_status, overall_error, apk_version_name, + local_apk_dir, local_apk_file_count, local_apk_is_stale + FROM app_download_record + WHERE local_apk_is_stale = 1 + ORDER BY downloads DESC + """ + ).fetchall() + return [dict(row) for row in rows] + + def get_download_record(self, package_name: str, + country_code: str = "") -> Optional[Dict[str, Any]]: + """按包名和国家码查询下载记录。""" + package_name = str(package_name or "").strip() + country_code = str(country_code or "").strip() + with self._connect() as conn: + row = conn.execute( + """ + SELECT * FROM app_download_record + WHERE package_name = ? AND country_code = ? + """, + (package_name, country_code), + ).fetchone() + return dict(row) if row else None + + def list_download_records(self, country_code: str = "", + overall_status: str = "") -> List[Dict[str, Any]]: + """列出下载记录，可按国家码和/或整体状态过滤。""" + with self._connect() as conn: + conditions = [] + params: List[Any] = [] + if country_code: + conditions.append("country_code = ?") + params.append(country_code) + if overall_status: + conditions.append("overall_status = ?") + params.append(overall_status) + where = "" + if conditions: + where = "WHERE " + " AND ".join(conditions) + rows = conn.execute( + f""" + SELECT * FROM app_download_record + {where} + ORDER BY downloads DESC + """, + params, + ).fetchall() + return [dict(row) for row in rows] + + def cleanup_stale_local(self, package_name: str) -> bool: + """删除本地过期的 APK 目录。返回是否成功删除。""" + pkg_dir = os.path.join(self.storage_dir, package_name) + if not os.path.isdir(pkg_dir): + return True + import shutil + try: + shutil.rmtree(pkg_dir, ignore_errors=True) + self.mark_unavailable(package_name) + return True + except OSError: + return False + + # ── local disk refresh (enhanced) ─────────────────────────────── + + def refresh_from_local_disk(self) -> int: + if not os.path.isdir(self.storage_dir): + return 0 + found = 0 + for entry_name in os.listdir(self.storage_dir): + pkg_dir = os.path.join(self.storage_dir, entry_name) + if not os.path.isdir(pkg_dir): + continue + manifest_path = os.path.join(pkg_dir, f"{entry_name}.json") + if os.path.isfile(manifest_path): + try: + with open(manifest_path, "r", encoding="utf-8") as fh: + manifest = json.load(fh) + except (json.JSONDecodeError, IOError): + manifest = None + else: + manifest = None + + if manifest: + package_name = manifest.get("package_name", entry_name) + download_date = manifest.get("download_date", "") + version_name = manifest.get("version_name", "") + apk_files = manifest.get("files") + source = manifest.get("source", "local_disk") + else: + package_name = entry_name + apk_files_info, earliest_ctime = self._scan_apk_files(pkg_dir) + if not apk_files_info: + continue + download_date = datetime.fromtimestamp(earliest_ctime).strftime("%Y-%m-%d %H:%M:%S") + version_name = "" + apk_files = apk_files_info + source = "local_disk_ctime" + + self.mark_available( + package_name=package_name, + download_date=download_date, + local_dir=self.storage_dir, + version_name=version_name, + apk_files=apk_files, + source=source, + ) + found += 1 + return found + + @staticmethod + def _scan_apk_files(pkg_dir: str) -> tuple: + """扫描目录下的 APK 文件，返回 (文件信息列表, 最早创建时间)。 + + 用于没有 manifest JSON 时的本地磁盘恢复。 + """ + apk_extensions = (".apk", ".xapk", ".apkm") + file_infos: List[Dict[str, Any]] = [] + earliest_ctime: Optional[float] = None + for name in sorted(os.listdir(pkg_dir)): + fp = os.path.join(pkg_dir, name) + if not os.path.isfile(fp): + continue + if not name.lower().endswith(apk_extensions): + continue + try: + st = os.stat(fp) + except OSError: + continue + ctime = st.st_ctime + if earliest_ctime is None or ctime < earliest_ctime: + earliest_ctime = ctime + file_infos.append({ + "filename": name, + "size": st.st_size, + }) + return file_infos, earliest_ctime or 0.0 + + # ── internal ─────────────────────────────────────────────────── + + @staticmethod + def _row_to_entry(row: sqlite3.Row) -> Dict[str, Any]: + """将 SQLite Row 转换为与旧 JSON 格式兼容的字典。""" + apk_files_raw = row["apk_files_json"] or "[]" + try: + apk_files = json.loads(apk_files_raw) + except (json.JSONDecodeError, TypeError): + apk_files = [] + return { + "download_date": row["download_date"], + "download_time": float(row["download_time"] or 0.0), + "local_dir": row["local_dir"] or "", + "smb_dir": row["smb_dir"] or "", + "version_name": row["version_name"] or "", + "source": row["source"] or "", + "apk_files": apk_files, + } + + +def main() -> int: + registry = ApkRegistry() + + # Test basic operations + registry.mark_available( + "com.test.app1", download_date="2026-05-19", + local_dir="/tmp/app1", version_name="1.0", source="test", + ) + entry = registry.get("com.test.app1") + assert entry is not None + assert entry["download_date"] == "2026-05-19" + assert entry["source"] == "test" + assert entry["smb_dir"].startswith("\\\\") + print(" [OK] mark_available + get") + + # Test freshness check + assert registry.is_fresh_enough("com.test.app1", "2026-05-18") + assert registry.is_fresh_enough("com.test.app1", "2026-05-19") + assert not registry.is_fresh_enough("com.test.app1", "2026-05-20") + assert not registry.is_fresh_enough("com.nonexistent", "2026-01-01") + print(" [OK] is_fresh_enough") + + # Test SMB path + smb = registry.get_smb_apk_dir("com.test.app1") + assert smb is not None and "mumu_apk" in smb + print(f" [OK] get_smb_apk_dir = {smb}") + + # Test pending download + registry.mark_pending_download("com.test.need_apk") + assert registry.is_pending_download("com.test.need_apk") + assert not registry.is_pending_download("com.test.app1") + assert "com.test.need_apk" in registry.get_pending_download_packages() + print(" [OK] pending download tracking") + + # Test that mark_available clears pending + registry.mark_available( + "com.test.need_apk", download_date="2026-05-19", + local_dir="/tmp/app2", + ) + assert not registry.is_pending_download("com.test.need_apk") + print(" [OK] mark_available clears pending") + + # Test pending_that_are_available + registry.mark_pending_download("com.test.app1") + ready = registry.get_pending_that_are_available() + assert "com.test.app1" in ready # already has local_dir from mark_available + print(f" [OK] pending_that_are_available: {ready}") + + # Test stale detection + catalog = { + "com.test.app1": "2026-05-19", + "com.test.app2": "2026-05-25", + } + stale = registry.get_stale_packages(catalog) + assert "com.test.app2" in stale + assert "com.test.app1" not in stale + print(" [OK] get_stale_packages") + + # Test list_all_packages + all_pkgs = registry.list_all_packages() + assert "com.test.app1" in all_pkgs + assert "com.test.need_apk" in all_pkgs + print(f" [OK] list_all_packages: {all_pkgs}") + + # Test get_download_time + dt = registry.get_download_time("com.test.app1") + assert dt > 0 + print(f" [OK] get_download_time: {dt}") + + # ── download record tests ── + print("\n Download record tests ...") + + # Test ensure_download_record (lazy init) + registry.ensure_download_record( + package_name="com.test.dl", + country_code="US", + app_name="Test DL App", + last_updated="2026-05-20", + downloads=5000000, + ) + record = registry.get_download_record("com.test.dl", "US") + assert record is not None + assert record["app_name"] == "Test DL App" + assert record["last_updated"] == "2026-05-20" + assert record["downloads"] == 5000000 + assert record["overall_status"] == "pending" + print(" [OK] ensure_download_record + get_download_record") + + # Test ensure_download_record updates existing + registry.ensure_download_record( + package_name="com.test.dl", + country_code="US", + app_name="Test DL App Updated", + last_updated="2026-05-25", + downloads=6000000, + ) + record = registry.get_download_record("com.test.dl", "US") + assert record["app_name"] == "Test DL App Updated" + assert record["last_updated"] == "2026-05-25" + assert record["downloads"] == 6000000 + print(" [OK] ensure_download_record update") + + # Test update_source_result + registry.update_source_result("com.test.dl", "US", "us_play", "success", + download_date="2026-05-21 10:00:00") + registry.update_source_result("com.test.dl", "US", "us_aurora", "not_attempted") + record = registry.get_download_record("com.test.dl", "US") + assert record["us_play_status"] == "success" + assert record["us_play_date"] == "2026-05-21 10:00:00" + assert record["us_aurora_status"] == "not_attempted" + assert record["overall_status"] == "success" + print(" [OK] update_source_result (success)") + + # Test update_source_result (failure) + registry.update_source_result("com.test.dl", "US", "worker_play", "failed", + error="network timeout") + record = registry.get_download_record("com.test.dl", "US") + assert record["worker_play_status"] == "failed" + assert record["worker_play_error"] == "network timeout" + assert record["overall_status"] == "success" # still overall success b/c us_play succeeded + print(" [OK] update_source_result (failure, but overall still success)") + + # Test all_failed + registry.update_source_result("com.test.dl", "US", "us_play", "failed", + error="region restricted") + record = registry.get_download_record("com.test.dl", "US") + assert record["us_play_status"] == "failed" + assert record["overall_status"] == "all_failed" + assert "region restricted" in (record["overall_error"] or "") + print(" [OK] overall_status = all_failed") + + # Test record_download_from_result (MinIO result) + result = { + "status": "ok", + "download_date": "2026-05-21 11:00:00", + "version_name": "2.0", + "files": [{"filename": "base.apk", "size": 5000000}], + "source_details": { + "us_play": {"status": "success", "error": None, "date": "2026-05-21 11:00:00"}, + "us_aurora": {"status": "not_attempted", "error": None, "date": None}, + "worker_play": {"status": "not_attempted", "error": None, "date": None}, + }, + } + registry.record_download_from_result("com.test.dl", "US", result) + record = registry.get_download_record("com.test.dl", "US") + assert record["us_play_status"] == "success" + assert record["overall_status"] == "success" + assert record["apk_version_name"] == "2.0" + assert record["local_apk_file_count"] == 1 + print(" [OK] record_download_from_result") + + # Test record_download_from_result (failed) + result_failed = { + "status": "failed", + "source_details": { + "us_play": {"status": "failed", "error": "region restricted", "date": None}, + "us_aurora": {"status": "failed", "error": "app not found", "date": None}, + }, + } + registry.record_download_from_result("com.test.dl", "US", result_failed) + record = registry.get_download_record("com.test.dl", "US") + assert record["us_play_status"] == "failed" + assert record["us_aurora_status"] == "failed" + assert record["overall_status"] == "all_failed" + print(" [OK] record_download_from_result (failed)") + + # Test list_download_records + all_records = registry.list_download_records() + assert len(all_records) >= 1 + us_records = registry.list_download_records(country_code="US") + assert len(us_records) == 1 + failed_records = registry.list_download_records(overall_status="all_failed") + assert len(failed_records) == 1 + print(f" [OK] list_download_records: total={len(all_records)}, US={len(us_records)}, failed={len(failed_records)}") + + # Test get_stale_records + # register an APK older than last_updated to trigger stale + registry.mark_available( + "com.test.dl", download_date="2026-05-01", + local_dir="/tmp/app_dl", version_name="1.0", source="test", + ) + stale_records = registry.get_stale_records() + assert any(r["package_name"] == "com.test.dl" for r in stale_records), \ + f"Expected com.test.dl to be stale, got: {[r['package_name'] for r in stale_records]}" + print(f" [OK] get_stale_records: {len(stale_records)} stale packages") + + # Test cleanup_stale_local + # This removes the apk_registry entry (simulating deletion on disk) + # Since /tmp/app_dl doesn't exist, it should return True + cleaned = registry.cleanup_stale_local("com.test.dl") + print(f" [OK] cleanup_stale_local: {cleaned}") + + # Cleanup download record test data + with registry._connect() as conn: + conn.execute("DELETE FROM app_download_record WHERE package_name = ?", ("com.test.dl",)) + print(" [OK] download record cleanup") + + # Cleanup + registry.mark_unavailable("com.test.app1") + registry.clear_pending_download("com.test.app1") + registry.clear_pending_download("com.test.need_apk") + registry.mark_unavailable("com.test.need_apk") + registry.mark_unavailable("com.test.dl") + assert registry.get("com.test.app1") is None + + print("\nAll tests passed!") + print(f"\nDB path: {registry.db_path}") + print(f"Storage dir: {registry.storage_dir}") + print(f"SMB dir: {registry.smb_dir}") + print(f"Packages: {registry.list_all_packages()}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apk_cloud/storage.py b/apk_cloud/storage.py new file mode 100644 index 0000000..19c4c21 --- /dev/null +++ b/apk_cloud/storage.py @@ -0,0 +1,524 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import hashlib +import io +import json +import os +import sys +from typing import Any, Dict, List, Optional + +import urllib3 +from minio import Minio +from minio.error import S3Error + +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if PROJECT_ROOT not in sys.path: + sys.path.insert(0, PROJECT_ROOT) + +from config import ( + MINIO_ENABLED, + MINIO_ENDPOINT, + MINIO_ACCESS_KEY, + MINIO_SECRET_KEY, + MINIO_BUCKET, + MINIO_SECURE, + APK_MINIO_MAX_BYTES, +) + +DOWNLOAD_QUEUE_KEY = "download-queue/current_batch.json" +RESULTS_PREFIX = "download-results/" +APKS_PREFIX = "apks/" + + +class MinioStorage: + """ + Minio 对象存储客户端 + + 注意：此类依赖 Minio 服务，仅在 MINIO_ENABLED=True 时可用。 + 跨团队部署时，如果没有 Minio 服务，请在配置中设置 MINIO_ENABLED=False。 + + 降级方案： + - 当 MINIO_ENABLED=False 时，直接实例化会抛出 RuntimeError + - 推荐使用 create() 工厂方法，它会在禁用时返回 None + - 调用方应检查返回值并使用直接下载模式（APK_DOWNLOAD_MODE='direct'） + """ + + def __init__( + self, + endpoint: str = MINIO_ENDPOINT, + access_key: str = MINIO_ACCESS_KEY, + secret_key: str = MINIO_SECRET_KEY, + bucket: str = MINIO_BUCKET, + secure: bool = MINIO_SECURE, + ): + if not MINIO_ENABLED: + raise RuntimeError( + "Minio 功能已禁用（MINIO_ENABLED=False）。\n" + "如需使用 Minio，请在配置文件中设置 MINIO_ENABLED=True 并配置 Minio 连接信息。\n" + "跨团队部署时，推荐使用 create() 工厂方法和 APK_DOWNLOAD_MODE='direct' 直接下载模式。" + ) + + if not endpoint or not access_key or not secret_key: + raise ValueError( + "Minio 配置不完整。请在配置文件中设置：\n" + " - MINIO_ENDPOINT\n" + " - MINIO_ACCESS_KEY\n" + " - MINIO_SECRET_KEY" + ) + + self.endpoint = endpoint + self.bucket = bucket + http_client = urllib3.PoolManager( + timeout=urllib3.Timeout(connect=10, read=60), + maxsize=30, + retries=urllib3.Retry(total=3, backoff_factor=0.3), + ) + self.client = Minio(endpoint, access_key=access_key, secret_key=secret_key, + secure=secure, http_client=http_client) + self._ensure_bucket() + + @classmethod + def create( + cls, + endpoint: str = MINIO_ENDPOINT, + access_key: str = MINIO_ACCESS_KEY, + secret_key: str = MINIO_SECRET_KEY, + bucket: str = MINIO_BUCKET, + secure: bool = MINIO_SECURE, + ) -> Optional["MinioStorage"]: + """ + 工厂方法：创建 MinioStorage 实例，如果 MINIO_ENABLED=False 则返回 None。 + + 用法示例： + storage = MinioStorage.create() + if storage is None: + # 降级到直接下载模式 + print("Minio 已禁用，使用直接下载模式") + return + + 返回: + MinioStorage 实例或 None（当 MINIO_ENABLED=False 时） + """ + if not MINIO_ENABLED: + return None + return cls(endpoint, access_key, secret_key, bucket, secure) + + @staticmethod + def is_enabled() -> bool: + """ + 检查 Minio 功能是否启用。 + + 返回: + bool: True 表示 Minio 已启用并可用 + """ + return MINIO_ENABLED + + def _ensure_bucket(self) -> None: + if not self.client.bucket_exists(self.bucket): + self.client.make_bucket(self.bucket) + + # ── low-level ops ────────────────────────────────────────────── + + def upload_file(self, local_path: str, object_name: str) -> str: + with open(local_path, "rb") as fh: + fh.seek(0, os.SEEK_END) + size = fh.tell() + fh.seek(0) + result = self.client.put_object( + self.bucket, object_name, fh, length=size, + ) + return result.etag or "" + + def upload_json(self, data: Any, object_name: str) -> str: + payload = json.dumps(data, ensure_ascii=False, indent=2).encode("utf-8") + self.client.put_object(self.bucket, object_name, io.BytesIO(payload), + length=len(payload), content_type="application/json") + return object_name + + def download_file(self, object_name: str, local_path: str, + *, overwrite: bool = True) -> str: + os.makedirs(os.path.dirname(local_path) or ".", exist_ok=True) + self.client.fget_object(self.bucket, object_name, local_path) + return local_path + + def download_json(self, object_name: str, retries: int = 3) -> Optional[Any]: + import time as _time + for attempt in range(retries): + try: + response = self.client.get_object(self.bucket, object_name) + data = response.read() + response.close() + response.release_conn() + return json.loads(data.decode("utf-8")) + except S3Error as exc: + if exc.code == "NoSuchKey": + return None + if attempt < retries - 1: + _time.sleep(1 * (attempt + 1)) + else: + raise + except Exception as exc: + if attempt < retries - 1: + _time.sleep(1 * (attempt + 1)) + else: + raise + return None + + def list_objects(self, prefix: str = "") -> List[Dict[str, Any]]: + objects = self.client.list_objects(self.bucket, prefix=prefix, recursive=True) + return [ + {"name": obj.object_name, "size": obj.size or 0, + "last_modified": obj.last_modified.isoformat() if obj.last_modified else ""} + for obj in objects + ] + + def delete_object(self, object_name: str) -> None: + self.client.remove_object(self.bucket, object_name) + + def object_exists(self, object_name: str) -> bool: + try: + self.client.stat_object(self.bucket, object_name) + return True + except S3Error as exc: + if exc.code == "NoSuchKey": + return False + raise + + # ── download queue ───────────────────────────────────────────── + + def push_download_queue(self, tasks: List[Dict[str, str]]) -> None: + tasks = [t for t in tasks if self._is_valid_queue_entry(t)] + self.upload_json(tasks, DOWNLOAD_QUEUE_KEY) + + def prepend_to_download_queue(self, tasks: List[Dict[str, str]]) -> None: + existing = self.pull_download_queue() + existing = [t for t in existing if self._is_valid_queue_entry(t)] + existing_packages = {t.get("package_name", "") for t in existing} + for task in reversed(tasks): + pkg = task.get("package_name", "") + if pkg in existing_packages: + existing = [t for t in existing if t.get("package_name") != pkg] + existing.insert(0, task) + self.upload_json(existing, DOWNLOAD_QUEUE_KEY) + + def pull_download_queue(self) -> List[Dict[str, str]]: + import time as _time + for attempt in range(3): + try: + data = self.download_json(DOWNLOAD_QUEUE_KEY) + if not isinstance(data, list): + return [] + return [t for t in data if self._is_valid_queue_entry(t)] + except Exception: + if attempt == 2: + return [] + _time.sleep(1 * (attempt + 1)) + + @staticmethod + def _is_valid_queue_entry(task: Dict[str, str]) -> bool: + pkg = str(task.get("package_name") or "").strip() + if not pkg: + return False + import re + if not re.match(r'^[a-zA-Z][a-zA-Z0-9_.]{2,}$', pkg): + return False + return True + + # ── APK upload / download ────────────────────────────────────── + + def upload_apk(self, package_name: str, local_dir: str, + download_date: str, version_code: str = "", + delete_after_upload: bool = False) -> Dict[str, Any]: + manifest: Dict[str, Any] = { + "package_name": package_name, + "download_date": download_date, + "version_code": version_code, + "files": [], + } + pkg_dir = os.path.join(local_dir, package_name) + if not os.path.isdir(pkg_dir): + raise FileNotFoundError(f"APK directory not found: {pkg_dir}") + + for filename in sorted(os.listdir(pkg_dir)): + filepath = os.path.join(pkg_dir, filename) + if not os.path.isfile(filepath): + continue + if not filename.endswith((".apk", ".xapk", ".apkm")): + continue + + md5 = _file_md5(filepath) + filesize = os.path.getsize(filepath) + remote_name = f"{APKS_PREFIX}{package_name}/{download_date}_{version_code}/{filename}" + self.upload_file(filepath, remote_name) + + manifest["files"].append({ + "remote_name": remote_name, + "filename": filename, + "size": filesize, + "md5": md5, + }) + + return manifest + + def download_apk(self, package_name: str, manifest: Dict[str, Any], + dest_dir: str) -> List[str]: + local_paths: List[str] = [] + pkg_dir = os.path.join(dest_dir, package_name) + + for old_file in os.listdir(pkg_dir) if os.path.isdir(pkg_dir) else []: + if old_file.endswith(('.apk', '.xapk', '.apkm')): + try: + os.remove(os.path.join(pkg_dir, old_file)) + except OSError: + pass + os.makedirs(pkg_dir, exist_ok=True) + + # 优先使用 uploaded_files（包含 remote_name），回退到 files + file_list = manifest.get("uploaded_files") or manifest.get("files") or [] + for file_info in file_list: + remote_name = file_info["remote_name"] + filename = file_info.get("filename", os.path.basename(remote_name)) + local_path = os.path.join(pkg_dir, filename) + + self.download_file(remote_name, local_path) + + actual_md5 = _file_md5(local_path) + expected_md5 = file_info.get("md5", "") + if expected_md5 and actual_md5 != expected_md5: + os.unlink(local_path) + raise IOError( + f"MD5 mismatch for {filename}: expected {expected_md5}, got {actual_md5}" + ) + + local_paths.append(local_path) + + return local_paths + + # ── results ──────────────────────────────────────────────────── + + def write_download_result(self, package_name: str, + result: Dict[str, Any]) -> None: + key = f"{RESULTS_PREFIX}{package_name}.json" + self.upload_json(result, key) + + def read_download_result(self, package_name: str) -> Optional[Dict[str, Any]]: + key = f"{RESULTS_PREFIX}{package_name}.json" + return self.download_json(key) + + def list_download_results(self) -> Dict[str, Dict[str, Any]]: + results: Dict[str, Dict[str, Any]] = {} + for obj in self.list_objects(RESULTS_PREFIX): + if not obj["name"].endswith(".json"): + continue + package_name = obj["name"][len(RESULTS_PREFIX):-len(".json")] + if not package_name: + continue + data = self.download_json(obj["name"]) + if isinstance(data, dict): + results[package_name] = data + return results + + def delete_download_result(self, package_name: str) -> None: + key = f"{RESULTS_PREFIX}{package_name}.json" + self.delete_object(key) + + # ── old APK cleanup ──────────────────────────────────────────── + + def list_apk_versions(self, package_name: str) -> List[str]: + prefix = f"{APKS_PREFIX}{package_name}/" + objects = self.list_objects(prefix) + dirs = set() + for obj in objects: + rel = obj["name"][len(prefix):] + parts = rel.split("/") + if len(parts) >= 2: + dirs.add(parts[0]) + return sorted(dirs) + + def delete_apk_version(self, package_name: str, version_dir: str) -> None: + prefix = f"{APKS_PREFIX}{package_name}/{version_dir}/" + for obj in self.list_objects(prefix): + self.delete_object(obj["name"]) + + def cleanup_old_versions(self, package_name: str, keep: int = 1) -> int: + versions = self.list_apk_versions(package_name) + if len(versions) <= keep: + return 0 + deleted = 0 + for version in versions[:-keep]: + self.delete_apk_version(package_name, version) + deleted += 1 + return deleted + + # ── storage limit enforcement ─────────────────────────────────── + + def list_all_versions_by_age(self) -> List[Dict[str, Any]]: + """按版本聚合列出所有 APK 版本，按最早修改时间排序。 + + 每个版本条目包含该版本下所有文件的总大小和最晚 last_modified。 + 这确保了 enforce_storage_limit 能正确统计每个版本的完整大小。 + """ + # 先按文件遍历，按 (package_name, version_dir) 聚合 + version_map: Dict[str, Dict[str, Any]] = {} + for obj in self.list_objects(APKS_PREFIX): + name = obj["name"] + if not name.endswith((".apk", ".xapk", ".apkm")): + continue + rel = name[len(APKS_PREFIX):] + parts = rel.split("/", 2) + if len(parts) < 3: + continue + package_name = parts[0] + version_dir = parts[1] + key = f"{package_name}/{version_dir}" + file_size = obj.get("size", 0) + file_modified = obj.get("last_modified", "") + + if key not in version_map: + version_map[key] = { + "package_name": package_name, + "version_dir": version_dir, + "size": file_size, + "last_modified": file_modified, + } + else: + # 累加文件大小，取最晚的 last_modified + version_map[key]["size"] += file_size + if file_modified > version_map[key]["last_modified"]: + version_map[key]["last_modified"] = file_modified + + versions = list(version_map.values()) + versions.sort(key=lambda v: v.get("last_modified", "")) + return versions + + def enforce_storage_limit(self, max_bytes: int = 0) -> dict: + limit = max_bytes or APK_MINIO_MAX_BYTES + # list_all_versions_by_age 已按版本聚合，size 为该版本所有文件的总大小 + all_versions = self.list_all_versions_by_age() + current_size = sum(v["size"] for v in all_versions) + if current_size <= limit: + return {"before_bytes": current_size, "after_bytes": current_size, + "deleted_versions": 0, "deleted_bytes": 0} + if not all_versions: + return {"before_bytes": current_size, "after_bytes": current_size, + "deleted_versions": 0, "deleted_bytes": 0} + + # 清理到 limit 的 80% + target = int(limit * 0.8) + deleted_bytes = 0 + to_delete: List[Dict[str, Any]] = [] + + for ver in all_versions: + if current_size - deleted_bytes <= target: + break + to_delete.append(ver) + # ver["size"] 已是该版本所有文件的总大小，无需额外聚合 + deleted_bytes += ver["size"] + + for ver in to_delete: + self.delete_apk_version(ver["package_name"], ver["version_dir"]) + + after_versions = self.list_all_versions_by_age() + after_size = sum(v["size"] for v in after_versions) + return { + "before_bytes": current_size, + "after_bytes": after_size, + "deleted_versions": len(to_delete), + "deleted_bytes": deleted_bytes, + } + + +# ── helpers ──────────────────────────────────────────────────────── + +def _file_md5(filepath: str) -> str: + h = hashlib.md5() + with open(filepath, "rb") as fh: + for chunk in iter(lambda: fh.read(65536), b""): + h.update(chunk) + return h.hexdigest() + + +# ── test ─────────────────────────────────────────────────────────── + +def main() -> int: + print(f"MinIO enabled: {MINIO_ENABLED}") + if not MINIO_ENABLED: + print("MinIO is disabled. Set MINIO_ENABLED=True to run this test.") + return 0 + + print(f"MinIO endpoint: {MINIO_ENDPOINT}") + print(f"MinIO bucket: {MINIO_BUCKET}") + print(f"MinIO secure: {MINIO_SECURE}") + + # 使用工厂方法创建实例 + storage = MinioStorage.create() + if storage is None: + print("Failed to create MinioStorage instance (MINIO_ENABLED=False)") + return 1 + + # 1. Test push/pull download queue + test_tasks = [ + {"package_name": "com.test.app1", "last_updated": "2026-05-01", "app_name": "Test App 1"}, + {"package_name": "com.test.app2", "last_updated": "2026-05-15", "app_name": "Test App 2"}, + ] + print("\n[1] push_download_queue ...") + storage.push_download_queue(test_tasks) + print(" pushed OK") + + pulled = storage.pull_download_queue() + print(f" pulled: {len(pulled)} tasks") + for t in pulled: + print(f" {t['package_name']}") + + # 2. Test upload/download a dummy file + print("\n[2] upload / download file ...") + dummy_dir = "/tmp/apk_cloud_test/com.test.app1" + os.makedirs(dummy_dir, exist_ok=True) + dummy_path = os.path.join(dummy_dir, "base.apk") + with open(dummy_path, "w") as fh: + fh.write("fake apk content for testing\n") + manifest = storage.upload_apk( + "com.test.app1", "/tmp/apk_cloud_test", download_date="2026-05-19", version_code="1", + ) + print(f" uploaded manifest: {json.dumps(manifest, indent=2)}") + + # 3. Test write / read download result + result = { + "status": "ok", + "download_date": "2026-05-19T10:00:00", + "version_code": "1", + "files": manifest["files"], + } + storage.write_download_result("com.test.app1", result) + read_back = storage.read_download_result("com.test.app1") + print(f"\n[3] write/read result: status={read_back.get('status') if read_back else 'NONE'}") + + # 4. Test download APK + print("\n[4] download APK ...") + os.makedirs("/tmp/apk_cloud_test_dl", exist_ok=True) + local_paths = storage.download_apk("com.test.app1", manifest, "/tmp/apk_cloud_test_dl") + print(f" downloaded: {local_paths}") + + # 5. List results + print("\n[5] list_download_results ...") + all_results = storage.list_download_results() + for pkg, res in all_results.items(): + print(f" {pkg}: {res.get('status')}") + + # Cleanup test data + print("\n[6] cleanup ...") + storage.delete_download_result("com.test.app1") + storage.delete_object(DOWNLOAD_QUEUE_KEY) + storage.delete_apk_version("com.test.app1", "2026-05-19_1") + import shutil + shutil.rmtree("/tmp/apk_cloud_test", ignore_errors=True) + shutil.rmtree("/tmp/apk_cloud_test_dl", ignore_errors=True) + print(" done") + + print("\nAll tests passed!") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apk_cloud/us_download_worker.py b/apk_cloud/us_download_worker.py new file mode 100644 index 0000000..99bfdb2 --- /dev/null +++ b/apk_cloud/us_download_worker.py @@ -0,0 +1,717 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +美国下载PC定时脚本 + +每 N 分钟从 MinIO 拉取下载任务列表，批量下载 → 导出APK → 上传到 MinIO， +并写入下载结果供中控读取。下载失败的包记录详细原因，汇总为报告定期上传。 + +多设备下载策略： +- 第一次在模拟器A（默认 Google Play）上尝试下载 +- 如果失败，自动回退到模拟器B（Aurora Store）再次尝试 +- 如果仍然失败，最终回退到 APKPure HTTP 直连下载 + +用法: + python us_download_worker.py + python us_download_worker.py --export-dir D:\\mumu_apk_export --interval 600 + python us_download_worker.py --once +""" + +from __future__ import annotations + +import json +import logging +import os +import shutil +import subprocess +import sys +import threading +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from datetime import datetime +from typing import Any, Dict, List, Optional + +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if PROJECT_ROOT not in sys.path: + sys.path.insert(0, PROJECT_ROOT) + +from config import ( + APK_EMULATOR_CLEANUP_ENABLED, + APK_US_DEVICE_A_SERIAL, + APK_US_DEVICE_A_VM_INDEX, + APK_US_DEVICE_B_SERIAL, + APK_US_DEVICE_B_VM_INDEX, + APK_US_EXPORT_DIR, + MUMU_MANAGER_PATH, +) +from apk_cloud.storage import MinioStorage +from apk_cloud.android_google_play_downloader import ( + GooglePlayDownloader, + STORE_GOOGLE_PLAY, + STORE_AURORA, +) +from apk_cloud.apkpure_downloader import download_from_apkpure, DOWNLOAD_DELAY_SECONDS + +logger = logging.getLogger("us_download_worker") + +_DEFAULT_POLL_SECONDS = 60 +_EMULATOR_CLEANUP_INTERVAL = 3600 +_FAILURE_REPORT_KEY = "download-reports/failures.json" +_HISTORY_FILE = "download_history.json" + +_MAX_PIPELINE_WORKERS = 8 +_APKPURE_MAX_CONCURRENT = 3 + +_device_a_lock = threading.Lock() +_device_b_lock = threading.Lock() +_apkpure_semaphore = threading.Semaphore(_APKPURE_MAX_CONCURRENT) +_history_lock = threading.Lock() + +US_DEVICE_A_SERIAL = APK_US_DEVICE_A_SERIAL +US_DEVICE_B_SERIAL = APK_US_DEVICE_B_SERIAL + +# 设备串号 → MuMu VM 编号 +_US_SERIAL_VM_MAP = { + US_DEVICE_A_SERIAL: APK_US_DEVICE_A_VM_INDEX, + US_DEVICE_B_SERIAL: APK_US_DEVICE_B_VM_INDEX, +} +_US_EMULATOR_RESTART_LOCK = threading.Lock() +_US_EMULATOR_RESTART_COOLDOWN: Dict[str, float] = {} + + +def _is_adb_device_offline_error(error_text: str) -> bool: + if not error_text: + return False + lowered = error_text.lower() + return "device" in lowered and "not found" in lowered + + +def _restart_emulator_if_needed(serial: str, reason: str) -> bool: + vm_index = _US_SERIAL_VM_MAP.get(serial) + if not vm_index: + return False + if not _is_adb_device_offline_error(reason): + return False + now = time.time() + with _US_EMULATOR_RESTART_LOCK: + last = _US_EMULATOR_RESTART_COOLDOWN.get(serial, 0) + if now - last < 120: + logger.info("[Emulator] %s 在冷却期内(%.0fs前刚重启)，跳过本次重启", + serial, now - last) + return False + _US_EMULATOR_RESTART_COOLDOWN[serial] = now + + logger.info("[Emulator] 检测到 ADB 断链 (%s)，正在重启模拟器 VM-%d ...", serial, vm_index) + manager = MUMU_MANAGER_PATH.strip() + if not os.path.isfile(manager): + logger.warning("[Emulator] MuMuManager.exe 不存在: %s", manager) + return False + + try: + subprocess.run( + [manager, "control", "-v", str(vm_index), "shutdown"], + capture_output=True, timeout=30, + ) + time.sleep(10) + subprocess.run( + [manager, "control", "-v", str(vm_index), "restart"], + capture_output=True, timeout=30, + ) + time.sleep(30) + subprocess.run(["adb", "disconnect", serial.split(":")[0]], + capture_output=True, timeout=10) + time.sleep(5) + subprocess.run(["adb", "connect", serial], + capture_output=True, timeout=15) + time.sleep(10) + logger.info("[Emulator] VM-%d 重启完成，ADB 已重连", vm_index) + return True + except subprocess.TimeoutExpired as exc: + logger.warning("[Emulator] VM-%d 重启命令超时: %s", vm_index, exc) + return False + except Exception as exc: + logger.warning("[Emulator] VM-%d 重启失败: %s", vm_index, exc) + return False + + +# ── Download History ──────────────────────────────────────────────── + +def _load_history(export_dir: str) -> Dict[str, Dict[str, str]]: + path = os.path.join(export_dir, _HISTORY_FILE) + if not os.path.isfile(path): + return {} + try: + with open(path, "r", encoding="utf-8") as f: + data = json.load(f) + return data if isinstance(data, dict) else {} + except (json.JSONDecodeError, IOError): + return {} + + +def _save_history(export_dir: str, history: Dict[str, Dict[str, str]]) -> None: + path = os.path.join(export_dir, _HISTORY_FILE) + tmp = path + ".tmp" + with open(tmp, "w", encoding="utf-8") as f: + json.dump(history, f, ensure_ascii=False, indent=2) + os.replace(tmp, path) + + +def _is_history_fresh(history: Dict[str, Dict[str, str]], + package_name: str, last_updated: str) -> bool: + entry = history.get(package_name) + if not entry: + return False + hist_updated = str(entry.get("last_updated") or "") + hist_date = str(entry.get("download_date") or "") + if not hist_date: + return False + if hist_updated != last_updated: + return False + return True + + +def _record_history(history: Dict[str, Dict[str, str]], + package_name: str, last_updated: str, + download_date: str = "", version_name: str = "") -> None: + history[package_name] = { + "last_updated": last_updated, + "download_date": download_date or datetime.now().strftime("%Y-%m-%d %H:%M:%S"), + "version_name": version_name or "", + } + + +def download_and_export(package_name: str, export_dir: str, + downloader: GooglePlayDownloader) -> tuple: + """返回 (result_dict_or_None, failure_reason_or_None)。""" + try: + if downloader.is_installed(package_name): + logger.info("Already installed, skip download: %s", package_name) + return _export_only(package_name, export_dir, downloader), None + + logger.info("Downloading: %s (device=%s)", package_name, + downloader._serial or "default") + success, message = downloader.start(package_name) + if not success: + reason = str(message or "unknown download error").strip() + logger.warning("Download failed: %s reason=%s", package_name, reason) + return None, reason + + logger.info("Download succeeded: %s", package_name) + result = _export_only(package_name, export_dir, downloader) + if result is None: + return None, "export failed after download" + return result, None + finally: + # 每次下载任务完成后（无论成功失败）重置设备状态 + # 清理商店页面、系统弹窗等残留，确保不影响下一个包的下载 + try: + downloader.reset_state() + except Exception as exc: + logger.warning("reset_state failed: %s", exc) + + +def download_with_fallback(package_name: str, export_dir: str) -> tuple: + """多设备下载 pipeline：先尝试模拟器A (Google Play)，失败则回退模拟器B (Aurora Store)， + 最终回退到 APKPure HTTP 直连下载。 + + 各阶段通过 device lock 协调并发： + - 一个任务释放 Device A → 下一个任务立即抢占 Device A + - Device A / Device B / APKPure 可同时处理互不干扰 + + 返回 (result_dict_or_None, failure_reason_or_None, source_details_dict)， + 其中 result_dict 包含 source_details 记录每个来源的尝试结果。 + """ + source_details = {} + + # ── 第一级：Google Play (Device A) ── + logger.info("[Device A / Google Play] Trying: %s", package_name) + with _device_a_lock: + device_a = GooglePlayDownloader(serial=US_DEVICE_A_SERIAL, store_type=STORE_GOOGLE_PLAY) + result, reason = download_and_export(package_name, export_dir, device_a) + if result: + source_details["us_play"] = { + "status": "success", + "error": None, + "date": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), + } + source_details["us_aurora"] = { + "status": "not_attempted", + "error": None, + "date": None, + } + source_details["apkpure"] = { + "status": "not_attempted", + "error": None, + "date": None, + } + result["source_details"] = source_details + return result, None, source_details + + source_details["us_play"] = { + "status": "failed", + "error": reason, + "date": None, + } + logger.warning("[Device A] Failed (%s), falling back to Device B / Aurora Store: %s", + reason, package_name) + + if _restart_emulator_if_needed(US_DEVICE_A_SERIAL, reason): + logger.info("[Device A] 模拟器已重启，重试一次: %s", package_name) + with _device_a_lock: + device_a = GooglePlayDownloader(serial=US_DEVICE_A_SERIAL, store_type=STORE_GOOGLE_PLAY) + result, reason = download_and_export(package_name, export_dir, device_a) + if result: + source_details["us_play"] = { + "status": "success", + "error": None, + "date": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), + } + source_details["us_aurora"] = {"status": "not_attempted", "error": None, "date": None} + source_details["apkpure"] = {"status": "not_attempted", "error": None, "date": None} + result["source_details"] = source_details + return result, None, source_details + + # ── 第二级：Aurora Store (Device B) ── + with _device_b_lock: + device_b = GooglePlayDownloader(serial=US_DEVICE_B_SERIAL, store_type=STORE_AURORA) + result2, reason2 = download_and_export(package_name, export_dir, device_b) + if result2: + source_details["us_aurora"] = { + "status": "success", + "error": None, + "date": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), + } + source_details["apkpure"] = { + "status": "not_attempted", + "error": None, + "date": None, + } + result2["source_details"] = source_details + return result2, None, source_details + + source_details["us_aurora"] = { + "status": "failed", + "error": reason2, + "date": None, + } + logger.warning( + "[Device A+B] Both failed (A=%s, B=%s), falling back to APKPure: %s", + reason, reason2, package_name, + ) + + if _restart_emulator_if_needed(US_DEVICE_B_SERIAL, reason2): + logger.info("[Device B] 模拟器已重启，重试一次: %s", package_name) + with _device_b_lock: + device_b = GooglePlayDownloader(serial=US_DEVICE_B_SERIAL, store_type=STORE_AURORA) + result2, reason2 = download_and_export(package_name, export_dir, device_b) + if result2: + source_details["us_aurora"] = { + "status": "success", + "error": None, + "date": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), + } + source_details["apkpure"] = {"status": "not_attempted", "error": None, "date": None} + result2["source_details"] = source_details + return result2, None, source_details + + # ── 第三级（最终回退）：APKPure HTTP 直连下载 ── + with _apkpure_semaphore: + result3, reason3 = download_from_apkpure(package_name, export_dir) + if result3: + source_details["apkpure"] = { + "status": "success", + "error": None, + "date": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), + } + result3["source_details"] = source_details + return result3, None, source_details + + source_details["apkpure"] = { + "status": "failed", + "error": reason3, + "date": None, + } + return ( + None, + f"all sources failed: Play={reason}, Aurora={reason2}, APKPure={reason3}", + source_details, + ) + + +def _export_only(package_name: str, export_dir: str, + downloader: GooglePlayDownloader, + uninstall_after: bool = True) -> Optional[Dict[str, Any]]: + version = downloader.get_apk_version(package_name) + pkg_export_dir = os.path.join(export_dir, package_name) + + for old_file in os.listdir(pkg_export_dir) if os.path.isdir(pkg_export_dir) else []: + if old_file.endswith(('.apk', '.xapk', '.apkm')): + try: + os.remove(os.path.join(pkg_export_dir, old_file)) + except OSError: + pass + os.makedirs(pkg_export_dir, exist_ok=True) + + apk_files = downloader.export_apk(package_name, pkg_export_dir) + if not apk_files: + logger.warning("APK export empty: %s", package_name) + return None + + if uninstall_after and APK_EMULATOR_CLEANUP_ENABLED: + try: + downloader._adb.run( + ["shell", "am", "force-stop", package_name], + check=False, timeout=10, + ) + except Exception: + pass + try: + downloader._adb.run( + ["shell", "pm", "clear", package_name], + check=False, timeout=15, + ) + except Exception: + pass + try: + downloader._adb.run( + ["shell", "pm", "uninstall", "--user", "0", package_name], + check=True, timeout=30, + ) + logger.info("Uninstalled after export: %s", package_name) + except Exception as exc: + logger.warning("Uninstall failed: %s: %s", package_name, exc) + try: + downloader._adb.run( + ["shell", "pm", "uninstall", package_name], + check=False, timeout=15, + ) + except Exception: + pass + + download_date = datetime.now().strftime("%Y-%m-%d %H:%M:%S") + file_infos = [] + for fp in apk_files: + size = os.path.getsize(fp) if os.path.isfile(fp) else 0 + file_infos.append({ + "filename": os.path.basename(fp), + "size": size, + }) + + return { + "package_name": package_name, + "download_date": download_date, + "version_name": version, + "files": file_infos, + "local_dir": export_dir, + } + + +def _process_single_task(task: Dict[str, str], export_dir: str, + storage: MinioStorage, + history: Dict[str, Dict[str, str]]) -> Dict[str, Any]: + """处理单个下载任务：pipeline 下载 → 上传 MinIO → 记录结果。 + + 返回 {"success": bool, "failures": [...]} 供 run_cycle 聚合统计。 + """ + package_name = str(task.get("package_name") or "").strip() + last_updated = str(task.get("last_updated") or "").strip() + if not package_name: + return {"success": False, "failures": []} + + try: + result, fail_reason, source_details = download_with_fallback(package_name, export_dir) + except Exception as exc: + logger.exception("download_with_fallback panicked: %s", package_name) + fail_info = { + "package_name": package_name, + "reason": f"pipeline error: {exc}", + "failed_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), + } + fallback = { + "status": "failed", + "package_name": package_name, + "download_date": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), + "reason": str(exc), + "source_details": {}, + } + try: + storage.write_download_result(package_name, fallback) + except Exception: + pass + with _history_lock: + _record_history(history, package_name, last_updated, download_date="none", version_name="") + _save_history(export_dir, history) + return {"success": False, "failures": [fail_info]} + + if result: + result["status"] = "ok" + pkg_dir = os.path.join(export_dir, package_name) + try: + version_code = result.get("version_name", "").replace(".", "_") or "v0" + manifest = storage.upload_apk( + package_name, export_dir, + download_date=result["download_date"][:10], + version_code=version_code, + ) + result["uploaded_files"] = manifest.get("files", []) + result["download_date"] = manifest.get("download_date", result.get("download_date", "")) + result["version_code"] = manifest.get("version_code", "") + storage.cleanup_old_versions(package_name) + + with _history_lock: + _record_history(history, package_name, last_updated, + download_date=result.get("download_date", ""), + version_name=result.get("version_name", "")) + _save_history(export_dir, history) + except Exception as exc: + logger.error("Upload failed for %s: %s", package_name, exc) + result["status"] = "download_ok_upload_failed" + result["upload_error"] = str(exc) + try: + storage.write_download_result(package_name, result) + except Exception: + pass + return { + "success": False, + "failures": [{ + "package_name": package_name, + "reason": f"upload failed: {exc}", + "failed_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), + }], + } + finally: + if os.path.isdir(pkg_dir): + shutil.rmtree(pkg_dir, ignore_errors=True) + logger.info("[Upload] 已清理本地文件: %s", pkg_dir) + + storage.write_download_result(package_name, result) + return {"success": True, "failures": []} + else: + reason = fail_reason or "download_or_export_failed" + result = { + "status": "failed", + "package_name": package_name, + "download_date": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), + "reason": reason, + "source_details": source_details, + } + storage.write_download_result(package_name, result) + with _history_lock: + _record_history(history, package_name, last_updated, download_date="none", version_name="") + _save_history(export_dir, history) + return { + "success": False, + "failures": [{ + "package_name": package_name, + "reason": reason, + "failed_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), + }], + } + + +def run_cycle(storage: MinioStorage, export_dir: str, + history: Dict[str, Dict[str, str]] = None) -> Dict[str, Any]: + history = history if history is not None else {} + + tasks = storage.pull_download_queue() + if not tasks: + logger.info("No tasks in download queue, sleeping") + return {"processed": 0, "uploaded": 0, "failures": []} + + original_count = len(tasks) + + filtered = [] + skipped = 0 + for task in tasks: + package_name = str(task.get("package_name") or "").strip() + last_updated = str(task.get("last_updated") or "").strip() + if not package_name or ' ' in package_name or '&&' in package_name: + logger.warning("Skipping invalid task entry: %s", task) + continue + if not task.get("force") and _is_history_fresh(history, package_name, last_updated): + skipped += 1 + continue + filtered.append(task) + + if skipped: + logger.info("History filter: %d/%d tasks skipped (already downloaded)", + skipped, original_count) + + tasks = filtered + logger.info("%d tasks to process (max_workers=%d, priority ordered)", len(tasks), _MAX_PIPELINE_WORKERS) + + processed = 0 + uploaded = 0 + failures: List[Dict[str, Any]] = [] + + if not tasks: + return {"processed": 0, "uploaded": 0, "failures": []} + + max_workers = min(_MAX_PIPELINE_WORKERS, max(len(tasks), 1)) + with ThreadPoolExecutor(max_workers=max_workers) as executor: + future_map = {} + for task in tasks: + future = executor.submit(_process_single_task, task, export_dir, storage, history) + future_map[future] = task + for future in as_completed(future_map): + task_result = future.result() + if task_result.get("success"): + processed += 1 + uploaded += 1 + failures.extend(task_result.get("failures", [])) + + return {"processed": processed, "uploaded": uploaded, "failures": failures} + + +def upload_failure_report(storage: MinioStorage, failures: List[Dict[str, Any]]) -> None: + if not failures: + return + try: + existing = storage.download_json(_FAILURE_REPORT_KEY) or {} + except Exception: + existing = {} + + existing.setdefault("failures", []) + existing_failures = existing.get("failures", []) + if not isinstance(existing_failures, list): + existing_failures = [] + existing["failures"] = existing_failures + + existing_failures.extend(failures) + existing["generated_at"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S") + existing["worker"] = "us-downloader" + existing["total_failures"] = len(existing_failures) + + try: + storage.upload_json(existing, _FAILURE_REPORT_KEY) + logger.info("Failure report uploaded: %d total failures", len(existing_failures)) + except Exception as exc: + logger.warning("Failed to upload failure report: %s", exc) + + +def cleanup_emulator_third_party_apps(downloader: GooglePlayDownloader) -> int: + result = downloader._adb.run( + ["shell", "pm", "list", "packages", "-3"], + check=False, timeout=30, + ) + if result.returncode != 0: + return 0 + packages = [] + for line in (result.stdout or "").splitlines(): + line = line.strip() + if line.startswith("package:"): + pkg = line.split(":", 1)[1].strip() + if pkg and pkg not in ("com.android.vending", "com.aurora.store"): + packages.append(pkg) + + uninstalled = 0 + for pkg in packages: + try: + r = downloader._adb.run( + ["shell", "pm", "uninstall", pkg], + check=False, timeout=15, + ) + if r.returncode == 0: + uninstalled += 1 + logger.info("Emulator cleanup uninstalled: %s", pkg) + except Exception: + pass + return uninstalled + + +def main_loop(export_dir: str, poll_seconds: int = _DEFAULT_POLL_SECONDS): + storage = MinioStorage() + downloader_for_cleanup = GooglePlayDownloader(serial=US_DEVICE_B_SERIAL, store_type=STORE_AURORA) + last_emulator_cleanup = time.time() + accumulated_failures: List[Dict[str, Any]] = [] + + history = _load_history(export_dir) + history_count = len(history) + logger.info("US Download Worker started, poll_interval=%ds, export_dir=%s", + poll_seconds, export_dir) + logger.info("Device A (Google Play): %s, Device B (Aurora): %s", + US_DEVICE_A_SERIAL, US_DEVICE_B_SERIAL) + logger.info("MinIO: %s / %s", storage.endpoint, storage.bucket) + logger.info("History: %d previously downloaded packages", history_count) + + while True: + start_time = time.time() + try: + stats = run_cycle(storage, export_dir, history) + elapsed = time.time() - start_time + logger.info("Cycle done: processed=%d uploaded=%d failed=%d elapsed=%.1fs", + stats["processed"], stats["uploaded"], + len(stats.get("failures", [])), elapsed) + + failures = stats.get("failures", []) + if failures: + accumulated_failures.extend(failures) + upload_failure_report(storage, accumulated_failures) + + if stats["uploaded"] > 0: + try: + limit_result = storage.enforce_storage_limit() + if limit_result.get("deleted_versions"): + logger.info( + "MinIO storage limit: %.1fGB -> %.1fGB (deleted %d versions)", + limit_result["before_bytes"] / (1024 ** 3), + limit_result["after_bytes"] / (1024 ** 3), + limit_result["deleted_versions"], + ) + except Exception as exc: + logger.warning("MinIO storage limit check failed: %s", exc) + + if APK_EMULATOR_CLEANUP_ENABLED and (time.time() - last_emulator_cleanup >= _EMULATOR_CLEANUP_INTERVAL): + try: + count = cleanup_emulator_third_party_apps(downloader_for_cleanup) + if count > 0: + logger.info("Emulator cleanup: uninstalled %d apps", count) + except Exception as exc: + logger.warning("Emulator cleanup failed: %s", exc) + last_emulator_cleanup = time.time() + except Exception: + logger.exception("Cycle error") + + sleep_left = max(0, poll_seconds - (time.time() - start_time)) + if sleep_left > 0: + logger.info("Sleeping %.0fs until next poll...", sleep_left) + time.sleep(sleep_left) + + +def main() -> int: + import argparse + + parser = argparse.ArgumentParser(description="US Download Worker") + parser.add_argument("--export-dir", default=APK_US_EXPORT_DIR, + help=f"APK导出目录 (default: {APK_US_EXPORT_DIR})") + parser.add_argument("--interval", type=int, default=_DEFAULT_POLL_SECONDS, + help=f"轮询间隔秒数 (default: {_DEFAULT_POLL_SECONDS})") + parser.add_argument("--once", action="store_true", + help="只执行一次循环后退出") + args = parser.parse_args() + + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(message)s", + datefmt="%H:%M:%S", + ) + + export_dir = str(args.export_dir or "").strip() or APK_US_EXPORT_DIR + os.makedirs(export_dir, exist_ok=True) + + if args.once: + storage = MinioStorage() + history = _load_history(export_dir) + stats = run_cycle(storage, export_dir, history) + failures = stats.get("failures", []) + if failures: + upload_failure_report(storage, failures) + print(f"done: processed={stats['processed']} uploaded={stats['uploaded']} failed={len(failures)}") + return 0 + + main_loop(export_dir, poll_seconds=args.interval) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/batch_rewrite_plan.py b/batch_rewrite_plan.py new file mode 100644 index 0000000..1d78687 --- /dev/null +++ b/batch_rewrite_plan.py @@ -0,0 +1,92 @@ +#!/usr/bin/env python3 +""" +批量替换 analytics.py 中所有 app_collect_summary 查询为新架构 + +策略： +1. 简单的 SELECT * FROM app_collect_summary → LEFT JOIN app_catalog + v_collection_latest +2. UPDATE app_collect_summary → UPDATE app_catalog +3. INSERT INTO app_collect_summary → INSERT INTO app_catalog + collection_task +""" + +import re + +def generate_new_query_patterns(): + """生成新架构的查询模式""" + + # 基础 JOIN 模式 + BASE_JOIN = """ + FROM app_catalog ac + LEFT JOIN v_collection_latest vl ON ac.package_name = vl.package_name + """ + + # 字段映射 + FIELD_MAPPING = { + 'catalog_active': 'ac.is_active', + 'incremental_batch_tag': 'ac.batch_tags', + 'latest_status': 'vl.execution_status', + 'latest_test_time': "CAST(strftime('%s', vl.completed_at) AS REAL)", + 'latest_worker_id': 'vl.worker_id', + 'latest_task_key': 'vl.task_key', + 'latest_failure_type': "COALESCE(vl.error_category, '') || '/' || COALESCE(vl.error_code, 0)", + 'restriction_status': """ + CASE + WHEN vl.execution_status = 'success' THEN + CASE WHEN COALESCE(vl.num_nodes, 0) >= 5 AND COALESCE(vl.self_ratio, 0) >= 30 + THEN 'success' + ELSE 'light_restricted' END + ELSE 'severe_restricted' + END""", + 'retryability': """ + CASE + WHEN vl.execution_status = 'success' THEN 'retryable' + WHEN vl.error_category = 'DOWNLOAD_ERROR' + AND vl.error_code IN (1,5,6,10,403,404) THEN 'non_retryable' + ELSE 'retryable' + END""", + 'collection_status': """ + CASE + WHEN vl.execution_status = 'success' + AND COALESCE(vl.num_nodes, 0) >= 5 + AND COALESCE(vl.self_ratio, 0) >= 30 THEN 'qualified' + ELSE 'restricted' + END""", + } + + return BASE_JOIN, FIELD_MAPPING + + +# 需要改写的函数列表（按优先级） +FUNCTIONS_TO_REWRITE = [ + # 批量查询 + 'list_pending_collection_tasks', # SELECT FROM app_collect_summary WHERE catalog_active=1 + 'list_qualified_apps', # SELECT FROM app_collect_summary WHERE restriction_status='success' + 'list_all_catalog_apps', # SELECT FROM app_collect_summary + 'list_model_eligible_apps', # SELECT FROM app_collect_summary WHERE model_eligible=1 + 'has_successful_related_magic_label', # SELECT FROM app_collect_summary WHERE app_magic_label=? + + # 目录管理 + 'upsert_catalog_entries', # INSERT INTO app_collect_summary + 'replace_catalog_entries', # INSERT + DELETE FROM app_collect_summary + 'update_collection_statuses', # UPDATE app_collect_summary SET collection_status + 'upsert_high_priority_app', # INSERT INTO app_collect_summary + + # 批次管理 + 'clear_incremental_batch_tags', # UPDATE app_collect_summary SET incremental_batch_tag + 'set_incremental_batch_tag_for_packages', # UPDATE app_collect_summary + + # 复杂聚合 + 'get_overview', # SELECT COUNT(*) FROM app_collect_summary GROUP BY + 'list_apps', # SELECT FROM app_collect_summary LIMIT OFFSET + 'get_app_detail', # SELECT FROM app_collect_summary WHERE package_name=? + + # 清理操作 + 'delete_snapshots_not_in', # DELETE FROM app_collect_summary + 'clear_all_snapshots', # UPDATE app_collect_summary SET latest_*=NULL +] + +print("=" * 60) +print("需要改写的函数：") +for i, func in enumerate(FUNCTIONS_TO_REWRITE, 1): + print(f"{i:2d}. {func}") +print(f"\n总计：{len(FUNCTIONS_TO_REWRITE)} 个函数") +print("=" * 60) diff --git a/calculate/count.py b/calculate/count.py new file mode 100644 index 0000000..3d465a7 --- /dev/null +++ b/calculate/count.py @@ -0,0 +1,462 @@ +import os +import csv +from collections import defaultdict + + +def _decode_mixed_csv_line(raw_line): + for encoding in ('utf-8-sig', 'utf-8', 'gb18030'): + try: + return raw_line.decode(encoding) + except UnicodeDecodeError: + continue + return raw_line.decode('gb18030', errors='replace') + + +def _read_csv_lines(csv_path): + with open(csv_path, 'rb') as f: + return [_decode_mixed_csv_line(line) for line in f] + + +def _csv_value(row, index): + if index >= len(row): + return '' + return str(row[index] or '').strip() + + +def _latest_task_key(task_time, seq): + task_time = str(task_time or '').strip() + return (1 if task_time else 0, task_time, seq) + + +def _load_latest_task_rows(task_paths, package_idx, time_idx, allowed_packages=None): + rows_by_package = {} + best_keys = {} + seq = 0 + + for task_path in task_paths: + if not os.path.exists(task_path): + continue + + try: + reader = csv.reader(_read_csv_lines(task_path)) + next(reader, None) + for row in reader: + package_name = _csv_value(row, package_idx) + if allowed_packages is not None and package_name not in allowed_packages: + continue + if not package_name: + continue + + seq += 1 + key = _latest_task_key(_csv_value(row, time_idx), seq) + if package_name in best_keys and key <= best_keys[package_name]: + continue + + best_keys[package_name] = key + rows_by_package[package_name] = row + except Exception as e: + print(f"读取 {task_path} 时出错: {e}") + + return rows_by_package + + +def _normalize_success_detail(detail_info): + detail_info = str(detail_info or '').strip() + if '|' not in detail_info: + return detail_info + + parts = detail_info.split('|', 1) + if len(parts) > 1: + return parts[1].strip() + return detail_info + + +def extract_second_level_domain(domain): + if not domain: + return "" + domain = str(domain).lower().replace('http://', '').replace('https://', '').strip().strip('"\'') + parts = domain.split('.') + if len(parts) <= 2: + return domain + return "*." + ".".join(parts[-2:]) + +def load_package_mapping(current_dir): + package_to_app = {} + app_to_package = {} + package_list_path = os.path.join(current_dir, 'package_list.csv') + if not os.path.exists(package_list_path): + print("未找到 package_list.csv，将使用包名作为应用名。") + return None, None + + try: + with open(package_list_path, 'r', encoding='utf-8-sig', newline='') as f: + reader = csv.DictReader(f) + for row in reader: + app_name = (row.get('app_name') or '').strip() + package_name = (row.get('package_name') or '').strip() + if package_name and app_name: + package_to_app[package_name] = app_name + app_to_package[app_name] = package_name + except Exception as e: + print(f"读取 package_list.csv 时出错: {e}") + return None, None + + return package_to_app, app_to_package + +def load_traffic_profile(current_dir, allowed_packages=None): + self_ratio_map = {} + traffic_profile_path = os.path.join(current_dir, 'app_traffic_profile_20260324.csv') + if not os.path.exists(traffic_profile_path): + print("未找到 app_traffic_profile_20260324.csv，将不使用 self_ratio。") + return self_ratio_map + + try: + with open(traffic_profile_path, 'r', encoding='utf-8-sig') as f: + reader = csv.DictReader(f) + for row in reader: + package_name = row.get('app_id', '').strip() + if allowed_packages is not None and package_name not in allowed_packages: + continue + if package_name: + try: + self_ratio = float(row.get('self_ratio', 0) or 0) + self_ratio_map[package_name] = self_ratio + except ValueError: + pass + except Exception as e: + print(f"读取 app_traffic_profile_20260224.csv 时出错: {e}") + + return self_ratio_map + +def load_success_tasks(current_dir, allowed_packages=None): + import glob + success_packages = set() + success_info = {} + success_time = {} + success_worker = {} + + success_task_files = sorted(glob.glob(os.path.join(current_dir, 'success_tasks*.csv'))) + + if not success_task_files: + return success_packages, success_info, success_time, success_worker + + for package_name, row in _load_latest_task_rows( + success_task_files, + package_idx=3, + time_idx=0, + allowed_packages=allowed_packages, + ).items(): + success_packages.add(package_name) + success_info[package_name] = _normalize_success_detail(_csv_value(row, 6)) + success_time[package_name] = _csv_value(row, 0) + success_worker[package_name] = _csv_value(row, 4) + + return success_packages, success_info, success_time, success_worker + + +def load_non_success_tasks(current_dir, allowed_packages=None): + failure_info = {} + failure_type = {} + failure_time = {} + failure_worker = {} + non_success_rows = {} + failed_rows = _load_latest_task_rows( + [os.path.join(current_dir, 'failed_tasks.csv')], + package_idx=3, + time_idx=0, + allowed_packages=allowed_packages, + ) + retry_rows = _load_latest_task_rows( + [os.path.join(current_dir, 'retry_tasks.csv')], + package_idx=3, + time_idx=0, + allowed_packages=allowed_packages, + ) + + for package_name, row in failed_rows.items(): + non_success_rows[package_name] = { + 'detail': _csv_value(row, 7) if len(row) >= 8 else _csv_value(row, 6), + 'failure_type': _csv_value(row, 5), + 'time': _csv_value(row, 0), + 'worker': _csv_value(row, 4), + } + + for package_name, row in retry_rows.items(): + if package_name in non_success_rows: + continue + non_success_rows[package_name] = { + 'detail': _csv_value(row, 6), + 'failure_type': _csv_value(row, 5), + 'time': _csv_value(row, 0), + 'worker': _csv_value(row, 4), + } + + for package_name, row in non_success_rows.items(): + failure_info[package_name] = row['detail'] + failure_type[package_name] = row['failure_type'] + failure_time[package_name] = row['time'] + failure_worker[package_name] = row['worker'] + + return failure_info, failure_type, failure_time, failure_worker + +def parse_utg_js(js_path, target_package=None): + import re + import json + result = { + 'num_nodes': 0, + 'num_reached_activities': 0, + 'app_num_total_activities': 0 + } + + try: + with open(js_path, 'r', encoding='utf-8-sig') as f: + content = f.read() + + json_match = re.search(r'var\s+utg\s*=\s*(\{.*\})\s*;?\s*$', content, re.DOTALL) + if json_match: + try: + utg_data = json.loads(json_match.group(1)) + + app_package = utg_data.get('app_package', '') + filter_package = target_package if target_package else app_package + + nodes = utg_data.get('nodes', []) + if filter_package: + filtered_count = sum(1 for node in nodes if node.get('package') == filter_package) + result['num_nodes'] = filtered_count + else: + result['num_nodes'] = len(nodes) + + num_reached_match = utg_data.get('num_reached_activities', 0) + result['num_reached_activities'] = num_reached_match if isinstance(num_reached_match, int) else 0 + + app_total_match = utg_data.get('app_num_total_activities', 0) + result['app_num_total_activities'] = app_total_match if isinstance(app_total_match, int) else 0 + + except json.JSONDecodeError as e: + print(f"解析 UTG JSON 时出错: {e}") + num_nodes_match = re.search(r'"num_nodes"\s*:\s*(\d+)', content) + if num_nodes_match: + result['num_nodes'] = int(num_nodes_match.group(1)) + + num_reached_match = re.search(r'"num_reached_activities"\s*:\s*(\d+)', content) + if num_reached_match: + result['num_reached_activities'] = int(num_reached_match.group(1)) + + app_total_match = re.search(r'"app_num_total_activities"\s*:\s*(\d+)', content) + if app_total_match: + result['app_num_total_activities'] = int(app_total_match.group(1)) + else: + num_nodes_match = re.search(r'"num_nodes"\s*:\s*(\d+)', content) + if num_nodes_match: + result['num_nodes'] = int(num_nodes_match.group(1)) + + num_reached_match = re.search(r'"num_reached_activities"\s*:\s*(\d+)', content) + if num_reached_match: + result['num_reached_activities'] = int(num_reached_match.group(1)) + + app_total_match = re.search(r'"app_num_total_activities"\s*:\s*(\d+)', content) + if app_total_match: + result['app_num_total_activities'] = int(app_total_match.group(1)) + + except Exception as e: + print(f"解析 UTG 文件 {js_path} 时出错: {e}") + + return result + +def process_all_data(): + app_data = defaultdict(lambda: { + 'domains': set(), + 'second_level_domains': set(), + 'droidbot_steps': 0, + 'gui_agent_steps': 0, + 'duration_seconds': 0, + 'num_nodes': 0, + 'num_reached_activities': 0, + 'app_num_total_activities': 0 + }) + + current_dir = os.getcwd() + package_to_app, app_to_package = load_package_mapping(current_dir) + allowed_packages = set(package_to_app.keys()) if package_to_app is not None else None + + batch_result_files = [] + + print("开始遍历文件收集数据...") + for root, dirs, files in os.walk(current_dir): + for file in files: + if file.endswith('.txt'): + file_path = os.path.join(root, file) + try: + with open(file_path, 'r', encoding='utf-8-sig') as f: + for line in f: + parts = line.strip().split(',') + + if len(parts) >= 3: + package_name = parts[0].strip() + app_name = parts[1].strip() + domain_name = parts[2].strip() + if allowed_packages is not None and package_name not in allowed_packages: + continue + + if domain_name != "model_data:" and package_name: + app_data[package_name]['domains'].add(domain_name) + second_level_domain = extract_second_level_domain(domain_name) + if second_level_domain: + app_data[package_name]['second_level_domains'].add(second_level_domain) + + except Exception as e: + print(f"处理文件 {file_path} 时出错: {e}") + + if file.endswith('.csv') and file.startswith('batch_result'): + file_path = os.path.join(root, file) + batch_result_files.append(file_path) + + if file.endswith('_utg.js'): + file_path = os.path.join(root, file) + package_name = file.replace('_utg.js', '') + if allowed_packages is not None and package_name not in allowed_packages: + continue + utg_result = parse_utg_js(file_path, package_name) + app_data[package_name]['num_nodes'] = max( + app_data[package_name]['num_nodes'], utg_result['num_nodes'] + ) + app_data[package_name]['num_reached_activities'] = max( + app_data[package_name]['num_reached_activities'], utg_result['num_reached_activities'] + ) + app_data[package_name]['app_num_total_activities'] = max( + app_data[package_name]['app_num_total_activities'], utg_result['app_num_total_activities'] + ) + + batch_result_files.sort(key=lambda x: os.path.getmtime(x)) + + print(f"找到 {len(batch_result_files)} 个 batch_result 文件，按修改时间排序处理...") + for file_path in batch_result_files: + try: + with open(file_path, 'r', encoding='utf-8-sig') as f: + reader = csv.DictReader(f) + for row in reader: + package_name = row.get('package_name', '').strip() + if allowed_packages is not None and package_name not in allowed_packages: + continue + if not package_name: + continue + + try: + droidbot_steps = int(row.get('droidbot_steps', 0) or 0) + app_data[package_name]['droidbot_steps'] = droidbot_steps + except ValueError: + pass + + try: + guiagent_steps = int(row.get('guiagent_steps', 0) or 0) + app_data[package_name]['gui_agent_steps'] = guiagent_steps + except ValueError: + pass + + try: + duration_seconds = float(row.get('duration_seconds', 0) or 0) + app_data[package_name]['duration_seconds'] = duration_seconds + except ValueError: + pass + + except Exception as e: + print(f"处理文件 {file_path} 时出错: {e}") + + found_packages = set(app_data.keys()) + print(f"从文件中共找到 {len(found_packages)} 个包的数据。") + + success_packages, success_info, success_time, success_worker = load_success_tasks( + current_dir, + allowed_packages=allowed_packages, + ) + failure_info, failure_type, failure_time, failure_worker = load_non_success_tasks( + current_dir, + allowed_packages=allowed_packages, + ) + self_ratio_map = load_traffic_profile(current_dir, allowed_packages=allowed_packages) + + if success_packages: + print(f"已加载 {len(success_packages)} 个成功任务。") + if failure_info: + print(f"已加载 {len(failure_info)} 个非成功任务。") + if self_ratio_map: + print(f"已加载 {len(self_ratio_map)} 个应用的流量配置。") + + final_packages = set() + + if package_to_app is not None: + final_packages = found_packages.intersection(allowed_packages) + print(f"根据 package_list.csv 筛选后，剩余 {len(final_packages)} 个包。") + print(f"过滤掉了 {len(found_packages) - len(final_packages)} 个不在列表中的包。") + else: + final_packages = found_packages + print("未进行筛选，使用所有找到的包。") + + output_file = 'app_domain_summary.csv' + headers = ['app_name', 'package_name', 'test_success', 'test_time', 'test_host', 'self_ratio', 'unique_domain_count', 'unique_domain_names', + 'unique_second_level_domain_count', 'unique_second_level_domains', + 'droidbot_steps', 'gui_agent_steps', 'duration_seconds', + 'num_nodes', 'num_reached_activities', 'app_num_total_activities', + 'task_detail', 'failure_type'] + + try: + with open(output_file, 'w', newline='', encoding='utf-8-sig') as csvfile: + writer = csv.writer(csvfile) + writer.writerow(headers) + + count = 0 + for package_name in sorted(final_packages): + info = app_data[package_name] + + if package_to_app: + app_name = package_to_app.get(package_name, package_name) + else: + app_name = package_name + + domains = sorted(list(info['domains'])) + second_level_domains = sorted(list(info['second_level_domains'])) + + is_success = package_name in success_packages + if is_success: + test_time = success_time.get(package_name, '') + test_host = success_worker.get(package_name, '') + task_detail = success_info.get(package_name, '') + package_failure_type = '' + else: + test_time = failure_time.get(package_name, '') + test_host = failure_worker.get(package_name, '') + task_detail = failure_info.get(package_name, '') + package_failure_type = failure_type.get(package_name, '') + self_ratio = self_ratio_map.get(package_name, '') + + writer.writerow([ + app_name, + package_name, + is_success, + test_time, + test_host, + self_ratio, + len(domains), + ", ".join(domains), + len(second_level_domains), + ", ".join(second_level_domains), + info['droidbot_steps'], + info['gui_agent_steps'], + info['duration_seconds'], + info['num_nodes'], + info['num_reached_activities'], + info['app_num_total_activities'], + task_detail, + package_failure_type + ]) + count += 1 + + print(f"处理完成！最终统计了 {count} 个 App。") + print(f"结果已保存至: {output_file}") + + except Exception as e: + print(f"写入 CSV 时出错: {e}") + +if __name__ == "__main__": + process_all_data() diff --git a/calculate/dpi_database_test.py b/calculate/dpi_database_test.py new file mode 100644 index 0000000..59d53b0 --- /dev/null +++ b/calculate/dpi_database_test.py @@ -0,0 +1,412 @@ +import ahocorasick +import pandas as pd +import os +import sys +from pathlib import Path +from collections import defaultdict + +# ================= 配置区域 ================= +TARGET_DATE = "20260331" +LIB_DATE = "20260224" +# 路径配置 +SCRIPT_DIR = Path(__file__).resolve().parent +PROJECT_ROOT = SCRIPT_DIR.parent + +APP_LIST_PATH = PROJECT_ROOT / f"Lib/{LIB_DATE}/input/TPDPI_app_list_{LIB_DATE}.csv" +URL_LIST_PATH = PROJECT_ROOT / f"Lib/{LIB_DATE}/input/TPDPI_url_lib_{LIB_DATE}.csv" +TRAFFIC_DATA_PATH = PROJECT_ROOT / f"TrafficData/traffic_summary/{TARGET_DATE}/traffic_summary.csv" + +OUTPUT_DIR = PROJECT_ROOT / f"TrafficData/traffic_summary/{TARGET_DATE}/result" + +# 输出文件定义 +OUTPUT_APP_PROFILE = OUTPUT_DIR / f"app_traffic_profile_{TARGET_DATE}.csv" +OUTPUT_UNMATCHED_FILE = OUTPUT_DIR / f"unmatched_domains_{TARGET_DATE}.csv" +OUTPUT_DETAIL_FILE = OUTPUT_DIR / f"traffic_detail_log_{TARGET_DATE}.csv" + + +# =========================================== + +def format_bytes(size): + """将字节转换为易读格式 (B, KB, MB, GB)，保留2位小数""" + power = 2 ** 10 + n = 0 + power_labels = {0: 'B', 1: 'KB', 2: 'MB', 3: 'GB', 4: 'TB'} + while size > power: + size /= power + n += 1 + if n == 0: + return f"{int(size)} B" + return f"{size:.2f} {power_labels[n]}" + + +class DpiDatabaseTester: + def __init__(self): + self.automaton = ahocorasick.Automaton() + + # 基础查找表 + self.tp_packages = {} + self.tp_pkg_string = {} + self.pkg_to_tp_mark = {} + + # 核心统计容器: Key = Source App ID + self.source_app_stats = defaultdict(lambda: { + 'app_label': '', + 'total_bytes': 0, # 包含 IP 流量 + 域名流量 + 'total_domain_bytes': 0, # 仅包含域名流量 + 'self_bytes': 0, + 'server_bytes': 0, + 'unrec_bytes': 0, + 'self_count': 0, + 'server_count': 0, + 'unrec_count': 0, + 'components': defaultdict(lambda: {'bytes': 0, 'is_self': False}) + }) + + # 存储未匹配记录 (用于 unmatched 文件) + self.unmatched_records = [] + + # 存储所有流水的详细记录 (用于 detail_log 文件) + self.all_detail_records = [] + + def check_paths(self): + files = [APP_LIST_PATH, URL_LIST_PATH, TRAFFIC_DATA_PATH] + missing = [f for f in files if not f.exists()] + if missing: + print(f"[Error] 缺少文件: {[f.name for f in missing]}") + sys.exit(1) + if not OUTPUT_DIR.exists(): + os.makedirs(OUTPUT_DIR) + + def load_app_list(self): + print(f"[1/5] 加载应用定义 (TPMark): {APP_LIST_PATH.name}") + try: + df = pd.read_csv(APP_LIST_PATH) + for row in df.itertuples(): + tp_mark = str(row.app_name).strip() + pkg_raw = str(row.package_name) if pd.notna(row.package_name) else "" + + self.tp_pkg_string[tp_mark] = pkg_raw + pkg_list = [p.strip() for p in pkg_raw.split(',') if p.strip()] + + if tp_mark not in self.tp_packages: + self.tp_packages[tp_mark] = set() + self.tp_packages[tp_mark].update(pkg_list) + + for pkg in pkg_list: + self.pkg_to_tp_mark[pkg] = tp_mark + + print(f" - 已加载 {len(self.tp_packages)} 个应用定义。") + except Exception as e: + print(f"[Error] 加载应用列表失败: {e}") + sys.exit(1) + + def load_url_list(self): + print(f"[2/5] 构建 AC 自动机: {URL_LIST_PATH.name}") + try: + df = pd.read_csv(URL_LIST_PATH) + df.columns = [c.strip().lower() for c in df.columns] + count = 0 + for row in df.itertuples(): + suffix = str(row.url).strip().lower() + tp_mark = str(row.app).strip() + if suffix: + self.automaton.add_word(suffix, (suffix, tp_mark)) + count += 1 + self.automaton.make_automaton() + print(f" - 加载了 {count} 条 URL 规则。") + except Exception as e: + print(f"[Error] 加载 URL 失败: {e}") + sys.exit(1) + + def search_domain(self, domain): + """返回 (TPMark, Pattern)""" + domain_lower = domain.lower() + valid_matches = [] + # iter 返回 (end_index, (pattern, tp_mark)) + for end_index, (pattern, tp_mark) in self.automaton.iter(domain_lower): + if end_index != len(domain_lower) - 1: continue + start_index = end_index - len(pattern) + 1 + if start_index == 0 or domain_lower[start_index - 1] == '.': + valid_matches.append((len(pattern), tp_mark, pattern)) + + if not valid_matches: return None, None + # 取最长匹配 + best = max(valid_matches, key=lambda x: x[0]) + return best[1], best[2] + + def process_traffic(self): + print(f"[3/5] 分析流量归属: {TRAFFIC_DATA_PATH.name}") + try: + df = pd.read_csv(TRAFFIC_DATA_PATH) + df.columns = [c.strip().replace(' ', '_').replace('(', '_').replace(')', '') for c in df.columns] + + processed = 0 + for row in df.itertuples(): + source_app_id = str(row.App_ID).strip() + + # Filter: 跳过 unknown 和 root + if source_app_id.lower() in ['unknown', 'root']: + continue + + domain = str(row.Domain).strip() + size = float(row.Traffic_Size_Bytes) + source_app_label = str(row.App_Name).strip() + + stats = self.source_app_stats[source_app_id] + stats['app_label'] = source_app_label + stats['total_bytes'] += size + + # === 临时变量，用于构建 Detail Record === + detail_tp_result = "" + detail_tp_url = "" + detail_is_self = False + detail_is_ip = False + + # === 判断是否为无域名IP流量 (model_data) === + is_ip_flow = domain.startswith("model_data:") + + if is_ip_flow: + detail_is_ip = True + detail_tp_result = "IP Flow" # 或保持空，根据需要 + continue + stats['unrec_bytes'] += size + stats['unrec_count'] += 1 + + # IP 流量也记入 unmatched_records 供 unmatched文件使用 + self.unmatched_records.append({ + 'app_id': source_app_id, + 'app_label': source_app_label, + 'input_url': domain, + 'traffic_size': size, + 'is_ip_flow': True + }) + + else: + # === 普通域名流量处理 === + stats['total_domain_bytes'] += size + + # 匹配域名 + matched_tp_mark, matched_pattern = self.search_domain(domain) + + if matched_tp_mark: + # 记录详情 + detail_tp_result = matched_tp_mark + detail_tp_url = matched_pattern + + owner_packages = self.tp_packages.get(matched_tp_mark, set()) + is_self = source_app_id in owner_packages + detail_is_self = is_self + + if is_self: + stats['self_bytes'] += size + stats['self_count'] += 1 + else: + stats['server_bytes'] += size + stats['server_count'] += 1 + + comp = stats['components'][matched_tp_mark] + comp['bytes'] += size + comp['is_self'] = is_self + + else: + # 域名未匹配 + detail_tp_result = "[|Unmatched|]" + + stats['unrec_bytes'] += size + stats['unrec_count'] += 1 + + self.unmatched_records.append({ + 'app_id': source_app_id, + 'app_label': source_app_label, + 'input_url': domain, + 'traffic_size': size, + 'is_ip_flow': False + }) + if detail_tp_result == "[|Unmatched|]": + continue + # === 保存详细流水记录 === + # 注意：此时还没计算 Ratio，因为 total_bytes 还在累加中 + # 存下原始数据，在导出时计算 Ratio + self.all_detail_records.append({ + 'tp_result': detail_tp_result if detail_tp_result else "", + 'app_name': source_app_label, + 'app_label': source_app_label, + 'app_id': source_app_id, + 'input_url': domain, + 'tp_url': detail_tp_url, + 'self': detail_is_self, + 'traffic_count_bytes': size, + 'is_ip_flow': detail_is_ip + }) + + processed += 1 + if processed % 50000 == 0: + print(f" - 已处理 {processed} 行...") + + except Exception as e: + print(f"[Error] 处理流量失败: {e}") + import traceback + traceback.print_exc() + sys.exit(1) + + def export_app_profile(self): + print(f"[4/5] 生成画像报告: {OUTPUT_APP_PROFILE.name}") + + data_rows = [] + for app_id, stats in self.source_app_stats.items(): + total = stats['total_bytes'] + if total == 0: continue + + tp_mark = self.pkg_to_tp_mark.get(app_id, "") + app_label = stats['app_label'] + app_name = app_label + + self_ratio = (stats['self_bytes'] / total) * 100 + recog_ratio = ((stats['self_bytes'] + stats['server_bytes']) / total) * 100 + + analysis_list = [] + sorted_comps = sorted(stats['components'].items(), key=lambda x: x[1]['bytes'], reverse=True) + + for comp_tp_mark, comp_data in sorted_comps: + comp_bytes = comp_data['bytes'] + comp_percent = (comp_bytes / total) * 100 + is_self_flag = '0' if comp_data['is_self'] else '1' + comp_pkg_str = self.tp_pkg_string.get(comp_tp_mark, "") + + analysis_item = [ + comp_tp_mark, comp_pkg_str, is_self_flag, f"{comp_percent:.2f}%", int(comp_bytes) + ] + analysis_list.append(analysis_item) + + data_rows.append({ + 'app_id': app_id, + 'app_label': app_label, + 'app_name': app_name, + 'tp_mark': tp_mark, + 'total_traffic(Bytes)': int(total), + 'self_traffic_count(Bytes)': format_bytes(stats['self_bytes']), + 'server_traffic_count(Bytes)': format_bytes(stats['server_bytes']), + 'unrecognized_traffic_count(Bytes)': format_bytes(stats['unrec_bytes']), + 'self_label_count': stats['self_count'], + 'server_label_count': stats['server_count'], + 'unrecognized_label_count': stats['unrec_count'], + 'self_ratio': f"{self_ratio:.2f}", + 'recognition_ratio': f"{recog_ratio:.2f}", + 'traffic_analysis': str(analysis_list) + }) + + df = pd.DataFrame(data_rows) + cols = [ + 'app_id', 'app_label', 'app_name', 'tp_mark', + 'total_traffic(Bytes)', 'self_traffic_count(Bytes)', 'server_traffic_count(Bytes)', + 'unrecognized_traffic_count(Bytes)', + 'self_label_count', 'server_label_count', 'unrecognized_label_count', + 'self_ratio', 'recognition_ratio', 'traffic_analysis' + ] + if not df.empty: + df = df[cols] + df.sort_values(by='total_traffic(Bytes)', ascending=False, inplace=True) + df.to_csv(OUTPUT_APP_PROFILE, index=False) + print(f" - 完成: {OUTPUT_APP_PROFILE.name}") + + def export_unmatched_list(self): + print(f"[5/5] 生成未匹配域名列表: {OUTPUT_UNMATCHED_FILE.name}") + if not self.unmatched_records: + print(" - 无未匹配记录。") + else: + export_rows = [] + for record in self.unmatched_records: + app_id = record['app_id'] + size = record['traffic_size'] + is_ip = record['is_ip_flow'] + + stats = self.source_app_stats[app_id] + total_bytes = stats['total_bytes'] + total_domain_bytes = stats['total_domain_bytes'] + + traffic_ratio = (size / total_bytes * 100) if total_bytes > 0 else 0 + + if is_ip: + domain_ratio_str = "0.00%" + else: + domain_ratio = (size / total_domain_bytes * 100) if total_domain_bytes > 0 else 0 + domain_ratio_str = f"{domain_ratio:.2f}%" + + export_rows.append({ + 'app_name': record['app_label'], + 'app_label': record['app_label'], + 'app_id': app_id, + 'input_url': record['input_url'], + 'traffic_count(Bytes)': int(size), + 'traffic_ratio': f"{traffic_ratio:.2f}%", + 'domain_traffic_ratio': domain_ratio_str, + 'organization': '' + }) + + df = pd.DataFrame(export_rows) + if not df.empty: + df.sort_values(by='traffic_count(Bytes)', ascending=False, inplace=True) + df.to_csv(OUTPUT_UNMATCHED_FILE, index=False) + print(f" - 完成: {OUTPUT_UNMATCHED_FILE.name}") + + def export_detail_log(self): + print(f"[Bonus] 生成详细流水日志: {OUTPUT_DETAIL_FILE.name}") + + if not self.all_detail_records: + print(" - 无流水记录。") + return + + export_rows = [] + # 批量处理，提升速度 + for r in self.all_detail_records: + app_id = r['app_id'] + size = r['traffic_count_bytes'] + is_ip = r['is_ip_flow'] + + # 获取该 App 的统计数据以计算 Ratio + stats = self.source_app_stats[app_id] + total_bytes = stats['total_bytes'] + total_domain_bytes = stats['total_domain_bytes'] + + # 1. Total Ratio + ratio = (size / total_bytes * 100) if total_bytes > 0 else 0 + + # 2. Domain Ratio + if is_ip: + domain_ratio_str = "0.00%" + else: + d_ratio = (size / total_domain_bytes * 100) if total_domain_bytes > 0 else 0 + domain_ratio_str = f"{d_ratio:.2f}%" + + export_rows.append({ + 'tp_result': r['tp_result'], + 'app_name': r['app_name'], + 'app_label': r['app_label'], + 'app_id': app_id, + 'input_url': r['input_url'], + 'tp_url': r['tp_url'], + 'self': str(r['self']), # 转字符串 + 'traffic_count(Bytes)': int(size), + 'traffic_ratio': f"{ratio:.2f}%", + 'domain_traffic_ratio': domain_ratio_str, + 'organization': '' + }) + + df = pd.DataFrame(export_rows) + # 排序建议：先按 App ID 聚类，再按流量大小降序 + if not df.empty: + df.sort_values(by=['app_id', 'traffic_count(Bytes)'], ascending=[True, False], inplace=True) + + df.to_csv(OUTPUT_DETAIL_FILE, index=False) + print(f" - 完成: {OUTPUT_DETAIL_FILE.name}") + + +if __name__ == "__main__": + tester = DpiDatabaseTester() + tester.check_paths() + tester.load_app_list() + tester.load_url_list() + tester.process_traffic() + tester.export_app_profile() + tester.export_unmatched_list() + tester.export_detail_log() diff --git a/calculate/get_traffic_summary.py b/calculate/get_traffic_summary.py new file mode 100644 index 0000000..83a0fd4 --- /dev/null +++ b/calculate/get_traffic_summary.py @@ -0,0 +1,255 @@ +#!/usr/bin/env python3 +""" +Merge all traffic_count*.txt files from the traffic_data directory into a single +traffic_summary.csv file. + +Directory structure: + \\192.168.2.75\dpi-sync\autool_config\data\traffic_data\ + └── / + └── / + └── traffic_count*.txt + +Input line format (comma-separated): + App_ID, App_Name, Domain, Flow(IP:Port-IP:Port-Protocol), Transport_Protocol, App_Protocol, Traffic_Size, [Organization] + +Output CSV format: + App ID, App Name, Domain, Traffic Size(Bytes), Flow Count, Traffic Ratio, Domain Traffic Ratio, Organization +""" + +import os +import re +import csv +import glob +import argparse +from collections import defaultdict +from datetime import datetime + + +def parse_traffic_size(size_str): + """Parse traffic size string like '128 B', '1.5 KB', '2.3 MB' into bytes.""" + size_str = size_str.strip() + if not size_str: + return 0 + + # Match number and optional unit + match = re.match(r'([\d.]+)\s*([KMGT]?B?)', size_str, re.IGNORECASE) + if not match: + return 0 + + value = float(match.group(1)) + unit = match.group(2).upper().strip() + + multipliers = { + '': 1, 'B': 1, + 'KB': 1024, 'K': 1024, + 'MB': 1024 ** 2, 'M': 1024 ** 2, + 'GB': 1024 ** 3, 'G': 1024 ** 3, + 'TB': 1024 ** 4, 'T': 1024 ** 4, + } + + return int(value * multipliers.get(unit, 1)) + + +def determine_domain_or_model(fields): + """ + Determine the domain field value. + If App_Protocol is recognized (e.g., DNS, TLS, HTTP, QUIC...), use the domain field directly. + Otherwise, generate a model_data identifier from the flow info. + """ + domain = fields[2].strip() if len(fields) > 2 else '' + # app_protocol = fields[5].strip() if len(fields) > 5 else '' + flow_info = fields[3].strip() if len(fields) > 3 else '' + + # # Known application protocols that have domain info + # known_protocols = { + # 'DNS', 'TLS', 'HTTP', 'HTTPS', 'QUIC', 'HTTP/S', 'SSL', + # 'NTP', 'STUN', 'DTLS', 'MQTT', 'MDNS', 'SSDP', 'LLMNR', + # } + + # if domain and app_protocol.upper() in known_protocols: + # return domain + # elif domain: + # return domain + if domain != 'model_data:': + return domain + else: + # Use flow info as model_data identifier + return f'model_data:{flow_info}' + + +def parse_traffic_file(filepath): + """ + Parse a single traffic_count file. + Returns a list of tuples: (app_id, app_name, domain, traffic_bytes, organization) + """ + records = [] + try: + with open(filepath, 'r', encoding='utf-8', errors='ignore') as f: + for line_num, line in enumerate(f, 1): + line = line.strip() + if not line: + continue + + # Split by comma, but be careful with trailing comma + fields = line.split(',') + + # Need at least 7 fields: app_id, app_name, domain, flow, transport, app_protocol, size + if len(fields) < 7: + continue + + app_id = fields[0].strip() + app_name = fields[1].strip() + domain = determine_domain_or_model(fields) + traffic_size_str = fields[6].strip() + organization = fields[7].strip() if len(fields) > 7 else '' + + traffic_bytes = parse_traffic_size(traffic_size_str) + + records.append((app_id, app_name, domain, traffic_bytes, organization)) + except Exception as e: + print(f" [WARNING] Error reading {filepath}: {e}") + + return records + + +def merge_traffic_data(base_dir, output_file): + """ + Scan all traffic_count*.txt files under base_dir and merge them into output_file. + """ + print(f"Scanning directory: {base_dir}") + print(f"Output file: {output_file}") + print() + + # Find all traffic_count*.txt files + pattern = os.path.join(base_dir, '**', 'traffic_count*.txt') + traffic_files = glob.glob(pattern, recursive=True) + + if not traffic_files: + print(f"No traffic_count files found in {base_dir}") + return + + print(f"Found {len(traffic_files)} traffic_count files") + print() + + # Aggregate data: key = (app_id, app_name, domain), value = {bytes, flow_count, organization} + aggregated = defaultdict(lambda: {'bytes': 0, 'flow_count': 0, 'organization': ''}) + + total_files_processed = 0 + total_records = 0 + + for filepath in sorted(traffic_files): + records = parse_traffic_file(filepath) + if records: + total_files_processed += 1 + total_records += len(records) + rel_path = os.path.relpath(filepath, base_dir) + # print(f" Processed: {rel_path} ({len(records)} records)") + + for app_id, app_name, domain, traffic_bytes, organization in records: + key = (app_id, app_name, domain) + aggregated[key]['bytes'] += traffic_bytes + aggregated[key]['flow_count'] += 1 + # Keep the first non-empty organization + if organization and not aggregated[key]['organization']: + aggregated[key]['organization'] = organization + + print() + print(f"Total files processed: {total_files_processed}") + print(f"Total raw records: {total_records}") + print(f"Total aggregated entries: {len(aggregated)}") + + # Pre-compute traffic totals per app (all entries and domain-only entries) + app_total_bytes = defaultdict(int) + app_domain_total_bytes = defaultdict(int) + print(f"start calculate traffic ratios per app") + for (app_id, app_name, domain), data in aggregated.items(): + app_total_bytes[(app_id, app_name)] += data['bytes'] + if not domain.startswith('model_data:'): + app_domain_total_bytes[(app_id, app_name)] += data['bytes'] + + # Sort by app_id, then by traffic bytes descending within each app + sorted_entries = sorted( + aggregated.items(), + key=lambda x: (x[0][0], -x[1]['bytes']) + ) + print(f"end calculate traffic ratios per app") + + # Write output CSV + os.makedirs(os.path.dirname(output_file), exist_ok=True) + print(f"start write output CSV") + with open(output_file, 'w', newline='', encoding='utf-8') as f: + writer = csv.writer(f) + writer.writerow([ + 'App ID', 'App Name', 'Domain', + 'Traffic Size(Bytes)', 'Flow Count', + 'Traffic Ratio', 'Domain Traffic Ratio', 'Organization' + ]) + + for (app_id, app_name, domain), data in sorted_entries: + total_bytes = data['bytes'] + flow_count = data['flow_count'] + app_total = app_total_bytes[(app_id, app_name)] + + # Traffic Ratio = this entry's bytes / total bytes for this app + if app_total > 0: + traffic_ratio = total_bytes / app_total * 100 + traffic_ratio_str = f"{traffic_ratio:.2f}%" + else: + traffic_ratio_str = "0.00%" + + # Domain Traffic Ratio: same as traffic ratio for domain entries, + # empty for model_data entries + if domain.startswith('model_data:'): + domain_traffic_ratio_str = '' + else: + app_domain_total = app_domain_total_bytes[(app_id, app_name)] + if app_domain_total > 0: + domain_ratio = total_bytes / app_domain_total * 100 + domain_traffic_ratio_str = f"{domain_ratio:.2f}%" + else: + domain_traffic_ratio_str = "0.00%" + + writer.writerow([ + app_id, app_name, domain, + total_bytes, flow_count, + traffic_ratio_str, domain_traffic_ratio_str, + data['organization'] + ]) + + print(f"\nOutput written to: {output_file}") + + +def main(): + parser = argparse.ArgumentParser( + description='Merge traffic_count files into traffic_summary.csv' + ) + parser.add_argument( + '--input-dir', + default=r'\\192.168.2.75\dpi-sync\autool_config\data\traffic_data', + help='Root directory containing PC subdirectories with traffic data' + ) + parser.add_argument( + '--output', + default=None, + help='Output CSV file path (default: TrafficData/traffic_summary//traffic_summary.csv)' + ) + + args = parser.parse_args() + + # Default output path with date + if args.output is None: + script_dir = os.path.dirname(os.path.abspath(__file__)) + project_dir = os.path.dirname(script_dir) + date_str = datetime.now().strftime('%Y%m%d') + # date_str = '20260224' + output_file = os.path.join( + project_dir, 'TrafficData', 'traffic_summary', date_str, 'traffic_summary.csv' + ) + else: + output_file = args.output + + merge_traffic_data(args.input_dir, output_file) + + +if __name__ == '__main__': + main() diff --git a/config.example.yaml b/config.example.yaml new file mode 100644 index 0000000..12cb244 --- /dev/null +++ b/config.example.yaml @@ -0,0 +1,208 @@ +# 配置文件示例（脱敏模板） +# 使用方式：将此文件复制为 config.yaml 并根据实际环境填写配置项 +# 敏感信息值留空，实际部署时填写 + +# ============================================================ +# 环境标识 +# ============================================================ +environment: + name: prod # 或 test + +instance_name: main + +# ============================================================ +# Redis 连接 +# ============================================================ +redis: + host: "127.0.0.1" + port: 6379 + db: 0 + channel_namespace: main + max_connections: 50 + +# ============================================================ +# 任务配置 +# ============================================================ +tasks: + max_retry_count: 3 + csv_file: package_list.csv + +# ============================================================ +# Worker 管理 +# ============================================================ +worker: + stale_timeout: 9000 + reinit_timeout: 600 + init_monitor_window: 1800 + action_max_parallel: 50 + managed_only_can_dispatch: false + allowed_ip_patterns: + - "10\\.\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}" + - "172\\.(1[6-9]|2\\d|3[0-1])\\.\\d{1,3}\\.\\d{1,3}" + - "192\\.168\\.\\d{1,3}\\.\\d{1,3}" + +# ============================================================ +# SSH / Git +# ============================================================ +ssh: + connect_timeout: 10 + action_timeout: 1800 + default_user: admin + default_password: "" + default_port: 22 + psexec_session_id: 1 + +git: + repo_url: "" # 如 //your-file-server/share/autool.git + +# ============================================================ +# MuMu 模拟器 +# ============================================================ +mumu: + manager_path: "C:\\Program Files\\Netease\\MuMu\\nx_main\\MuMuManager.exe" + vm_index: 2 + max_vm_index: 10 + restart_settle_seconds: 15 + recover_script_path: "" + recover_ahk_exe: "" + clean_backup_path: "" + local_import_dir: "D:\\mumu_backups" + network_bridge_card: "Realtek PCIe GbE Family Controller" + adb_ip_offset: 100 + bridge_ip_offset: 100 + network_script_path: "" + network_gateways: + "192.168.1": "192.168.1.1" + "192.168.2": "192.168.2.1" + +# ============================================================ +# 网络共享 +# ============================================================ +share: + smb_user: "" + smb_password: "" + +# ============================================================ +# 日志 +# ============================================================ +logging: + file: dispatcher.log + max_bytes: 52428800 + backup_count: 3 + date_format: "%Y-%m-%d %H:%M:%S" + +# ============================================================ +# 通知（微信 / 企业微信） +# ============================================================ +notifications: + wechat_tokens: {} + wecom_tokens: {} + +# ============================================================ +# 告警策略 +# ============================================================ +alert: + window_seconds: 600 + threshold: 30 + cooldown_seconds: 1800 + +# ============================================================ +# 监控 +# ============================================================ +monitoring: + status_push_interval: 10800 + timezone: Asia/Shanghai + timeline_bucket_minutes: 15 + +# ============================================================ +# 流水线 +# ============================================================ +pipeline: + run_option_keys: + - stop_worker + - clone_if_missing + - git_pull + - setup + - pull_pcap_files + - recover_mumu + - restart_mumu + - start_worker + default_options: + stop_worker: true + clone_if_missing: false + git_pull: true + setup: false + pull_pcap_files: true + recover_mumu: false + restart_mumu: true + start_worker: true + auto_reboot_recovery_delay: 120 + legacy_end_worker_ips: + - "192.168.1.51" + - "192.168.1.61" + - "192.168.2.101" + force_kill_images: + - OpenConsole.exe + - WindowsTerminal.exe + - conhost.exe + - powershell.exe + +# ============================================================ +# Dashboard +# ============================================================ +dashboard: + host: "0.0.0.0" + port: 8890 + +# ============================================================ +# 数据分析 +# ============================================================ +analytics: + traffic_root: "" + traffic_root_block: "" + traversal_root: "" + tpdpi_app_list_file: TPDPI_app_list.csv + tpdpi_url_lib_file: TPDPI_url_lib.csv + artifact_wait_seconds: 180 + artifact_file_wait_seconds: 60 + job_poll_seconds: 1 + model_traffic_threshold: 50 + +# ============================================================ +# MinIO / APK 存储 +# ============================================================ +minio: + enabled: false + endpoint: "" + access_key: "" + secret_key: "" + bucket: autool-apk + secure: false + +apk: + download_mode: direct # "direct" 或 "smb" + local_storage_dir: "" + smb_dir: "" + us_export_dir: "C:\\mumu_apks" + us_device_a_serial: "127.0.0.1:7555" + us_device_b_serial: "127.0.0.1:7556" + us_device_a_vm_index: 2 + us_device_b_vm_index: 1 + prefetch_poll_interval: 120 + download_queue_interval: 60 + minio_max_bytes: 107374182400 + emulator_cleanup_enabled: true + +# ============================================================ +# 任务路由 / Worker 清单 +# ============================================================ +task_routing_rules: + package_name: {} + task_key: {} + +worker_inventory: + - worker_id: "192.168.1.10" + ssh_target: "192.168.1.10" + repo_dir: "D:/autool" + python_exe: python + tags: [] diff --git a/config.py b/config.py new file mode 100644 index 0000000..e47d87b --- /dev/null +++ b/config.py @@ -0,0 +1,557 @@ +# -*- encoding=utf8 -*- +""" +Task dispatcher configuration. +""" + +import json +import os +import re +from typing import Any, Dict, List, Optional, Tuple + +try: + import yaml +except ImportError: # pragma: no cover - exercised only when dependency is absent. + yaml = None + + +BASE_DIR = os.path.dirname(os.path.abspath(__file__)) +CONFIG_DIR = os.path.join(BASE_DIR, "config") +CONFIG_PATH_ENV = os.environ.get("AUTOOL_DISPATCHER_CONFIG") +CONFIG_PATH = CONFIG_PATH_ENV or os.path.join(CONFIG_DIR, "config.json") +CONFIG_YAML_PATH = os.path.join(BASE_DIR, "config.yaml") +CONFIG_EXAMPLE_PATH = os.path.join(CONFIG_DIR, "config.example.json") +CONFIG_EXAMPLE_YAML_PATH = os.path.join(BASE_DIR, "config.example.yaml") +CURRENT_ENV_PATH = os.path.join(CONFIG_DIR, "current_env.txt") +LOCAL_CONFIG_PATH = os.path.join(CONFIG_DIR, "local.json") +ENV_CONFIG_PATHS = { + "prod": os.path.join(CONFIG_DIR, "prod.json"), + "test": os.path.join(CONFIG_DIR, "test.json"), +} +ENV_ALIASES = { + "": "prod", + "prod": "prod", + "production": "prod", + "main": "prod", + "master": "prod", + "test": "test", + "testing": "test", +} + + +def _load_config_payload(path: str, *, required: bool) -> Any: + if not os.path.exists(path): + if required: + raise FileNotFoundError(f"config file not found: {path}") + return {} + + with open(path, "r", encoding="utf-8") as handle: + if path.endswith((".yaml", ".yml")): + if yaml is None: + raise ImportError("PyYAML is required to read YAML config files. Install: pip install PyYAML") + payload = yaml.safe_load(handle) or {} + else: + payload = json.load(handle) + + return payload + + +def _load_config_dict(path: str, *, required: bool) -> Dict[str, Any]: + payload = _load_config_payload(path, required=required) + + if not isinstance(payload, dict): + raise ValueError(f"config payload must be an object: {path}") + return payload + + +def _load_json_dict(path: str, *, required: bool) -> Dict[str, Any]: + return _load_config_dict(path, required=required) + + +def _merge_dict(base: Dict[str, Any], override: Dict[str, Any]) -> Dict[str, Any]: + merged = dict(base) + for key, value in override.items(): + current = merged.get(key) + if isinstance(current, dict) and isinstance(value, dict): + merged[key] = _merge_dict(current, value) + continue + merged[key] = value + return merged + + +def _read_current_env() -> str: + if not os.path.exists(CURRENT_ENV_PATH): + return "prod" + with open(CURRENT_ENV_PATH, "r", encoding="utf-8") as handle: + return handle.read().strip().lower() + + +def normalize_env_name(env_name: Optional[str] = None) -> str: + raw = str("prod" if env_name is None else env_name).strip().lower() + normalized = ENV_ALIASES.get(raw) + if normalized: + return normalized + raise ValueError(f"unsupported config env: {env_name or raw}") + + +def get_runtime_env_name() -> str: + return CONFIG_ENV + + +def _load_legacy_layered_config(env_name: Optional[str] = None) -> Dict[str, Any]: + normalized_env = normalize_env_name(env_name) + config = _load_config_dict(ENV_CONFIG_PATHS[normalized_env], required=True) + config = _merge_dict(config, _load_config_dict(LOCAL_CONFIG_PATH, required=False)) + return config + + +# 嵌套小写 → 扁平大写键名映射 +_NESTED_TO_FLAT_MAP = { + ("environment", "name"): "CONFIG_ENV", + ("instance_name",): "INSTANCE_NAME", + ("redis", "host"): "REDIS_HOST", + ("redis", "port"): "REDIS_PORT", + ("redis", "db"): "REDIS_DB", + ("redis", "channel_namespace"): "CHANNEL_NAMESPACE", + ("redis", "max_connections"): "REDIS_MAX_CONNECTIONS_DISPATCHER", + ("tasks", "max_retry_count"): "MAX_RETRY_COUNT", + ("tasks", "csv_file"): "TASK_CSV_FILE", + ("worker", "stale_timeout"): "WORKER_STALE_TIMEOUT", + ("worker", "reinit_timeout"): "WORKER_REINIT_TIMEOUT", + ("worker", "init_monitor_window"): "WORKER_INIT_MONITOR_WINDOW", + ("worker", "action_max_parallel"): "WORKER_ACTION_MAX_PARALLEL", + ("worker", "managed_only_can_dispatch"): "ONLY_MANAGED_WORKERS_CAN_DISPATCH", + ("worker", "allowed_ip_patterns"): "WORKER_ALLOWED_IP_PATTERNS", + ("ssh", "connect_timeout"): "SSH_CONNECT_TIMEOUT_SECONDS", + ("ssh", "action_timeout"): "SSH_ACTION_TIMEOUT_SECONDS", + ("ssh", "default_user"): "SSH_DEFAULT_USER", + ("ssh", "default_password"): "SSH_DEFAULT_PASSWORD", + ("ssh", "default_port"): "SSH_DEFAULT_PORT", + ("ssh", "psexec_session_id"): "PSEXEC_DEFAULT_SESSION_ID", + ("git", "repo_url"): "AUTOOL_GIT_REPO_URL", + ("mumu", "manager_path"): "MUMU_MANAGER_PATH", + ("mumu", "vm_index"): "MUMU_VM_INDEX", + ("mumu", "max_vm_index"): "MUMU_MAX_VM_INDEX", + ("mumu", "restart_settle_seconds"): "MUMU_RESTART_SETTLE_SECONDS", + ("mumu", "recover_script_path"): "MUMU_RECOVER_SCRIPT_PATH", + ("mumu", "recover_ahk_exe"): "MUMU_RECOVER_AHK_EXE", + ("mumu", "clean_backup_path"): "CLEAN_BACKUP_PATH", + ("mumu", "local_import_dir"): "LOCAL_IMPORT_DIR", + ("mumu", "network_bridge_card"): "NET_BRIDGE_CARD", + ("mumu", "adb_ip_offset"): "EMULATOR_ADB_IP_OFFSET", + ("mumu", "bridge_ip_offset"): "MUMU_BRIDGE_IP_OFFSET", + ("mumu", "network_script_path"): "MUMU_NETWORK_SCRIPT_PATH", + ("mumu", "network_gateways"): "MUMU_NETWORK_GATEWAYS", + ("share", "smb_user"): "SHARE_SMB_USER", + ("share", "smb_password"): "SHARE_SMB_PASSWORD", + ("logging", "file"): "LOG_FILE", + ("logging", "max_bytes"): "LOG_MAX_BYTES", + ("logging", "backup_count"): "LOG_BACKUP_COUNT", + ("logging", "date_format"): "LOG_DATE_FORMAT", + ("notifications", "wechat_tokens"): "WECHAT_TOKENS", + ("notifications", "wecom_tokens"): "WECOM_TOKENS", + ("alert", "window_seconds"): "ALERT_WINDOW_SECONDS", + ("alert", "threshold"): "ALERT_THRESHOLD", + ("alert", "cooldown_seconds"): "ALERT_COOLDOWN_SECONDS", + ("monitoring", "status_push_interval"): "WORKER_STATUS_PUSH_INTERVAL", + ("monitoring", "timezone"): "MONITORING_TIMEZONE", + ("monitoring", "timeline_bucket_minutes"): "MONITORING_TIMELINE_BUCKET_MINUTES", + ("pipeline", "run_option_keys"): "RUN_PIPELINE_OPTION_KEYS", + ("pipeline", "default_options"): "RUN_PIPELINE_DEFAULT_OPTIONS", + ("pipeline", "auto_reboot_recovery_delay"): "AUTO_REBOOT_RECOVERY_DELAY_SECONDS", + ("pipeline", "legacy_end_worker_ips"): "RUNBATCH_LEGACY_END_WORKER_IPS", + ("pipeline", "force_kill_images"): "RUNBATCH_FORCE_KILL_IMAGES", + ("dashboard", "host"): "DASHBOARD_HOST", + ("dashboard", "port"): "DASHBOARD_PORT", + ("analytics", "traffic_root"): "ANALYTICS_TRAFFIC_ROOT", + ("analytics", "traffic_root_block"): "ANALYTICS_TRAFFIC_ROOT_BLOCK", + ("analytics", "traversal_root"): "ANALYTICS_TRAVERSAL_ROOT", + ("analytics", "tpdpi_app_list_file"): "ANALYTICS_TPDPI_APP_LIST_FILE", + ("analytics", "tpdpi_url_lib_file"): "ANALYTICS_TPDPI_URL_LIB_FILE", + ("analytics", "artifact_wait_seconds"): "ANALYTICS_ARTIFACT_WAIT_SECONDS", + ("analytics", "artifact_file_wait_seconds"): "ANALYTICS_ARTIFACT_FILE_WAIT_SECONDS", + ("analytics", "job_poll_seconds"): "ANALYTICS_JOB_POLL_SECONDS", + ("analytics", "model_traffic_threshold"): "MODEL_TRAFFIC_THRESHOLD", + ("minio", "enabled"): "MINIO_ENABLED", + ("minio", "endpoint"): "MINIO_ENDPOINT", + ("minio", "access_key"): "MINIO_ACCESS_KEY", + ("minio", "secret_key"): "MINIO_SECRET_KEY", + ("minio", "bucket"): "MINIO_BUCKET", + ("minio", "secure"): "MINIO_SECURE", + ("apk", "download_mode"): "APK_DOWNLOAD_MODE", + ("apk", "local_storage_dir"): "APK_LOCAL_STORAGE_DIR", + ("apk", "smb_dir"): "APK_SMB_DIR", + ("apk", "us_export_dir"): "APK_US_EXPORT_DIR", + ("apk", "us_device_a_serial"): "APK_US_DEVICE_A_SERIAL", + ("apk", "us_device_b_serial"): "APK_US_DEVICE_B_SERIAL", + ("apk", "us_device_a_vm_index"): "APK_US_DEVICE_A_VM_INDEX", + ("apk", "us_device_b_vm_index"): "APK_US_DEVICE_B_VM_INDEX", + ("apk", "prefetch_poll_interval"): "APK_PREFETCH_POLL_INTERVAL", + ("apk", "download_queue_interval"): "APK_DOWNLOAD_QUEUE_INTERVAL", + ("apk", "minio_max_bytes"): "APK_MINIO_MAX_BYTES", + ("apk", "emulator_cleanup_enabled"): "APK_EMULATOR_CLEANUP_ENABLED", + ("task_routing_rules",): "TASK_ROUTING_RULES", + ("worker_inventory",): "WORKER_INVENTORY", +} + + +def _is_nested_format(config: Dict[str, Any]) -> bool: + """检测是否为嵌套小写格式：顶层有 environment 键且无 CONFIG_ENV""" + return "environment" in config and "CONFIG_ENV" not in config + + +def _dig(config: Dict[str, Any], *keys: str) -> Any: + """从嵌套字典中按路径取值""" + value = config + for key in keys: + if not isinstance(value, dict): + return None + value = value.get(key) + return value + + +def _normalize_nested_config(config: Dict[str, Any]) -> Dict[str, Any]: + """将嵌套小写格式转为扁平大写格式，已是扁平格式则直接返回""" + if not _is_nested_format(config): + return config + + flat = {} + for nested_keys, flat_key in _NESTED_TO_FLAT_MAP.items(): + value = _dig(config, *nested_keys) + if value is not None: + flat[flat_key] = value + return flat + + +def load_dispatcher_config( + env_name: Optional[str] = None, + config_path: Optional[str] = None, +) -> Dict[str, Any]: + if config_path: + return _normalize_nested_config(_load_config_dict(os.path.abspath(config_path), required=True)) + + if env_name is not None: + return _load_legacy_layered_config(env_name) + + if CONFIG_PATH_ENV: + return _normalize_nested_config(_load_config_dict(os.path.abspath(CONFIG_PATH_ENV), required=True)) + + if os.path.exists(CONFIG_YAML_PATH): + return _normalize_nested_config(_load_config_dict(CONFIG_YAML_PATH, required=True)) + + if os.path.exists(CONFIG_PATH): + return _normalize_nested_config(_load_config_dict(CONFIG_PATH, required=True)) + + if os.path.exists(CONFIG_EXAMPLE_YAML_PATH): + return _normalize_nested_config(_load_config_dict(CONFIG_EXAMPLE_YAML_PATH, required=True)) + + if os.path.exists(CONFIG_EXAMPLE_PATH): + return _normalize_nested_config(_load_config_dict(CONFIG_EXAMPLE_PATH, required=True)) + + legacy_env = _read_current_env() + return _load_legacy_layered_config(legacy_env) + + +CONFIG = load_dispatcher_config() +CONFIG_ENV = normalize_env_name(CONFIG.get("CONFIG_ENV") or CONFIG.get("INSTANCE_NAME") or "prod") + + +# ==================== Instance ==================== +INSTANCE_NAME = str(CONFIG["INSTANCE_NAME"]).strip() +INSTANCE_DIR = os.path.join(BASE_DIR, "runtime", INSTANCE_NAME) +REPORT_DIR = os.path.join(INSTANCE_DIR, "reports") + + +# ==================== Redis ==================== +REDIS_HOST = str(CONFIG["REDIS_HOST"]).strip() +REDIS_PORT = int(CONFIG["REDIS_PORT"]) +REDIS_DB = int(CONFIG["REDIS_DB"]) +CHANNEL_NAMESPACE = str(CONFIG["CHANNEL_NAMESPACE"]).strip() +REDIS_MAX_CONNECTIONS_DISPATCHER = int(CONFIG["REDIS_MAX_CONNECTIONS_DISPATCHER"]) + +# ==================== Tasks ==================== +MAX_RETRY_COUNT = int(CONFIG["MAX_RETRY_COUNT"]) +TASK_CSV_PATH = os.path.join(BASE_DIR, str(CONFIG["TASK_CSV_FILE"]).strip()) + +# ==================== Worker ==================== +WORKER_STALE_TIMEOUT = int(CONFIG["WORKER_STALE_TIMEOUT"]) +WORKER_REINIT_TIMEOUT = int(CONFIG["WORKER_REINIT_TIMEOUT"]) +WORKER_INIT_MONITOR_WINDOW = int(CONFIG["WORKER_INIT_MONITOR_WINDOW"]) +WORKER_ACTION_MAX_PARALLEL = int(CONFIG["WORKER_ACTION_MAX_PARALLEL"]) +SSH_CONNECT_TIMEOUT_SECONDS = int(CONFIG["SSH_CONNECT_TIMEOUT_SECONDS"]) +SSH_ACTION_TIMEOUT_SECONDS = int(CONFIG["SSH_ACTION_TIMEOUT_SECONDS"]) +AUTOOL_GIT_REPO_URL = str(CONFIG["AUTOOL_GIT_REPO_URL"]).strip() +MUMU_MANAGER_PATH = str(CONFIG["MUMU_MANAGER_PATH"]).strip() +MUMU_VM_INDEX = int(CONFIG["MUMU_VM_INDEX"]) +MUMU_MAX_VM_INDEX = int(CONFIG["MUMU_MAX_VM_INDEX"]) +CLEAN_BACKUP_PATH = str(CONFIG["CLEAN_BACKUP_PATH"]).strip() +LOCAL_IMPORT_DIR = str(CONFIG["LOCAL_IMPORT_DIR"]).strip() +NET_BRIDGE_CARD = str(CONFIG["NET_BRIDGE_CARD"]).strip() +EMULATOR_ADB_IP_OFFSET = int(CONFIG.get("EMULATOR_ADB_IP_OFFSET", 100)) +MUMU_BRIDGE_IP_OFFSET = int(CONFIG.get("MUMU_BRIDGE_IP_OFFSET", 100)) +MUMU_NETWORK_GATEWAYS = dict( + CONFIG.get("MUMU_NETWORK_GATEWAYS") + or {"192.168.1": "192.168.1.1", "192.168.2": "192.168.2.1"} +) +WORKER_ALLOWED_IP_PATTERNS = tuple(str(item).strip() for item in CONFIG.get("WORKER_ALLOWED_IP_PATTERNS", []) if str(item).strip()) +SSH_DEFAULT_USER = str(CONFIG["SSH_DEFAULT_USER"]).strip() +SSH_DEFAULT_PASSWORD = str(CONFIG["SSH_DEFAULT_PASSWORD"]).strip() +SSH_DEFAULT_PORT = int(CONFIG["SSH_DEFAULT_PORT"]) +PSEXEC_DEFAULT_SESSION_ID = int(CONFIG["PSEXEC_DEFAULT_SESSION_ID"]) +MUMU_RESTART_SETTLE_SECONDS = int(CONFIG.get("MUMU_RESTART_SETTLE_SECONDS", 15)) +RUNBATCH_LEGACY_END_WORKER_IPS = frozenset( + str(item).strip() + for item in CONFIG.get("RUNBATCH_LEGACY_END_WORKER_IPS", ["192.168.1.51", "192.168.1.61", "192.168.2.101"]) + if str(item).strip() +) +RUNBATCH_FORCE_KILL_IMAGES = tuple( + str(item).strip() + for item in CONFIG.get("RUNBATCH_FORCE_KILL_IMAGES", ["OpenConsole.exe", "WindowsTerminal.exe", "conhost.exe", "powershell.exe"]) + if str(item).strip() +) +MUMU_RECOVER_SCRIPT_PATH = str(CONFIG.get("MUMU_RECOVER_SCRIPT_PATH", r"\\your-file-server\share\recover_mumu_image.cmd")).strip() +MUMU_NETWORK_SCRIPT_PATH = str(CONFIG.get("MUMU_NETWORK_SCRIPT_PATH", r"\\your-file-server\share\recover_mumu_network.cmd")).strip() +MUMU_RECOVER_AHK_EXE = str(CONFIG.get("MUMU_RECOVER_AHK_EXE", r"\\your-file-server\share\mumu_recover.exe")).strip() + +# ==================== Backward compatibility ==================== +SHARE_SMB_USER = str(CONFIG.get("SHARE_SMB_USER") or "").strip() +SHARE_SMB_PASSWORD = str(CONFIG.get("SHARE_SMB_PASSWORD") or "").strip() + + +def get_repo_url(worker_ip: str) -> str: + del worker_ip + return AUTOOL_GIT_REPO_URL + + +def get_clean_backup_path(worker_ip: str) -> str: + del worker_ip + return CLEAN_BACKUP_PATH + + +def get_share_smb_target(worker_ip: str) -> str: + del worker_ip + path = CLEAN_BACKUP_PATH + if not path.startswith("\\\\"): + return "" + parts = [part for part in path.split("\\") if part] + if not parts: + return "" + return f"\\\\{parts[0]}" + +# ==================== Output files ==================== +FAILED_TASKS_CSV = os.path.join(REPORT_DIR, "failed_tasks.csv") +RETRY_TASKS_CSV = os.path.join(REPORT_DIR, "retry_tasks.csv") +SUCCESS_TASKS_CSV = os.path.join(REPORT_DIR, "success_tasks.csv") +FINAL_REPORT_CSV = os.path.join(REPORT_DIR, "final_report.csv") +WORKER_REPORT_CSV = os.path.join(REPORT_DIR, "worker_report.csv") + +# ==================== Logging ==================== +LOG_DIR = os.path.join(INSTANCE_DIR, "logs") +LOG_FILE = str(CONFIG["LOG_FILE"]).strip() +LOG_MAX_BYTES = int(CONFIG["LOG_MAX_BYTES"]) +LOG_BACKUP_COUNT = int(CONFIG["LOG_BACKUP_COUNT"]) +LOG_FORMAT = f"%(asctime)s [{INSTANCE_NAME}] [%(levelname)s] %(message)s" +LOG_DATE_FORMAT = str(CONFIG["LOG_DATE_FORMAT"]).strip() + +# ==================== Notifications ==================== +WECHAT_TOKENS = dict(CONFIG.get("WECHAT_TOKENS") or {}) +WECOM_TOKENS = dict(CONFIG.get("WECOM_TOKENS") or {}) + +# ==================== Alert strategy ==================== +ALERT_WINDOW_SECONDS = int(CONFIG["ALERT_WINDOW_SECONDS"]) +ALERT_THRESHOLD = int(CONFIG["ALERT_THRESHOLD"]) +ALERT_COOLDOWN_SECONDS = int(CONFIG["ALERT_COOLDOWN_SECONDS"]) + +# ==================== Status push ==================== +WORKER_STATUS_PUSH_INTERVAL = int(CONFIG["WORKER_STATUS_PUSH_INTERVAL"]) + +RUN_PIPELINE_OPTION_KEYS = tuple(CONFIG.get("RUN_PIPELINE_OPTION_KEYS") or ()) +RUN_PIPELINE_DEFAULT_OPTIONS = dict(CONFIG.get("RUN_PIPELINE_DEFAULT_OPTIONS") or {}) +AUTO_REBOOT_RECOVERY_DELAY_SECONDS = int(CONFIG["AUTO_REBOOT_RECOVERY_DELAY_SECONDS"]) + + +# ==================== Worker registry ==================== +WORKER_INVENTORY_CONFIG_KEY = "WORKER_INVENTORY" +WORKER_INVENTORY_PATH = CONFIG_PATH + + +def _resolve_ssh_endpoint(ssh_target: str, ssh_port: int) -> Tuple[str, int]: + """直连模式：解析 SSH 目标地址，验证合法性后直接返回""" + host = str(ssh_target).split("@")[-1].strip() + ip = extract_worker_ip(host) + if ip != host: + return host, ssh_port + if WORKER_ALLOWED_IP_PATTERNS and not any(re.fullmatch(pattern, ip) for pattern in WORKER_ALLOWED_IP_PATTERNS): + raise ValueError(f"worker ip is not allowed by WORKER_ALLOWED_IP_PATTERNS: {ssh_target}") + return host, ssh_port + + +def extract_worker_ip(value: Any) -> str: + text = str(value or "").strip() + match = re.search(r"(? str: + return extract_worker_ip(ip_address) or extract_worker_ip(worker_id) or str(worker_id or "").strip() + + +def _normalize_worker_entry(raw: Dict[str, Any]) -> Dict[str, Any]: + ssh_target = str(raw["ssh_target"]).strip() + worker_id = normalize_worker_id(raw.get("worker_id", ""), ssh_target) + if not worker_id: + raise ValueError("worker inventory entry missing worker_id") + if not ssh_target: + raise ValueError(f"worker inventory entry missing ssh_target: {worker_id}") + ssh_user = str(raw.get("ssh_user") or SSH_DEFAULT_USER).strip() + ssh_password = str(raw.get("ssh_password") or "").strip() + ssh_port = int(raw.get("ssh_port") or SSH_DEFAULT_PORT) + ssh_host, ssh_port = _resolve_ssh_endpoint(ssh_target, ssh_port) + repo_dir = str(raw["repo_dir"]).strip() + if not repo_dir: + raise ValueError(f"worker inventory entry missing repo_dir: {worker_id}") + python_exe = str(raw["python_exe"]).strip() + if not python_exe: + raise ValueError(f"worker inventory entry missing python_exe: {worker_id}") + psexec_session_id = int(raw.get("psexec_session_id") or PSEXEC_DEFAULT_SESSION_ID) + tags = [] + seen = set() + for tag in raw.get("tags", []) or []: + candidate = str(tag).strip() + if candidate and candidate not in seen: + seen.add(candidate) + tags.append(candidate) + + item = { + "worker_id": worker_id, + "ssh_target": ssh_target, + "ssh_host": ssh_host, + "ssh_port": ssh_port, + "ssh_user": ssh_user, + "repo_dir": repo_dir, + "python_exe": python_exe, + "psexec_session_id": psexec_session_id, + "tags": tags, + } + if ssh_password: + item["ssh_password"] = ssh_password + return item + + +def _worker_matches_env(worker: Dict[str, Any], env_name: str) -> bool: + if env_name == "test": + return "test" in worker.get("tags", []) + return True + + +def load_worker_inventory(path: Optional[str] = None, env_name: Optional[str] = None) -> List[Dict[str, Any]]: + normalized_env = normalize_env_name(CONFIG_ENV if env_name is None else env_name) + target_path = path + if target_path: + if not os.path.exists(target_path): + return [] + payload = _load_config_payload(target_path, required=True) + if isinstance(payload, dict): + payload = payload.get(WORKER_INVENTORY_CONFIG_KEY, []) + else: + payload = CONFIG.get(WORKER_INVENTORY_CONFIG_KEY, []) + + if not isinstance(payload, list): + raise ValueError(f"worker inventory must be a list: {target_path}") + + normalized: List[Dict[str, Any]] = [] + seen = set() + for item in payload: + if not isinstance(item, dict): + raise ValueError(f"worker inventory entry must be an object: {item!r}") + normalized_item = _normalize_worker_entry(item) + worker_id = normalized_item["worker_id"] + if worker_id in seen: + raise ValueError(f"duplicate worker_id in worker inventory: {worker_id}") + seen.add(worker_id) + if not _worker_matches_env(normalized_item, normalized_env): + continue + normalized.append(normalized_item) + return normalized + + +WORKER_INVENTORY = load_worker_inventory() +MANAGED_WORKER_IDS = [item["worker_id"] for item in WORKER_INVENTORY] +ONLY_MANAGED_WORKERS_CAN_DISPATCH = bool(CONFIG.get("ONLY_MANAGED_WORKERS_CAN_DISPATCH", False)) + +# Exact-match routing rules for targeted dispatch. +# Values must be worker_id strings; worker IDs are normalized to the worker IP. +TASK_ROUTING_RULES = dict(CONFIG.get("TASK_ROUTING_RULES") or {"package_name": {}, "task_key": {}}) + +# ==================== Dashboard ==================== +DASHBOARD_HOST = str(CONFIG["DASHBOARD_HOST"]).strip() +DASHBOARD_PORT = int(CONFIG["DASHBOARD_PORT"]) +WORKER_ONLINE_TIMEOUT_SECONDS = WORKER_STALE_TIMEOUT + +# ==================== Monitoring ==================== +MONITORING_DB_PATH = os.path.join(INSTANCE_DIR, "monitoring.sqlite3") +MONITORING_TIMEZONE = str(CONFIG["MONITORING_TIMEZONE"]).strip() +MONITORING_TIMELINE_BUCKET_MINUTES = int(CONFIG["MONITORING_TIMELINE_BUCKET_MINUTES"]) + +# ==================== Analytics ==================== +ANALYTICS_TRAFFIC_ROOT = str(CONFIG["ANALYTICS_TRAFFIC_ROOT"]).strip() +ANALYTICS_TRAFFIC_ROOT_BLOCK = str(CONFIG["ANALYTICS_TRAFFIC_ROOT_BLOCK"]).strip() +ANALYTICS_TRAVERSAL_ROOT = str(CONFIG["ANALYTICS_TRAVERSAL_ROOT"]).strip() +ANALYTICS_TPDPI_APP_LIST = os.path.join( + BASE_DIR, + "calculate", + str(CONFIG["ANALYTICS_TPDPI_APP_LIST_FILE"]).strip(), +) +ANALYTICS_TPDPI_URL_LIB = os.path.join( + BASE_DIR, + "calculate", + str(CONFIG["ANALYTICS_TPDPI_URL_LIB_FILE"]).strip(), +) +ANALYTICS_ARTIFACT_WAIT_SECONDS = int(CONFIG["ANALYTICS_ARTIFACT_WAIT_SECONDS"]) +ANALYTICS_ARTIFACT_FILE_WAIT_SECONDS = int(CONFIG["ANALYTICS_ARTIFACT_FILE_WAIT_SECONDS"]) +ANALYTICS_JOB_POLL_SECONDS = int(CONFIG["ANALYTICS_JOB_POLL_SECONDS"]) +MODEL_TRAFFIC_THRESHOLD = int(CONFIG["MODEL_TRAFFIC_THRESHOLD"]) + + +def channel_name(name: str) -> str: + return f"{CHANNEL_NAMESPACE}:{name}" if CHANNEL_NAMESPACE else name + + +# ==================== MinIO / APK Cloud Storage ==================== +# 功能开关：是否启用 Minio 上传下载功能 +# - True: 使用 Minio 进行 APK 上传和分发（原团队内网模式） +# - False: 禁用 Minio，使用直接下载模式（跨团队部署推荐） +MINIO_ENABLED = bool(CONFIG.get("MINIO_ENABLED", True)) + +# APK 下载模式（当 MINIO_ENABLED=False 时生效） +# - "direct": Worker 直接从下载链接获取 APK（跨团队部署推荐） +# - "smb": 从 SMB 共享目录读取 APK（需要内网访问权限） +APK_DOWNLOAD_MODE = str(CONFIG.get("APK_DOWNLOAD_MODE", "direct")).strip().lower() + +# Minio 连接配置（仅当 MINIO_ENABLED=True 时必需） +MINIO_ENDPOINT = str(CONFIG.get("MINIO_ENDPOINT", "")).strip() +MINIO_ACCESS_KEY = str(CONFIG.get("MINIO_ACCESS_KEY", "")).strip() +MINIO_SECRET_KEY = str(CONFIG.get("MINIO_SECRET_KEY", "")).strip() +MINIO_BUCKET = str(CONFIG.get("MINIO_BUCKET", "autool-apk")).strip() +MINIO_SECURE = bool(CONFIG.get("MINIO_SECURE", False)) + +# APK 存储路径配置 +APK_LOCAL_STORAGE_DIR = str(CONFIG.get("APK_LOCAL_STORAGE_DIR", "")).strip() +APK_SMB_DIR = str(CONFIG.get("APK_SMB_DIR", "")).strip() +APK_US_EXPORT_DIR = str(CONFIG.get("APK_US_EXPORT_DIR", "D:\\mumu_apks")).strip() + +# US 拓扑设备配置（仅当使用 US 拓扑模式时需要） +APK_US_DEVICE_A_SERIAL = str(CONFIG.get("APK_US_DEVICE_A_SERIAL", "127.0.0.1:7555")).strip() +APK_US_DEVICE_B_SERIAL = str(CONFIG.get("APK_US_DEVICE_B_SERIAL", "127.0.0.1:7556")).strip() +APK_US_DEVICE_A_VM_INDEX = int(CONFIG.get("APK_US_DEVICE_A_VM_INDEX", 2)) +APK_US_DEVICE_B_VM_INDEX = int(CONFIG.get("APK_US_DEVICE_B_VM_INDEX", 1)) + +# APK 管理配置 +APK_PREFETCH_POLL_INTERVAL = int(CONFIG.get("APK_PREFETCH_POLL_INTERVAL", 120)) +APK_DOWNLOAD_QUEUE_INTERVAL = int(CONFIG.get("APK_DOWNLOAD_QUEUE_INTERVAL", 60)) +APK_MINIO_MAX_BYTES = int(CONFIG.get("APK_MINIO_MAX_BYTES", 107374182400)) # 100GB +APK_EMULATOR_CLEANUP_ENABLED = bool(CONFIG.get("APK_EMULATOR_CLEANUP_ENABLED", True)) diff --git a/config/config.example.json b/config/config.example.json new file mode 100644 index 0000000..296ce70 --- /dev/null +++ b/config/config.example.json @@ -0,0 +1,132 @@ +{ + "CONFIG_ENV": "prod", + "INSTANCE_NAME": "main", + "REDIS_HOST": "127.0.0.1", + "REDIS_PORT": 6379, + "REDIS_DB": 0, + "CHANNEL_NAMESPACE": "main", + "REDIS_MAX_CONNECTIONS_DISPATCHER": 50, + "MAX_RETRY_COUNT": 3, + "TASK_CSV_FILE": "package_list.csv", + "WORKER_STALE_TIMEOUT": 9000, + "WORKER_REINIT_TIMEOUT": 600, + "WORKER_INIT_MONITOR_WINDOW": 1800, + "WORKER_ACTION_MAX_PARALLEL": 50, + "SSH_CONNECT_TIMEOUT_SECONDS": 10, + "SSH_ACTION_TIMEOUT_SECONDS": 1800, + "AUTOOL_GIT_REPO_URL": "//your-file-server/share/autool.git", + "MUMU_MANAGER_PATH": "C:\\Program Files\\Netease\\MuMu\\nx_main\\MuMuManager.exe", + "MUMU_VM_INDEX": 2, + "MUMU_MAX_VM_INDEX": 10, + "MUMU_RESTART_SETTLE_SECONDS": 15, + "MUMU_RECOVER_SCRIPT_PATH": "\\\\your-file-server\\share\\recover_mumu_image.cmd", + "MUMU_RECOVER_AHK_EXE": "\\\\your-file-server\\share\\mumu_recover.exe", + "CLEAN_BACKUP_PATH": "\\\\your-file-server\\share\\autool_config\\taskagent.mumudata", + "LOCAL_IMPORT_DIR": "D:\\mumu_backups", + "NET_BRIDGE_CARD": "Realtek PCIe GbE Family Controller", + "EMULATOR_ADB_IP_OFFSET": 100, + "MUMU_BRIDGE_IP_OFFSET": 100, + "MUMU_NETWORK_SCRIPT_PATH": "\\\\your-file-server\\share\\recover_mumu_network.cmd", + "MUMU_NETWORK_GATEWAYS": { + "192.168.1": "192.168.1.1", + "192.168.2": "192.168.2.1" + }, + "WORKER_ALLOWED_IP_PATTERNS": [ + "10\\.\\d{1,3}\\.\\d{1,3}\\.\\d{1,3}", + "172\\.(1[6-9]|2\\d|3[0-1])\\.\\d{1,3}\\.\\d{1,3}", + "192\\.168\\.\\d{1,3}\\.\\d{1,3}" + ], + "RUNBATCH_LEGACY_END_WORKER_IPS": [ + "192.168.1.51", + "192.168.1.61", + "192.168.2.101" + ], + "RUNBATCH_FORCE_KILL_IMAGES": [ + "OpenConsole.exe", + "WindowsTerminal.exe", + "conhost.exe", + "powershell.exe" + ], + "SSH_DEFAULT_USER": "admin", + "SSH_DEFAULT_PASSWORD": "", + "SHARE_SMB_USER": "", + "SHARE_SMB_PASSWORD": "", + "SSH_DEFAULT_PORT": 22, + "PSEXEC_DEFAULT_SESSION_ID": 1, + "LOG_FILE": "dispatcher.log", + "LOG_MAX_BYTES": 52428800, + "LOG_BACKUP_COUNT": 3, + "LOG_DATE_FORMAT": "%Y-%m-%d %H:%M:%S", + "WECHAT_TOKENS": {}, + "WECOM_TOKENS": {}, + "ALERT_WINDOW_SECONDS": 600, + "ALERT_THRESHOLD": 30, + "ALERT_COOLDOWN_SECONDS": 1800, + "WORKER_STATUS_PUSH_INTERVAL": 10800, + "RUN_PIPELINE_OPTION_KEYS": [ + "stop_worker", + "clone_if_missing", + "git_pull", + "setup", + "pull_pcap_files", + "recover_mumu", + "restart_mumu", + "start_worker" + ], + "RUN_PIPELINE_DEFAULT_OPTIONS": { + "stop_worker": true, + "clone_if_missing": false, + "git_pull": true, + "setup": false, + "pull_pcap_files": true, + "recover_mumu": false, + "restart_mumu": true, + "start_worker": true + }, + "AUTO_REBOOT_RECOVERY_DELAY_SECONDS": 120, + "ONLY_MANAGED_WORKERS_CAN_DISPATCH": false, + "TASK_ROUTING_RULES": { + "package_name": {}, + "task_key": {} + }, + "WORKER_INVENTORY": [ + { + "worker_id": "192.168.1.10", + "ssh_target": "192.168.1.10", + "repo_dir": "D:/autool", + "python_exe": "python", + "tags": [] + } + ], + "DASHBOARD_HOST": "0.0.0.0", + "DASHBOARD_PORT": 8890, + "MONITORING_TIMEZONE": "Asia/Shanghai", + "MONITORING_TIMELINE_BUCKET_MINUTES": 15, + "ANALYTICS_TRAFFIC_ROOT": "\\\\your-file-server\\share\\autool_config\\data\\traffic_data", + "ANALYTICS_TRAFFIC_ROOT_BLOCK": "\\\\your-file-server\\share\\autool_config\\data\\traffic_data_block", + "ANALYTICS_TRAVERSAL_ROOT": "\\\\your-file-server\\share\\autool_config\\data\\traversal_log", + "ANALYTICS_TPDPI_APP_LIST_FILE": "TPDPI_app_list.csv", + "ANALYTICS_TPDPI_URL_LIB_FILE": "TPDPI_url_lib.csv", + "ANALYTICS_ARTIFACT_WAIT_SECONDS": 180, + "ANALYTICS_ARTIFACT_FILE_WAIT_SECONDS": 60, + "ANALYTICS_JOB_POLL_SECONDS": 1, + "MODEL_TRAFFIC_THRESHOLD": 50, + "MINIO_ENABLED": false, + "APK_DOWNLOAD_MODE": "direct", + "MINIO_ENDPOINT": "", + "MINIO_ACCESS_KEY": "", + "MINIO_SECRET_KEY": "", + "MINIO_BUCKET": "autool-apk", + "MINIO_SECURE": false, + "APK_LOCAL_STORAGE_DIR": "", + "APK_SMB_DIR": "", + "APK_US_EXPORT_DIR": "C:\\mumu_apks", + "APK_US_DEVICE_A_SERIAL": "127.0.0.1:7555", + "APK_US_DEVICE_A_VM_INDEX": 2, + "APK_US_DEVICE_B_SERIAL": "127.0.0.1:7556", + "APK_US_DEVICE_B_VM_INDEX": 1, + "APK_PREFETCH_POLL_INTERVAL": 120, + "APK_DOWNLOAD_QUEUE_INTERVAL": 60, + "APK_MINIO_MAX_BYTES": 107374182400, + "APK_EMULATOR_CLEANUP_ENABLED": true +} diff --git a/config_new.py b/config_new.py new file mode 100644 index 0000000..1cadaec --- /dev/null +++ b/config_new.py @@ -0,0 +1,78 @@ +# -*- encoding=utf8 -*- +""" +Simplified configuration loader for dispatcher. +统一使用 config.yaml，向后兼容旧的 JSON 格式。 +""" + +import json +import os +from typing import Any, Dict + +try: + import yaml + YAML_AVAILABLE = True +except ImportError: + YAML_AVAILABLE = False + + +BASE_DIR = os.path.dirname(os.path.abspath(__file__)) + +# 新配置路径 +CONFIG_YAML_PATH = os.path.join(BASE_DIR, "config.yaml") +CONFIG_EXAMPLE_YAML_PATH = os.path.join(BASE_DIR, "config.example.yaml") + + +def load_config_file(path: str) -> Dict[str, Any]: + """加载配置文件（支持 JSON 和 YAML）""" + if not os.path.exists(path): + return {} + + with open(path, "r", encoding="utf-8") as f: + if path.endswith(('.yaml', '.yml')): + if not YAML_AVAILABLE: + raise ImportError("PyYAML is required. Install: pip install pyyaml") + return yaml.safe_load(f) or {} + else: + return json.load(f) + + +def load_dispatcher_config() -> Dict[str, Any]: + """ + 加载 dispatcher 配置 + + 优先级： + 1. config.yaml（新格式，推荐） + 2. config/config.example.json（向后兼容） + """ + # 尝试新格式 + if os.path.exists(CONFIG_YAML_PATH): + print(f"[Config] Loading from {CONFIG_YAML_PATH}") + return load_config_file(CONFIG_YAML_PATH) + + # 尝试示例配置 + if os.path.exists(CONFIG_EXAMPLE_YAML_PATH): + print(f"[Config] Loading from {CONFIG_EXAMPLE_YAML_PATH}") + return load_config_file(CONFIG_EXAMPLE_YAML_PATH) + + # 向后兼容：尝试旧的 JSON 配置 + old_config_path = os.path.join(BASE_DIR, "config", "config.example.json") + if os.path.exists(old_config_path): + print(f"[Config] Loading from {old_config_path} (legacy)") + print("[Config] Consider migrating to config.yaml") + return load_config_file(old_config_path) + + raise FileNotFoundError( + "No configuration file found. Expected one of:\n" + f" - {CONFIG_YAML_PATH}\n" + f" - {CONFIG_EXAMPLE_YAML_PATH}\n" + "Create config.yaml from config.example.yaml" + ) + + +# 测试 +if __name__ == "__main__": + config = load_dispatcher_config() + print(f"\n✓ Config loaded successfully") + print(f" Instance: {config.get('INSTANCE_NAME')}") + print(f" Redis: {config.get('REDIS_HOST')}:{config.get('REDIS_PORT')}") + print(f" Minio: {config.get('MINIO_ENABLED')}") diff --git a/dashboard_main.py b/dashboard_main.py new file mode 100644 index 0000000..a3f0781 --- /dev/null +++ b/dashboard_main.py @@ -0,0 +1,4931 @@ +import json +import logging +import traceback +from concurrent.futures import ThreadPoolExecutor, as_completed +import threading +import time +import uuid +from datetime import datetime +from typing import Any, Dict, List, Optional + +from config import ( + DASHBOARD_HOST, + DASHBOARD_PORT, + INSTANCE_NAME, + MONITORING_TIMEZONE, + REDIS_DB, + RUN_PIPELINE_DEFAULT_OPTIONS, + WORKER_ACTION_MAX_PARALLEL, +) +from redis_task_distribute import RedisTaskDispatcher + +try: + from flask import Flask, jsonify, render_template_string, request +except ImportError: # pragma: no cover - runtime dependency check + Flask = None + jsonify = None + render_template_string = None + request = None + + +def silence_flask_runtime_logs(): + try: + import flask.cli as flask_cli + + flask_cli.show_server_banner = lambda *args, **kwargs: None + except Exception: + pass + + try: + import werkzeug.serving as werkzeug_serving + + werkzeug_serving._log = lambda *args, **kwargs: None + except Exception: + pass + + logging.getLogger("werkzeug").disabled = True + + +class WorkerActionJobManager: + def __init__(self, dispatcher): + self.dispatcher = dispatcher + self._lock = threading.Lock() + self._jobs: Dict[str, Dict[str, Any]] = {} + + def submit(self, worker_ids: List[str], action: str, options: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: + job_id = uuid.uuid4().hex + now_ts = time.time() + job = { + "job_id": job_id, + "action": action, + "worker_ids": list(worker_ids), + "options": dict(options or {}), + "status": "running", + "started_at": now_ts, + "finished_at": 0.0, + "error": "", + "results": [ + { + "worker_id": worker_id, + "ok": None, + "message": "等待执行", + "steps": [], + } + for worker_id in worker_ids + ], + } + with self._lock: + self._jobs[job_id] = job + thread = threading.Thread( + target=self._run_job, + args=(job_id, list(worker_ids), action, dict(options or {})), + daemon=True, + ) + thread.start() + return self.get(job_id) + + def get(self, job_id: str) -> Optional[Dict[str, Any]]: + with self._lock: + job = self._jobs.get(job_id) + if not job: + return None + return json.loads(json.dumps(job)) + + def _run_job(self, job_id: str, worker_ids: List[str], action: str, options: Dict[str, Any]) -> None: + results: List[Dict[str, Any]] = [] + try: + max_workers = min(max(len(worker_ids), 1), WORKER_ACTION_MAX_PARALLEL) + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = { + executor.submit( + self.dispatcher.run_dashboard_action, + [worker_id], + action, + options=options, + progress_callback=lambda current_worker_id, payload: self._record_progress(job_id, current_worker_id, payload), + ): worker_id + for worker_id in worker_ids + } + for future in as_completed(futures): + worker_id = futures[future] + try: + worker_results = future.result() + except Exception as exc: + detail = traceback.format_exc().strip() + worker_results = [{ + "worker_id": worker_id, + "action": action, + "ok": False, + "message": detail or str(exc), + "stderr": detail, + }] + results.extend(worker_results) + except Exception as exc: + with self._lock: + job = self._jobs.get(job_id) + if not job: + return + job["status"] = "failed" + job["finished_at"] = time.time() + job["error"] = traceback.format_exc().strip() or str(exc) + return + + with self._lock: + job = self._jobs.get(job_id) + if not job: + return + by_worker = {item["worker_id"]: item for item in job["results"]} + for result in results: + item = by_worker.setdefault( + result.get("worker_id", ""), + {"worker_id": result.get("worker_id", ""), "ok": None, "message": "", "steps": []}, + ) + item["ok"] = result.get("ok") + item["message"] = self._result_message(result) + if isinstance(result.get("steps"), list): + item["steps"] = [ + { + "step": str(step.get("step", "")), + "ok": step.get("ok"), + "message": str(step.get("message", "")), + } + for step in result.get("steps", []) + if isinstance(step, dict) + ] + item["stdout"] = result.get("stdout", "") + item["stderr"] = result.get("stderr", "") + item["command"] = result.get("command", "") + job["results"] = list(by_worker.values()) + job["finished_at"] = time.time() + job["status"] = "failed" if any(item.get("ok") is False for item in job["results"]) else "completed" + + @staticmethod + def _result_message(result: Dict[str, Any]) -> str: + base = str(result.get("message", "") or "").strip() + prefixes = [] + failed_step = str(result.get("failed_step", "") or "").strip() + stage = str(result.get("stage", "") or "").strip() + if failed_step: + prefixes.append(f"step={failed_step}") + if stage: + prefixes.append(f"stage={stage}") + if prefixes: + prefix = "[" + " ".join(prefixes) + "]" + base = f"{prefix} {base}".strip() + if result.get("ok") is not False: + return base + parts = [base] if base else [] + for label, key in (("stdout", "stdout"), ("stderr", "stderr")): + value = str(result.get(key, "") or "").strip() + if not value or value == base: + continue + parts.append(f"{label}:\n{value}") + return "\n\n".join(parts) + + def _record_progress(self, job_id: str, worker_id: str, payload: Dict[str, Any]) -> None: + with self._lock: + job = self._jobs.get(job_id) + if not job: + return + result = None + for item in job["results"]: + if item["worker_id"] == worker_id: + result = item + break + if result is None: + result = {"worker_id": worker_id, "ok": None, "message": "", "steps": []} + job["results"].append(result) + result["message"] = str(payload.get("message", result.get("message", ""))) + step_name = str(payload.get("step", "")).strip() + if not step_name: + return + status = payload.get("status") + step_ok = None + if status in {"success", "ok"}: + step_ok = True + elif status in {"failed", "error"}: + step_ok = False + steps = result.setdefault("steps", []) + current_step = None + for step in steps: + if step.get("step") == step_name: + current_step = step + break + if current_step is None: + current_step = {"step": step_name, "ok": step_ok, "message": result["message"]} + steps.append(current_step) + else: + current_step["message"] = result["message"] + current_step["ok"] = step_ok + + +TEMPLATE = """ + + + + + + {{ instance_name }} Monitoring Dashboard + + + + + + + +
+
+
Control Center / {{ instance_name }}
+
+
+

Worker Monitoring

+
+
+
+ + + + +
+
+
+
+ 实例 {{ instance_name }} / Redis DB {{ redis_db }} + 日期 {{ default_date }} + 范围 全部 Worker + 今日视图自动刷新 +
+
+ +
+
+

运行快照

+
+
+
+
+
+ + + +
+
+

整体效率

+
+
+
+
+

今日采集概况

+
--
+
+
+
+
+
+
+
+
+
+

整体运行时间构成

+
--
+
+
+
+ +
+
+

合格任务阶段耗时

+
--
+
+
+
+ +
+
+

运行状态时间分布

+
--
+
+
+
+ +
+
+

各 Worker 采集占比

+
--
+
+
+
+
+
+ + + + + + + + +
+ +
+ + + + +""" + + +def create_app(dispatcher=None): + if Flask is None: + raise RuntimeError("Flask is not installed. Run `pip install -r requirements.txt` first.") + + dispatcher = dispatcher or RedisTaskDispatcher() + action_jobs = WorkerActionJobManager(dispatcher) + app = Flask(__name__) + + def _get_date_arg() -> str: + return request.args.get("date") or datetime.now().strftime("%Y-%m-%d") + + def _get_int_arg(name, default): + value = request.args.get(name, default) + try: + return int(value) + except (TypeError, ValueError): + return default + + def _parse_worker_action_request(): + payload = request.get_json(silent=True) or {} + worker_ids = payload.get("worker_ids", []) + action = str(payload.get("action", "")).strip() + allowed_actions = { + "run_pipeline", + "start", + "drain", + "recover_and_start", + "enable", + "disable", + "clone", + "pull", + "pull_pcap_files", + "setup", + "restart_mumu", + "stop_worker", + "reboot", + "status", + "execute_command", + "fix_adb_connection", + "configure_mumu_network", + "recover_mumu_full", + "force_reapk", + } + if not isinstance(worker_ids, list) or not worker_ids: + return None, (jsonify({"ok": False, "error": "worker_ids is required"}), 400) + if not action: + return None, (jsonify({"ok": False, "error": "action is required"}), 400) + if action not in allowed_actions: + return None, (jsonify({"ok": False, "error": f"unsupported action: {action}"}), 400) + return { + "worker_ids": worker_ids, + "action": action, + "options": payload.get("options"), + }, None + + @app.get("/") + def index(): + return render_template_string( + TEMPLATE, + instance_name=INSTANCE_NAME, + redis_db=REDIS_DB, + timezone=MONITORING_TIMEZONE, + run_pipeline_defaults=json.dumps(RUN_PIPELINE_DEFAULT_OPTIONS, ensure_ascii=True), + default_date=datetime.now().strftime("%Y-%m-%d"), + ) + + @app.get("/api/monitor/overview") + def monitor_overview(): + return jsonify( + dispatcher.get_monitoring_overview( + date_str=_get_date_arg(), + worker_id=request.args.get("worker_id"), + ) + ) + + @app.get("/api/monitor/timeline") + def monitor_timeline(): + return jsonify( + dispatcher.get_monitoring_timeline( + date_str=_get_date_arg(), + bucket_minutes=_get_int_arg("bucket_minutes", 15), + worker_id=request.args.get("worker_id"), + ) + ) + + @app.get("/api/monitor/distributions") + def monitor_distributions(): + return jsonify( + dispatcher.get_monitoring_distributions( + date_str=_get_date_arg(), + worker_id=request.args.get("worker_id"), + ) + ) + + @app.get("/api/monitor/workers") + def monitor_workers(): + return jsonify(dispatcher.get_monitoring_workers(date_str=_get_date_arg())) + + @app.get("/api/workers/control") + def worker_control(): + return jsonify( + dispatcher.get_worker_control_rows( + date_str=_get_date_arg(), + worker_id=request.args.get("worker_id"), + ) + ) + + @app.post("/api/workers/actions") + def worker_actions(): + try: + parsed, error_response = _parse_worker_action_request() + if error_response is not None: + return error_response + results = dispatcher.run_dashboard_action(parsed["worker_ids"], parsed["action"], options=parsed["options"]) + return jsonify({"ok": True, "results": results}) + except Exception as e: + logger = logging.getLogger("worker-actions") + logger.exception("Failed to run worker action") + return jsonify({"ok": False, "error": str(e)}), 500 + + @app.post("/api/workers/actions/submit") + def worker_actions_submit(): + try: + parsed, error_response = _parse_worker_action_request() + if error_response is not None: + return error_response + job = action_jobs.submit(parsed["worker_ids"], parsed["action"], options=parsed["options"]) + return jsonify({"ok": True, "job": job}) + except Exception as e: + logger = logging.getLogger("worker-actions-submit") + logger.exception("Failed to submit worker action") + return jsonify({"ok": False, "error": str(e)}), 500 + + @app.get("/api/workers/actions/status") + def worker_actions_status(): + job_id = str(request.args.get("job_id", "")).strip() + if not job_id: + return jsonify({"ok": False, "error": "job_id is required"}), 400 + job = action_jobs.get(job_id) + if not job: + return jsonify({"ok": False, "error": "job not found"}), 404 + return jsonify({"ok": True, "job": job}) + + @app.get("/api/monitor/tasks") + def monitor_tasks(): + return jsonify( + dispatcher.get_monitoring_tasks( + date_str=_get_date_arg(), + worker_id=request.args.get("worker_id"), + limit=_get_int_arg("limit", 60), + ) + ) + + @app.post("/api/monitor/reset") + def monitor_reset(): + try: + return jsonify(dispatcher.reset_monitoring_dashboard()) + except Exception as e: + logger = logging.getLogger("monitor-reset") + logger.exception("Failed to reset monitoring data") + return jsonify({"ok": False, "error": str(e)}), 500 + + @app.get("/api/analytics/overview") + def analytics_overview(): + return jsonify( + dispatcher.get_analytics_overview( + incremental_batch_tag=str(request.args.get("incremental_batch_tag", "")).strip(), + top_n=_get_int_arg("top_n", 3000), + ) + ) + + @app.get("/api/analytics/apps") + def analytics_apps(): + return jsonify( + dispatcher.get_analytics_apps( + q=str(request.args.get("q", "")).strip(), + latest_status=str(request.args.get("latest_status", "")).strip(), + artifact_status=str(request.args.get("artifact_status", "")).strip(), + collection_status=str(request.args.get("collection_status", "")).strip(), + restriction_status=str(request.args.get("restriction_status", "")).strip(), + retryability=str(request.args.get("retryability", "")).strip(), + incremental_batch_tag=str(request.args.get("incremental_batch_tag", "")).strip(), + top_n=_get_int_arg("top_n", 3000), + sort=str(request.args.get("sort", "updated_at")).strip(), + order=str(request.args.get("order", "desc")).strip(), + page=_get_int_arg("page", 1), + page_size=_get_int_arg("page_size", 50), + ) + ) + + @app.get("/api/analytics/apps/") + def analytics_app_detail(package_name): + payload = dispatcher.get_analytics_app_detail(package_name) + if payload is None: + return jsonify({"ok": False, "error": "package not found"}), 404 + return jsonify(payload) + + @app.post("/api/analytics/apps//rebuild") + def analytics_app_rebuild(package_name): + if not str(package_name or "").strip(): + return jsonify({"ok": False, "error": "package_name is required"}), 400 + job = dispatcher.enqueue_analytics_rebuild(str(package_name).strip()) + return jsonify({"ok": True, "job": job}) + + @app.get("/api/analytics/jobs") + def analytics_jobs(): + return jsonify( + dispatcher.get_analytics_jobs( + job_type=str(request.args.get("job_type", "")).strip(), + status=str(request.args.get("status", "")).strip(), + package_name=str(request.args.get("package_name", "")).strip(), + limit=_get_int_arg("limit", 50), + ) + ) + + @app.post("/api/analytics/backfill") + def analytics_backfill(): + payload = request.get_json(silent=True) or {} + scope = str(payload.get("scope", "all") or "all").strip() + packages = payload.get("packages") or [] + if scope not in {"all", "packages"}: + return jsonify({"ok": False, "error": "scope must be all or packages"}), 400 + if scope == "packages": + if not isinstance(packages, list) or not packages: + return jsonify({"ok": False, "error": "packages is required when scope=packages"}), 400 + packages = [str(item).strip() for item in packages if str(item).strip()] + if not packages: + return jsonify({"ok": False, "error": "packages is required when scope=packages"}), 400 + else: + packages = None + job = dispatcher.enqueue_analytics_backfill(packages=packages) + return jsonify({"ok": True, "job": job}) + + return app + + +def run_dashboard(dispatcher=None, host=None, port=None): + silence_flask_runtime_logs() + app = create_app(dispatcher=dispatcher) + app.logger.disabled = True + app.run( + host=host or DASHBOARD_HOST, + port=port or DASHBOARD_PORT, + debug=False, + use_reloader=False, + threaded=True, + ) + + +def main(): + run_dashboard() + + +if __name__ == "__main__": + main() diff --git a/dispatcher_main.py b/dispatcher_main.py new file mode 100755 index 0000000..b1f9d7e --- /dev/null +++ b/dispatcher_main.py @@ -0,0 +1,690 @@ +import csv +import json +import os +import threading +import time +import ctypes +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime + +from config import ( + channel_name, + DASHBOARD_HOST, + DASHBOARD_PORT, + CHANNEL_NAMESPACE, + FINAL_REPORT_CSV, + INSTANCE_NAME, + REDIS_HOST, + REDIS_DB, + REDIS_PORT, + REPORT_DIR, + WORKER_REPORT_CSV, + WORKER_STALE_TIMEOUT, + WORKER_STATUS_PUSH_INTERVAL, + APK_DOWNLOAD_QUEUE_INTERVAL, + APK_PREFETCH_POLL_INTERVAL, +) +from dashboard_main import run_dashboard +from log_manager import logger +from redis_task_distribute import RedisTaskDispatcher + + +def _set_console_title(): + title = ( + f"AUTOOL Dispatcher [{INSTANCE_NAME}] " + f"Redis {REDIS_PORT}/{REDIS_DB} Dashboard {DASHBOARD_PORT} " + f"Channel {CHANNEL_NAMESPACE or 'default'}" + ) + if os.name != "nt": + return + try: + ctypes.windll.kernel32.SetConsoleTitleW(title) + except Exception: + pass + + +class DispatcherService: + def __init__(self, redis_host=REDIS_HOST, redis_port=REDIS_PORT, redis_db=REDIS_DB): + self.dispatcher = self._wait_for_redis(redis_host, redis_port, redis_db) + self.redis = self.dispatcher.redis + version = os.environ.get("AUTOOL_DISPATCHER_VERSION", INSTANCE_NAME) + self.dispatcher.monitor.start_controller_session(version=version) + self.running = True + self.pubsub_list = [] + self._threads = [] + self._lock = threading.Lock() + self._report_executor = ThreadPoolExecutor(max_workers=32, thread_name_prefix="worker-report") + + def _wait_for_redis(self, redis_host, redis_port, redis_db, retry_interval=5): + while True: + try: + dispatcher = RedisTaskDispatcher(redis_host=redis_host, redis_port=redis_port, redis_db=redis_db) + dispatcher.redis.ping() + logger.info(f"[Redis] 成功连接到 {redis_host}:{redis_port}/{redis_db}") + return dispatcher + except Exception as e: + logger.warning(f"[Redis] 连接失败: {e}，{retry_interval}秒后重试...") + time.sleep(retry_interval) + + def listen_worker_init(self): + pubsub = self.redis.pubsub() + with self._lock: + self.pubsub_list.append(pubsub) + pubsub.subscribe(channel_name("worker:init")) + + logger.info("[监听] Worker 初始化频道已启动") + + try: + while self.running: + message = pubsub.get_message(timeout=1) + if message is None or message["type"] != "message": + continue + + try: + request = json.loads(message["data"]) + task = self.dispatcher.worker_init( + request["worker_id"], + request["ip_address"], + request["mac_address"], + request["hostname"], + request["platform"], + request.get("device_type", ""), + ) + self.redis.publish( + channel_name(f"worker:init:response:{request['worker_id']}"), + json.dumps({"task": task}), + ) + except Exception as e: + logger.error(f"处理 Worker 初始化请求时出错: {e}") + finally: + pubsub.unsubscribe() + pubsub.close() + + def listen_worker_report(self): + pubsub = self.redis.pubsub() + with self._lock: + self.pubsub_list.append(pubsub) + pubsub.subscribe(channel_name("worker:report")) + + logger.info("[监听] Worker 上报频道已启动") + + try: + while self.running: + message = pubsub.get_message(timeout=1) + if message is None or message["type"] != "message": + continue + + try: + request = json.loads(message["data"]) + self._report_executor.submit(self._handle_worker_report_request, request) + except Exception as e: + logger.error(f"处理 Worker 上报请求时出错: {e}") + finally: + pubsub.unsubscribe() + pubsub.close() + + def _handle_worker_report_request(self, request): + worker_id = str(request.get("worker_id", "")).strip() + if not worker_id: + logger.error("处理 Worker 上报请求时出错: missing worker_id") + return + response_channel = channel_name(f"worker:report:response:{worker_id}") + try: + next_task = self.dispatcher.worker_report( + worker_id, + request["previous_task_key"], + request.get("report_data", {}), + ) + payload = {"task": next_task} + except Exception as e: + logger.exception(f"处理 Worker 上报请求时出错: worker={worker_id}") + payload = {"task": None, "error": str(e)} + self.redis.publish(response_channel, json.dumps(payload)) + + def listen_worker_event(self): + pubsub = self.redis.pubsub() + with self._lock: + self.pubsub_list.append(pubsub) + pubsub.subscribe(channel_name("worker:event")) + + logger.info("[监听] Worker 事件频道已启动") + + try: + while self.running: + message = pubsub.get_message(timeout=1) + if message is None or message["type"] != "message": + continue + + try: + request = json.loads(message["data"]) + if str(request.get("event_type") or "").strip() == "artifacts_synced": + self.dispatcher.handle_analytics_worker_event(request) + else: + self.dispatcher.monitor.handle_worker_event(request) + except Exception as e: + logger.error(f"处理 Worker 事件时出错: {e}") + finally: + pubsub.unsubscribe() + pubsub.close() + + def listen_worker_retry(self): + pubsub = self.redis.pubsub() + with self._lock: + self.pubsub_list.append(pubsub) + pubsub.subscribe(channel_name("worker:retry")) + + logger.info("[监听] Worker 重试频道已启动") + + try: + while self.running: + message = pubsub.get_message(timeout=1) + if message is None or message["type"] != "message": + continue + + try: + request = json.loads(message["data"]) + success = self.dispatcher.worker_retry( + request["worker_id"], + request["current_task_key"], + ) + self.redis.publish( + channel_name(f"worker:retry:response:{request['worker_id']}"), + json.dumps({"success": success}), + ) + except Exception as e: + logger.error(f"处理 Worker 重试请求时出错: {e}") + finally: + pubsub.unsubscribe() + pubsub.close() + + def listen_catalog_update(self): + pubsub = self.redis.pubsub() + with self._lock: + self.pubsub_list.append(pubsub) + pubsub.subscribe(channel_name("catalog:updated")) + + logger.info("[监听] 应用目录更新频道已启动") + + try: + while self.running: + message = pubsub.get_message(timeout=1) + if message is None or message["type"] != "message": + continue + + try: + data = json.loads(message["data"]) + added = int(data.get("added") or 0) + priority = str(data.get("priority") or "") + if priority == "high_block": + logger.info(f"[目录同步] 收到 block 任务加载通知 (added={added})，任务已直接入队，跳过 DB 刷新") + continue + logger.info(f"[目录同步] 收到 catalog:updated 通知 (added={added})，刷新任务队列...") + new_count = self.dispatcher.refresh_tasks_from_app_summary() + if new_count > 0: + logger.info(f"[目录同步] 已从数据库加载 {new_count} 个新任务") + else: + logger.info("[目录同步] 未发现新任务") + except Exception as e: + logger.error(f"[目录同步] 处理 catalog:updated 消息时出错: {e}") + finally: + pubsub.unsubscribe() + pubsub.close() + + def start_listeners(self): + threads = [] + + init_thread = threading.Thread(target=self.listen_worker_init, name="worker-init-listener") + init_thread.start() + threads.append(init_thread) + + report_thread = threading.Thread(target=self.listen_worker_report, name="worker-report-listener") + report_thread.start() + threads.append(report_thread) + + event_thread = threading.Thread(target=self.listen_worker_event, name="worker-event-listener") + event_thread.start() + threads.append(event_thread) + + retry_thread = threading.Thread(target=self.listen_worker_retry, name="worker-retry-listener") + retry_thread.start() + threads.append(retry_thread) + + catalog_thread = threading.Thread(target=self.listen_catalog_update, name="catalog-update-listener") + catalog_thread.start() + threads.append(catalog_thread) + + self._threads = threads + return threads + + def stop(self, clear_data=True): + logger.info("[停止] 正在停止服务...") + self.running = False + self.dispatcher.monitor.close_controller_session( + stop_reason="clear_data" if clear_data else "completed" + ) + + for thread in self._threads: + thread.join(timeout=3) + + with self._lock: + for pubsub in self.pubsub_list: + try: + pubsub.unsubscribe() + pubsub.close() + except Exception: + pass + self.pubsub_list.clear() + self._report_executor.shutdown(wait=False, cancel_futures=True) + self.dispatcher.close() + + if clear_data: + logger.info("[清理] 正在清空任务列表和 Worker 列表...") + self._clear_all_data() + logger.info("[清理] 清理完成") + + def _clear_all_data(self): + try: + self.redis.delete("task:queue:high") + self.redis.delete("task:queue:default") + self.redis.delete("task:queue:low") + self.redis.delete("task:queue") + self.redis.delete("task:status") + self.redis.delete("task:details") + self.redis.delete("task:completed") + self.redis.delete("task:failed") + self.redis.delete("workers:info") + self.redis.delete("worker:tasks") + self.redis.delete("workers:idle") + self.redis.delete("workers:busy") + self.redis.delete("workers:meta") + for key in self.redis.scan_iter(match="alert:*"): + self.redis.delete(key) + for key in self.redis.scan_iter(match="cooldown:*"): + self.redis.delete(key) + for key in self.redis.scan_iter(match="worker:recovery_observe:*"): + self.redis.delete(key) + for key in self.redis.scan_iter(match="worker:recovery_pending:*"): + self.redis.delete(key) + + logger.info("[清理] 已清空: task:queue:high/default/low, task:status, task:details, task:completed, task:failed") + logger.info("[清理] 已清空: workers:info, worker:tasks, workers:idle, workers:busy, workers:meta") + except Exception as e: + logger.error(f"[清理] 清理数据时出错: {e}") + + +def start_dashboard_thread(dispatcher): + def _target(): + try: + logger.info(f"[Dashboard:{INSTANCE_NAME}] 已启动: http://{DASHBOARD_HOST}:{DASHBOARD_PORT}") + run_dashboard(dispatcher=dispatcher, host=DASHBOARD_HOST, port=DASHBOARD_PORT) + except Exception as e: + logger.warning(f"[Dashboard] 未启动: {e}") + + thread = threading.Thread(target=_target, name="dashboard-server", daemon=True) + thread.start() + return thread + + +def generate_final_report(dispatcher): + try: + os.makedirs(os.path.dirname(FINAL_REPORT_CSV) or ".", exist_ok=True) + with open(FINAL_REPORT_CSV, mode="w", newline="", encoding="utf-8") as handle: + writer = csv.writer(handle) + writer.writerow(["任务键", "应用名称", "包名", "状态", "Worker ID", "IP地址", "开始时间", "结束时间"]) + + all_status = dispatcher.redis.hgetall("task:status") + for task_key, status_json in sorted(all_status.items()): + status = json.loads(status_json) + worker_id = status.get("worker_id", "N/A") + + task_details_json = dispatcher.redis.hget("task:details", task_key) + task_details = json.loads(task_details_json) if task_details_json else {} + + ip_address = "N/A" + if worker_id != "N/A": + worker_info_json = dispatcher.redis.hget("workers:info", worker_id) + if worker_info_json: + worker_info = json.loads(worker_info_json) + ip_address = worker_info.get("ip_address", "N/A") + + writer.writerow([ + task_key, + task_details.get("app_name", "N/A"), + task_details.get("package_name", "N/A"), + status.get("status", "unknown"), + worker_id, + ip_address, + time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(status.get("start_time", 0))), + time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(status.get("end_time", 0))), + ]) + + logger.info(f"最终报告已保存到: {FINAL_REPORT_CSV}") + except Exception as e: + logger.error(f"生成最终报告时出错: {e}") + + +def generate_block_test_report(dispatcher, report_dir=None): + try: + output_dir = report_dir or REPORT_DIR + block_report_csv = os.path.join(output_dir, "block_test_report.csv") + os.makedirs(os.path.dirname(block_report_csv) or ".", exist_ok=True) + with open(block_report_csv, mode="w", newline="", encoding="utf-8") as handle: + writer = csv.writer(handle) + writer.writerow([ + "任务键", "应用名称", "包名", "状态", "Worker ID", "IP地址", + "开始时间", "结束时间", "失败原因", "失败类型", "重试次数", + ]) + + all_details = dispatcher.redis.hgetall("task:details") + all_status = dispatcher.redis.hgetall("task:status") + + success_count = 0 + failed_count = 0 + retry_count = 0 + + for task_key in sorted(all_details.keys()): + if not task_key.endswith("_block"): + continue + task_details_json = all_details.get(task_key) + task_details = json.loads(task_details_json) if task_details_json else {} + if not task_details.get("is_block_task"): + continue + + status_json = all_status.get(task_key) + status = json.loads(status_json) if status_json else {} + + task_status = status.get("status", "unknown") + worker_id = status.get("worker_id", "N/A") + + ip_address = "N/A" + if worker_id != "N/A": + worker_info_json = dispatcher.redis.hget("workers:info", worker_id) + if worker_info_json: + worker_info = json.loads(worker_info_json) + ip_address = worker_info.get("ip_address", "N/A") + + start_time = status.get("start_time", 0) + end_time = status.get("end_time", 0) + start_str = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(start_time)) if start_time else "" + end_str = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(end_time)) if end_time else "" + + writer.writerow([ + task_key, + task_details.get("app_name", "N/A"), + task_details.get("package_name", "N/A"), + task_status, + worker_id, + ip_address, + start_str, + end_str, + status.get("last_fail_message", ""), + status.get("last_fail_type", ""), + status.get("retry_count", 0), + ]) + + if task_status in ("success", "completed", "qualified"): + success_count += 1 + elif task_status in ("failed",): + failed_count += 1 + else: + retry_count += 1 + + writer.writerow([]) + writer.writerow(["统计"]) + writer.writerow(["成功", str(success_count)]) + writer.writerow(["失败", str(failed_count)]) + writer.writerow(["重试中/待分发", str(retry_count)]) + + logger.info(f"[Block测试] Block 测试报告已保存到: {block_report_csv}") + except Exception as e: + logger.error(f"生成 Block 测试报告时出错: {e}") + + +def generate_worker_report(dispatcher): + try: + all_workers = dispatcher.get_registered_workers() + os.makedirs(os.path.dirname(WORKER_REPORT_CSV) or ".", exist_ok=True) + with open(WORKER_REPORT_CSV, mode="w", newline="", encoding="utf-8") as handle: + writer = csv.writer(handle) + writer.writerow(["Worker ID", "IP地址", "MAC地址", "主机名", "平台", "状态", "当前任务", "注册时间", "最后更新时间"]) + + for worker in all_workers: + writer.writerow([ + worker.get("worker_id", "N/A"), + worker.get("ip_address", "N/A"), + worker.get("mac_address", "N/A"), + worker.get("hostname", "N/A"), + worker.get("platform", "N/A"), + worker.get("status", "N/A"), + worker.get("current_task", "N/A"), + time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(worker.get("register_time", 0))), + time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(worker.get("last_update", 0))), + ]) + + logger.info(f"Worker 使用报告已保存到: {WORKER_REPORT_CSV}") + except Exception as e: + logger.error(f"生成 Worker 报告时出错: {e}") + + +def push_worker_status(dispatcher): + try: + workers = dispatcher.get_dashboard_workers() + summary = dispatcher.get_dashboard_summary() + online_workers = [worker for worker in workers if worker.get("online")] + idle_workers = [worker for worker in online_workers if worker.get("status") == "idle"] + busy_workers = [worker for worker in online_workers if worker.get("status") == "busy"] + + message_lines = [ + "Worker 状态报告", + f"时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", + "", + f"任务统计: 待分发 {summary['pending']} | 运行中 {summary['running']} | 成功 {summary['completed']} | 失败 {summary['failed']}", + f"Worker 统计: 在线 {len(online_workers)} | 空闲 {len(idle_workers)} | 忙碌 {len(busy_workers)}", + ] + + if online_workers: + message_lines.append("") + message_lines.append("在线 Worker:") + for worker in online_workers: + state = "idle" if worker.get("status") == "idle" else "busy" + message_lines.append(f"- {worker.get('ip_address', 'N/A')} ({worker.get('hostname', 'N/A')}) [{state}]") + + dispatcher.notifier.send_weCom_alert("SYSTEM", "\n".join(message_lines)) + logger.info(f"[状态推送] 已推送 Worker 状态，在线 {len(online_workers)} 台") + except Exception as e: + logger.error(f"[状态推送] 推送 Worker 状态时出错: {e}") + + +def main(): + _set_console_title() + logger.info("=" * 60) + logger.info(f"Redis 任务分发器 [{INSTANCE_NAME}]") + logger.info(f"分发服务: {REDIS_HOST}:{REDIS_PORT}/{REDIS_DB}") + logger.info(f"频道命名空间: {CHANNEL_NAMESPACE or 'default'}") + logger.info(f"Dashboard: http://{DASHBOARD_HOST}:{DASHBOARD_PORT}") + logger.info("=" * 60) + + service = DispatcherService(redis_host=REDIS_HOST, redis_port=REDIS_PORT, redis_db=REDIS_DB) + start_dashboard_thread(service.dispatcher) + + logger.info("[1] 同步 Worker 注册表...") + service.dispatcher.sync_worker_registry() + + logger.info("[2] 从应用全表加载待采集任务...") + task_count = service.dispatcher.load_tasks_from_app_summary() + if task_count == 0: + logger.warning("应用全表中没有待采集任务，请先运行应用名单同步脚本后再启动中控") + return + + logger.info("[3] 启动监听服务...") + logger.info("等待 Worker 连接...") + logger.info("按 Ctrl+C 停止程序") + service.start_listeners() + + # 启动后立即同步一次下载队列 + 轮询一次APK结果 + initial_queue_count = service.dispatcher.sync_download_queue() + if initial_queue_count > 0: + logger.info(f"[APK队列] 启动时已推送 {initial_queue_count} 个任务到下载队列") + else: + logger.info("[APK队列] 启动同步：当前无需下载的新APK任务") + + initial_poll_count = service.dispatcher.poll_apk_results() + if initial_poll_count > 0: + logger.info(f"[APK轮询] 启动轮询：导入 {initial_poll_count} 个APK") + else: + logger.info("[APK轮询] 启动轮询：无新APK结果") + + last_status_time = time.time() + last_worker_status_time = time.time() + last_cleanup_time = time.time() + last_recovery_check_time = time.time() + last_worker_push_time = time.time() + last_monitor_heartbeat_time = time.time() + last_task_refresh_time = time.time() + last_apk_queue_time = time.time() + last_apk_poll_time = time.time() + last_apk_cleanup_time = time.time() + last_block_report_time = time.time() + tasks_done_logged = False + + try: + while True: + current_time = time.time() + + if current_time - last_cleanup_time > 60: + cleaned_count = service.dispatcher.cleanup_stale_workers(timeout=WORKER_STALE_TIMEOUT) + if cleaned_count > 0: + logger.info(f"[清理] 已清理 {cleaned_count} 个超时 Worker") + last_cleanup_time = current_time + + if current_time - last_recovery_check_time > 60: + reset_workers = service.dispatcher.check_recovery_windows() + if reset_workers: + logger.info(f"[恢复窗口] 已重置自动恢复标记: {reset_workers}") + last_recovery_check_time = current_time + + if current_time - last_worker_push_time >= WORKER_STATUS_PUSH_INTERVAL: + push_worker_status(service.dispatcher) + last_worker_push_time = current_time + + if current_time - last_task_refresh_time > 120: + new_count = service.dispatcher.refresh_tasks_from_app_summary() + if new_count > 0: + logger.info(f"[任务刷新] 定期检查：从数据库加载了 {new_count} 个新任务") + last_task_refresh_time = current_time + + if current_time - last_block_report_time >= 300: + generate_block_test_report(service.dispatcher) + last_block_report_time = current_time + + if current_time - last_apk_queue_time >= APK_DOWNLOAD_QUEUE_INTERVAL: + queue_count = service.dispatcher.sync_download_queue() + if queue_count > 0: + logger.info(f"[APK队列] 已推送 {queue_count} 个任务到下载队列") + last_apk_queue_time = current_time + + if current_time - last_apk_poll_time >= APK_PREFETCH_POLL_INTERVAL: + imported = service.dispatcher.poll_apk_results() + if imported > 0: + logger.info(f"[APK轮询] 本轮导入 {imported} 个APK文件") + last_apk_poll_time = current_time + + if current_time - last_apk_cleanup_time >= 3600: + service.dispatcher._enforce_apk_storage_limits() + service.dispatcher._manage_stale_apks() + last_apk_cleanup_time = current_time + + # 心跳间隔从15秒增大到30秒，减少SQLite写入频率 + if current_time - last_monitor_heartbeat_time >= 30: + service.dispatcher.monitor.heartbeat_controller_session() + last_monitor_heartbeat_time = current_time + + # 状态日志间隔从15秒增大到60秒，减少日志I/O和Redis全量查询频率 + if current_time - last_status_time > 60: + stats = service.dispatcher.get_statistics() + logger.info( + f"[状态] 待分发: {stats['pending']} | 运行中: {stats['running']} | " + f"已完成: {stats['completed']} | 失败: {stats['failed']}" + ) + + running_tasks = service.redis.hgetall("worker:tasks") + if running_tasks: + logger.info("[运行中任务]") + for worker_id, task_key in sorted(running_tasks.items()): + worker_info_json = service.redis.hget("workers:info", worker_id) + if worker_info_json: + worker_info = json.loads(worker_info_json) + logger.info(f" - Worker {worker_id} ({worker_info.get('ip_address', 'N/A')}): {task_key}") + + last_status_time = current_time + + # Worker状态日志间隔从30秒增大到120秒，减少日志I/O + if current_time - last_worker_status_time > 120: + all_workers = service.dispatcher.get_dashboard_workers() + idle_workers = [worker for worker in all_workers if worker.get("online") and worker.get("status") == "idle"] + busy_workers = [worker for worker in all_workers if worker.get("online") and worker.get("status") == "busy"] + + logger.info("[Worker 状态]") + logger.info(f" 总数: {len(all_workers)} | 空闲: {len(idle_workers)} | 忙碌: {len(busy_workers)}") + + if idle_workers: + logger.info(" 空闲 Worker:") + for worker in idle_workers: + logger.info(f" - {worker.get('worker_id', 'N/A')} ({worker.get('ip_address', 'N/A')})") + + if busy_workers: + logger.info(" 忙碌 Worker:") + for worker in busy_workers: + logger.info( + f" - {worker.get('worker_id', 'N/A')} ({worker.get('ip_address', 'N/A')}) " + f"-> {worker.get('current_task', 'N/A')}" + ) + + last_worker_status_time = current_time + + # 使用轻量级Redis命令检测任务是否全部完成，避免每次循环都调用全量get_statistics() + queue_len = service.dispatcher.get_pending_queue_length() + running_count = service.redis.hlen("worker:tasks") + if queue_len == 0 and running_count == 0: + stats = service.dispatcher.get_statistics() + if stats["pending"] == 0 and stats["running"] == 0: + if not tasks_done_logged: + logger.info("=" * 60) + logger.info("所有任务已完成，静默等待新任务...") + logger.info(f"总计: 成功 {stats['completed']} | 失败 {stats['failed']}") + logger.info("=" * 60) + tasks_done_logged = True + else: + tasks_done_logged = False + else: + tasks_done_logged = False + + # 循环间隔从2秒增大到5秒，减少轮询频率 + time.sleep(5) + + except KeyboardInterrupt: + logger.info("检测到中断信号，正在清理数据并退出...") + service.stop(clear_data=True) + logger.info("程序已中断退出") + return + except Exception as e: + logger.error(f"程序发生错误: {e}", exc_info=True) + service.stop(clear_data=True) + return + + service.stop(clear_data=False) + + logger.info("[4] 生成最终报告...") + generate_final_report(service.dispatcher) + + logger.info("[5] 生成 Worker 使用报告...") + generate_worker_report(service.dispatcher) + + logger.info("[6] 生成 Block 测试报告...") + generate_block_test_report(service.dispatcher) + + logger.info("程序结束") + + +if __name__ == "__main__": + main() diff --git a/log_manager.py b/log_manager.py new file mode 100644 index 0000000..c93aee5 --- /dev/null +++ b/log_manager.py @@ -0,0 +1,55 @@ +# -*- encoding=utf8 -*- +""" +日志管理模块 +提供统一的日志记录，同时输出到控制台和文件，支持日志轮转 +""" +import os +import logging +from logging.handlers import RotatingFileHandler +from config import LOG_DIR, LOG_FILE, LOG_MAX_BYTES, LOG_BACKUP_COUNT, LOG_FORMAT, LOG_DATE_FORMAT + + +def setup_logger(name='dispatcher', level=logging.INFO): + """创建并配置 logger + + Args: + name: logger 名称 + level: 日志级别 + + Returns: + logging.Logger 实例 + """ + logger = logging.getLogger(name) + + # 避免重复添加 handler + if logger.handlers: + return logger + + logger.setLevel(level) + + formatter = logging.Formatter(LOG_FORMAT, datefmt=LOG_DATE_FORMAT) + + # 控制台 handler（保留原有的控制台输出） + console_handler = logging.StreamHandler() + console_handler.setLevel(level) + console_handler.setFormatter(formatter) + logger.addHandler(console_handler) + + # 文件 handler（带日志轮转） + os.makedirs(LOG_DIR, exist_ok=True) + log_path = os.path.join(LOG_DIR, LOG_FILE) + file_handler = RotatingFileHandler( + log_path, + maxBytes=LOG_MAX_BYTES, + backupCount=LOG_BACKUP_COUNT, + encoding='utf-8' + ) + file_handler.setLevel(level) + file_handler.setFormatter(formatter) + logger.addHandler(file_handler) + + return logger + + +# 全局 logger 实例，直接导入使用 +logger = setup_logger() diff --git a/migrate_to_new_db.py b/migrate_to_new_db.py new file mode 100644 index 0000000..4e2c64e --- /dev/null +++ b/migrate_to_new_db.py @@ -0,0 +1,134 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +新数据库初始化/校验脚本。 + +运行时代码不再兼容旧库，也不创建 app_collect_summary / v_collection_latest +这类兼容视图。本脚本只负责用 schema_new.sql 初始化新库，并校验核心新表。 +""" + +import argparse +import os +import sqlite3 +import sys +import time + +from analytics import AnalyticsRepository + + +BASE_DIR = os.path.dirname(os.path.abspath(__file__)) +RUNTIME_MAIN = os.path.join(BASE_DIR, "runtime", "main") +DEFAULT_DB = os.path.join(RUNTIME_MAIN, "monitoring_new.sqlite3") +SCHEMA_SQL = os.path.join(BASE_DIR, "schema_new.sql") + +REQUIRED_TABLES = { + "collection_task", + "app_catalog", + "apk_registry", + "worker_activity_summary", + "analytics_job", + "analytics_source_file", + "app_domain_traffic", + "app_traffic_component", +} + +FORBIDDEN_OBJECTS = { + "app_collect_summary", + "v_collection_latest", + "v_collection_stats_by_tag", +} + + +def log(message: str) -> None: + print(f"[new-db] {message}", flush=True) + + +def _remove_sqlite_files(db_path: str) -> None: + for suffix in ("", "-wal", "-shm"): + path = f"{db_path}{suffix}" + if os.path.exists(path): + os.remove(path) + + +def initialize_new_db(db_path: str, *, replace: bool = False) -> None: + if not os.path.exists(SCHEMA_SQL): + raise FileNotFoundError(f"schema file not found: {SCHEMA_SQL}") + if os.path.exists(db_path): + if not replace: + log(f"数据库已存在，跳过初始化: {db_path}") + return + _remove_sqlite_files(db_path) + os.makedirs(os.path.dirname(db_path) or ".", exist_ok=True) + with open(SCHEMA_SQL, "r", encoding="utf-8") as handle: + schema_sql = handle.read() + with sqlite3.connect(db_path) as connection: + connection.executescript(schema_sql) + AnalyticsRepository(db_path=db_path) + log(f"✅ 新库 schema 初始化完成: {db_path}") + + +def validate_new_db(db_path: str) -> bool: + if not os.path.exists(db_path): + log(f"❌ 数据库不存在: {db_path}") + return False + ok = True + with sqlite3.connect(db_path) as connection: + connection.row_factory = sqlite3.Row + objects = { + str(row["name"]) + for row in connection.execute( + "SELECT name FROM sqlite_master WHERE type IN ('table', 'view')" + ).fetchall() + } + missing = sorted(REQUIRED_TABLES - objects) + forbidden = sorted(FORBIDDEN_OBJECTS & objects) + if missing: + ok = False + log("❌ 缺少新 schema 表: " + ", ".join(missing)) + if forbidden: + ok = False + log("❌ 存在旧兼容对象: " + ", ".join(forbidden)) + if ok: + log("✅ schema 对象校验通过") + + for table_name in sorted(REQUIRED_TABLES & objects): + row = connection.execute(f"SELECT COUNT(*) AS total FROM {table_name}").fetchone() + log(f"{table_name}: {int(row['total'] or 0)} 行") + + task_state = connection.execute( + """ + SELECT task_status, COUNT(*) AS total + FROM collection_task + GROUP BY task_status + ORDER BY total DESC + """ + ).fetchall() if "collection_task" in objects else [] + if task_state: + log( + "collection_task.task_status 分布: " + + str({str(row["task_status"] or "NULL"): int(row["total"] or 0) for row in task_state}) + ) + return ok + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Initialize and validate the new analytics database schema.") + parser.add_argument("--db-path", default=DEFAULT_DB, help="Target sqlite path.") + parser.add_argument("--replace", action="store_true", help="Delete the target sqlite files before initializing.") + parser.add_argument("--validate-only", action="store_true", help="Only validate an existing database.") + return parser + + +def main() -> int: + args = build_parser().parse_args() + db_path = str(args.db_path or "").strip() or DEFAULT_DB + started_at = time.time() + if not args.validate_only: + initialize_new_db(db_path, replace=bool(args.replace)) + ok = validate_new_db(db_path) + log(f"完成，用时 {time.time() - started_at:.1f}s") + return 0 if ok else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/monitoring.py b/monitoring.py new file mode 100644 index 0000000..cd82cf5 --- /dev/null +++ b/monitoring.py @@ -0,0 +1,2463 @@ +import json +import os +import socket +import sqlite3 +import threading +import time +from collections import Counter, defaultdict +from contextlib import contextmanager +from datetime import datetime, timedelta +from typing import Any, Dict, Iterable, List, Optional, Tuple +from zoneinfo import ZoneInfo + +from config import ( + MANAGED_WORKER_IDS, + MONITORING_DB_PATH, + MONITORING_TIMEZONE, + MONITORING_TIMELINE_BUCKET_MINUTES, + normalize_worker_id, +) +from log_manager import logger +from result_codes import ( + APP_ERROR_DESC, + BUSINESS_ERROR_DESC, + DOWNLOAD_ERROR_DESC, + INFRA_ERROR_DESC, + AppError, + BusinessError, + DownloadError, + InfraError, +) + + +SHANGHAI_TZ = ZoneInfo(MONITORING_TIMEZONE) +RUNNING_STATES = {"running_download", "running_collect"} +ONLINE_STATES = {"idle_waiting_task", "running_download", "running_collect", "disabled"} +DISPLAY_STATE_LABELS = { + "idle_waiting_task": "idle", + "running_download": "download", + "running_collect": "collect", + "disabled": "disabled", + "offline": "offline", +} +RUNTIME_STATE_ORDER = [ + "running_download", + "running_collect", + "idle_waiting_task", + "disabled", + "offline", +] +RUNTIME_STATE_LABELS = { + "running_download": "下载中", + "running_collect": "采集中", + "idle_waiting_task": "空闲等待", + "disabled": "已禁用", + "offline": "离线", +} +STAGE_LABELS = { + "download": "下载", + "collect": "采集", +} +STAGE_ORDER = ["download", "collect"] +FAILURE_DOMAIN_LABELS = { + "download_network": "下载网络", + "download_other": "下载其他", + "infra": "基础设施", + "app": "应用问题", + "business": "业务限制", + "other": "其他", +} +DOWNLOAD_NETWORK_ERROR_CODES = { + int(DownloadError.DOWNLOAD_TIMEOUT), + int(DownloadError.NETWORK_ERROR), +} +DOWNLOAD_NETWORK_MESSAGE_HINTS = ( + "timeout", + "overtime", + "jump failed", + "app page load failed", + "page load failed", + "network error", + "connection error", + "连接超时", + "网络错误", + "超时", +) + + +def _safe_json_dumps(payload: Dict[str, Any]) -> str: + return json.dumps(payload, ensure_ascii=False, sort_keys=True) + + +def _safe_json_loads(payload: Optional[str], fallback: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: + if not payload: + return fallback or {} + try: + value = json.loads(payload) + return value if isinstance(value, dict) else (fallback or {}) + except (TypeError, ValueError): + return fallback or {} + + +def _date_str_from_ts(ts: float) -> str: + return datetime.fromtimestamp(ts, SHANGHAI_TZ).strftime("%Y-%m-%d") + + +def _day_bounds(date_str: str, now_ts: Optional[float] = None) -> Tuple[float, float, bool]: + local_day = datetime.strptime(date_str, "%Y-%m-%d").replace(tzinfo=SHANGHAI_TZ) + day_start = local_day.timestamp() + day_end = (local_day + timedelta(days=1)).timestamp() + now_ts = now_ts or time.time() + is_today = _date_str_from_ts(now_ts) == date_str + return day_start, min(day_end, now_ts) if is_today else day_end, is_today + + +def _interval_overlap(start: float, end: float, bucket_start: float, bucket_end: float) -> bool: + return max(start, bucket_start) < min(end, bucket_end) + + +def _build_histogram(values: List[float], edges: List[float]) -> List[Dict[str, Any]]: + buckets: List[Dict[str, Any]] = [] + previous = 0.0 + for edge in edges: + buckets.append({"label": _format_bucket_label(previous, edge), "start": previous, "end": edge, "count": 0}) + previous = edge + buckets.append({"label": _format_bucket_label(previous, None), "start": previous, "end": None, "count": 0}) + for value in values: + if value is None: + continue + placed = False + for bucket in buckets: + bucket_end = bucket["end"] + if bucket_end is None: + bucket["count"] += 1 + placed = True + break + if bucket["start"] <= value < bucket_end: + bucket["count"] += 1 + placed = True + break + if not placed and buckets: + buckets[-1]["count"] += 1 + return buckets + + +def _percentile(values: List[float], ratio: float) -> Optional[float]: + if not values: + return None + ordered = sorted(values) + index = int(round((len(ordered) - 1) * ratio)) + return round(ordered[index], 2) + + +def _format_minutes(seconds: Optional[float]) -> str: + seconds = float(seconds or 0.0) + minutes = seconds / 60.0 + if minutes >= 10: + return f"{round(minutes):.0f}m" + if minutes >= 1: + return f"{minutes:.1f}m" + if seconds <= 0: + return "0m" + return f"{minutes:.1f}m" + + +def _format_bucket_label(start_seconds: float, end_seconds: Optional[float]) -> str: + if end_seconds is None: + return f"{_format_minutes(start_seconds)}+" + return f"{_format_minutes(start_seconds)}-{_format_minutes(end_seconds)}" + + +def _humanize_failure_reason(error_type: Optional[str], error_message: Optional[str]) -> str: + if error_message: + return str(error_message) + if not error_type: + return "未知原因" + try: + category_name, code_str = str(error_type).split("/", 1) + code = int(code_str) + except (ValueError, TypeError): + return str(error_type) + + mapping = { + "INFRA_ERROR": INFRA_ERROR_DESC, + "DOWNLOAD_ERROR": DOWNLOAD_ERROR_DESC, + "APP_ERROR": APP_ERROR_DESC, + "BUSINESS_ERROR": BUSINESS_ERROR_DESC, + }.get(category_name) + if not mapping: + return str(error_type) + return mapping.get(code, str(error_type)) + + +def _task_row_total_duration(row: sqlite3.Row) -> float: + return float(row["total_duration_seconds"] or row["failed_total_duration_seconds"] or 0.0) + + +def _build_failure_task_sample(row: sqlite3.Row) -> Dict[str, Any]: + task_started_at = float(row["task_started_at"]) if row["task_started_at"] is not None else None + task_ended_at = float(row["task_ended_at"]) if row["task_ended_at"] is not None else None + failure_domain = row["failure_domain"] + failure_subtype = row["failure_subtype"] + diagnostics = _extract_task_diagnostics(row) + if (row["status"] or "unknown") != "success" and (not failure_domain or not failure_subtype): + failure_domain, failure_subtype = _classify_failure(row["error_type"], row["error_message"]) + return { + "execution_id": row["execution_id"], + "task_key": row["task_key"], + "app_name": row["app_name"] or "--", + "package_name": row["package_name"] or "--", + "worker_id": row["worker_id"], + "status": row["status"] or "unknown", + "failure_domain": failure_domain or "", + "failure_domain_label": _humanize_failure_domain(failure_domain), + "failure_subtype": failure_subtype or "", + "failure_subtype_label": _humanize_failure_subtype( + failure_domain, + failure_subtype, + row["error_type"], + row["error_message"], + ), + "failure_reason": _humanize_failure_reason(row["error_type"], row["error_message"]), + "error_type": row["error_type"] or "", + "error_message": row["error_message"] or "", + "failed_stage": row["failed_stage"] or "--", + "download_duration_seconds": row["download_duration_seconds"] or 0.0, + "collect_duration_seconds": row["collect_duration_seconds"] or 0.0, + "total_duration_seconds": row["total_duration_seconds"] or 0.0, + "failed_total_duration_seconds": row["failed_total_duration_seconds"] or 0.0, + "diagnostics": diagnostics, + "diagnostics_note": _format_diagnostics_note(diagnostics, int(row["retry_count"] or 0)), + "task_started_at": datetime.fromtimestamp(task_started_at, SHANGHAI_TZ).strftime("%Y-%m-%d %H:%M:%S") + if task_started_at + else "--", + "task_ended_at": datetime.fromtimestamp(task_ended_at, SHANGHAI_TZ).strftime("%Y-%m-%d %H:%M:%S") + if task_ended_at + else "--", + } + + +def _parse_error_type(error_type: Optional[str]) -> Tuple[Optional[str], Optional[int]]: + if not error_type: + return None, None + try: + category_name, code_str = str(error_type).split("/", 1) + return category_name, int(code_str) + except (TypeError, ValueError): + return str(error_type), None + + +def _download_error_message_is_network(message: Optional[str]) -> bool: + if not message: + return False + lowered = str(message).lower() + return any(token in lowered for token in DOWNLOAD_NETWORK_MESSAGE_HINTS) + + +def _normalize_download_errors(download_errors: Any) -> Dict[str, Dict[str, Any]]: + if isinstance(download_errors, dict): + return download_errors + return {} + + +def _all_download_errors_are_network(download_errors: Any) -> bool: + errors = _normalize_download_errors(download_errors) + if not errors: + return False + for detail in errors.values(): + if not isinstance(detail, dict): + return False + code = detail.get("code") + try: + if code is not None and int(code) in DOWNLOAD_NETWORK_ERROR_CODES: + continue + except (TypeError, ValueError): + pass + if _download_error_message_is_network(detail.get("message")): + continue + return False + return True + + +def _error_code_name(category_name: Optional[str], code: Optional[int]) -> str: + if category_name == "DOWNLOAD_ERROR" and code is not None: + try: + return DownloadError(int(code)).name + except ValueError: + return f"DOWNLOAD_{code}" + if category_name == "INFRA_ERROR" and code is not None: + try: + return InfraError(int(code)).name + except ValueError: + return f"INFRA_{code}" + if category_name == "APP_ERROR" and code is not None: + try: + return AppError(int(code)).name + except ValueError: + return f"APP_{code}" + if category_name == "BUSINESS_ERROR" and code is not None: + try: + return BusinessError(int(code)).name + except ValueError: + return f"BUSINESS_{code}" + if category_name: + return str(category_name) + return "OTHER" + + +def _classify_failure( + error_type: Optional[str], + error_message: Optional[str], + download_errors: Any = None, +) -> Tuple[str, str]: + category_name, code = _parse_error_type(error_type) + subtype = _error_code_name(category_name, code) + if category_name == "DOWNLOAD_ERROR": + if code in DOWNLOAD_NETWORK_ERROR_CODES: + return "download_network", subtype + if code == int(DownloadError.ALL_SOURCES_FAILED) and _all_download_errors_are_network(download_errors): + return "download_network", "ALL_SOURCES_NETWORK_FAILED" + if _download_error_message_is_network(error_message): + return "download_network", subtype + return "download_other", subtype + if category_name == "INFRA_ERROR": + return "infra", subtype + if category_name == "APP_ERROR": + return "app", subtype + if category_name == "BUSINESS_ERROR": + return "business", subtype + return "other", subtype + + +def _humanize_failure_domain(domain: Optional[str]) -> str: + return FAILURE_DOMAIN_LABELS.get(domain or "", "其他") + + +def _humanize_failure_subtype( + domain: Optional[str], + subtype: Optional[str], + error_type: Optional[str], + error_message: Optional[str], +) -> str: + if domain == "download_network" and subtype == "ALL_SOURCES_NETWORK_FAILED": + return "所有下载源网络失败" + reason = _humanize_failure_reason(error_type, None) + if reason and reason != str(error_type): + return reason + if error_message and not subtype: + return str(error_message) + if subtype: + return str(subtype) + return "未知原因" + + +def _extract_task_diagnostics(row: sqlite3.Row) -> Dict[str, Any]: + trace_json = row["trace_json"] if "trace_json" in row.keys() else None + trace = _safe_json_loads(trace_json) + diagnostics = trace.get("diagnostics") + return diagnostics if isinstance(diagnostics, dict) else {} + + +def _format_diagnostics_note(diagnostics: Dict[str, Any], retry_count: int) -> str: + labels = [] + if retry_count > 0: + labels.append(f"重试{retry_count}次") + if diagnostics.get("emulator_restarted"): + labels.append("模拟器重启") + if diagnostics.get("airtest_reinitialized"): + labels.append("Airtest重连") + if diagnostics.get("download_recovered"): + labels.append("下载恢复") + if diagnostics.get("collect_recovered"): + labels.append("采集恢复") + return " / ".join(labels) if labels else "--" + + +class MonitoringRepository: + def __init__(self, db_path: str = MONITORING_DB_PATH, initialize_schema: bool = True): + self.db_path = db_path + self._write_lock = threading.Lock() + if initialize_schema: + self._initialize_schema() + + @contextmanager + def _connect(self): + """每次操作使用独立连接，避免跨线程共享连接导致游标状态串扰。""" + connection = sqlite3.connect(self.db_path, timeout=30) + connection.row_factory = sqlite3.Row + connection.execute("PRAGMA busy_timeout=30000") + connection.execute("PRAGMA synchronous=NORMAL") + connection.execute("PRAGMA wal_autocheckpoint=1000") + connection.execute("PRAGMA cache_size=-8192") + connection.execute("PRAGMA temp_store=MEMORY") + try: + yield connection + connection.commit() + except Exception: + try: + connection.rollback() + except Exception: + pass + raise + finally: + connection.close() + + def _table_exists(self, connection: sqlite3.Connection, table_name: str) -> bool: + row = connection.execute( + """ + SELECT 1 + FROM sqlite_master + WHERE type = 'table' AND name = ? + """, + (table_name,), + ).fetchone() + return bool(row) + + def close(self): + """保留接口兼容；独立连接模式下无需额外关闭。""" + return + + def _initialize_schema(self): + os.makedirs(os.path.dirname(self.db_path) or ".", exist_ok=True) + with self._connect() as connection: + connection.execute("PRAGMA journal_mode=WAL") + connection.executescript( + """ + CREATE TABLE IF NOT EXISTS controller_session ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + started_at REAL NOT NULL, + ended_at REAL, + host TEXT, + version TEXT, + stop_reason TEXT, + last_heartbeat_at REAL + ); + + CREATE TABLE IF NOT EXISTS worker_state_event ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + worker_id TEXT NOT NULL, + session_id INTEGER, + event_time REAL NOT NULL, + state TEXT NOT NULL, + reason TEXT, + task_key TEXT + ); + CREATE INDEX IF NOT EXISTS idx_worker_state_event_worker_time + ON worker_state_event(worker_id, event_time); + + CREATE TABLE IF NOT EXISTS task_stage_event ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + execution_id TEXT NOT NULL, + task_key TEXT NOT NULL, + worker_id TEXT NOT NULL, + session_id INTEGER, + event_time REAL NOT NULL, + event_type TEXT NOT NULL, + stage TEXT, + status TEXT, + reason TEXT, + metrics_json TEXT + ); + CREATE INDEX IF NOT EXISTS idx_task_stage_event_exec_time + ON task_stage_event(execution_id, event_time); + + CREATE TABLE IF NOT EXISTS task_execution ( + execution_id TEXT PRIMARY KEY, + task_key TEXT NOT NULL, + app_name TEXT, + package_name TEXT, + last_updated TEXT, + collection_task_type TEXT, + worker_id TEXT NOT NULL, + session_id INTEGER, + stat_date TEXT NOT NULL, + task_started_at REAL, + task_ended_at REAL, + status TEXT, + failed_stage TEXT, + failure_bucket TEXT, + retry_count INTEGER DEFAULT 0, + download_duration_seconds REAL, + collect_duration_seconds REAL, + failed_total_duration_seconds REAL, + total_duration_seconds REAL, + download_source TEXT, + error_type TEXT, + error_message TEXT, + result_detail TEXT, + droidbot_steps INTEGER, + guiagent_steps INTEGER, + total_steps INTEGER, + num_nodes INTEGER, + num_reached_activities INTEGER, + app_num_total_activities INTEGER, + failure_domain TEXT, + failure_subtype TEXT, + trace_json TEXT, + last_updated_at REAL + ); + CREATE INDEX IF NOT EXISTS idx_task_execution_stat_date + ON task_execution(stat_date, worker_id, task_started_at); + """ + ) + columns = { + row["name"] + for row in connection.execute("PRAGMA table_info(task_execution)").fetchall() + } + if "failure_bucket" not in columns: + connection.execute("ALTER TABLE task_execution ADD COLUMN failure_bucket TEXT") + if "failure_domain" not in columns: + connection.execute("ALTER TABLE task_execution ADD COLUMN failure_domain TEXT") + if "failure_subtype" not in columns: + connection.execute("ALTER TABLE task_execution ADD COLUMN failure_subtype TEXT") + if "trace_json" not in columns: + connection.execute("ALTER TABLE task_execution ADD COLUMN trace_json TEXT") + if "result_detail" not in columns: + connection.execute("ALTER TABLE task_execution ADD COLUMN result_detail TEXT") + if "droidbot_steps" not in columns: + connection.execute("ALTER TABLE task_execution ADD COLUMN droidbot_steps INTEGER") + if "guiagent_steps" not in columns: + connection.execute("ALTER TABLE task_execution ADD COLUMN guiagent_steps INTEGER") + if "total_steps" not in columns: + connection.execute("ALTER TABLE task_execution ADD COLUMN total_steps INTEGER") + if "num_nodes" not in columns: + connection.execute("ALTER TABLE task_execution ADD COLUMN num_nodes INTEGER") + if "num_reached_activities" not in columns: + connection.execute("ALTER TABLE task_execution ADD COLUMN num_reached_activities INTEGER") + if "app_num_total_activities" not in columns: + connection.execute("ALTER TABLE task_execution ADD COLUMN app_num_total_activities INTEGER") + if "last_updated" not in columns: + connection.execute("ALTER TABLE task_execution ADD COLUMN last_updated TEXT") + if "collection_task_type" not in columns: + connection.execute("ALTER TABLE task_execution ADD COLUMN collection_task_type TEXT") + + # ── APK 注册表（由 apk_cloud.registry.ApkRegistry 使用）── + connection.executescript( + """ + CREATE TABLE IF NOT EXISTS apk_registry ( + package_name TEXT PRIMARY KEY, + download_date TEXT NOT NULL, + download_time REAL NOT NULL, + local_dir TEXT, + smb_dir TEXT, + version_name TEXT, + source TEXT, + apk_files_json TEXT, + updated_at REAL NOT NULL + ); + + CREATE TABLE IF NOT EXISTS apk_pending_download ( + package_name TEXT PRIMARY KEY, + created_at REAL NOT NULL + ); + + CREATE TABLE IF NOT EXISTS app_download_record ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + package_name TEXT NOT NULL, + country_code TEXT NOT NULL DEFAULT '', + app_name TEXT DEFAULT '', + last_updated TEXT NOT NULL DEFAULT '', + downloads INTEGER DEFAULT 0, + + worker_play_status TEXT NOT NULL DEFAULT 'pending', + worker_play_error TEXT, + worker_play_date TEXT, + + us_play_status TEXT NOT NULL DEFAULT 'pending', + us_play_error TEXT, + us_play_date TEXT, + + us_aurora_status TEXT NOT NULL DEFAULT 'pending', + us_aurora_error TEXT, + us_aurora_date TEXT, + + overall_status TEXT NOT NULL DEFAULT 'pending', + overall_error TEXT, + apk_version_name TEXT DEFAULT '', + + local_apk_dir TEXT DEFAULT '', + local_apk_file_count INTEGER DEFAULT 0, + local_apk_is_stale INTEGER DEFAULT 1, + + created_at REAL NOT NULL, + updated_at REAL NOT NULL, + + UNIQUE(package_name, country_code) + ); + CREATE INDEX IF NOT EXISTS idx_download_record_pkg + ON app_download_record(package_name); + CREATE INDEX IF NOT EXISTS idx_download_record_country + ON app_download_record(country_code); + CREATE INDEX IF NOT EXISTS idx_download_record_overall + ON app_download_record(overall_status); + CREATE INDEX IF NOT EXISTS idx_download_record_stale + ON app_download_record(local_apk_is_stale); + """ + ) + + def clear_monitoring_tables(self): + table_names = [ + "controller_session", + "worker_state_event", + "task_stage_event", + "task_execution", + ] + with self._write_lock, self._connect() as connection: + for table_name in table_names: + if self._table_exists(connection, table_name): + connection.execute(f"DELETE FROM {table_name}") + if self._table_exists(connection, "sqlite_sequence"): + for table_name in table_names: + connection.execute("DELETE FROM sqlite_sequence WHERE name = ?", (table_name,)) + + def create_controller_session(self, host: str, version: str, started_at: float) -> int: + with self._write_lock, self._connect() as connection: + cursor = connection.execute( + """ + INSERT INTO controller_session ( + started_at, host, version, last_heartbeat_at + ) VALUES (?, ?, ?, ?) + """, + (started_at, host, version, started_at), + ) + return int(cursor.lastrowid) + + def heartbeat_controller_session(self, session_id: int, heartbeat_at: float): + with self._write_lock, self._connect() as connection: + connection.execute( + """ + UPDATE controller_session + SET last_heartbeat_at = ? + WHERE id = ? + """, + (heartbeat_at, session_id), + ) + + def close_controller_session(self, session_id: int, ended_at: float, stop_reason: str): + with self._write_lock, self._connect() as connection: + connection.execute( + """ + UPDATE controller_session + SET ended_at = ?, stop_reason = ?, last_heartbeat_at = COALESCE(last_heartbeat_at, ?) + WHERE id = ? + """, + (ended_at, stop_reason, ended_at, session_id), + ) + + def repair_open_sessions(self, repaired_at: float): + with self._write_lock, self._connect() as connection: + rows = connection.execute( + """ + SELECT id, last_heartbeat_at + FROM controller_session + WHERE ended_at IS NULL + """ + ).fetchall() + for row in rows: + ended_at = float(row["last_heartbeat_at"] or repaired_at) + connection.execute( + """ + UPDATE controller_session + SET ended_at = ?, stop_reason = COALESCE(stop_reason, 'recovered_after_restart') + WHERE id = ? + """, + (ended_at, row["id"]), + ) + + def insert_worker_state_event( + self, + worker_id: str, + session_id: Optional[int], + event_time: float, + state: str, + reason: Optional[str], + task_key: Optional[str], + ): + with self._write_lock, self._connect() as connection: + connection.execute( + """ + INSERT INTO worker_state_event ( + worker_id, session_id, event_time, state, reason, task_key + ) VALUES (?, ?, ?, ?, ?, ?) + """, + (worker_id, session_id, event_time, state, reason, task_key), + ) + + def insert_task_stage_event( + self, + execution_id: str, + task_key: str, + worker_id: str, + session_id: Optional[int], + event_time: float, + event_type: str, + stage: Optional[str], + status: Optional[str], + reason: Optional[str], + metrics: Optional[Dict[str, Any]], + ): + with self._write_lock, self._connect() as connection: + connection.execute( + """ + INSERT INTO task_stage_event ( + execution_id, task_key, worker_id, session_id, event_time, + event_type, stage, status, reason, metrics_json + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + execution_id, + task_key, + worker_id, + session_id, + event_time, + event_type, + stage, + status, + reason, + _safe_json_dumps(metrics or {}), + ), + ) + + def upsert_task_execution( + self, + execution_id: str, + task_key: str, + worker_id: str, + session_id: Optional[int], + stat_date: str, + task_started_at: float, + app_name: Optional[str] = None, + package_name: Optional[str] = None, + last_updated: Optional[str] = None, + collection_task_type: Optional[str] = None, + retry_count: int = 0, + ): + with self._write_lock, self._connect() as connection: + connection.execute( + """ + INSERT INTO task_execution ( + execution_id, task_key, app_name, package_name, last_updated, collection_task_type, worker_id, session_id, + stat_date, task_started_at, retry_count, last_updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(execution_id) DO UPDATE SET + app_name = COALESCE(excluded.app_name, task_execution.app_name), + package_name = COALESCE(excluded.package_name, task_execution.package_name), + last_updated = COALESCE(excluded.last_updated, task_execution.last_updated), + collection_task_type = COALESCE(excluded.collection_task_type, task_execution.collection_task_type), + retry_count = MAX(task_execution.retry_count, excluded.retry_count), + worker_id = excluded.worker_id, + session_id = excluded.session_id, + last_updated_at = excluded.last_updated_at + """, + ( + execution_id, + task_key, + app_name, + package_name, + last_updated, + collection_task_type, + worker_id, + session_id, + stat_date, + task_started_at, + retry_count, + time.time(), + ), + ) + + def update_task_execution_stage( + self, + execution_id: str, + stage: str, + duration_seconds: Optional[float], + download_source: Optional[str] = None, + ): + if duration_seconds is None: + return + with self._write_lock, self._connect() as connection: + if stage == "download": + connection.execute( + """ + UPDATE task_execution + SET download_duration_seconds = ?, + download_source = COALESCE(?, download_source), + last_updated_at = ? + WHERE execution_id = ? + """, + (duration_seconds, download_source, time.time(), execution_id), + ) + elif stage == "collect": + connection.execute( + """ + UPDATE task_execution + SET collect_duration_seconds = ?, + last_updated_at = ? + WHERE execution_id = ? + """, + (duration_seconds, time.time(), execution_id), + ) + + def finish_task_execution( + self, + execution_id: str, + task_ended_at: float, + status: str, + failed_stage: Optional[str], + failure_bucket: Optional[str], + total_duration_seconds: Optional[float], + failed_total_duration_seconds: Optional[float], + error_type: Optional[str], + error_message: Optional[str], + result_detail: Optional[str], + droidbot_steps: Optional[int], + guiagent_steps: Optional[int], + total_steps: Optional[int], + num_nodes: Optional[int], + num_reached_activities: Optional[int], + app_num_total_activities: Optional[int], + failure_domain: Optional[str], + failure_subtype: Optional[str], + trace_payload: Optional[Dict[str, Any]], + retry_count: Optional[int] = None, + ): + with self._write_lock, self._connect() as connection: + connection.execute( + """ + UPDATE task_execution + SET task_ended_at = ?, + status = ?, + failed_stage = ?, + failure_bucket = COALESCE(?, failure_bucket), + total_duration_seconds = COALESCE(?, total_duration_seconds), + failed_total_duration_seconds = COALESCE(?, failed_total_duration_seconds), + error_type = COALESCE(?, error_type), + error_message = COALESCE(?, error_message), + result_detail = COALESCE(?, result_detail), + droidbot_steps = COALESCE(?, droidbot_steps), + guiagent_steps = COALESCE(?, guiagent_steps), + total_steps = COALESCE(?, total_steps), + num_nodes = COALESCE(?, num_nodes), + num_reached_activities = COALESCE(?, num_reached_activities), + app_num_total_activities = COALESCE(?, app_num_total_activities), + failure_domain = COALESCE(?, failure_domain), + failure_subtype = COALESCE(?, failure_subtype), + trace_json = COALESCE(?, trace_json), + retry_count = COALESCE(?, retry_count), + last_updated_at = ? + WHERE execution_id = ? + """, + ( + task_ended_at, + status, + failed_stage, + failure_bucket, + total_duration_seconds, + failed_total_duration_seconds, + error_type, + error_message, + result_detail, + droidbot_steps, + guiagent_steps, + total_steps, + num_nodes, + num_reached_activities, + app_num_total_activities, + failure_domain, + failure_subtype, + _safe_json_dumps(trace_payload) if trace_payload else None, + retry_count, + time.time(), + execution_id, + ), + ) + + def list_task_executions(self, stat_date: str, worker_id: Optional[str] = None, limit: int = 200) -> List[sqlite3.Row]: + with self._connect() as connection: + if worker_id: + return connection.execute( + """ + SELECT * + FROM task_execution + WHERE stat_date = ? AND worker_id = ? + ORDER BY COALESCE(task_started_at, last_updated_at, 0) DESC + LIMIT ? + """, + (stat_date, worker_id, limit), + ).fetchall() + return connection.execute( + """ + SELECT * + FROM task_execution + WHERE stat_date = ? + ORDER BY COALESCE(task_started_at, last_updated_at, 0) DESC + LIMIT ? + """, + (stat_date, limit), + ).fetchall() + + def list_app_collect_statuses(self, package_names: Iterable[str]) -> Dict[str, Dict[str, str]]: + """ + 查询应用的采集状态分类（新架构）。 + + 直接从 collection_task 查询最新执行记录，基于 execution_status, + error_category, error_code, num_nodes 和流量比例动态计算状态。 + """ + normalized = [str(item or "").strip() for item in package_names if str(item or "").strip()] + if not normalized: + return {} + + with self._connect() as connection: + rows = [] + for package_name in normalized: + row = connection.execute( + """ + SELECT + package_name, + execution_status, + error_category, + error_code, + num_nodes, + CASE WHEN total_traffic_bytes > 0 + THEN self_traffic_bytes * 100.0 / total_traffic_bytes + ELSE 0 END AS self_ratio + FROM collection_task + WHERE package_name = ? + ORDER BY + datetime(COALESCE(NULLIF(completed_at, ''), NULLIF(started_at, ''), created_at)) DESC, + attempt DESC, + batch_tag DESC, + run_kind DESC + LIMIT 1 + """, + (package_name,), + ).fetchone() + if row: + rows.append(row) + + result = {} + # 不可重试的 DOWNLOAD_ERROR 错误码 + NON_RETRYABLE_DOWNLOAD_CODES = {1, 5, 6, 10, 403, 404} + + for row in rows: + pkg = row["package_name"] + if not pkg: + continue + + exec_status = (row["execution_status"] or "").strip() + error_cat = (row["error_category"] or "").strip() + error_code = row["error_code"] + nodes = row["num_nodes"] or 0 + self_ratio = row["self_ratio"] or 0 + + # 动态计算 restriction_status + if exec_status == "success": + # 成功任务：根据质量指标判断 + if nodes >= 5 and self_ratio >= 30: + restriction_status = "success" + else: + restriction_status = "light_restricted" + else: + # 失败任务：严重受限 + restriction_status = "severe_restricted" + + # 动态计算 retryability + if exec_status == "success": + retryability = "retryable" + elif error_cat == "DOWNLOAD_ERROR" and error_code in NON_RETRYABLE_DOWNLOAD_CODES: + retryability = "non_retryable" + else: + retryability = "retryable" + + # 动态计算 collection_status（业务层分类） + if restriction_status == "success": + collection_status = "qualified" + else: + collection_status = "restricted" + + result[pkg] = { + "collection_status": collection_status, + "restriction_status": restriction_status, + "retryability": retryability, + } + + return result + + def get_task_execution_totals(self, worker_id: Optional[str] = None) -> Dict[str, int]: + with self._connect() as connection: + if worker_id: + row = connection.execute( + """ + SELECT + COUNT(*) AS task_total, + COALESCE(SUM(CASE WHEN status = 'success' THEN 1 ELSE 0 END), 0) AS task_success, + COALESCE(SUM(CASE WHEN status != 'success' THEN 1 ELSE 0 END), 0) AS task_non_success, + COALESCE(SUM(CASE WHEN status = 'failed' THEN 1 ELSE 0 END), 0) AS task_failed + FROM task_execution + WHERE worker_id = ? + """, + (worker_id,), + ).fetchone() + else: + row = connection.execute( + """ + SELECT + COUNT(*) AS task_total, + COALESCE(SUM(CASE WHEN status = 'success' THEN 1 ELSE 0 END), 0) AS task_success, + COALESCE(SUM(CASE WHEN status != 'success' THEN 1 ELSE 0 END), 0) AS task_non_success, + COALESCE(SUM(CASE WHEN status = 'failed' THEN 1 ELSE 0 END), 0) AS task_failed + FROM task_execution + """ + ).fetchone() + return { + "task_total": int(row["task_total"] or 0), + "task_success": int(row["task_success"] or 0), + "task_non_success": int(row["task_non_success"] or 0), + "task_failed": int(row["task_failed"] or 0), + } + + def get_task_overview_metrics(self, stat_date: str, worker_id: Optional[str] = None) -> Dict[str, Any]: + with self._connect() as connection: + where_clause = "WHERE stat_date = ?" + params: List[Any] = [stat_date] + if worker_id: + where_clause += " AND worker_id = ?" + params.append(worker_id) + + summary_row = connection.execute( + f""" + SELECT + COUNT(*) AS task_total, + COALESCE(SUM(CASE WHEN status = 'success' THEN 1 ELSE 0 END), 0) AS task_success, + COALESCE(SUM(CASE WHEN status != 'success' THEN 1 ELSE 0 END), 0) AS task_non_success, + COALESCE(SUM(CASE WHEN status = 'failed' THEN 1 ELSE 0 END), 0) AS task_failed, + COALESCE(AVG(download_duration_seconds), 0) AS avg_download_duration_seconds, + COALESCE(AVG(collect_duration_seconds), 0) AS avg_collect_duration_seconds, + COALESCE(SUM(CASE WHEN download_duration_seconds IS NOT NULL THEN 1 ELSE 0 END), 0) AS download_value_count, + COALESCE(SUM(CASE WHEN collect_duration_seconds IS NOT NULL THEN 1 ELSE 0 END), 0) AS collect_value_count + FROM task_execution + {where_clause} + """, + params, + ).fetchone() + + def percentile_for(column: str, count: int, ratio: float) -> float: + if count <= 0: + return 0.0 + offset = int(round((count - 1) * ratio)) + row = connection.execute( + f""" + SELECT {column} AS value + FROM task_execution + {where_clause} AND {column} IS NOT NULL + ORDER BY {column} + LIMIT 1 OFFSET ? + """, + [*params, offset], + ).fetchone() + if not row or row["value"] is None: + return 0.0 + return round(float(row["value"]), 2) + + download_count = int(summary_row["download_value_count"] or 0) + collect_count = int(summary_row["collect_value_count"] or 0) + return { + "task_total": int(summary_row["task_total"] or 0), + "task_success": int(summary_row["task_success"] or 0), + "task_non_success": int(summary_row["task_non_success"] or 0), + "task_failed": int(summary_row["task_failed"] or 0), + "avg_download_duration_seconds": round(float(summary_row["avg_download_duration_seconds"] or 0.0), 2), + "avg_collect_duration_seconds": round(float(summary_row["avg_collect_duration_seconds"] or 0.0), 2), + "p50_download_duration_seconds": percentile_for("download_duration_seconds", download_count, 0.5), + "p90_download_duration_seconds": percentile_for("download_duration_seconds", download_count, 0.9), + "p50_collect_duration_seconds": percentile_for("collect_duration_seconds", collect_count, 0.5), + "p90_collect_duration_seconds": percentile_for("collect_duration_seconds", collect_count, 0.9), + } + + def list_state_events_before(self, day_end: float, worker_ids: Optional[Iterable[str]] = None) -> List[sqlite3.Row]: + with self._connect() as connection: + if worker_ids: + worker_ids = list(worker_ids) + placeholders = ",".join("?" for _ in worker_ids) + query = f""" + SELECT worker_id, event_time, state, reason, task_key + FROM worker_state_event + WHERE event_time < ? AND worker_id IN ({placeholders}) + ORDER BY worker_id, event_time + """ + return connection.execute(query, [day_end, *worker_ids]).fetchall() + return connection.execute( + """ + SELECT worker_id, event_time, state, reason, task_key + FROM worker_state_event + WHERE event_time < ? + ORDER BY worker_id, event_time + """, + (day_end,), + ).fetchall() + + def list_worker_ids_up_to(self, day_end: float) -> List[str]: + with self._connect() as connection: + rows = connection.execute( + """ + SELECT DISTINCT worker_id + FROM worker_state_event + WHERE event_time < ? + """, + (day_end,), + ).fetchall() + return [row["worker_id"] for row in rows] + + def list_sessions_overlapping(self, day_start: float, day_end: float) -> List[sqlite3.Row]: + with self._connect() as connection: + return connection.execute( + """ + SELECT * + FROM controller_session + WHERE started_at < ? AND COALESCE(ended_at, last_heartbeat_at, ?) > ? + ORDER BY started_at + """, + (day_end, day_end, day_start), + ).fetchall() + + def update_failure_bucket(self, execution_id: str, failure_bucket: str): + with self._write_lock, self._connect() as connection: + connection.execute( + """ + UPDATE task_execution + SET failure_bucket = ?, + last_updated_at = ? + WHERE execution_id = ? + """, + (failure_bucket, time.time(), execution_id), + ) + + def update_failure_bucket_by_task_key(self, task_key: str, failure_bucket: str): + with self._write_lock, self._connect() as connection: + connection.execute( + """ + UPDATE task_execution + SET failure_bucket = ?, + last_updated_at = ? + WHERE execution_id = ( + SELECT execution_id + FROM task_execution + WHERE task_key = ? + ORDER BY COALESCE(task_ended_at, last_updated_at, task_started_at) DESC + LIMIT 1 + ) + """, + (failure_bucket, time.time(), task_key), + ) + + def get_latest_execution_id(self, task_key: str) -> Optional[str]: + with self._connect() as connection: + row = connection.execute( + """ + SELECT execution_id + FROM task_execution + WHERE task_key = ? + ORDER BY COALESCE(task_ended_at, last_updated_at, task_started_at) DESC + LIMIT 1 + """, + (task_key,), + ).fetchone() + return row["execution_id"] if row else None + + def list_stage_finish_events(self, stat_date: str) -> List[sqlite3.Row]: + with self._connect() as connection: + return connection.execute( + """ + SELECT + te.execution_id, + te.status, + te.failure_bucket, + te.total_duration_seconds, + te.failed_total_duration_seconds, + te.error_type, + te.error_message, + e.stage, + e.metrics_json + FROM task_stage_event e + JOIN task_execution te ON te.execution_id = e.execution_id + WHERE te.stat_date = ? + AND e.event_type = 'stage_finished' + ORDER BY e.event_time + """, + (stat_date,), + ).fetchall() + + +class MonitoringService: + def __init__(self, redis_client, managed_worker_ids: Optional[List[str]] = None): + self.redis = redis_client + self.repo = MonitoringRepository() + self.managed_worker_ids = list(managed_worker_ids or MANAGED_WORKER_IDS or []) + self._lock = threading.RLock() + self._session_id: Optional[int] = None + self._version = "unknown" + self.repo.repair_open_sessions(time.time()) + + def start_controller_session(self, version: str = "unknown") -> int: + with self._lock: + if self._session_id is not None: + return self._session_id + started_at = time.time() + self._version = version + self._session_id = self.repo.create_controller_session( + host=socket.gethostname(), + version=version, + started_at=started_at, + ) + self.redis.set("controller:current_session", str(self._session_id)) + return self._session_id + + def heartbeat_controller_session(self): + with self._lock: + session_id = self.current_session_id + if session_id is None: + return + self.repo.heartbeat_controller_session(session_id, time.time()) + + def close_controller_session(self, stop_reason: str = "stopped"): + with self._lock: + session_id = self._session_id + if session_id is None: + return + self.repo.close_controller_session(session_id, time.time(), stop_reason) + self.redis.delete("controller:current_session") + self._session_id = None + + def _classify_rows_by_collection_quality( + self, + rows: List[sqlite3.Row], + ) -> Tuple[List[sqlite3.Row], List[sqlite3.Row], Dict[str, int]]: + package_map = self.repo.list_app_collect_statuses( + row["package_name"] for row in rows if row["package_name"] + ) + qualified_items: List[sqlite3.Row] = [] + severe_items: List[sqlite3.Row] = [] + counters = { + "qualified_task_count": 0, + "qualified_success_task_count": 0, + "qualified_light_task_count": 0, + "non_success_task_count": 0, + "non_retryable_task_count": 0, + } + for row in rows: + package_name = str(row["package_name"] or "").strip() + row_state = package_map.get(package_name, {}) + restriction_status = row_state.get("restriction_status", "") + retryability = row_state.get("retryability", "") + if restriction_status == "severe_restricted": + severe_items.append(row) + counters["non_success_task_count"] += 1 + if retryability == "non_retryable": + counters["non_retryable_task_count"] += 1 + else: + qualified_items.append(row) + counters["qualified_task_count"] += 1 + if restriction_status == "light_restricted": + counters["qualified_light_task_count"] += 1 + else: + counters["qualified_success_task_count"] += 1 + return qualified_items, severe_items, counters + + def _split_rows_by_quality(self, rows: List[sqlite3.Row]) -> Tuple[List[sqlite3.Row], List[sqlite3.Row]]: + qualified_items, severe_items, _ = self._classify_rows_by_collection_quality(rows) + return qualified_items, severe_items + + @property + def current_session_id(self) -> Optional[int]: + if self._session_id is not None: + return self._session_id + raw = self.redis.get("controller:current_session") + if raw: + try: + self._session_id = int(raw) + except (TypeError, ValueError): + self._session_id = None + return self._session_id + + def _execution_key(self, task_key: str) -> str: + return f"monitor:execution:{task_key}" + + def _current_state_key(self, worker_id: str) -> str: + return f"worker:current_state:{worker_id}" + + def _get_current_state(self, worker_id: str) -> Dict[str, Any]: + return _safe_json_loads(self.redis.get(self._current_state_key(worker_id))) + + def _set_current_state(self, worker_id: str, payload: Dict[str, Any]): + self.redis.set(self._current_state_key(worker_id), _safe_json_dumps(payload)) + + def _delete_monitoring_redis_state(self): + self.redis.delete("controller:current_session") + patterns = [ + "worker:current_state:*", + "worker:last_event:*", + "monitor:execution:*", + ] + keys_to_delete: List[str] = [] + for pattern in patterns: + keys_to_delete.extend(list(self.redis.scan_iter(match=pattern))) + if keys_to_delete: + self.redis.delete(*keys_to_delete) + + def _normalize_reseed_state(self, worker: Dict[str, Any]) -> Optional[str]: + if not worker.get("online", True): + return None + state = str(worker.get("state") or "").strip() + if state in {"idle_waiting_task", "running_download", "running_collect", "disabled"}: + return state + if worker.get("status") == "busy": + return "running_collect" + if worker.get("status") == "idle": + return "idle_waiting_task" + return "idle_waiting_task" + + def reset_monitoring_data(self, worker_snapshots: Optional[List[Dict[str, Any]]] = None) -> Dict[str, Any]: + with self._lock: + reset_time = time.time() + self.close_controller_session(stop_reason="dashboard_reset") + self.repo.clear_monitoring_tables() + self._delete_monitoring_redis_state() + self._session_id = None + session_id = self.start_controller_session(version=self._version) + + reseeded_workers = 0 + for worker in worker_snapshots or []: + worker_id = worker.get("worker_id") + if not worker_id: + continue + state = self._normalize_reseed_state(worker) + if not state: + continue + self.set_worker_state( + worker_id, + state, + reason="dashboard_reset_reseed", + task_key=worker.get("current_task"), + event_time=reset_time, + ) + current_task = worker.get("current_task") + if current_task: + task_metadata = self._resolve_task_metadata(current_task) + retry_count = 0 + status_json = self.redis.hget("task:status", current_task) + if status_json: + status_payload = _safe_json_loads(status_json) + retry_count = int(status_payload.get("retry_count") or 0) + execution_id = f"{current_task}:{int(reset_time * 1000)}" + self.repo.upsert_task_execution( + execution_id=execution_id, + task_key=current_task, + worker_id=worker_id, + session_id=self.current_session_id, + stat_date=_date_str_from_ts(reset_time), + task_started_at=reset_time, + app_name=task_metadata.get("app_name"), + package_name=task_metadata.get("package_name"), + last_updated=task_metadata.get("last_updated"), + collection_task_type=task_metadata.get("collection_task_type"), + retry_count=retry_count, + ) + self.redis.set(self._execution_key(current_task), execution_id) + reseeded_workers += 1 + + return { + "ok": True, + "reset_at": datetime.fromtimestamp(reset_time, SHANGHAI_TZ).strftime("%Y-%m-%d %H:%M:%S"), + "session_id": session_id, + "reseeded_workers": reseeded_workers, + } + + def _update_worker_info_state( + self, + worker_id: str, + state: str, + task_key: Optional[str], + reason: Optional[str], + event_time: float, + ): + info_json = self.redis.hget("workers:info", worker_id) + if not info_json: + return + info = _safe_json_loads(info_json) + info["state"] = state + info["state_since"] = event_time + info["pause_reason"] = reason + if task_key is not None: + info["current_task"] = task_key + self.redis.hset("workers:info", worker_id, _safe_json_dumps(info)) + + def set_worker_state( + self, + worker_id: str, + state: str, + reason: Optional[str] = None, + task_key: Optional[str] = None, + event_time: Optional[float] = None, + ): + with self._lock: + event_time = event_time or time.time() + current = self._get_current_state(worker_id) + if ( + current.get("state") == state + and current.get("reason") == reason + and current.get("task_key") == task_key + ): + return + payload = {"state": state, "reason": reason, "task_key": task_key, "since": event_time} + self._set_current_state(worker_id, payload) + self._update_worker_info_state(worker_id, state, task_key, reason, event_time) + self.repo.insert_worker_state_event( + worker_id=worker_id, + session_id=self.current_session_id, + event_time=event_time, + state=state, + reason=reason, + task_key=task_key, + ) + + def register_worker(self, worker_id: str, task: Optional[Dict[str, Any]] = None, event_time: Optional[float] = None): + with self._lock: + event_time = event_time or time.time() + if task and task.get("task_key"): + self._update_worker_info_state(worker_id, "idle_waiting_task", task.get("task_key"), None, event_time) + else: + self.set_worker_state(worker_id, "idle_waiting_task", reason="waiting_task", event_time=event_time) + + def handle_task_assigned(self, worker_id: str, task_key: str, event_time: Optional[float] = None): + with self._lock: + event_time = event_time or time.time() + info_json = self.redis.hget("workers:info", worker_id) + if not info_json: + return + info = _safe_json_loads(info_json) + info["current_task"] = task_key + info["last_update"] = event_time + self.redis.hset("workers:info", worker_id, _safe_json_dumps(info)) + # 任务分发后将 worker state 标记为 running（等 stage_started 事件进一步细化为 download/collect） + self.set_worker_state(worker_id, "running_download", task_key=task_key, event_time=event_time) + if not self.redis.get(self._execution_key(task_key)): + status_payload = _safe_json_loads(self.redis.hget("task:status", task_key)) + retry_count = int(status_payload.get("retry_count") or 0) + task_metadata = self._resolve_task_metadata(task_key) + execution_id = f"{task_key}:{int(event_time * 1000)}" + self.repo.upsert_task_execution( + execution_id=execution_id, + task_key=task_key, + worker_id=worker_id, + session_id=self.current_session_id, + stat_date=_date_str_from_ts(event_time), + task_started_at=event_time, + app_name=task_metadata.get("app_name"), + package_name=task_metadata.get("package_name"), + last_updated=task_metadata.get("last_updated"), + collection_task_type=task_metadata.get("collection_task_type"), + retry_count=retry_count, + ) + self.redis.set(self._execution_key(task_key), execution_id) + + def _resolve_task_details(self, task_key: str) -> Tuple[Optional[str], Optional[str]]: + details = self._resolve_task_metadata(task_key) + return details.get("app_name"), details.get("package_name") + + def _resolve_task_metadata(self, task_key: str) -> Dict[str, Any]: + raw = self.redis.hget("task:details", task_key) + if not raw: + return {} + details = _safe_json_loads(raw) + return details if isinstance(details, dict) else {} + + def _ensure_execution_id( + self, + worker_id: str, + task_key: str, + event_time: float, + retry_count: int = 0, + ) -> str: + with self._lock: + execution_id = self.redis.get(self._execution_key(task_key)) + if execution_id: + return execution_id + execution_id = f"{task_key}:{int(event_time * 1000)}" + task_metadata = self._resolve_task_metadata(task_key) + self.repo.upsert_task_execution( + execution_id=execution_id, + task_key=task_key, + worker_id=worker_id, + session_id=self.current_session_id, + stat_date=_date_str_from_ts(event_time), + task_started_at=event_time, + app_name=task_metadata.get("app_name"), + package_name=task_metadata.get("package_name"), + last_updated=task_metadata.get("last_updated"), + collection_task_type=task_metadata.get("collection_task_type"), + retry_count=retry_count, + ) + self.redis.set(self._execution_key(task_key), execution_id) + return execution_id + + def _resolve_report_execution_id( + self, + worker_id: str, + task_key: str, + event_time: float, + retry_count: int = 0, + ) -> str: + execution_id = self.redis.get(self._execution_key(task_key)) + if execution_id: + return execution_id + execution_id = self.repo.get_latest_execution_id(task_key) + if execution_id: + return execution_id + return self._ensure_execution_id(worker_id, task_key, event_time, retry_count=retry_count) + + def _resolve_failure_fields( + self, + error_type: Optional[str], + error_message: Optional[str], + download_errors: Any = None, + trace_payload: Optional[Dict[str, Any]] = None, + ) -> Tuple[Optional[str], Optional[str], Optional[Dict[str, Any]]]: + trace_payload = dict(trace_payload) if isinstance(trace_payload, dict) else {} + normalized_download_errors = _normalize_download_errors(download_errors) + if normalized_download_errors: + trace_payload["download_errors"] = normalized_download_errors + failure_domain = trace_payload.get("failure_domain") + failure_subtype = trace_payload.get("failure_subtype") + has_failure = bool(error_type or error_message or normalized_download_errors) + if has_failure and (not failure_domain or not failure_subtype): + failure_domain, failure_subtype = _classify_failure( + error_type, + error_message, + download_errors=normalized_download_errors, + ) + elif not has_failure: + failure_domain = None + failure_subtype = None + if not trace_payload: + trace_payload = None + elif has_failure and not trace_payload.get("failure_domain"): + trace_payload["failure_domain"] = failure_domain + trace_payload["failure_subtype"] = failure_subtype + return failure_domain, failure_subtype, trace_payload + + def handle_worker_event(self, payload: Dict[str, Any]): + with self._lock: + worker_id = normalize_worker_id(payload.get("worker_id")) + task_key = payload.get("task_key") + event_type = payload.get("event_type") + if not worker_id or not event_type: + return + event_time = float(payload.get("event_time") or time.time()) + stage = payload.get("stage") + state = payload.get("state") + status = payload.get("status") + reason = payload.get("reason") + metrics = payload.get("metrics") or {} + + if event_type == "worker_state_changed" and state: + self.set_worker_state(worker_id, state, reason=reason, task_key=task_key, event_time=event_time) + return + + if not task_key: + return + + retry_count = int(metrics.get("retry_count") or 0) + execution_id = self._ensure_execution_id(worker_id, task_key, event_time, retry_count=retry_count) + self.repo.insert_task_stage_event( + execution_id=execution_id, + task_key=task_key, + worker_id=worker_id, + session_id=self.current_session_id, + event_time=event_time, + event_type=event_type, + stage=stage, + status=status, + reason=reason, + metrics=metrics, + ) + + if event_type == "task_started": + task_metadata = self._resolve_task_metadata(task_key) + self.repo.upsert_task_execution( + execution_id=execution_id, + task_key=task_key, + worker_id=worker_id, + session_id=self.current_session_id, + stat_date=_date_str_from_ts(event_time), + task_started_at=event_time, + app_name=task_metadata.get("app_name"), + package_name=task_metadata.get("package_name"), + last_updated=task_metadata.get("last_updated"), + collection_task_type=task_metadata.get("collection_task_type"), + retry_count=retry_count, + ) + self.handle_task_assigned(worker_id, task_key, event_time) + return + + if event_type == "stage_started": + if stage == "download": + self.set_worker_state(worker_id, "running_download", task_key=task_key, event_time=event_time) + elif stage == "collect": + self.set_worker_state(worker_id, "running_collect", task_key=task_key, event_time=event_time) + return + + if event_type == "stage_finished": + self.repo.update_task_execution_stage( + execution_id=execution_id, + stage=stage or "", + duration_seconds=metrics.get("duration_seconds"), + download_source=metrics.get("download_source"), + ) + return + + if event_type == "task_finished": + error_type = metrics.get("error_type") + error_message = metrics.get("error_message") or reason + total_duration = metrics.get("total_duration_seconds") + failed_total_duration = metrics.get("failed_total_duration_seconds") + failed_stage = metrics.get("failed_stage") + failure_domain, failure_subtype, trace_payload = self._resolve_failure_fields( + error_type, + error_message, + download_errors=metrics.get("download_errors"), + trace_payload=metrics.get("trace"), + ) + self.repo.finish_task_execution( + execution_id=execution_id, + task_ended_at=event_time, + status=status or "unknown", + failed_stage=failed_stage, + failure_bucket="failed" if (status or "unknown") == "failed" else status, + total_duration_seconds=total_duration, + failed_total_duration_seconds=failed_total_duration, + error_type=error_type, + error_message=error_message, + result_detail=error_message, + droidbot_steps=None, + guiagent_steps=None, + total_steps=None, + num_nodes=None, + num_reached_activities=None, + app_num_total_activities=None, + failure_domain=failure_domain, + failure_subtype=failure_subtype, + trace_payload=trace_payload, + retry_count=retry_count, + ) + self.redis.delete(self._execution_key(task_key)) + self.set_worker_state(worker_id, "idle_waiting_task", reason="waiting_task", event_time=event_time) + + def handle_worker_report( + self, + worker_id: str, + previous_task_key: str, + report_data: Dict[str, Any], + current_time: Optional[float] = None, + retry_count: int = 0, + resolved_message: Optional[str] = None, + ): + worker_id = normalize_worker_id(worker_id) + with self._lock: + if not previous_task_key: + return + current_time = current_time or time.time() + execution_id = self._resolve_report_execution_id( + worker_id, + previous_task_key, + current_time, + retry_count=retry_count, + ) + statistics = report_data.get("statistics") or {} + error_data = report_data.get("error") or {} + trace_payload = report_data.get("trace") + status = report_data.get("status", "unknown") + current_state = self._get_current_state(worker_id).get("state") + failed_stage = "download" if not statistics.get("download_source") else "collect" + if status == "success": + failed_stage = None + elif current_state == "running_download": + failed_stage = "download" + elif current_state == "running_collect": + failed_stage = "collect" + error_type = ( + f"{error_data.get('category')}/{error_data.get('code')}" + if error_data.get("category") is not None and error_data.get("code") is not None + else None + ) + error_message = resolved_message or error_data.get("reason") or statistics.get("error_reason") + failure_domain, failure_subtype, trace_payload = self._resolve_failure_fields( + error_type, + error_message, + download_errors=report_data.get("download_errors"), + trace_payload=trace_payload, + ) + self.repo.finish_task_execution( + execution_id=execution_id, + task_ended_at=current_time, + status=status, + failed_stage=failed_stage, + failure_bucket=status if status != "failed" else "failed", + total_duration_seconds=statistics.get("duration_seconds"), + failed_total_duration_seconds=statistics.get("duration_seconds") if status != "success" else None, + error_type=error_type, + error_message=error_message, + result_detail=resolved_message or error_message, + droidbot_steps=statistics.get("droidbot_steps"), + guiagent_steps=statistics.get("guiagent_steps"), + total_steps=statistics.get("total_steps"), + num_nodes=statistics.get("num_nodes"), + num_reached_activities=statistics.get("num_reached_activities"), + app_num_total_activities=statistics.get("app_num_total_activities"), + failure_domain=failure_domain, + failure_subtype=failure_subtype, + trace_payload=trace_payload, + retry_count=retry_count, + ) + self.redis.delete(self._execution_key(previous_task_key)) + + def mark_worker_disabled(self, worker_id: str, reason: str): + self.set_worker_state(worker_id, "disabled", reason=reason, event_time=time.time()) + + def mark_worker_offline(self, worker_id: str, reason: str): + self.set_worker_state(worker_id, "offline", reason=reason, event_time=time.time()) + + def set_failure_bucket(self, task_key: str, failure_bucket: str): + with self._lock: + execution_id = self.redis.get(self._execution_key(task_key)) + if execution_id: + self.repo.update_failure_bucket(execution_id, failure_bucket) + return + self.repo.update_failure_bucket_by_task_key(task_key, failure_bucket) + + def _controller_up_seconds(self, date_str: str, now_ts: Optional[float] = None) -> float: + day_start, day_end, _ = _day_bounds(date_str, now_ts) + total = 0.0 + for row in self.repo.list_sessions_overlapping(day_start, day_end): + start = max(float(row["started_at"]), day_start) + end = min(float(row["ended_at"] or row["last_heartbeat_at"] or day_end), day_end) + total += max(0.0, end - start) + return round(total, 2) + + def _resolve_worker_ids( + self, + date_str: str, + dispatcher=None, + worker_id: Optional[str] = None, + ) -> List[str]: + if worker_id: + return [worker_id] + _, day_end, is_today = _day_bounds(date_str) + # 以 managed_worker_ids（inventory.json）为准 + worker_ids = set(self.managed_worker_ids) + if not is_today: + # 仅在查看历史日期时合并 SQLite 中的 worker（兼容查看旧数据） + worker_ids.update(self.repo.list_worker_ids_up_to(day_end)) + if dispatcher is not None and is_today: + for worker in dispatcher.get_dashboard_workers(): + worker_ids.add(worker.get("worker_id")) + worker_ids.discard(None) + worker_ids.discard("") + return sorted(worker_ids) + + def _build_worker_options(self, date_str: str, dispatcher=None) -> List[Dict[str, Any]]: + worker_ids = self._resolve_worker_ids(date_str, dispatcher=dispatcher) + current_workers: Dict[str, Dict[str, Any]] = {} + if dispatcher is not None: + current_workers = {worker.get("worker_id"): worker for worker in dispatcher.get_dashboard_workers()} + + options: List[Dict[str, Any]] = [] + for worker_id in worker_ids: + current = current_workers.get(worker_id, {}) + ip_address = current.get("ip_address") or worker_id.split("_", 1)[0] + hostname = current.get("hostname") or "" + label = f"{ip_address} · {hostname}" if hostname else ip_address + options.append( + { + "worker_id": worker_id, + "ip_address": ip_address, + "hostname": hostname, + "label": label, + } + ) + return options + + def _build_worker_intervals( + self, + date_str: str, + dispatcher=None, + worker_id: Optional[str] = None, + ) -> Dict[str, List[Tuple[float, float, str]]]: + day_start, day_end, _ = _day_bounds(date_str) + worker_ids = self._resolve_worker_ids(date_str, dispatcher=dispatcher, worker_id=worker_id) + rows = self.repo.list_state_events_before(day_end, worker_ids=worker_ids) + grouped: Dict[str, List[sqlite3.Row]] = defaultdict(list) + for row in rows: + grouped[row["worker_id"]].append(row) + + intervals: Dict[str, List[Tuple[float, float, str]]] = {} + for worker_id in worker_ids: + events = grouped.get(worker_id, []) + prev_state = "offline" + prev_time = day_start + worker_intervals: List[Tuple[float, float, str]] = [] + for row in events: + event_time = float(row["event_time"]) + if event_time < day_start: + prev_state = row["state"] + continue + if event_time > day_end: + break + if event_time > prev_time: + worker_intervals.append((prev_time, event_time, prev_state)) + prev_state = row["state"] + prev_time = event_time + if day_end > prev_time: + worker_intervals.append((prev_time, day_end, prev_state)) + intervals[worker_id] = worker_intervals + return intervals + + def _summarize_worker_seconds( + self, + date_str: str, + dispatcher=None, + worker_id: Optional[str] = None, + ) -> Dict[str, Counter]: + intervals = self._build_worker_intervals(date_str, dispatcher=dispatcher, worker_id=worker_id) + summary: Dict[str, Counter] = {} + for worker_id, worker_intervals in intervals.items(): + counts = Counter() + for start, end, state in worker_intervals: + counts[state] += max(0.0, end - start) + summary[worker_id] = counts + return summary + + def _runtime_state_breakdown( + self, + date_str: str, + dispatcher=None, + worker_id: Optional[str] = None, + ) -> List[Dict[str, Any]]: + totals = Counter() + for counts in self._summarize_worker_seconds(date_str, dispatcher=dispatcher, worker_id=worker_id).values(): + for state in RUNTIME_STATE_ORDER: + totals[state] += float(counts.get(state, 0.0)) + total_seconds = sum(totals.values()) + return [ + { + "state": state, + "label": RUNTIME_STATE_LABELS[state], + "duration_seconds": round(float(totals[state]), 2), + "share_percent": round((float(totals[state]) / total_seconds) * 100, 2) if total_seconds > 0 else 0.0, + } + for state in RUNTIME_STATE_ORDER + ] + + def get_overview( + self, + date_str: Optional[str] = None, + dispatcher=None, + worker_id: Optional[str] = None, + ) -> Dict[str, Any]: + now_ts = time.time() + date_str = date_str or _date_str_from_ts(now_ts) + day_start, day_end, is_today = _day_bounds(date_str, now_ts=now_ts) + up_seconds = self._controller_up_seconds(date_str, now_ts=now_ts) + down_seconds = round(max(0.0, day_end - day_start - up_seconds), 2) + metrics = self.repo.get_task_overview_metrics(date_str, worker_id=worker_id) + overall_totals = self.repo.get_task_execution_totals(worker_id=worker_id) + overall_success_total = overall_totals["task_success"] + # The global dashboard should reflect the dispatcher's in-memory success set, + # including tasks preloaded from the current active catalog queue. + if dispatcher is not None and not worker_id: + overall_success_total = int(dispatcher.get_statistics().get("completed") or 0) + + state_counts = { + "online_workers": 0, + "running_workers": 0, + "running_download_workers": 0, + "running_collect_workers": 0, + "idle_workers": 0, + "offline_workers": 0, + "disabled_workers": 0, + "pending_tasks": 0, + } + if is_today and dispatcher is not None: + dashboard_workers = dispatcher.get_dashboard_workers() + if worker_id: + dashboard_workers = [worker for worker in dashboard_workers if worker.get("worker_id") == worker_id] + state_counts["pending_tasks"] = 0 if worker_id else dispatcher.get_statistics().get("pending", 0) + for worker in dashboard_workers: + current_worker_id = worker.get("worker_id", "") + state = worker.get("state") or self._get_current_state(current_worker_id).get("state") or "offline" + if not worker.get("online"): + state = "offline" + if state in ONLINE_STATES: + state_counts["online_workers"] += 1 + if state in RUNNING_STATES: + state_counts["running_workers"] += 1 + if state == "running_download": + state_counts["running_download_workers"] += 1 + elif state == "running_collect": + state_counts["running_collect_workers"] += 1 + elif state == "idle_waiting_task": + state_counts["idle_workers"] += 1 + elif state == "disabled": + state_counts["disabled_workers"] += 1 + else: + state_counts["offline_workers"] += 1 + else: + timeline = self.get_timeline(date_str=date_str, dispatcher=dispatcher, worker_id=worker_id) + last = timeline["points"][-1] if timeline["points"] else {} + state_counts = { + "online_workers": last.get("online_workers", 0), + "running_workers": last.get("running_workers", 0), + "running_download_workers": last.get("running_download_workers", 0), + "running_collect_workers": last.get("running_collect_workers", 0), + "idle_workers": last.get("idle_workers", 0), + "offline_workers": last.get("offline_workers", 0), + "disabled_workers": last.get("disabled_workers", 0), + } + + return { + "date": date_str, + "is_today": is_today, + "controller_status": "running" if (is_today and self.current_session_id) else "stopped", + "controller_up_seconds": up_seconds, + "controller_down_seconds": down_seconds, + "task_total": metrics["task_total"], + "task_success": metrics["task_success"], + "task_non_success": metrics["task_non_success"], + "task_failed": metrics["task_failed"], + "task_total_all": overall_totals["task_total"], + "task_success_all": overall_success_total, + "task_non_success_all": overall_totals["task_non_success"], + "task_failed_all": overall_totals["task_failed"], + "avg_download_duration_seconds": metrics["avg_download_duration_seconds"], + "avg_collect_duration_seconds": metrics["avg_collect_duration_seconds"], + "p50_download_duration_seconds": metrics["p50_download_duration_seconds"], + "p90_download_duration_seconds": metrics["p90_download_duration_seconds"], + "p50_collect_duration_seconds": metrics["p50_collect_duration_seconds"], + "p90_collect_duration_seconds": metrics["p90_collect_duration_seconds"], + **state_counts, + "selected_worker_id": worker_id or "", + "worker_options": self._build_worker_options(date_str, dispatcher=dispatcher), + "generated_at": datetime.fromtimestamp(now_ts, SHANGHAI_TZ).strftime("%Y-%m-%d %H:%M:%S"), + } + + def get_distributions( + self, + date_str: Optional[str] = None, + worker_id: Optional[str] = None, + dispatcher=None, + ) -> Dict[str, Any]: + date_str = date_str or _date_str_from_ts(time.time()) + task_rows = self.repo.list_task_executions(date_str, worker_id=worker_id, limit=100000) + _, severe_items = self._split_rows_by_quality(task_rows) + download_values = [float(row["download_duration_seconds"]) for row in task_rows if row["download_duration_seconds"] is not None] + collect_values = [float(row["collect_duration_seconds"]) for row in task_rows if row["collect_duration_seconds"] is not None] + total_values = [float(row["total_duration_seconds"]) for row in task_rows if row["total_duration_seconds"] is not None] + severe_values = [ + float(row["failed_total_duration_seconds"] or row["total_duration_seconds"] or 0.0) + for row in severe_items + if (row["failed_total_duration_seconds"] is not None or row["total_duration_seconds"] is not None) + ] + bucket_edges = [60, 180, 300, 600, 1200, 1800, 3600] + return { + "date": date_str, + "download_duration": _build_histogram(download_values, bucket_edges), + "collect_duration": _build_histogram(collect_values, bucket_edges), + "total_duration": _build_histogram(total_values, bucket_edges), + "non_success_total_duration": _build_histogram(severe_values, bucket_edges), + **self.get_stage_analysis(date_str, worker_id=worker_id, dispatcher=dispatcher, rows=task_rows), + } + + def get_stage_analysis( + self, + date_str: Optional[str] = None, + worker_id: Optional[str] = None, + dispatcher=None, + rows: Optional[List[sqlite3.Row]] = None, + ) -> Dict[str, Any]: + date_str = date_str or _date_str_from_ts(time.time()) + rows = rows if rows is not None else self.repo.list_task_executions(date_str, worker_id=worker_id, limit=100000) + worker_options = self._build_worker_options(date_str, dispatcher=dispatcher) + worker_label_map = { + item["worker_id"]: item.get("label") or item["worker_id"] + for item in worker_options + } + qualified_items, severe_items, collection_quality = self._classify_rows_by_collection_quality(rows) + + def row_total(row: sqlite3.Row) -> float: + return _task_row_total_duration(row) + + def summarize(items: List[sqlite3.Row]) -> List[Dict[str, Any]]: + if not items: + return [] + total_sum = sum(max(row_total(item), 0.0) for item in items) + result = [] + for stage in STAGE_ORDER: + key = f"{stage}_duration_seconds" + stage_sum = sum(float(item[key] or 0.0) for item in items) + result.append( + { + "stage": stage, + "label": STAGE_LABELS.get(stage, stage), + "avg_duration_seconds": round(stage_sum / len(items), 2), + "share_percent": round((stage_sum / total_sum) * 100, 2) if total_sum > 0 else 0.0, + "total_duration_seconds": round(stage_sum, 2), + } + ) + return result + + qualified_total = round(sum(row_total(item) for item in qualified_items), 2) + severe_total = round(sum(row_total(item) for item in severe_items), 2) + total_runtime = qualified_total + severe_total + time_composition = [ + { + "bucket": "qualified", + "label": "合格", + "duration_seconds": qualified_total, + "share_percent": round((qualified_total / total_runtime) * 100, 2) if total_runtime > 0 else 0.0, + }, + { + "bucket": "severe_restricted", + "label": "严重受限", + "duration_seconds": severe_total, + "share_percent": round((severe_total / total_runtime) * 100, 2) if total_runtime > 0 else 0.0, + }, + ] + + domain_buckets: Dict[str, Dict[str, Any]] = {} + for row in severe_items: + failure_domain = row["failure_domain"] + failure_subtype = row["failure_subtype"] + if not failure_domain or not failure_subtype: + failure_domain, failure_subtype = _classify_failure(row["error_type"], row["error_message"]) + duration_seconds = row_total(row) + reason_label = _humanize_failure_reason(row["error_type"], row["error_message"]) + message_label = row["error_message"] or reason_label + sample = _build_failure_task_sample(row) + + domain_entry = domain_buckets.setdefault( + failure_domain, + { + "domain": failure_domain, + "label": _humanize_failure_domain(failure_domain), + "duration_seconds": 0.0, + "task_count": 0, + "_subtypes": {}, + "_reasons": {}, + "_messages": {}, + "_tasks": [], + }, + ) + domain_entry["duration_seconds"] += duration_seconds + domain_entry["task_count"] += 1 + domain_entry["_tasks"].append(sample) + + domain_reason = domain_entry["_reasons"].setdefault(reason_label, {"label": reason_label, "duration_seconds": 0.0, "task_count": 0}) + domain_reason["duration_seconds"] += duration_seconds + domain_reason["task_count"] += 1 + + domain_message = domain_entry["_messages"].setdefault(message_label, {"label": message_label, "duration_seconds": 0.0, "task_count": 0}) + domain_message["duration_seconds"] += duration_seconds + domain_message["task_count"] += 1 + + subtype_entry = domain_entry["_subtypes"].setdefault( + failure_subtype, + { + "subtype": failure_subtype, + "label": _humanize_failure_subtype(failure_domain, failure_subtype, row["error_type"], row["error_message"]), + "duration_seconds": 0.0, + "task_count": 0, + "_reasons": {}, + "_messages": {}, + "_tasks": [], + }, + ) + subtype_entry["duration_seconds"] += duration_seconds + subtype_entry["task_count"] += 1 + subtype_entry["_tasks"].append(sample) + + subtype_reason = subtype_entry["_reasons"].setdefault(reason_label, {"label": reason_label, "duration_seconds": 0.0, "task_count": 0}) + subtype_reason["duration_seconds"] += duration_seconds + subtype_reason["task_count"] += 1 + + subtype_message = subtype_entry["_messages"].setdefault(message_label, {"label": message_label, "duration_seconds": 0.0, "task_count": 0}) + subtype_message["duration_seconds"] += duration_seconds + subtype_message["task_count"] += 1 + + failure_domain_breakdown = [] + failure_subtype_breakdown = [] + for domain_entry in sorted(domain_buckets.values(), key=lambda item: item["duration_seconds"], reverse=True): + domain_duration = float(domain_entry["duration_seconds"]) + subtypes = [] + for subtype_entry in sorted(domain_entry["_subtypes"].values(), key=lambda item: item["duration_seconds"], reverse=True): + subtype_duration = float(subtype_entry["duration_seconds"]) + subtype_reasons = sorted(subtype_entry["_reasons"].values(), key=lambda item: item["duration_seconds"], reverse=True) + subtype_messages = sorted(subtype_entry["_messages"].values(), key=lambda item: item["duration_seconds"], reverse=True) + subtype_tasks = sorted( + subtype_entry["_tasks"], + key=lambda item: ( + float(item.get("failed_total_duration_seconds") or item.get("total_duration_seconds") or 0.0), + item.get("task_ended_at", ""), + ), + reverse=True, + ) + reasons = [ + { + "label": item["label"], + "duration_seconds": round(float(item["duration_seconds"]), 2), + "task_count": int(item["task_count"]), + "overall_share_percent": round((float(item["duration_seconds"]) / total_runtime) * 100, 2) if total_runtime > 0 else 0.0, + "error_share_percent": round((float(item["duration_seconds"]) / severe_total) * 100, 2) if severe_total > 0 else 0.0, + "domain_share_percent": round((float(item["duration_seconds"]) / domain_duration) * 100, 2) if domain_duration > 0 else 0.0, + "subtype_share_percent": round((float(item["duration_seconds"]) / subtype_duration) * 100, 2) if subtype_duration > 0 else 0.0, + } + for item in subtype_reasons + ] + messages = [ + { + "label": item["label"], + "message": item["label"], + "duration_seconds": round(float(item["duration_seconds"]), 2), + "task_count": int(item["task_count"]), + "overall_share_percent": round((float(item["duration_seconds"]) / total_runtime) * 100, 2) if total_runtime > 0 else 0.0, + "error_share_percent": round((float(item["duration_seconds"]) / severe_total) * 100, 2) if severe_total > 0 else 0.0, + "domain_share_percent": round((float(item["duration_seconds"]) / domain_duration) * 100, 2) if domain_duration > 0 else 0.0, + "subtype_share_percent": round((float(item["duration_seconds"]) / subtype_duration) * 100, 2) if subtype_duration > 0 else 0.0, + } + for item in subtype_messages + ] + subtype_item = { + "domain": domain_entry["domain"], + "subtype": subtype_entry["subtype"], + "label": subtype_entry["label"], + "duration_seconds": round(subtype_duration, 2), + "task_count": int(subtype_entry["task_count"]), + "avg_duration_seconds": round(subtype_duration / subtype_entry["task_count"], 2) if subtype_entry["task_count"] else 0.0, + "overall_share_percent": round((subtype_duration / total_runtime) * 100, 2) if total_runtime > 0 else 0.0, + "error_share_percent": round((subtype_duration / severe_total) * 100, 2) if severe_total > 0 else 0.0, + "domain_share_percent": round((subtype_duration / domain_duration) * 100, 2) if domain_duration > 0 else 0.0, + "reasons": reasons, + "messages": messages, + "tasks": subtype_tasks, + } + subtypes.append(subtype_item) + failure_subtype_breakdown.append(subtype_item) + + domain_reasons = sorted(domain_entry["_reasons"].values(), key=lambda item: item["duration_seconds"], reverse=True) + domain_messages = sorted(domain_entry["_messages"].values(), key=lambda item: item["duration_seconds"], reverse=True) + domain_tasks = sorted( + domain_entry["_tasks"], + key=lambda item: ( + float(item.get("failed_total_duration_seconds") or item.get("total_duration_seconds") or 0.0), + item.get("task_ended_at", ""), + ), + reverse=True, + ) + failure_domain_breakdown.append( + { + "domain": domain_entry["domain"], + "label": domain_entry["label"], + "duration_seconds": round(domain_duration, 2), + "task_count": int(domain_entry["task_count"]), + "overall_share_percent": round((domain_duration / total_runtime) * 100, 2) if total_runtime > 0 else 0.0, + "error_share_percent": round((domain_duration / severe_total) * 100, 2) if severe_total > 0 else 0.0, + "share_percent": round((domain_duration / severe_total) * 100, 2) if severe_total > 0 else 0.0, + "subtypes": subtypes, + "reasons": [ + { + "label": item["label"], + "duration_seconds": round(float(item["duration_seconds"]), 2), + "task_count": int(item["task_count"]), + "overall_share_percent": round((float(item["duration_seconds"]) / total_runtime) * 100, 2) if total_runtime > 0 else 0.0, + "error_share_percent": round((float(item["duration_seconds"]) / severe_total) * 100, 2) if severe_total > 0 else 0.0, + "domain_share_percent": round((float(item["duration_seconds"]) / domain_duration) * 100, 2) if domain_duration > 0 else 0.0, + } + for item in domain_reasons + ], + "messages": [ + { + "label": item["label"], + "message": item["label"], + "duration_seconds": round(float(item["duration_seconds"]), 2), + "task_count": int(item["task_count"]), + "overall_share_percent": round((float(item["duration_seconds"]) / total_runtime) * 100, 2) if total_runtime > 0 else 0.0, + "error_share_percent": round((float(item["duration_seconds"]) / severe_total) * 100, 2) if severe_total > 0 else 0.0, + "domain_share_percent": round((float(item["duration_seconds"]) / domain_duration) * 100, 2) if domain_duration > 0 else 0.0, + } + for item in domain_messages + ], + "tasks": domain_tasks, + } + ) + + failure_reasons = Counter() + for row in severe_items: + label = _humanize_failure_reason(row["error_type"], row["error_message"]) + failure_reasons[label] += 1 + + if len(time_composition) > 1: + time_composition[1]["details"] = [ + { + "domain": item["domain"], + "label": item["label"], + "duration_seconds": item["duration_seconds"], + "overall_share_percent": item["overall_share_percent"], + "error_share_percent": item["error_share_percent"], + "task_count": item["task_count"], + } + for item in failure_domain_breakdown + ] + + worker_stage_buckets: Dict[str, Dict[str, Any]] = {} + total_collect_seconds = 0.0 + for row in rows: + current_worker_id = row["worker_id"] + if not current_worker_id: + continue + download_seconds = float(row["download_duration_seconds"] or 0.0) + collect_seconds = float(row["collect_duration_seconds"] or 0.0) + if download_seconds <= 0 and collect_seconds <= 0: + continue + stage_total_seconds = download_seconds + collect_seconds + bucket = worker_stage_buckets.setdefault( + current_worker_id, + { + "worker_id": current_worker_id, + "label": worker_label_map.get(current_worker_id, current_worker_id), + "download_duration_seconds": 0.0, + "collect_duration_seconds": 0.0, + "stage_total_seconds": 0.0, + "task_count": 0, + "success_count": 0, + "non_success_count": 0, + }, + ) + bucket["download_duration_seconds"] += download_seconds + bucket["collect_duration_seconds"] += collect_seconds + bucket["stage_total_seconds"] += stage_total_seconds + bucket["task_count"] += 1 + if row["status"] == "success": + bucket["success_count"] += 1 + else: + bucket["non_success_count"] += 1 + total_collect_seconds += collect_seconds + + worker_collect_share_breakdown = [ + { + "worker_id": item["worker_id"], + "label": item["label"], + "download_duration_seconds": round(float(item["download_duration_seconds"]), 2), + "collect_duration_seconds": round(float(item["collect_duration_seconds"]), 2), + "stage_total_seconds": round(float(item["stage_total_seconds"]), 2), + "collect_share_percent": round((float(item["collect_duration_seconds"]) / float(item["stage_total_seconds"])) * 100, 2) + if item["stage_total_seconds"] > 0 + else 0.0, + "overall_collect_share_percent": round((float(item["collect_duration_seconds"]) / total_collect_seconds) * 100, 2) + if total_collect_seconds > 0 + else 0.0, + "task_count": int(item["task_count"]), + "success_count": int(item["success_count"]), + "non_success_count": int(item["non_success_count"]), + } + for item in sorted( + worker_stage_buckets.values(), + key=lambda current: ( + round((float(current["collect_duration_seconds"]) / float(current["stage_total_seconds"])) * 100, 6) + if current["stage_total_seconds"] > 0 + else 0.0, + float(current["collect_duration_seconds"]), + ), + reverse=True, + ) + ] + + return { + "success_stage_breakdown": summarize(qualified_items), + "non_success_stage_breakdown": summarize(severe_items), + "time_composition": time_composition, + "daily_collection_overview": { + "qualified_task_count": int(collection_quality["qualified_task_count"]), + "qualified_success_task_count": int(collection_quality["qualified_success_task_count"]), + "qualified_light_task_count": int(collection_quality["qualified_light_task_count"]), + "non_success_task_count": int(collection_quality["non_success_task_count"]), + "non_retryable_task_count": int(collection_quality["non_retryable_task_count"]), + }, + "runtime_state_breakdown": self._runtime_state_breakdown( + date_str, + dispatcher=dispatcher, + worker_id=worker_id, + ), + "failure_domain_breakdown": failure_domain_breakdown, + "failure_subtype_breakdown": failure_subtype_breakdown, + "failure_reasons": [ + {"label": label, "count": count} + for label, count in failure_reasons.most_common() + ], + "success_task_count": len(qualified_items), + "non_success_task_count": len(severe_items), + "worker_collect_share_breakdown": worker_collect_share_breakdown, + "total_collect_duration_seconds": round(total_collect_seconds, 2), + "selected_worker_id": worker_id or "", + "worker_options": worker_options, + } + + def get_timeline( + self, + date_str: Optional[str] = None, + dispatcher=None, + bucket_minutes: int = MONITORING_TIMELINE_BUCKET_MINUTES, + worker_id: Optional[str] = None, + ) -> Dict[str, Any]: + date_str = date_str or _date_str_from_ts(time.time()) + day_start, day_end, _ = _day_bounds(date_str) + intervals = self._build_worker_intervals(date_str, dispatcher=dispatcher, worker_id=worker_id) + bucket_seconds = max(1, bucket_minutes) * 60 + buckets: List[Dict[str, Any]] = [] + bucket_start = day_start + while bucket_start < day_end: + buckets.append( + { + "timestamp": bucket_start, + "label": datetime.fromtimestamp(bucket_start, SHANGHAI_TZ).strftime("%H:%M"), + "online_workers": 0, + "running_workers": 0, + "running_download_workers": 0, + "running_collect_workers": 0, + "idle_workers": 0, + "offline_workers": 0, + "disabled_workers": 0, + } + ) + bucket_start += bucket_seconds + + for worker_intervals in intervals.values(): + for interval_start, interval_end, state in worker_intervals: + for bucket in buckets: + bucket_start = bucket["timestamp"] + bucket_end = bucket_start + bucket_seconds + overlap_start = max(interval_start, bucket_start) + overlap_end = min(interval_end, bucket_end) + if overlap_start >= overlap_end: + continue + weight = (overlap_end - overlap_start) / bucket_seconds + if state in ONLINE_STATES: + bucket["online_workers"] += weight + if state in RUNNING_STATES: + bucket["running_workers"] += weight + if state == "running_download": + bucket["running_download_workers"] += weight + elif state == "running_collect": + bucket["running_collect_workers"] += weight + elif state == "idle_waiting_task": + bucket["idle_workers"] += weight + elif state == "disabled": + bucket["disabled_workers"] += weight + else: + bucket["offline_workers"] += weight + for bucket in buckets: + for key in ( + "online_workers", + "running_workers", + "running_download_workers", + "running_collect_workers", + "idle_workers", + "offline_workers", + "disabled_workers", + ): + bucket[key] = round(float(bucket[key]), 2) + return {"date": date_str, "bucket_minutes": bucket_minutes, "points": buckets} + + def get_worker_stats( + self, + date_str: Optional[str] = None, + dispatcher=None, + ) -> List[Dict[str, Any]]: + now_ts = time.time() + date_str = date_str or _date_str_from_ts(now_ts) + is_today = _date_str_from_ts(now_ts) == date_str + seconds_by_worker = self._summarize_worker_seconds(date_str, dispatcher=dispatcher) + task_rows = self.repo.list_task_executions(date_str, limit=100000) + task_stats: Dict[str, Dict[str, Any]] = defaultdict( + lambda: { + "task_count": 0, + "success_count": 0, + "failed_count": 0, + "download_values": [], + "collect_values": [], + "last_task_at": None, + } + ) + for row in task_rows: + entry = task_stats[row["worker_id"]] + entry["task_count"] += 1 + if row["status"] == "success": + entry["success_count"] += 1 + elif row["status"] == "failed": + entry["failed_count"] += 1 + if row["download_duration_seconds"] is not None: + entry["download_values"].append(float(row["download_duration_seconds"])) + if row["collect_duration_seconds"] is not None: + entry["collect_values"].append(float(row["collect_duration_seconds"])) + last_task_at = row["task_ended_at"] or row["task_started_at"] + if last_task_at is not None: + if entry["last_task_at"] is None or float(last_task_at) > entry["last_task_at"]: + entry["last_task_at"] = float(last_task_at) + + current_workers = {} + if dispatcher is not None and is_today: + current_workers = {worker.get("worker_id"): worker for worker in dispatcher.get_dashboard_workers()} + + results: List[Dict[str, Any]] = [] + for worker_id in sorted(set(self._resolve_worker_ids(date_str, dispatcher=dispatcher))): + seconds = seconds_by_worker.get(worker_id, Counter()) + metrics = task_stats.get(worker_id, {}) + current_worker = current_workers.get(worker_id, {}) + current_state = ( + current_worker.get("state") + or self._get_current_state(worker_id).get("state") + or ("offline" if not current_worker.get("online") else "idle_waiting_task") + ) + if is_today and current_worker and not current_worker.get("online"): + current_state = "offline" + last_task_at = metrics.get("last_task_at") + results.append( + { + "worker_id": worker_id, + "ip_address": current_worker.get("ip_address") or worker_id.split("_", 1)[0], + "status": DISPLAY_STATE_LABELS.get(current_state, current_state), + "running_seconds": round(seconds.get("running_download", 0.0) + seconds.get("running_collect", 0.0), 2), + "idle_waiting_seconds": round(seconds.get("idle_waiting_task", 0.0), 2), + "offline_seconds": round(seconds.get("offline", 0.0), 2), + "disabled_seconds": round(seconds.get("disabled", 0.0), 2), + "task_count": metrics.get("task_count", 0), + "success_count": metrics.get("success_count", 0), + "failed_count": metrics.get("failed_count", 0), + "avg_download_duration_seconds": round(sum(metrics.get("download_values", [])) / len(metrics.get("download_values", [])), 2) + if metrics.get("download_values") + else 0.0, + "avg_collect_duration_seconds": round(sum(metrics.get("collect_values", [])) / len(metrics.get("collect_values", [])), 2) + if metrics.get("collect_values") + else 0.0, + "last_task_at": datetime.fromtimestamp(last_task_at, SHANGHAI_TZ).strftime("%Y-%m-%d %H:%M:%S") + if last_task_at + else "--", + } + ) + return results + + def get_task_rows( + self, + date_str: Optional[str] = None, + worker_id: Optional[str] = None, + limit: int = 200, + ) -> List[Dict[str, Any]]: + date_str = date_str or _date_str_from_ts(time.time()) + rows = self.repo.list_task_executions(date_str, worker_id=worker_id, limit=limit) + return [_build_failure_task_sample(row) for row in rows] diff --git a/notify_manager.py b/notify_manager.py new file mode 100644 index 0000000..c947313 --- /dev/null +++ b/notify_manager.py @@ -0,0 +1,62 @@ +import requests +import logging +import time + +logger = logging.getLogger(__name__) + +class Notifier: + def __init__(self, wechat_tokens, wecom_tokens=None): + self.wechat_tokens = wechat_tokens + self.wecom_tokens = wecom_tokens or {} + # Dispatcher 端不需要获取本机 IP，告警信息中通常包含 Worker ID + # self.ip = self._get_ip_address() + + def send_wechat_alert(self, worker_id, message): + """发送微信告警 + Args: + worker_id: 发生告警的 Worker ID + message: 告警内容 + """ + title = f"设备告警：{worker_id}" + if not self.wechat_tokens: + logger.warning("未配置微信推送 Token，跳过发送。") + return + + for name, token in self.wechat_tokens.items(): + url = f'https://wx.xtuis.cn/{token}.send' + try: + requests.post(url, data={'text': title, 'desp': f"设备 {worker_id} {message}"}, timeout=10) + logger.info(f"告警已发送给: {name}") + time.sleep(3) # 避免发送太快导致失败 + except Exception as e: + logger.warning(f"发送告警失败: {e}") + + def send_weCom_alert(self, worker_id, message): + """发送企业微信告警 + Args: + worker_id: 发生告警的 Worker ID + message: 告警内容 + """ + title = f"{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}, {worker_id}" + content = f"{title}\n\n{message}" + + if not self.wecom_tokens: + logger.warning("未配置企业微信 Webhook Key，跳过发送。") + return + + for name, key in self.wecom_tokens.items(): + url = f'https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key={key}' + payload = { + "msgtype": "text", + "text": { + "content": content + } + } + try: + response = requests.post(url, json=payload, timeout=10) + if response.status_code == 200: + logger.info(f"企业微信告警已发送给: {name}") + else: + logger.warning(f"发送企业微信告警到 {name} 失败: HTTP {response.status_code}, {response.text}") + except Exception as e: + logger.warning(f"发送企业微信告警到 {name} 失败: {e}") diff --git a/redis_task_distribute.py b/redis_task_distribute.py new file mode 100755 index 0000000..aad7998 --- /dev/null +++ b/redis_task_distribute.py @@ -0,0 +1,2921 @@ +import redis +from redis import ConnectionPool +import json +import os +import threading +import time +import traceback +from datetime import datetime +from typing import Optional, Dict, Any, List, Tuple +from config import ( + ANALYTICS_ARTIFACT_WAIT_SECONDS, + ANALYTICS_TRAFFIC_ROOT, + ANALYTICS_TRAFFIC_ROOT_BLOCK, + REDIS_HOST, REDIS_PORT, REDIS_DB, REDIS_MAX_CONNECTIONS_DISPATCHER, + WECHAT_TOKENS, WECOM_TOKENS, + ALERT_WINDOW_SECONDS, ALERT_THRESHOLD, ALERT_COOLDOWN_SECONDS, + WORKER_REINIT_TIMEOUT, WORKER_INIT_MONITOR_WINDOW, + RUN_PIPELINE_DEFAULT_OPTIONS, + AUTO_REBOOT_RECOVERY_DELAY_SECONDS, + WORKER_INVENTORY, + MANAGED_WORKER_IDS, TASK_ROUTING_RULES, + ONLY_MANAGED_WORKERS_CAN_DISPATCH, WORKER_ONLINE_TIMEOUT_SECONDS, + normalize_worker_id, + APK_DOWNLOAD_MODE, + APK_PREFETCH_POLL_INTERVAL, + APK_DOWNLOAD_QUEUE_INTERVAL, +) +from analytics import AnalyticsService, _make_task_key +from log_manager import logger +from monitoring import MonitoringService +from notify_manager import Notifier +from remote_worker_controller import RemoteWorkerController +from result_codes import ErrorInfo, ErrorCategory, AppError, BusinessError, InfraError, DownloadError + +class RedisTaskDispatcher: + HOT_APP_RETRY_DOWNLOADS_THRESHOLD = 1_000_000 + INFRA_REBOOT_RECOVERY_DELAY_SECONDS = 300 + TASK_QUEUE_NAMES = ("task:queue:high", "task:queue:default", "task:queue:low") + TASK_QUEUE_CLEANUP_NAMES = (*TASK_QUEUE_NAMES, "task:queue") + DIRECT_APK_SOURCES = ["google_play", "apkpure"] + APK_CACHE_SOURCES = ["google_play", "local", "apkpure"] + # 类级别的连接池（所有实例共享） + _connection_pool: Optional[ConnectionPool] = None + _connection_pool_config: Optional[Tuple[str, int, int, int]] = None + + def __init__( + self, + redis_host=REDIS_HOST, + redis_port=REDIS_PORT, + redis_db=REDIS_DB, + max_connections=REDIS_MAX_CONNECTIONS_DISPATCHER, + ): + """初始化任务分发器 + + Args: + redis_host: Redis服务器地址 + redis_port: Redis端口 + redis_db: Redis数据库编号 + max_connections: 连接池最大连接数（需要支持30+台PC并发） + """ + # 使用连接池提高性能 + pool_config = (redis_host, redis_port, redis_db, max_connections) + if ( + RedisTaskDispatcher._connection_pool is None + or RedisTaskDispatcher._connection_pool_config != pool_config + ): + RedisTaskDispatcher._connection_pool = ConnectionPool( + host=redis_host, + port=redis_port, + db=redis_db, + decode_responses=True, + max_connections=max_connections, + socket_timeout=30, # 连接超时时间 + socket_connect_timeout=10, # 连接建立超时 + retry_on_timeout=True # 超时时自动重试 + ) + RedisTaskDispatcher._connection_pool_config = pool_config + self.redis = redis.Redis(connection_pool=RedisTaskDispatcher._connection_pool) + self.worker_inventory = {item["worker_id"]: dict(item) for item in WORKER_INVENTORY} + self.managed_worker_ids = set(MANAGED_WORKER_IDS) + self.only_managed_workers_can_dispatch = ONLY_MANAGED_WORKERS_CAN_DISPATCH + self.worker_online_timeout = WORKER_ONLINE_TIMEOUT_SECONDS + self.task_routing_rules = self._normalize_routing_rules(TASK_ROUTING_RULES) + self.remote_controller = RemoteWorkerController(WORKER_INVENTORY) + + # 初始化告警推送器 + self.notifier = Notifier(WECHAT_TOKENS, WECOM_TOKENS) + self.monitor = MonitoringService(self.redis, managed_worker_ids=MANAGED_WORKER_IDS) + self.analytics = AnalyticsService( + db_path=self.monitor.repo.db_path, + artifact_wait_seconds=ANALYTICS_ARTIFACT_WAIT_SECONDS, + on_snapshot_callback=self._handle_analytics_snapshot, + traffic_root_block=ANALYTICS_TRAFFIC_ROOT_BLOCK, + ) + self._apk_registry = None # lazy init + self._minio_storage = None # lazy init + + def _normalize_worker_ids(self, worker_ids: Optional[List[str]]) -> List[str]: + normalized = [] + seen = set() + for worker_id in worker_ids or []: + candidate = normalize_worker_id(worker_id) + if candidate and candidate not in seen: + seen.add(candidate) + normalized.append(candidate) + return normalized + + def _migrate_worker_runtime_id(self, old_worker_id: str, worker_id: str) -> None: + if not old_worker_id or not worker_id or old_worker_id == worker_id: + return + old_info_json = self.redis.hget("workers:info", old_worker_id) + if old_info_json and not self.redis.hget("workers:info", worker_id): + try: + info = json.loads(old_info_json) + except (TypeError, ValueError): + info = {} + info["worker_id"] = worker_id + info.setdefault("ip_address", worker_id) + self.redis.hset("workers:info", worker_id, json.dumps(info)) + if old_info_json: + self.redis.hdel("workers:info", old_worker_id) + for set_key in ("workers:idle", "workers:busy"): + if self.redis.sismember(set_key, old_worker_id): + self.redis.srem(set_key, old_worker_id) + self.redis.sadd(set_key, worker_id) + old_task = self.redis.hget("worker:tasks", old_worker_id) + if old_task and not self.redis.hget("worker:tasks", worker_id): + self.redis.hset("worker:tasks", worker_id, old_task) + if old_task: + self.redis.hdel("worker:tasks", old_worker_id) + + @staticmethod + def _normalize_device_type(raw_device_type: Any) -> str: + candidate = str(raw_device_type or '').strip().lower() + if candidate in {'emulator', 'physical'}: + return candidate + return '' + + def _normalize_task_device_type(self, raw_device_type: Any) -> str: + candidate = str(raw_device_type or '').strip().lower() + if candidate in {'emulator', 'physical', 'any'}: + return candidate + return 'emulator' + + def _normalize_routing_rules(self, rules: Optional[Dict[str, Dict[str, List[str]]]]) -> Dict[str, Dict[str, List[str]]]: + normalized = { + 'package_name': {}, + 'task_key': {}, + } + for rule_type in normalized: + for key, worker_ids in (rules or {}).get(rule_type, {}).items(): + normalized[rule_type][str(key).strip()] = self._normalize_worker_ids(worker_ids) + return normalized + + def _resolve_task_targets(self, task_key: str, package_name: str) -> Tuple[List[str], str]: + task_targets = self.task_routing_rules.get('task_key', {}).get(task_key, []) + if task_targets: + return task_targets, 'task_key' + + package_targets = self.task_routing_rules.get('package_name', {}).get(package_name, []) + if package_targets: + return package_targets, 'package_name' + + return [], 'public' + + def _normalize_country_codes(self, raw_country_code: Any) -> List[str]: + if isinstance(raw_country_code, (list, tuple)): + raw_items = raw_country_code + else: + raw_items = str(raw_country_code or '').split(',') + country_codes = ['US'] + seen = {'US'} + for item in raw_items: + candidate = str(item or '').strip().upper() + if not candidate or candidate in seen: + continue + seen.add(candidate) + country_codes.append(candidate) + return country_codes + + def _build_task_info_from_payload(self, payload: Dict[str, Any]) -> Dict[str, Any]: + normalized_payload = dict(payload or {}) + app_name = str(normalized_payload.get('app_name') or '').strip() + package_name = str(normalized_payload.get('package_name') or '').strip() + task_key = str(normalized_payload.get('task_key') or _make_task_key(app_name, package_name)).strip() + target_worker_ids, routing_mode = self._resolve_task_targets(task_key, package_name) + country_code = str(normalized_payload.get('country_code', '') or '').strip() + country_codes = normalized_payload.get('country_codes') or country_code + original_row = normalized_payload.get('original_row') or {} + original_app_magic_label = original_row.get('app_magic_label') if isinstance(original_row, dict) else '' + original_last_updated = original_row.get('last_updated') if isinstance(original_row, dict) else '' + return { + 'app_name': app_name, + 'package_name': package_name, + 'country_code': country_code, + 'country_codes': self._normalize_country_codes(country_codes), + 'app_magic_label': str(normalized_payload.get('app_magic_label') or original_app_magic_label or '').strip(), + 'last_updated': str(normalized_payload.get('last_updated') or original_last_updated or '').strip(), + 'collection_task_type': str(normalized_payload.get('collection_task_type') or 'new_app').strip() or 'new_app', + 'last_update_interval_days': int(normalized_payload.get('last_update_interval_days') or 0), + 'downloads': int(normalized_payload.get('downloads') or 0), + 'device_type': self._normalize_task_device_type(normalized_payload.get('device_type', '')), + 'available_sources': normalized_payload.get('available_sources') or ['google_play', 'local'], + 'original_row': original_row, + 'target_worker_ids': target_worker_ids, + 'routing_mode': routing_mode, + } + + @staticmethod + def _normalize_task_queue_label(raw_queue: Any, default: str = "default") -> str: + candidate = str(raw_queue or default).strip().lower() + if candidate in {"high", "default", "low"}: + return candidate + return default + + def _should_auto_reroute_to_physical(self, task_details: Dict[str, Any], error_info: ErrorInfo) -> bool: + current_device_type = self._normalize_task_device_type(task_details.get('device_type', '')) + return ( + current_device_type != 'physical' + and error_info.category == ErrorCategory.APP_ERROR + and error_info.code in {AppError.CRASH, AppError.ROOT_MODE_UNSUPPORTED} + ) + + def _build_physical_reroute_task_details(self, task_details: Dict[str, Any]) -> Dict[str, Any]: + rerouted_task_details = dict(task_details or {}) + rerouted_task_details['device_type'] = 'physical' + original_row = rerouted_task_details.get('original_row') + if isinstance(original_row, dict): + original_row_copy = dict(original_row) + original_row_copy['device_type'] = 'physical' + rerouted_task_details['original_row'] = original_row_copy + return rerouted_task_details + + def _is_task_queued(self, task_key: str) -> bool: + try: + for queue_name in self.TASK_QUEUE_NAMES: + if self.redis.lpos(queue_name, task_key) is not None: + return True + return False + except Exception: + for queue_name in self.TASK_QUEUE_NAMES: + if task_key in set(self.redis.lrange(queue_name, 0, -1)): + return True + return False + + def _remove_task_from_queue(self, task_key: str) -> int: + removed = 0 + try: + for queue_name in self.TASK_QUEUE_CLEANUP_NAMES: + removed += int(self.redis.lrem(queue_name, 0, task_key) or 0) + except Exception: + pass + return removed + + def _get_queued_task_keys(self) -> List[str]: + task_keys = [] + seen = set() + for queue_name in self.TASK_QUEUE_NAMES: + for task_key in reversed(self.redis.lrange(queue_name, 0, -1)): + if task_key in seen: + continue + seen.add(task_key) + task_keys.append(task_key) + return task_keys + + def get_pending_queue_length(self) -> int: + return sum(int(self.redis.llen(queue_name) or 0) for queue_name in self.TASK_QUEUE_NAMES) + + def _is_task_running(self, task_key: str) -> bool: + status_json = self.redis.hget("task:status", task_key) + if status_json: + try: + status = json.loads(status_json) + if status.get("status") == "running": + return True + except (json.JSONDecodeError, TypeError): + pass + return task_key in set(self.redis.hvals("worker:tasks")) + + def _mark_task_pending_for_queue(self, task_key: str) -> None: + status_json = self.redis.hget("task:status", task_key) + retry_count = 0 + if status_json: + try: + retry_count = int((json.loads(status_json) or {}).get("retry_count") or 0) + except (json.JSONDecodeError, TypeError, ValueError): + retry_count = 0 + self.redis.srem("task:completed", task_key) + self.redis.srem("task:failed", task_key) + self.redis.hset("task:status", task_key, json.dumps({ + "status": "pending", + "end_time": time.time(), + "retry_count": retry_count, + })) + + def _is_hot_new_app_task(self, task_details: Dict[str, Any]) -> bool: + if str(task_details.get('collection_task_type') or 'new_app').strip() != 'new_app': + return False + try: + downloads = int(task_details.get('downloads') or 0) + except (TypeError, ValueError): + downloads = 0 + return downloads >= self.HOT_APP_RETRY_DOWNLOADS_THRESHOLD + + def _mark_retry_avoid_worker(self, task_key: str, task_details: Dict[str, Any], worker_id: str) -> Dict[str, Any]: + next_details = dict(task_details or {}) + excluded = [ + str(item or '').strip() + for item in (next_details.get('excluded_worker_ids') or []) + if str(item or '').strip() + ] + if worker_id and worker_id not in excluded: + excluded.append(worker_id) + next_details['excluded_worker_ids'] = excluded[-8:] + next_details['last_failed_worker_id'] = worker_id + self.redis.hset("task:details", task_key, json.dumps(next_details)) + return next_details + + def _queue_retry_task(self, task_key: str, task_details: Dict[str, Any], *, prioritize: bool = False): + self._remove_task_from_queue(task_key) + queue_name = self._resolve_task_queue(task_details) + if prioritize: + self.redis.rpush(queue_name, task_key) + else: + self.redis.lpush(queue_name, task_key) + + def _resolve_task_queue(self, task_details: Dict[str, Any]) -> str: + task_queue = self._normalize_task_queue_label(task_details.get("task_queue")) + if task_queue == "low": + return "task:queue:low" + if task_queue == "high": + return "task:queue:high" + return "task:queue:default" + + def _handle_analytics_snapshot(self, package_name: str, summary: Dict[str, Any]) -> None: + payload = dict(summary.get("task_payload") or {}) + app_name = str(payload.get("app_name") or summary.get("app_name") or package_name).strip() or package_name + task_key = str(summary.get("latest_task_key") or payload.get("task_key") or _make_task_key(app_name, package_name)).strip() + if not task_key: + return + existing_task_details_json = self.redis.hget("task:details", task_key) + existing_task_details = json.loads(existing_task_details_json) if existing_task_details_json else {} + if existing_task_details.get("is_block_task") or task_key.endswith("_block"): + return + if existing_task_details.get("task_queue") == "low": + if summary.get("model_eligible"): + self._remove_task_from_queue(task_key) + self.redis.srem("task:failed", task_key) + self.redis.srem("task:completed", task_key) + self.redis.hset("task:status", task_key, json.dumps({ + "status": "pending", + "end_time": time.time(), + "retry_count": 0, + })) + self.redis.lpush("task:queue:low", task_key) + logger.info(f"[模型流量] 任务 {task_key} 继续留在低优队列 (model_flow={summary.get('model_flow_count', 0)})") + else: + self._remove_task_from_queue(task_key) + self.redis.sadd("task:completed", task_key) + self.redis.hset("task:status", task_key, json.dumps({ + "status": "completed", + "end_time": time.time(), + })) + logger.info(f"[模型流量] 任务 {task_key} 已完成，退出低优队列 (model_flow={summary.get('model_flow_count', 0)})") + return + collection_status = str(summary.get("collection_status") or "pending").strip() or "pending" + task_info = self._build_task_info_from_payload( + payload or { + "app_name": app_name, + "package_name": package_name, + "country_code": summary.get("country_code", ""), + "device_type": summary.get("device_type", ""), + } + ) + if existing_task_details.get("excluded_worker_ids"): + task_info["excluded_worker_ids"] = existing_task_details.get("excluded_worker_ids") or [] + task_info["last_failed_worker_id"] = existing_task_details.get("last_failed_worker_id", "") + existing_queue = self._normalize_task_queue_label(existing_task_details.get("task_queue")) + task_info["task_queue"] = existing_queue if existing_queue in {"high", "default"} else "default" + self.redis.hset("task:details", task_key, json.dumps(task_info)) + prev_status_json = self.redis.hget("task:status", task_key) + task_status = json.loads(prev_status_json) if prev_status_json else {} + retry_count = int(task_status.get("retry_count") or 0) + if collection_status == "qualified": + self._remove_task_from_queue(task_key) + self.redis.srem("task:failed", task_key) + self.redis.sadd("task:completed", task_key) + task_status.update({ + "status": "qualified", + "end_time": time.time(), + "retry_count": retry_count, + }) + self.redis.hset("task:status", task_key, json.dumps(task_status)) + return + if collection_status == "failed_terminal": + self._remove_task_from_queue(task_key) + self.redis.srem("task:completed", task_key) + self.redis.sadd("task:failed", task_key) + task_status.update({ + "status": "failed", + "end_time": time.time(), + "retry_count": retry_count, + }) + self.redis.hset("task:status", task_key, json.dumps(task_status)) + return + self.redis.srem("task:completed", task_key) + self.redis.srem("task:failed", task_key) + if task_status.get("status") != "running": + task_status.update({ + "status": "pending", + "end_time": task_status.get("end_time", time.time()), + "retry_count": retry_count, + }) + self.redis.hset("task:status", task_key, json.dumps(task_status)) + if not self._is_task_queued(task_key): + self._queue_retry_task(task_key, task_info, prioritize=self._is_hot_new_app_task(task_info)) + + def _task_to_payload(self, task_key: str, task_details: Dict[str, Any]) -> Dict[str, Any]: + task_details = self._enrich_task_with_apk(task_details) + payload = { + 'task_key': task_key, + 'app_name': task_details['app_name'], + 'package_name': task_details['package_name'], + 'country_code': task_details.get('country_code', ''), + 'country_codes': self._normalize_country_codes(task_details.get('country_codes') or task_details.get('country_code', '')), + 'app_magic_label': str(task_details.get('app_magic_label') or '').strip(), + 'last_updated': str(task_details.get('last_updated') or '').strip(), + 'collection_task_type': str(task_details.get('collection_task_type') or 'new_app').strip() or 'new_app', + 'last_update_interval_days': int(task_details.get('last_update_interval_days') or 0), + 'device_type': self._normalize_task_device_type(task_details.get('device_type', '')), + 'available_sources': task_details.get('available_sources') or self._default_available_sources(), + 'local_apk_dir': task_details.get('local_apk_dir', ''), + 'local_apk_files': task_details.get('local_apk_files', []), + 'target_worker_ids': task_details.get('target_worker_ids', []), + 'routing_mode': task_details.get('routing_mode', 'public'), + 'keep_app_installed': bool(task_details.get('keep_app_installed')), + 'blocked': bool(task_details.get('is_block_task', False)), + 'traffic_root': ANALYTICS_TRAFFIC_ROOT_BLOCK if task_details.get('is_block_task') else ANALYTICS_TRAFFIC_ROOT, + } + if payload.get('local_apk_dir'): + logger.info("[下发] 任务 %s 附带本地APK: dir=%s files=%s", + task_key, payload['local_apk_dir'], ', '.join(f.get('filename', '') for f in payload['local_apk_files'])) + elif payload.get('available_sources') and 'local' in payload['available_sources']: + logger.warning("[下发] 任务 %s 标记了local源但无本地APK路径 (available_sources=%s)", + task_key, payload['available_sources']) + return payload + + @staticmethod + def _apk_download_mode() -> str: + return str(APK_DOWNLOAD_MODE or "direct").strip().lower() + + @classmethod + def _uses_central_apk_cache(cls) -> bool: + return cls._apk_download_mode() in {"minio", "smb"} + + @classmethod + def _default_available_sources(cls) -> List[str]: + if cls._uses_central_apk_cache(): + return list(cls.APK_CACHE_SOURCES) + return list(cls.DIRECT_APK_SOURCES) + + def _get_worker_info(self, worker_id: str) -> Optional[Dict[str, Any]]: + worker_id = normalize_worker_id(worker_id) + worker_info_json = self.redis.hget("workers:info", worker_id) + if not worker_info_json: + return None + return json.loads(worker_info_json) + + def _get_worker_ip(self, worker_id: str) -> str: + worker_id = normalize_worker_id(worker_id) + worker_info = self._get_worker_info(worker_id) + if worker_info: + return worker_info.get('ip_address', '') + inventory = self.worker_inventory.get(worker_id) + if inventory: + return inventory["ssh_target"].split("@")[-1] + return worker_id.split('_', 1)[0] if '_' in worker_id else '' + + def _workers_meta_key(self) -> str: + return "workers:meta" + + def _worker_recovery_observe_key(self, worker_id: str) -> str: + return f"worker:recovery_observe:{worker_id}" + + def _worker_recovery_pending_key(self, worker_id: str) -> str: + return f"worker:recovery_pending:{worker_id}" + + def _default_worker_meta(self, worker_id: str) -> Dict[str, Any]: + allow_dispatch = (not self.only_managed_workers_can_dispatch) or self._is_managed_worker_id(worker_id) + return { + 'policy': 'enabled' if allow_dispatch else 'disabled', + 'auto_recover_used': False, + 'last_action': '', + 'last_action_status': '', + 'last_action_message': '', + 'last_action_at': 0.0, + } + + def _get_worker_meta(self, worker_id: str) -> Dict[str, Any]: + payload = self.redis.hget(self._workers_meta_key(), worker_id) + if not payload: + return self._default_worker_meta(worker_id) + try: + meta = json.loads(payload) + except (TypeError, ValueError): + return self._default_worker_meta(worker_id) + if not isinstance(meta, dict): + return self._default_worker_meta(worker_id) + normalized = self._default_worker_meta(worker_id) + normalized.update(meta) + normalized['auto_recover_used'] = bool(normalized.get('auto_recover_used')) + return normalized + + def _set_worker_meta(self, worker_id: str, meta: Dict[str, Any]) -> Dict[str, Any]: + normalized = self._default_worker_meta(worker_id) + normalized.update(meta) + self.redis.hset(self._workers_meta_key(), worker_id, json.dumps(normalized)) + return normalized + + def _update_worker_meta(self, worker_id: str, **updates: Any) -> Dict[str, Any]: + meta = self._get_worker_meta(worker_id) + meta.update(updates) + return self._set_worker_meta(worker_id, meta) + + def _mark_worker_action(self, worker_id: str, action: str, result: Dict[str, Any]) -> Dict[str, Any]: + return self._update_worker_meta( + worker_id, + last_action=action, + last_action_status='success' if result.get('ok') else 'failed', + last_action_message=result.get('message', ''), + last_action_at=time.time(), + ) + + def _set_worker_policy(self, worker_id: str, policy: str) -> Dict[str, Any]: + return self._update_worker_meta(worker_id, policy=policy) + + @staticmethod + def _is_worker_recovery_action(action: str) -> bool: + return action in {'recover_and_start', 'run_pipeline'} + + def _schedule_post_reboot_pipeline(self, worker_id: str, delay_seconds: int = AUTO_REBOOT_RECOVERY_DELAY_SECONDS) -> None: + thread = threading.Thread( + target=self._run_post_reboot_pipeline_after_delay, + args=(worker_id, delay_seconds), + name=f"post-reboot-pipeline-{worker_id}", + daemon=True, + ) + thread.start() + + def _run_post_reboot_pipeline_after_delay(self, worker_id: str, delay_seconds: int = AUTO_REBOOT_RECOVERY_DELAY_SECONDS) -> None: + pending_key = self._worker_recovery_pending_key(worker_id) + observe_key = self._worker_recovery_observe_key(worker_id) + try: + if delay_seconds > 0: + logger.info( + "[Worker自动恢复] worker=%s 主机重启后等待%s秒，再执行默认恢复 pipeline", + worker_id, + delay_seconds, + ) + time.sleep(delay_seconds) + + self.redis.delete(observe_key) + self.redis.delete(pending_key) + pipeline_result = self.run_dashboard_action([worker_id], 'run_pipeline')[0] + self._mark_worker_action(worker_id, 'run_pipeline', pipeline_result) + + if pipeline_result.get('ok'): + self.redis.setex(pending_key, WORKER_REINIT_TIMEOUT, str(time.time())) + self._update_worker_meta(worker_id, auto_recover_used=True) + logger.warning(f"[Worker自动恢复] Worker {worker_id} 主机重启后已执行默认恢复 pipeline") + self.notifier.send_weCom_alert( + worker_id, + f"【Worker主机重启后自动恢复】\nWorker: {worker_id}\n已在主机重启后执行 dashboard 默认恢复链路，等待 Worker 重新 init" + ) + return + + self.redis.delete(pending_key) + self._update_worker_meta(worker_id, auto_recover_used=False) + logger.error(f"[Worker自动恢复失败] Worker {worker_id} 主机重启后执行默认恢复 pipeline 失败") + self.notifier.send_weCom_alert( + worker_id, + f"【Worker主机重启后自动恢复失败】\nWorker: {worker_id}\n恢复结果: {pipeline_result.get('message', '')}" + ) + except Exception: + detail = traceback.format_exc().strip() + self.redis.delete(pending_key) + self._update_worker_meta(worker_id, auto_recover_used=False) + self._mark_worker_action( + worker_id, + 'run_pipeline', + { + 'ok': False, + 'message': detail, + 'stderr': detail, + }, + ) + logger.exception(f"[Worker自动恢复失败] Worker {worker_id} 主机重启后恢复线程异常") + self.notifier.send_weCom_alert( + worker_id, + f"【Worker主机重启后自动恢复异常】\nWorker: {worker_id}\n异常信息: {detail}" + ) + + def _get_control_display_status(self, worker_id: str, current_worker: Optional[Dict[str, Any]] = None) -> str: + meta = self._get_worker_meta(worker_id) + if meta.get('policy') == 'disabled': + return 'disabled' + if meta.get('policy') == 'drain': + return 'drain' + if current_worker: + if not current_worker.get('online'): + return 'offline' + if current_worker.get('state') == 'disabled': + return 'disabled' + if current_worker.get('status') == 'busy': + return 'running' + return 'idle' + return 'offline' + + def _run_worker_action( + self, + worker_id: str, + action: str, + options: Optional[Dict[str, Any]] = None, + progress_callback=None, + ) -> Dict[str, Any]: + logger.info( + "[DashboardAction] start worker=%s action=%s options=%s", + worker_id, + action, + json.dumps(options or {}, ensure_ascii=False, sort_keys=True), + ) + result = self.remote_controller.run_action(worker_id, action, options=options, progress_callback=progress_callback) + self._mark_worker_action(worker_id, action, result) + logger.info( + "[DashboardAction] finish worker=%s action=%s ok=%s failed_step=%s stage=%s message=%s", + worker_id, + action, + result.get('ok'), + result.get('failed_step', ''), + result.get('stage', ''), + str(result.get('message', '')).replace('\r', ' ').replace('\n', ' | '), + ) + return result + + def get_worker_control_rows( + self, + date_str: Optional[str] = None, + worker_id: Optional[str] = None, + ) -> List[Dict[str, Any]]: + filter_worker_id = worker_id or "" + monitor_rows = { + row["worker_id"]: row + for row in self.monitor.get_worker_stats(date_str=date_str, dispatcher=self) + } + current_workers = {worker["worker_id"]: worker for worker in self.get_dashboard_workers()} + worker_ids = list(self.managed_worker_ids) or list(current_workers.keys()) + rows: List[Dict[str, Any]] = [] + for current_worker_id in sorted(set(worker_ids) | set(monitor_rows) | set(current_workers)): + if current_worker_id not in self.worker_inventory: + continue + if filter_worker_id and current_worker_id != filter_worker_id: + continue + inventory = self.worker_inventory.get(current_worker_id, {}) + monitor_row = monitor_rows.get(current_worker_id, {}) + current_worker = current_workers.get(current_worker_id, {}) + meta = self._get_worker_meta(current_worker_id) + current_status = current_worker.get('state', 'offline') + current_status = { + 'idle_waiting_task': 'idle', + 'running_download': 'download', + 'running_collect': 'collect', + }.get(current_status, current_status) + rows.append( + { + 'worker_id': current_worker_id, + 'ip_address': current_worker.get('ip_address') or monitor_row.get('ip_address') or self._get_worker_ip(current_worker_id), + 'device_type': self._normalize_device_type(current_worker.get('device_type', '')), + 'status': monitor_row.get('status', current_status), + 'control_status': self._get_control_display_status(current_worker_id, current_worker), + 'policy': meta.get('policy', 'enabled'), + 'dispatch_enabled': self._worker_can_dispatch(current_worker_id), + 'online': bool(current_worker.get('online')), + 'task_count': monitor_row.get('task_count', 0), + 'avg_download_duration_seconds': monitor_row.get('avg_download_duration_seconds', 0.0), + 'avg_collect_duration_seconds': monitor_row.get('avg_collect_duration_seconds', 0.0), + 'running_seconds': monitor_row.get('running_seconds', 0.0), + 'idle_waiting_seconds': monitor_row.get('idle_waiting_seconds', 0.0), + 'failed_count': monitor_row.get('failed_count', 0), + 'last_task_at': monitor_row.get('last_task_at', '--'), + 'current_task': current_worker.get('current_task'), + 'last_action': meta.get('last_action', ''), + 'last_action_status': meta.get('last_action_status', ''), + 'last_action_message': meta.get('last_action_message', ''), + 'last_action_at': meta.get('last_action_at', 0.0), + 'tags': inventory.get('tags', []), + } + ) + return rows + + def _is_physical_worker(self, worker_id: str) -> bool: + worker_info = self._get_worker_info(worker_id) + if not worker_info: + return False + return self._normalize_device_type(worker_info.get("device_type", "")) == "physical" + + def run_dashboard_action( + self, + worker_ids: List[str], + action: str, + options: Optional[Dict[str, Any]] = None, + progress_callback=None, + ) -> List[Dict[str, Any]]: + results = [] + if action == 'force_reapk': + pkg = str((options or {}).get("package_name", "")).strip() + if not pkg: + return [{"worker_id": "", "action": "force_reapk", "ok": False, "message": "package_name is required"}] + if not self._uses_central_apk_cache(): + return [{ + "worker_id": "", + "action": "force_reapk", + "ok": False, + "message": "当前为 direct 模式，未启用中心 APK 缓存，无需 force_reapk", + }] + n = self.force_reapk(pkg) + return [{"worker_id": "", "action": "force_reapk", "ok": True, "message": f"已重置 {pkg} ({n} 个任务)，等待中心 APK 缓存刷新"}] + + for worker_id in worker_ids: + try: + if action == 'enable': + self.redis.delete(self._worker_recovery_observe_key(worker_id)) + self.redis.delete(self._worker_recovery_pending_key(worker_id)) + meta = self._update_worker_meta(worker_id, policy='enabled', auto_recover_used=False) + worker_info = self._get_worker_info(worker_id) + if self._get_worker_online_state(worker_info): + self.monitor.set_worker_state(worker_id, 'idle_waiting_task', reason='manual_enable') + meta['ok'] = True + meta['message'] = 'worker enabled' + result = self._mark_worker_action(worker_id, 'enable', meta) + result.update({'worker_id': worker_id, 'action': 'enable', 'ok': True, 'message': 'worker enabled'}) + elif action == 'disable': + self.redis.delete(self._worker_recovery_observe_key(worker_id)) + self.redis.delete(self._worker_recovery_pending_key(worker_id)) + self.monitor.mark_worker_disabled(worker_id, 'manual_disable') + meta = self._set_worker_policy(worker_id, 'disabled') + meta['ok'] = True + meta['message'] = 'worker disabled' + result = self._mark_worker_action(worker_id, 'disable', meta) + result.update({'worker_id': worker_id, 'action': 'disable', 'ok': True, 'message': 'worker disabled'}) + elif action == 'drain': + meta = self._set_worker_policy(worker_id, 'drain') + meta['ok'] = True + meta['message'] = 'worker drain enabled' + result = self._mark_worker_action(worker_id, 'drain', meta) + result.update({'worker_id': worker_id, 'action': 'drain', 'ok': True, 'message': 'worker drain enabled'}) + elif action == 'start': + self._set_worker_policy(worker_id, 'enabled') + pcap_result = self._run_worker_action(worker_id, 'pull_pcap_files', progress_callback=progress_callback) + if not pcap_result.get('ok'): + self._mark_worker_action(worker_id, 'start', pcap_result) + results.append(pcap_result) + continue + result = self._run_worker_action(worker_id, 'start_worker', progress_callback=progress_callback) + self._mark_worker_action(worker_id, 'start', result) + elif action == 'run_pipeline': + pipeline_options = dict(RUN_PIPELINE_DEFAULT_OPTIONS) + if options: + pipeline_options.update(options) + if self._is_physical_worker(worker_id): + pipeline_options.pop("recover_mumu", None) + pipeline_options.pop("restart_mumu", None) + if pipeline_options.get('start_worker'): + self._set_worker_policy(worker_id, 'enabled') + result = self._run_worker_action(worker_id, 'run_pipeline', options=pipeline_options, progress_callback=progress_callback) + elif action == 'stop_worker': + result = self._run_worker_action(worker_id, 'stop_worker', progress_callback=progress_callback) + elif action == 'recover_and_start': + self._set_worker_policy(worker_id, 'enabled') + self._run_worker_action(worker_id, 'stop_worker', progress_callback=progress_callback) + pcap_result = self._run_worker_action(worker_id, 'pull_pcap_files', progress_callback=progress_callback) + if not pcap_result.get('ok'): + self._mark_worker_action(worker_id, 'recover_and_start', pcap_result) + results.append(pcap_result) + continue + if self._is_physical_worker(worker_id): + result = self._run_worker_action(worker_id, 'start_worker', progress_callback=progress_callback) + self._mark_worker_action(worker_id, 'recover_and_start', result) + results.append(result) + continue + recover_result = self._run_worker_action(worker_id, 'recover_mumu', progress_callback=progress_callback) + if not recover_result.get('ok'): + self._mark_worker_action(worker_id, 'recover_and_start', recover_result) + results.append(recover_result) + continue + restart_result = self._run_worker_action(worker_id, 'restart_mumu', progress_callback=progress_callback) + if not restart_result.get('ok'): + self._mark_worker_action(worker_id, 'recover_and_start', restart_result) + results.append(restart_result) + continue + self._run_worker_action(worker_id, 'configure_mumu_network', progress_callback=progress_callback) + self._run_worker_action(worker_id, 'fix_adb_connection', progress_callback=progress_callback) + result = self._run_worker_action(worker_id, 'start_worker', progress_callback=progress_callback) + self._mark_worker_action(worker_id, 'recover_and_start', result) + elif action == 'clone': + result = self._run_worker_action(worker_id, 'clone', progress_callback=progress_callback) + elif action == 'pull': + result = self._run_worker_action(worker_id, 'pull', progress_callback=progress_callback) + elif action == 'pull_pcap_files': + result = self._run_worker_action(worker_id, 'pull_pcap_files', progress_callback=progress_callback) + elif action == 'setup': + result = self._run_worker_action(worker_id, 'setup', progress_callback=progress_callback) + elif action == 'restart_mumu': + if self._is_physical_worker(worker_id): + result = {"worker_id": worker_id, "action": "restart_mumu", "ok": False, "message": "真机设备，跳过重启MuMu"} + results.append(result) + continue + result = self._run_worker_action(worker_id, 'restart_mumu', progress_callback=progress_callback) + elif action == 'recover_mumu': + if self._is_physical_worker(worker_id): + result = {"worker_id": worker_id, "action": "recover_mumu", "ok": False, "message": "真机设备，跳过MuMu镜像恢复"} + results.append(result) + continue + result = self._run_worker_action(worker_id, 'recover_mumu', progress_callback=progress_callback) + elif action == 'reboot': + result = self._run_worker_action(worker_id, 'reboot', progress_callback=progress_callback) + elif action == 'status': + result = self._run_worker_action(worker_id, 'status', progress_callback=progress_callback) + elif action == 'execute_command': + result = self._run_worker_action(worker_id, 'execute_command', options=options, progress_callback=progress_callback) + elif action == 'fix_adb_connection': + result = self._run_worker_action(worker_id, 'fix_adb_connection', progress_callback=progress_callback) + elif action == 'configure_mumu_network': + result = self._run_worker_action(worker_id, 'configure_mumu_network', progress_callback=progress_callback) + elif action == 'recover_mumu_full': + if self._is_physical_worker(worker_id): + result = {"worker_id": worker_id, "action": "recover_mumu_full", "ok": False, "message": "真机设备，跳过MuMu恢复"} + results.append(result) + continue + result = self._run_worker_action(worker_id, 'recover_mumu_full', progress_callback=progress_callback) + else: + raise ValueError(f'unsupported action: {action}') + except Exception as e: + detail = traceback.format_exc().strip() + result = { + 'worker_id': worker_id, + 'action': action, + 'ok': False, + 'message': detail or str(e), + 'stderr': detail, + } + self._mark_worker_action(worker_id, action, result) + results.append(result) + return results + + def check_recovery_windows(self) -> List[str]: + reset_workers = [] + for worker_id in self.managed_worker_ids: + meta = self._get_worker_meta(worker_id) + if not meta.get('auto_recover_used'): + continue + if self.redis.exists(self._worker_recovery_pending_key(worker_id)): + continue + observe_key = self._worker_recovery_observe_key(worker_id) + if self.redis.exists(observe_key): + continue + if self._is_worker_recovery_action(meta.get('last_action', '')) and meta.get('last_action_status') == 'success': + self._update_worker_meta(worker_id, auto_recover_used=False) + reset_workers.append(worker_id) + return reset_workers + + def _is_managed_worker_id(self, worker_id: str) -> bool: + worker_id = normalize_worker_id(worker_id) + return not self.managed_worker_ids or worker_id in self.managed_worker_ids + + def _worker_can_dispatch(self, worker_id: str) -> bool: + worker_id = normalize_worker_id(worker_id) + if not worker_id: + return False + meta = self._get_worker_meta(worker_id) + if meta.get('policy') != 'enabled': + return False + if self.only_managed_workers_can_dispatch: + return self._is_managed_worker_id(worker_id) + return True + + def _worker_matches_task_device_type( + self, + task_details: Dict[str, Any], + worker_info: Optional[Dict[str, Any]], + ) -> bool: + required_device_type = self._normalize_task_device_type(task_details.get('device_type', '')) + worker_device_type = self._normalize_device_type((worker_info or {}).get('device_type', '')) + if required_device_type == 'any': + return worker_device_type in {'emulator', 'physical'} + return worker_device_type == required_device_type + + def _task_matches_worker( + self, + task_details: Dict[str, Any], + worker_id: str, + worker_info: Optional[Dict[str, Any]] = None, + ) -> bool: + excluded_worker_ids = { + normalize_worker_id(item) + for item in (task_details.get('excluded_worker_ids') or []) + if normalize_worker_id(item) + } + if worker_id in excluded_worker_ids: + return False + target_worker_ids = self._normalize_worker_ids(task_details.get('target_worker_ids', [])) + if target_worker_ids: + if worker_id not in target_worker_ids: + return False + elif not self._worker_can_dispatch(worker_id): + return False + return self._worker_matches_task_device_type(task_details, worker_info) + + def _restore_unmatched_tasks(self, skipped: List[Tuple[str, str]]): + for task_key, queue_name in skipped: + self.redis.rpush(queue_name, task_key) + + def _get_worker_online_state(self, worker_info: Optional[Dict[str, Any]], current_time: Optional[float] = None) -> bool: + if not worker_info: + return False + current_time = current_time or time.time() + last_update = worker_info.get('last_update', worker_info.get('register_time', 0)) + return (current_time - last_update) <= self.worker_online_timeout + + def _get_dashboard_worker_lookup(self) -> Dict[str, Dict[str, Any]]: + current_time = time.time() + worker_lookup: Dict[str, Dict[str, Any]] = {} + + for worker in self.get_registered_workers(): + worker_copy = dict(worker) + meta = self._get_worker_meta(worker_copy.get('worker_id', '')) + inventory = self.worker_inventory.get(worker_copy.get('worker_id', ''), {}) + worker_copy['online'] = self._get_worker_online_state(worker_copy, current_time) + worker_copy['is_managed'] = self._is_managed_worker_id(worker_copy.get('worker_id', '')) + worker_copy['dispatch_enabled'] = self._worker_can_dispatch(worker_copy.get('worker_id', '')) + worker_copy['policy'] = meta.get('policy', 'enabled') + worker_copy['last_action'] = meta.get('last_action', '') + worker_copy['last_action_status'] = meta.get('last_action_status', '') + worker_copy['last_action_message'] = meta.get('last_action_message', '') + worker_copy['last_action_at'] = meta.get('last_action_at', 0.0) + worker_copy['device_type'] = self._normalize_device_type(worker_copy.get('device_type', '')) + worker_copy['tags'] = inventory.get('tags', []) + worker_copy['ssh_target'] = inventory.get('ssh_target', '') + worker_lookup[worker_copy['worker_id']] = worker_copy + + return worker_lookup + + def _build_pending_task_entry(self, task_key: str, worker_lookup: Dict[str, Dict[str, Any]]) -> Dict[str, Any]: + task_details_json = self.redis.hget("task:details", task_key) + task_details = json.loads(task_details_json) if task_details_json else {} + target_worker_ids = task_details.get('target_worker_ids', []) + routing_state = 'public_pending' + + if target_worker_ids: + online_workers = [ + worker for worker in worker_lookup.values() + if worker.get('online') and worker.get('worker_id') in target_worker_ids + ] + idle_workers = [worker for worker in online_workers if worker.get('status') == 'idle'] + if idle_workers: + routing_state = 'routable_pending' + elif online_workers: + routing_state = 'waiting_target_busy' + else: + routing_state = 'waiting_target_offline' + + status_json = self.redis.hget("task:status", task_key) + status = json.loads(status_json) if status_json else {} + return { + 'task_key': task_key, + 'app_name': task_details.get('app_name', ''), + 'package_name': task_details.get('package_name', ''), + 'status': status.get('status', 'pending'), + 'retry_count': status.get('retry_count', 0), + 'device_type': self._normalize_task_device_type(task_details.get('device_type', '')), + 'target_worker_ids': target_worker_ids, + 'routing_mode': task_details.get('routing_mode', 'public'), + 'routing_state': routing_state, + } + + def load_tasks_from_app_summary(self) -> int: + """从应用全表加载任务到三个优先级队列。""" + for queue_name in self.TASK_QUEUE_CLEANUP_NAMES: + self.redis.delete(queue_name) + self.redis.delete("task:status") + self.redis.delete("task:details") + self.redis.delete("task:completed") + self.redis.delete("task:failed") + logger.info("[清空] 已清空任务队列和相关数据") + + entries = self.analytics.list_pending_collection_tasks() + high_count = 0 + default_count = 0 + active_base_keys: set = set() + + for entry in entries: + payload = entry.get("task_payload") or {} + package_name = str(payload.get("package_name") or entry.get("package_name") or "").strip() + app_name = str(payload.get("app_name") or entry.get("app_name") or "").strip() + if not package_name or not app_name: + continue + task_key = _make_task_key(app_name, package_name) + task_info = self._build_task_info_from_payload(payload) + unmanaged_targets = [ + target_worker_id for target_worker_id in task_info.get('target_worker_ids', []) + if not self._is_managed_worker_id(target_worker_id) + ] + if unmanaged_targets: + logger.warning( + f"[定向配置] 任务 {task_key} 配置了未纳入名单的目标Worker: {unmanaged_targets}" + ) + task_queue = self._normalize_task_queue_label(entry.get("task_queue"), "default") + task_info["task_queue"] = "high" if task_queue == "high" else "default" + self.redis.hset("task:details", task_key, json.dumps(task_info)) + active_base_keys.add(task_key) + + if task_queue == "high": + self.redis.lpush("task:queue:high", task_key) + high_count += 1 + else: + self.redis.lpush("task:queue:default", task_key) + default_count += 1 + + # 加载低优先级模型流量任务（使用 _model 后缀 task_key 避免与采集任务冲突） + model_entries = self.analytics.list_model_eligible_apps() + low_count = 0 + for entry in model_entries: + payload = entry.get("task_payload") or {} + package_name = str(payload.get("package_name") or entry.get("package_name") or "").strip() + app_name = str(payload.get("app_name") or entry.get("app_name") or "").strip() + if not package_name or not app_name: + continue + base_task_key = _make_task_key(app_name, package_name) + if base_task_key in active_base_keys: + continue + task_key = f"{base_task_key}_model" + task_info = self._build_task_info_from_payload(payload) + task_info["keep_app_installed"] = True + task_info["task_queue"] = "low" + task_info["package_name"] = package_name + task_info["app_name"] = app_name + self.redis.hset("task:details", task_key, json.dumps(task_info)) + self.redis.lpush("task:queue:low", task_key) + low_count += 1 + + logger.info( + f"[加载完成] 高优 {high_count} | 默认 {default_count} | 低优(模型) {low_count}" + ) + return high_count + default_count + low_count + + def refresh_tasks_from_app_summary(self) -> int: + """增量刷新：从应用全表加载新增的 pending/eligible 任务到三个优先级队列。""" + entries = self.analytics.list_pending_collection_tasks() + added_count = 0 + active_base_keys: set = set() + for entry in entries: + payload = entry.get("task_payload") or {} + package_name = str(payload.get("package_name") or entry.get("package_name") or "").strip() + app_name = str(payload.get("app_name") or entry.get("app_name") or "").strip() + if not package_name or not app_name: + continue + task_key = _make_task_key(app_name, package_name) + active_base_keys.add(task_key) + + if self._is_task_queued(task_key): + continue + if self._is_task_running(task_key): + continue + + task_info = self._build_task_info_from_payload(payload) + unmanaged_targets = [ + target_worker_id for target_worker_id in task_info.get('target_worker_ids', []) + if not self._is_managed_worker_id(target_worker_id) + ] + if unmanaged_targets: + logger.warning( + f"[定向配置] 任务 {task_key} 配置了未纳入名单的目标Worker: {unmanaged_targets}" + ) + task_queue = self._normalize_task_queue_label(entry.get("task_queue"), "default") + task_info["task_queue"] = "high" if task_queue == "high" else "default" + self.redis.hset("task:details", task_key, json.dumps(task_info)) + target_queue = "task:queue:high" if task_queue == "high" else "task:queue:default" + self._remove_task_from_queue(task_key) + self._mark_task_pending_for_queue(task_key) + self.redis.lpush(target_queue, task_key) + added_count += 1 + + model_entries = self.analytics.list_model_eligible_apps() + for entry in model_entries: + payload = entry.get("task_payload") or {} + package_name = str(payload.get("package_name") or entry.get("package_name") or "").strip() + app_name = str(payload.get("app_name") or entry.get("app_name") or "").strip() + if not package_name or not app_name: + continue + base_task_key = _make_task_key(app_name, package_name) + if base_task_key in active_base_keys: + continue + task_key = f"{base_task_key}_model" + if self._is_task_queued(task_key): + continue + if self._is_task_running(task_key): + continue + + task_info = self._build_task_info_from_payload(payload) + task_info["keep_app_installed"] = True + task_info["task_queue"] = "low" + task_info["package_name"] = package_name + task_info["app_name"] = app_name + self.redis.hset("task:details", task_key, json.dumps(task_info)) + self._remove_task_from_queue(task_key) + self._mark_task_pending_for_queue(task_key) + self.redis.lpush("task:queue:low", task_key) + added_count += 1 + + if added_count > 0: + logger.info(f"[增量刷新] 从应用全表新增任务 {added_count} 个") + return added_count + + def _clear_block_tasks(self) -> int: + cleared = 0 + all_details = dict(self.redis.hgetall("task:details") or {}) + for task_key in all_details: + if not task_key.endswith("_block"): + continue + if self._is_task_running(task_key): + continue + self._remove_task_from_queue(task_key) + self.redis.hdel("task:details", task_key) + self.redis.hdel("task:status", task_key) + self.redis.srem("task:completed", task_key) + self.redis.srem("task:failed", task_key) + cleared += 1 + return cleared + + def load_block_tasks(self, packages: List[Dict[str, Any]]) -> int: + cleared = self._clear_block_tasks() + if cleared: + logger.info(f"[Block测试] 已清除 {cleared} 个旧的 block 任务") + count = 0 + for entry in packages: + payload = entry.get("task_payload") or {} + package_name = str(payload.get("package_name") or entry.get("package_name") or "").strip() + app_name = str(payload.get("app_name") or entry.get("app_name") or "").strip() + if not package_name or not app_name: + continue + base_task_key = _make_task_key(app_name, package_name) + task_key = f"{base_task_key}_block" + if self._is_task_running(task_key): + continue + task_info = self._build_task_info_from_payload(payload) + task_info["is_block_task"] = True + task_info["task_queue"] = "high" + self.redis.hset("task:details", task_key, json.dumps(task_info)) + self._remove_task_from_queue(task_key) + self._mark_task_pending_for_queue(task_key) + self.redis.lpush("task:queue:high", task_key) + count += 1 + logger.info(f"[Block测试] 已加载 {count} 个 block 任务到高优队列") + return count + + def load_all_qualified_block_tasks(self) -> int: + entries = self.analytics.repo.list_qualified_apps() + return self.load_block_tasks(entries) + + def worker_init(self, worker_id, ip_address, mac_address, hostname, platform, device_type=''): + """Worker初始化：注册Worker并领取第一个任务 + + Args: + worker_id: Worker标识符（IP_MAC） + ip_address: IP地址 + mac_address: MAC地址 + hostname: 主机名 + platform: 平台 + + Returns: + dict: 包含第一个任务信息或None（无任务） + """ + raw_worker_id = str(worker_id or "").strip() + worker_id = normalize_worker_id(worker_id, ip_address) + self._migrate_worker_runtime_id(raw_worker_id, worker_id) + existing_worker = self.redis.hget("workers:info", worker_id) + previous_task = None + meta = self._get_worker_meta(worker_id) + + if existing_worker: + # 获取之前分配的任务 + previous_task_key = self.redis.hget("worker:tasks", worker_id) + if previous_task_key: + logger.warning(f"[重复初始化] Worker {worker_id} 再次初始化，将重新分配之前的任务: {previous_task_key}") + # 获取任务详情 + task_details_json = self.redis.hget("task:details", previous_task_key) + if task_details_json: + task_details = json.loads(task_details_json) + previous_task = self._task_to_payload(previous_task_key, task_details) + else: + logger.warning(f"[重复初始化] Worker {worker_id} 再次初始化，之前无分配任务") + + # 注册/更新Worker信息 + worker_info = { + 'worker_id': worker_id, + 'ip_address': ip_address, + 'mac_address': mac_address, + 'hostname': hostname, + 'platform': platform, + 'device_type': self._normalize_device_type(device_type), + 'register_time': time.time(), + 'last_update': time.time(), + 'status': 'idle', + 'is_managed': self._is_managed_worker_id(worker_id), + 'dispatch_enabled': self._worker_can_dispatch(worker_id), + } + + self.redis.hset("workers:info", worker_id, json.dumps(worker_info)) + self.redis.sadd("workers:idle", worker_id) + self.monitor.register_worker(worker_id, event_time=worker_info['last_update']) + if ( + meta.get('auto_recover_used') + and self._is_worker_recovery_action(meta.get('last_action', '')) + and meta.get('last_action_status') == 'success' + and self.redis.exists(self._worker_recovery_pending_key(worker_id)) + ): + self.redis.delete(self._worker_recovery_pending_key(worker_id)) + self.redis.setex(self._worker_recovery_observe_key(worker_id), WORKER_INIT_MONITOR_WINDOW, str(time.time())) + + # 如果有之前的任务，直接返回该任务 + if previous_task: + self._update_worker_status(worker_id, 'busy', previous_task['task_key']) + self.monitor.handle_task_assigned(worker_id, previous_task['task_key']) + logger.info(f"[初始化] Worker {worker_id} 重新领取之前的任务: {previous_task['app_name']}") + return previous_task + + if not worker_info['dispatch_enabled']: + logger.info(f"[初始化] Worker {worker_id} 已注册，但当前未启用分发") + return None + + # 否则领取新任务 + task = self._assign_task(worker_id) + + if task: + logger.info(f"[初始化] Worker {worker_id} 注册成功，领取任务: {task['app_name']}") + else: + logger.info(f"[初始化] Worker {worker_id} 注册成功，暂无任务") + + return task + + def worker_report(self, worker_id, previous_task_key, report_data: Optional[Dict[str, Any]] = None, **legacy_kwargs): + """上报前一个任务的完成状态，并获取下一个任务 + + Args: + worker_id: Worker标识符 + previous_task_key: 前一个任务的键（应用名称_包名） + report_data: 上报数据字典，结构如下: + { + "status": "success|failed|stop", # 必填 + "error": { # 失败时必填 + "category": "INFRA_ERROR|DOWNLOAD_ERROR|APP_ERROR|BUSINESS_ERROR", + "code": 1, + "reason": "错误原因描述", + "details": "详细错误信息（可选）", + "crashed_source": "google_play" # APP_CRASH时必填，闪退的下载源 + }, + "metrics": { # 可选 + "login_count": 0, + "register_count": 0, + "stuck_reason_code": 0, + "guiagent_message": "GuiAgent消息", + "scenario_triggered": true + }, + "statistics": { # 可选 + "round": 1, + "exit_code": 0, + "error_reason": "...", + "droidbot_steps": 100, + "guiagent_steps": 50, + "total_steps": 150, + "duration_seconds": 120.5, + "is_retry": false, + "download_source": "google" + } + } + + Returns: + dict: 包含下一个任务信息或None（无任务/stop状态/无效Worker） + """ + VALID_STATUS = ['success', 'failed', 'stop'] + raw_worker_id = str(worker_id or "").strip() + worker_id = normalize_worker_id(worker_id) + self._migrate_worker_runtime_id(raw_worker_id, worker_id) + if not isinstance(report_data, dict): + report_data = { + 'status': legacy_kwargs.get('status', 'failed'), + 'error': {}, + 'metrics': {}, + 'statistics': {}, + } + message = legacy_kwargs.get('message') + error_type = legacy_kwargs.get('error_type') + if message: + report_data['error']['reason'] = message + report_data['error']['details'] = message + if error_type: + category, _, code_text = str(error_type).partition('/') + report_data['error']['category'] = category or 'INFRA_ERROR' + try: + report_data['error']['code'] = int(code_text) + except (TypeError, ValueError): + report_data['error']['code'] = 0 + if not report_data['error']: + report_data.pop('error') + + status = report_data.get('status', 'failed') + error_data = report_data.get('error', {}) + metrics = report_data.get('metrics', {}) + statistics = report_data.get('statistics', {}) + download_errors = report_data.get('download_errors', {}) + + error_info = ErrorInfo.from_report_dict(error_data) + crashed_source = error_data.get('crashed_source') if error_data else None + + login_count = metrics.get('login_count', 0) if metrics else 0 + register_count = metrics.get('register_count', 0) if metrics else 0 + guiagent_message = metrics.get('guiagent_message', '') if metrics else '' + scenario_triggered = metrics.get('scenario_triggered', False) if metrics else False + + message = error_info.reason + + if download_errors: + source_msgs = [f"{k}: {v['message']}" for k, v in download_errors.items()] + message = f"下载失败: {', '.join(source_msgs)}" + elif guiagent_message: + if guiagent_message != error_info.reason: + message = guiagent_message + elif error_info.details and error_info.details != error_info.reason: + message = error_info.details + + if status == 'success': + if login_count > 0 or register_count > 0: + message = f"成功 (登录{login_count}次, 注册{register_count}次)" + + if not isinstance(status, str) or status not in VALID_STATUS: + logger.warning(f"Worker {worker_id} 上报了无效的status: {status} (类型: {type(status).__name__})") + logger.warning(f"预期status值: {VALID_STATUS}") + return None + + worker_exists = self.redis.hget("workers:info", worker_id) + retry_count = 0 + if previous_task_key: + previous_status_json = self.redis.hget("task:status", previous_task_key) + if previous_status_json: + previous_status = json.loads(previous_status_json) + retry_count = previous_status.get('retry_count', 0) + self.monitor.handle_worker_report( + worker_id=worker_id, + previous_task_key=previous_task_key, + report_data=report_data, + current_time=time.time(), + retry_count=retry_count, + resolved_message=message, + ) + if not worker_exists: + logger.warning(f"[忽略] Worker {worker_id} 已不在Worker列表中，返回错误码") + return { + 'task_key': 'ERROR', + 'app_name': 'ERROR', + 'package_name': '-1', + 'country_code': '', + 'device_type': '', + 'error_code': -1, + 'error_message': 'Worker已从列表中移除' + } + + if previous_task_key: + if status in ['stop', 'failed'] and message: + logger.warning(f"[告警] Worker {worker_id} [{error_info.category.name}/{error_info.code}]: {message}") + + if status in ['stop', 'failed']: + self._record_failed_task( + previous_task_key, + worker_id, + error_info, + message, + statistics, + crashed_source, + download_errors=download_errors, + ) + if self.redis.hget("worker:tasks", worker_id) == previous_task_key: + self.redis.hdel("worker:tasks", worker_id) + self._update_worker_status(worker_id, 'idle', None) + self.monitor.set_worker_state(worker_id, 'idle_waiting_task', reason='waiting_task') + elif status == 'success': + self._complete_task(previous_task_key, worker_id, status, message, statistics) + task_details_json = self.redis.hget("task:details", previous_task_key) + package_name = "" + if task_details_json: + package_name = json.loads(task_details_json).get("package_name", "") + if package_name: + self.analytics.enqueue_incremental( + package_name=package_name, + task_key=previous_task_key, + worker_id=worker_id, + report_time=time.time(), + trigger_source="worker_report", + waiting_for_artifacts=True, + ) + + # stop状态：清除Worker并不再分配任务 + if status == 'stop': + self.monitor.mark_worker_offline(worker_id, "worker_stopped") + self._remove_worker(worker_id) + logger.info(f"[停止] Worker {worker_id} 请求停止，已从Worker列表移除") + return None + + # 获取下一个任务 + next_task = self._assign_task(worker_id) + + if next_task: + logger.info(f"[上报] Worker {worker_id} 上报任务 {previous_task_key}={status}，领取新任务: {next_task['app_name']}") + else: + self.monitor.set_worker_state(worker_id, 'idle_waiting_task', reason='waiting_task') + logger.info(f"[上报] Worker {worker_id} 上报任务 {previous_task_key}={status}，暂无新任务") + + return next_task + + def get_analytics_overview( + self, + *, + incremental_batch_tag: str = "", + top_n: int = 3000, + ) -> Dict[str, Any]: + return self.analytics.get_overview( + incremental_batch_tag=incremental_batch_tag, + top_n=top_n, + ) + + def get_analytics_apps( + self, + *, + q: str = "", + latest_status: str = "", + artifact_status: str = "", + collection_status: str = "", + restriction_status: str = "", + retryability: str = "", + incremental_batch_tag: str = "", + top_n: int = 3000, + sort: str = "updated_at", + order: str = "desc", + page: int = 1, + page_size: int = 50, + ) -> Dict[str, Any]: + return self.analytics.list_apps( + q=q, + latest_status=latest_status, + artifact_status=artifact_status, + collection_status=collection_status, + restriction_status=restriction_status, + retryability=retryability, + incremental_batch_tag=incremental_batch_tag, + top_n=top_n, + sort=sort, + order=order, + page=page, + page_size=page_size, + ) + + def get_analytics_app_detail(self, package_name: str) -> Optional[Dict[str, Any]]: + return self.analytics.get_app_detail(package_name) + + def get_analytics_jobs( + self, + *, + job_type: str = "", + status: str = "", + package_name: str = "", + limit: int = 50, + ) -> List[Dict[str, Any]]: + return self.analytics.list_jobs( + job_type=job_type or None, + status=status or None, + package_name=package_name or None, + limit=limit, + ) + + def enqueue_analytics_backfill(self, packages: Optional[List[str]] = None) -> Dict[str, Any]: + return self.analytics.enqueue_backfill(packages=packages) + + def enqueue_analytics_rebuild(self, package_name: str) -> Dict[str, Any]: + return self.analytics.enqueue_manual_rebuild(package_name) + + def handle_analytics_worker_event(self, payload: Dict[str, Any]) -> int: + return self.analytics.handle_worker_event(payload) + + def close(self): + try: + self.analytics.close() + except Exception: + logger.exception("Failed to close analytics service") + + def worker_retry(self, worker_id, current_task_key): + """重试：告知该任务需重新计时 + + Args: + worker_id: Worker标识符 + current_task_key: 当前任务的键（应用名称_包名） + + Returns: + bool: 是否成功更新 + """ + raw_worker_id = str(worker_id or "").strip() + worker_id = normalize_worker_id(worker_id) + self._migrate_worker_runtime_id(raw_worker_id, worker_id) + # 检查Worker是否已被清除 + worker_exists = self.redis.hget("workers:info", worker_id) + if not worker_exists: + logger.warning(f"[忽略] Worker {worker_id} 已不在Worker列表中，忽略其retry") + return False + + if not current_task_key: + return False + + try: + # 获取任务状态 + status_json = self.redis.hget("task:status", current_task_key) + if not status_json: + return False + + status = json.loads(status_json) + if status['status'] != 'running' or status.get('worker_id') != worker_id: + return False + + current_time = time.time() + # 更新重试时间（保留原始start_time，用于超时计算时以重试时间为准） + status['last_retry'] = current_time + self.redis.hset("task:status", current_task_key, json.dumps(status)) + + # 更新Worker最后更新时间 + worker_info_json = self.redis.hget("workers:info", worker_id) + if worker_info_json: + worker_info = json.loads(worker_info_json) + worker_info['last_update'] = current_time + self.redis.hset("workers:info", worker_id, json.dumps(worker_info)) + + logger.info(f"[重试] Worker {worker_id} 任务 {current_task_key} 重新计时成功") + return True + + except Exception as e: + logger.error(f"重试更新失败: {e}") + return False + + def _assign_task(self, worker_id) -> Optional[Dict[str, Any]]: + """分配任务给Worker（内部方法），按优先级：高优 → 默认 → 低优。""" + lock_key = "lock:assign:queue" + lock = self.redis.lock(lock_key, timeout=5, blocking_timeout=1) + + try: + with lock: + current_task = self.redis.hget("worker:tasks", worker_id) + if current_task: + return None + + if not self._worker_can_dispatch(worker_id): + return None + + worker_info = self._get_worker_info(worker_id) + + task_key = None + task_details = None + skipped: List[Tuple[str, str]] = [] + + for queue_name in self.TASK_QUEUE_NAMES: + queue_length = self.redis.llen(queue_name) + for _ in range(queue_length): + candidate_task_key = self.redis.rpop(queue_name) + if not candidate_task_key: + break + + task_details_json = self.redis.hget("task:details", candidate_task_key) + if not task_details_json: + logger.warning(f"[任务详情缺失] 任务 {candidate_task_key} 缺少详情，已跳过") + continue + + candidate_task_details = json.loads(task_details_json) + if self._task_matches_worker(candidate_task_details, worker_id, worker_info=worker_info): + task_key = candidate_task_key + task_details = candidate_task_details + break + + skipped.append((candidate_task_key, queue_name)) + + if task_key: + break + + if skipped: + self._restore_unmatched_tasks(skipped) + + if not task_key or not task_details: + return None + + current_time = time.time() + if task_details.get('excluded_worker_ids'): + task_details = dict(task_details) + task_details.pop('excluded_worker_ids', None) + self.redis.hset("task:details", task_key, json.dumps(task_details)) + + # 获取之前的重试次数 + prev_status_json = self.redis.hget("task:status", task_key) + retry_count = 0 + if prev_status_json: + prev_status = json.loads(prev_status_json) + retry_count = prev_status.get('retry_count', 0) + + # 更新任务状态 + self.redis.hset("task:status", task_key, json.dumps({ + 'status': 'running', + 'worker_id': worker_id, + 'start_time': current_time, + 'last_retry': current_time, + 'retry_count': retry_count + })) + + # 更新Worker任务映射 + self.redis.hset("worker:tasks", worker_id, task_key) + + # 更新Worker状态 + self._update_worker_status(worker_id, 'busy', task_key) + self.monitor.handle_task_assigned(worker_id, task_key, event_time=current_time) + + # 发布任务分发事件 + self.redis.publish(f"task:assign:{worker_id}", json.dumps({ + 'task_key': task_key, + 'task_details': task_details + })) + + logger.info( + f"[分配] 任务 {task_key} -> Worker {worker_id} " + f"(routing={task_details.get('routing_mode', 'public')}, " + f"device_type={task_details.get('device_type', '')}, " + f"targets={task_details.get('target_worker_ids', [])})" + ) + + return self._task_to_payload(task_key, task_details) + + except redis.exceptions.LockNotOwnedError: + # 锁已过期或被其他进程获取 + return None + except redis.exceptions.LockError: + # 无法获取锁 + return None + except Exception as e: + logger.error(f"分配任务到 Worker {worker_id} 时出错: {e}") + return None + + def _update_worker_status(self, worker_id: str, status: str, current_task: Optional[str] = None): + """更新Worker状态的辅助方法""" + current_time = time.time() + + if status == 'busy': + self.redis.sadd("workers:busy", worker_id) + self.redis.srem("workers:idle", worker_id) + else: # idle + self.redis.sadd("workers:idle", worker_id) + self.redis.srem("workers:busy", worker_id) + + worker_info_json = self.redis.hget("workers:info", worker_id) + if worker_info_json: + worker_info = json.loads(worker_info_json) + worker_info['status'] = status + worker_info['current_task'] = current_task + worker_info['last_update'] = current_time + self.redis.hset("workers:info", worker_id, json.dumps(worker_info)) + + def _remove_worker(self, worker_id: str): + """从Worker列表中移除Worker（内部方法） + + Args: + worker_id: Worker标识符 + """ + # 从idle/busy集合中移除 + self.redis.srem("workers:idle", worker_id) + self.redis.srem("workers:busy", worker_id) + + # 删除Worker信息 + self.redis.hdel("workers:info", worker_id) + + # 清除Worker的任务映射 + self.redis.hdel("worker:tasks", worker_id) + + def _complete_task(self, task_key, worker_id, status, message='', statistics=None): + """任务完成（内部方法）""" + current_time = time.time() + + prev_status_json = self.redis.hget("task:status", task_key) + retry_count = 0 + if prev_status_json: + prev_status = json.loads(prev_status_json) + retry_count = prev_status.get('retry_count', 0) + + task_status = { + 'status': status, + 'worker_id': worker_id, + 'end_time': current_time, + 'retry_count': retry_count + } + + if statistics: + task_status['statistics'] = statistics + + self.redis.hset("task:status", task_key, json.dumps(task_status)) + + if self.redis.hget("worker:tasks", worker_id) == task_key: + self.redis.hdel("worker:tasks", worker_id) + self._update_worker_status(worker_id, 'idle', None) + self.monitor.set_worker_state(worker_id, 'idle_waiting_task', reason='waiting_task') + + if status == 'success': + self.redis.sadd("task:completed", task_key) + else: + new_retry_count = retry_count + 1 + task_details_json = self.redis.hget("task:details", task_key) + task_details = json.loads(task_details_json) if task_details_json else {} + retry_task_details = self._mark_retry_avoid_worker(task_key, task_details, worker_id) + self.redis.hset("task:status", task_key, json.dumps({ + 'status': 'pending_retry', + 'worker_id': worker_id, + 'end_time': current_time, + 'retry_count': new_retry_count + })) + self.redis.srem("task:failed", task_key) + self._queue_retry_task( + task_key, + retry_task_details, + prioritize=self._is_hot_new_app_task(retry_task_details), + ) + logger.info(f"[重试] 任务 {task_key} 失败，已重新入队 (第{new_retry_count}次失败)") + + self.redis.publish("task:completed", json.dumps({ + 'task_key': task_key, + 'worker_id': worker_id, + 'status': status, + 'retry_count': retry_count, + 'statistics': statistics + })) + + def get_statistics(self): + """获取统计信息 + + Returns: + dict: 包含各项统计数据 + - pending: 待分发任务数（含重试中任务） + - running: 运行中任务数 + - completed: 成功任务数 + - failed: 永久失败任务数 + """ + all_status = self.redis.hgetall("task:status") + running_count = 0 + pending_entries = [] + worker_lookup = self._get_dashboard_worker_lookup() + queue_task_keys = self._get_queued_task_keys() + + for task_key, status_json in all_status.items(): + status = json.loads(status_json) + if status.get('status') == 'running': + running_count += 1 + + for task_key in queue_task_keys: + pending_entries.append(self._build_pending_task_entry(task_key, worker_lookup)) + + return { + 'pending': len(queue_task_keys), + 'running': running_count, + 'completed': self.redis.scard("task:completed"), + 'failed': self.redis.scard("task:failed"), + 'pending_targeted': sum(1 for item in pending_entries if item['target_worker_ids']), + 'waiting_target_busy': sum(1 for item in pending_entries if item['routing_state'] == 'waiting_target_busy'), + 'waiting_target_offline': sum(1 for item in pending_entries if item['routing_state'] == 'waiting_target_offline'), + } + + def get_registered_workers(self): + """获取所有已注册的Worker（以 worker_inventory.json 为准）""" + workers_info = self.redis.hgetall("workers:info") + workers = [] + for worker_id, inventory_data in self.worker_inventory.items(): + info_json = workers_info.get(worker_id) + if info_json: + info = json.loads(info_json) + else: + info = { + "worker_id": worker_id, + "ip_address": inventory_data.get("ssh_target", ""), + "hostname": "", + "platform": "", + "device_type": "", + "status": "offline", + "register_time": 0, + "last_update": 0, + "current_task": None + } + info.setdefault('dispatch_enabled', self._worker_can_dispatch(worker_id)) + info.setdefault('is_managed', True) + workers.append(info) + return workers + + def get_idle_workers(self): + """获取空闲的Worker列表""" + idle_workers = self.redis.smembers("workers:idle") + workers_info = self.redis.hgetall("workers:info") + + idle_worker_list = [] + for worker_id in idle_workers: + if worker_id in self.worker_inventory and worker_id in workers_info: + info = json.loads(workers_info[worker_id]) + idle_worker_list.append(info) + + return idle_worker_list + + def get_busy_workers(self): + """获取忙碌的Worker列表""" + busy_workers = self.redis.smembers("workers:busy") + workers_info = self.redis.hgetall("workers:info") + + busy_worker_list = [] + for worker_id in busy_workers: + if worker_id in self.worker_inventory and worker_id in workers_info: + info = json.loads(workers_info[worker_id]) + busy_worker_list.append(info) + + return busy_worker_list + + def get_dashboard_workers(self) -> List[Dict[str, Any]]: + current_time = time.time() + registered_by_id: Dict[str, Dict[str, Any]] = {} + extra_workers: List[Dict[str, Any]] = [] + + for worker in self.get_registered_workers(): + worker_copy = dict(worker) + worker_copy['online'] = self._get_worker_online_state(worker_copy, current_time) + worker_copy['dispatch_enabled'] = self._worker_can_dispatch(worker_copy.get('worker_id', '')) + worker_copy['is_managed'] = self._is_managed_worker_id(worker_copy.get('worker_id', '')) + worker_copy['state'] = worker_copy.get('state', 'idle_waiting_task' if worker_copy['online'] else 'offline') + registered_by_id[worker_copy.get('worker_id', '')] = worker_copy + + if not self.managed_worker_ids: + return sorted( + registered_by_id.values(), + key=lambda item: (item.get('ip_address', ''), item.get('worker_id', '')), + ) + + managed_workers: List[Dict[str, Any]] = [] + for managed_worker_id in sorted(self.managed_worker_ids): + worker = registered_by_id.get(managed_worker_id, {}) + meta = self._get_worker_meta(managed_worker_id) + inventory = self.worker_inventory.get(managed_worker_id, {}) + managed_workers.append({ + 'worker_id': managed_worker_id, + 'ip_address': worker.get('ip_address', self._get_worker_ip(managed_worker_id)), + 'hostname': worker.get('hostname', ''), + 'platform': worker.get('platform', ''), + 'device_type': self._normalize_device_type(worker.get('device_type', '')), + 'current_task': worker.get('current_task'), + 'status': worker.get('status', 'offline') if worker else 'offline', + 'state': worker.get('state', 'idle_waiting_task') if worker else 'offline', + 'state_since': worker.get('state_since'), + 'pause_reason': worker.get('pause_reason'), + 'register_time': worker.get('register_time'), + 'last_update': worker.get('last_update'), + 'online': self._get_worker_online_state(worker, current_time) if worker else False, + 'is_managed': True, + 'dispatch_enabled': self._worker_can_dispatch(managed_worker_id), + 'policy': meta.get('policy', 'enabled'), + 'last_action': meta.get('last_action', ''), + 'last_action_status': meta.get('last_action_status', ''), + 'last_action_message': meta.get('last_action_message', ''), + 'last_action_at': meta.get('last_action_at', 0.0), + 'tags': inventory.get('tags', []), + 'ssh_target': inventory.get('ssh_target', ''), + }) + + return managed_workers + + def get_dashboard_tasks( + self, + pending_offset: int = 0, + pending_limit: int = 50, + recent_limit: int = 20, + ) -> Dict[str, Any]: + worker_lookup = self._get_dashboard_worker_lookup() + running_tasks = [] + failed_tasks = [] + pending_tasks = [] + + all_status = self.redis.hgetall("task:status") + for task_key, status_json in all_status.items(): + status = json.loads(status_json) + task_details_json = self.redis.hget("task:details", task_key) + task_details = json.loads(task_details_json) if task_details_json else {} + task_entry = { + 'task_key': task_key, + 'app_name': task_details.get('app_name', ''), + 'package_name': task_details.get('package_name', ''), + 'worker_id': status.get('worker_id'), + 'status': status.get('status', 'unknown'), + 'retry_count': status.get('retry_count', 0), + 'device_type': self._normalize_task_device_type(task_details.get('device_type', '')), + 'target_worker_ids': task_details.get('target_worker_ids', []), + 'routing_mode': task_details.get('routing_mode', 'public'), + 'last_fail_type': status.get('last_fail_type'), + 'last_fail_message': status.get('last_fail_message'), + 'start_time': status.get('start_time'), + 'end_time': status.get('end_time'), + 'last_fail_time': status.get('last_fail_time'), + } + if status.get('status') == 'running': + running_tasks.append(task_entry) + elif status.get('status') == 'failed': + failed_tasks.append(task_entry) + + queue_task_keys = self._get_queued_task_keys() + pending_total = len(queue_task_keys) + safe_offset = max(pending_offset, 0) + safe_limit = max(pending_limit, 1) + pending_window = queue_task_keys[safe_offset:safe_offset + safe_limit] + for task_key in pending_window: + pending_tasks.append(self._build_pending_task_entry(task_key, worker_lookup)) + + running_tasks.sort(key=lambda item: item.get('start_time') or 0, reverse=True) + failed_tasks.sort(key=lambda item: item.get('last_fail_time') or 0, reverse=True) + + return { + 'running': running_tasks[:recent_limit], + 'pending': pending_tasks, + 'failed': failed_tasks[:recent_limit], + 'pending_total': pending_total, + 'pending_offset': safe_offset, + 'pending_limit': safe_limit, + 'pending_window_end': min(safe_offset + safe_limit, pending_total), + } + + def get_dashboard_summary(self) -> Dict[str, Any]: + stats = self.get_statistics() + workers = self.get_dashboard_workers() + total_tasks = stats['pending'] + stats['running'] + stats['completed'] + stats['failed'] + managed_workers = [worker for worker in workers if worker.get('is_managed')] + online_workers = [worker for worker in managed_workers if worker.get('online')] + idle_workers = [worker for worker in online_workers if worker.get('status') == 'idle'] + busy_workers = [worker for worker in online_workers if worker.get('status') == 'busy'] + + return { + **stats, + 'managed_workers_total': len(managed_workers), + 'managed_workers_online': len(online_workers), + 'managed_workers_idle': len(idle_workers), + 'managed_workers_busy': len(busy_workers), + 'completion_rate': round((stats['completed'] / total_tasks) * 100, 2) if total_tasks else 0.0, + 'generated_at': datetime.now().strftime('%Y-%m-%d %H:%M:%S'), + } + + def get_dashboard_snapshot( + self, + pending_offset: int = 0, + pending_limit: int = 50, + recent_limit: int = 20, + ) -> Dict[str, Any]: + return { + 'summary': self.get_dashboard_summary(), + 'workers': self.get_dashboard_workers(), + 'tasks': self.get_dashboard_tasks( + pending_offset=pending_offset, + pending_limit=pending_limit, + recent_limit=recent_limit, + ), + } + + def get_monitoring_overview(self, date_str: Optional[str] = None, worker_id: Optional[str] = None) -> Dict[str, Any]: + return self.monitor.get_overview(date_str=date_str, dispatcher=self, worker_id=worker_id) + + def get_monitoring_timeline( + self, + date_str: Optional[str] = None, + bucket_minutes: int = 15, + worker_id: Optional[str] = None, + ) -> Dict[str, Any]: + return self.monitor.get_timeline( + date_str=date_str, + dispatcher=self, + bucket_minutes=bucket_minutes, + worker_id=worker_id, + ) + + def get_monitoring_distributions(self, date_str: Optional[str] = None, worker_id: Optional[str] = None) -> Dict[str, Any]: + return self.monitor.get_distributions(date_str=date_str, worker_id=worker_id, dispatcher=self) + + def get_monitoring_workers(self, date_str: Optional[str] = None) -> List[Dict[str, Any]]: + return self.monitor.get_worker_stats(date_str=date_str, dispatcher=self) + + def get_monitoring_tasks( + self, + date_str: Optional[str] = None, + worker_id: Optional[str] = None, + limit: int = 200, + ) -> List[Dict[str, Any]]: + return self.monitor.get_task_rows(date_str=date_str, worker_id=worker_id, limit=limit) + + def reset_monitoring_dashboard(self) -> Dict[str, Any]: + return self.monitor.reset_monitoring_data(worker_snapshots=self.get_dashboard_workers()) + + def sync_worker_registry(self): + """启动时同步：以 worker_inventory.json 为准清理 Redis 中的残留 Worker 数据""" + valid_ids = set(self.worker_inventory.keys()) + + # 清理 workers:info 中不在 inventory 的历史 worker + all_info_ids = set(self.redis.hgetall("workers:info").keys()) + stale_info_ids = all_info_ids - valid_ids + for stale_id in stale_info_ids: + self.redis.hdel("workers:info", stale_id) + + # 清理 idle/busy 集合中不在 inventory 的条目 + for set_key in ("workers:idle", "workers:busy"): + for member in self.redis.smembers(set_key): + if member not in valid_ids: + self.redis.srem(set_key, member) + + # 清理 worker:tasks 映射中不在 inventory 的条目 + all_task_mappings = self.redis.hgetall("worker:tasks") + for wid in list(all_task_mappings.keys()): + if wid not in valid_ids: + self.redis.hdel("worker:tasks", wid) + + cleaned = len(stale_info_ids) + if cleaned: + logger.info(f"[启动同步] 已清理 {cleaned} 个不在 inventory 中的历史 Worker: {sorted(stale_info_ids)}") + else: + logger.info("[启动同步] Worker 注册表与 inventory 一致，无需清理") + + def cleanup_stale_workers(self, timeout=7200): + """清理超时的Worker + + Args: + timeout: Worker超时时间(秒)，默认7200秒(120分钟) + """ + current_time = time.time() + workers_info = self.redis.hgetall("workers:info") + cleaned_count = 0 + + for worker_id, info_json in workers_info.items(): + info = json.loads(info_json) + last_update = info.get('last_update', info.get('register_time', 0)) + + if current_time - last_update > timeout: + # Worker超时，清理 + self.monitor.mark_worker_offline(worker_id, "heartbeat_timeout") + self.redis.srem("workers:idle", worker_id) + self.redis.srem("workers:busy", worker_id) + self.redis.hdel("workers:info", worker_id) + + # 获取Worker的任务并标记为失败 + current_task = self.redis.hget("worker:tasks", worker_id) + if current_task: + # 记录任务失败并更新运行态/分析态 + worker_timeout_error = ErrorInfo.infra(InfraError.WORKER_TIMEOUT, "Worker heartbeat timeout") + self._record_failed_task(current_task, worker_id, worker_timeout_error, "Worker heartbeat timeout") + self.redis.hdel("worker:tasks", worker_id) + + cleaned_count += 1 + logger.info(f"[清理] Worker {worker_id} 超时，已清理") + + return cleaned_count + + def _handle_infra_error_recovery_alert(self, task_key: str, worker_id: str, error_type_str: str, message: str) -> bool: + alert_key = f"alert:{worker_id}:infra_errors" + pending_key = self._worker_recovery_pending_key(worker_id) + observe_key = self._worker_recovery_observe_key(worker_id) + current_time = time.time() + + pipe = self.redis.pipeline() + pipe.zremrangebyscore(alert_key, 0, current_time - ALERT_WINDOW_SECONDS) + pipe.zadd(alert_key, {str(current_time): current_time}) + pipe.zcard(alert_key) + pipe.expire(alert_key, ALERT_WINDOW_SECONDS + 60) + _, _, count, _ = pipe.execute() + + if count < ALERT_THRESHOLD: + return False + + self.redis.delete(alert_key) + self.redis.delete(observe_key) + self.redis.delete(pending_key) + logger.warning( + f"[Worker基础设施错误高频告警] Worker {worker_id} 高频基础设施错误，" + f"{count}次/{ALERT_WINDOW_SECONDS//60}分钟，自动重启已关闭，请手动处理" + ) + self.notifier.send_weCom_alert( + worker_id, + f"【Worker基础设施错误高频告警】\nWorker: {worker_id}\n任务: {task_key}\n" + f"频率: {count}次/{ALERT_WINDOW_SECONDS//60}分钟\n最新错误: {error_type_str}\n消息: {message}\n" + f"自动重启已关闭，请手动处理" + ) + return True + + def _dispatch_alert(self, task_key, worker_id, error_info: ErrorInfo, message, retry_count, download_errors=None): + """根据错误类型分发告警逻辑 + + Args: + task_key: 任务键 + worker_id: Worker ID + error_info: 错误信息对象 + message: 错误消息 + retry_count: 当前重试次数 + """ + error_type_str = f"{error_info.category.name}/{error_info.code}" + + infra_recovery_triggered = False + if error_info.category == ErrorCategory.INFRA_ERROR: + infra_recovery_triggered = self._handle_infra_error_recovery_alert( + task_key, + worker_id, + error_type_str, + message, + ) + + if error_info.category == ErrorCategory.INFRA_ERROR and error_info.code == InfraError.WORKER_TIMEOUT: + logger.error(f"[告警推送] Worker超时: {worker_id}") + if not infra_recovery_triggered: + self.notifier.send_weCom_alert( + worker_id, + f"【Worker超时】\nWorker: {worker_id}\n任务: {task_key}\n消息: {message}" + ) + return + + if error_info.is_no_retry(download_errors=download_errors): + logger.info(f"[直接失败] 任务 {task_key} 不可重试错误: {error_type_str}, Worker: {worker_id}") + return + + alert_key = f"alert:{worker_id}:all_errors" + observe_key = self._worker_recovery_observe_key(worker_id) + pending_key = self._worker_recovery_pending_key(worker_id) + current_time = time.time() + + pipe = self.redis.pipeline() + pipe.zremrangebyscore(alert_key, 0, current_time - ALERT_WINDOW_SECONDS) + pipe.zadd(alert_key, {str(current_time): current_time}) + pipe.zcard(alert_key) + pipe.expire(alert_key, ALERT_WINDOW_SECONDS + 60) + _, _, count, _ = pipe.execute() + + if count < ALERT_THRESHOLD or infra_recovery_triggered: + return + + meta = self._get_worker_meta(worker_id) + recovery_in_effect = bool(meta.get('auto_recover_used')) and ( + self.redis.exists(pending_key) or self.redis.exists(observe_key) + ) + self.redis.delete(alert_key) + + if recovery_in_effect: + self.redis.delete(observe_key) + self.redis.delete(pending_key) + self._update_worker_meta(worker_id, auto_recover_used=False) + + logger.warning( + f"[Worker高频错误告警] Worker {worker_id} 自动恢复后再次高频错误，" + f"自动重启已关闭，请手动处理" + ) + self.notifier.send_weCom_alert( + worker_id, + f"【Worker高频错误告警】\nWorker: {worker_id}\n频率: {count}次/{ALERT_WINDOW_SECONDS//60}分钟\n最新错误: {error_type_str}\n消息: {message}\n自动恢复后再次触发连续错误，自动重启已关闭，请手动处理" + ) + return + + self.redis.delete(observe_key) + self.redis.delete(pending_key) + pipeline_result = self.run_dashboard_action([worker_id], 'run_pipeline')[0] + self._mark_worker_action(worker_id, 'run_pipeline', pipeline_result) + + if pipeline_result.get('ok'): + self.redis.setex(pending_key, WORKER_REINIT_TIMEOUT, str(time.time())) + self._update_worker_meta(worker_id, auto_recover_used=True) + logger.warning(f"[Worker自动恢复] Worker {worker_id} 高频错误，已执行默认恢复 pipeline") + self.notifier.send_weCom_alert( + worker_id, + f"【Worker高频错误自动恢复】\nWorker: {worker_id}\n频率: {count}次/{ALERT_WINDOW_SECONDS//60}分钟\n最新错误: {error_type_str}\n消息: {message}\n已执行 dashboard 默认恢复链路" + ) + return + + self.redis.delete(pending_key) + self._update_worker_meta(worker_id, auto_recover_used=False) + logger.error(f"[Worker自动恢复失败] Worker {worker_id} 高频错误后执行默认恢复 pipeline 失败") + self.notifier.send_weCom_alert( + worker_id, + f"【Worker高频错误自动恢复失败】\nWorker: {worker_id}\n频率: {count}次/{ALERT_WINDOW_SECONDS//60}分钟\n最新错误: {error_type_str}\n消息: {message}\n恢复结果: {pipeline_result.get('message', '')}" + ) + + def _record_failed_task(self, task_key, worker_id, error_info: ErrorInfo, message='', statistics=None, crashed_source=None, download_errors=None): + """记录失败任务，分流重试和最终失败 + + - APP_ERROR.CRASH / APP_ERROR.ROOT_MODE_UNSUPPORTED: 若当前任务不是 physical，则自动改派到真机并重新入队 + - APP_ERROR / BUSINESS_ERROR: 默认直接标记为永久失败 + - 满足终态下载组合条件的 DOWNLOAD_ERROR（Google Play 终态类错误 + local 任意失败）→ 直接标记为永久失败 + - 可重试错误 → 等待应用画像决定是否重新入队 + + Args: + task_key: 任务键 + worker_id: Worker标识符 + error_info: 错误信息对象 + message: 附加消息/告警信息 + statistics: 统计数据字典 + crashed_source: 闪退的下载源（已废弃，保留参数兼容） + """ + if error_info.category == ErrorCategory.SUCCESS and error_info.code == 0: + logger.warning(f"[异常上报] Worker {worker_id} 上报失败但错误码为 SUCCESS/0，按成功处理: {task_key}") + self._complete_task(task_key, worker_id, 'success', message or 'worker reported failure with SUCCESS/0, treated as success') + return + + is_no_retry_error = error_info.is_no_retry(download_errors=download_errors) + is_app_crash = (error_info.category == ErrorCategory.APP_ERROR and error_info.code == AppError.CRASH) + is_terminal_local_fallback_failed = ( + error_info.category == ErrorCategory.DOWNLOAD_ERROR + and error_info.code in { + DownloadError.REGION_RESTRICTED, + DownloadError.APP_NOT_FOUND, + DownloadError.INCOMPATIBLE, + DownloadError.COUNTRY_LOCKED, + } + and is_no_retry_error + ) + + prev_status_json = self.redis.hget("task:status", task_key) + retry_count = 0 + if prev_status_json: + prev_status = json.loads(prev_status_json) + retry_count = prev_status.get('retry_count', 0) + + new_retry_count = retry_count + 1 + + task_details_json = self.redis.hget("task:details", task_key) + task_details = {} + app_name = "" + package_name = "" + if task_details_json: + task_details = json.loads(task_details_json) + app_name = task_details.get('app_name', '') + package_name = task_details.get('package_name', '') + app_magic_label = str(task_details.get('app_magic_label') or '').strip() + is_block_task = bool(task_details.get('is_block_task', False)) + + error_type_str = f"{error_info.category.name}/{error_info.code}" + task_status = { + 'status': 'failed', + 'worker_id': worker_id, + 'retry_count': new_retry_count, + 'last_fail_time': time.time(), + 'last_fail_type': error_type_str, + 'last_fail_message': message + } + + if statistics: + task_status['statistics'] = statistics + + is_download_error = ( + error_info.category == ErrorCategory.DOWNLOAD_ERROR + or (error_info.category == ErrorCategory.INFRA_ERROR and error_info.code == InfraError.DOWNLOAD_FAILED) + ) + + is_local_install_failed = ( + is_download_error + and message + and ("local: install failed" in message.lower() + or "local: Install" in message) + ) + if is_local_install_failed and self._uses_central_apk_cache(): + prev_status = json.loads(prev_status_json) if prev_status_json else {} + if prev_status.get("force_reapk_done"): + self._log_apk_install_failure(task_key, package_name, app_name, message) + logger.warning("[自动重下] %s 已重下过一次仍失败，不再重试: %s", task_key, message) + self.redis.hset("task:status", task_key, json.dumps(task_status)) + self.monitor.set_failure_bucket(task_key, "failed") + self.redis.sadd("task:failed", task_key) + return + logger.info("[自动重下] %s 本地APK安装失败，触发force_reapk后重试", task_key) + self.force_reapk(package_name) + self.redis.hset("task:status", task_key, json.dumps({ + "status": "pending_retry", + "worker_id": worker_id, + "retry_count": new_retry_count, + "last_fail_time": time.time(), + "last_fail_type": error_type_str, + "last_fail_message": message, + "force_reapk_done": True, + })) + self.monitor.set_failure_bucket(task_key, "retry") + self.redis.srem("task:failed", task_key) + return + + if is_block_task: + self.redis.hset("task:status", task_key, json.dumps(task_status)) + if is_download_error: + last_updated = str(task_details.get("last_updated") or "").strip() + registry = self._get_apk_registry() if self._uses_central_apk_cache() else None + if registry and registry.is_fresh_enough(package_name, last_updated): + self.monitor.set_failure_bucket(task_key, "retry") + self.redis.srem("task:failed", task_key) + retry_task_details = self._mark_retry_avoid_worker(task_key, task_details, worker_id) + retry_status = dict(task_status) + retry_status["status"] = "pending" + self.redis.hset("task:status", task_key, json.dumps(retry_status)) + self._queue_retry_task(task_key, retry_task_details, prioritize=True) + logger.info(f"[Block测试] 任务 {task_key} 下载失败({error_type_str})，本地已有APK，立即重试") + else: + self._mark_download_error_awaiting_apk(package_name, + app_name=app_name, + last_updated=last_updated) + self.monitor.set_failure_bucket(task_key, "retry") + self.redis.srem("task:failed", task_key) + retry_task_details = self._mark_retry_avoid_worker(task_key, task_details, worker_id) + retry_status = dict(task_status) + retry_status["status"] = "pending_retry" + self.redis.hset("task:status", task_key, json.dumps(retry_status)) + if self._uses_central_apk_cache(): + logger.info(f"[Block测试] 任务 {task_key} 下载失败({error_type_str})，等待中心 APK 缓存刷新后重试一次") + else: + logger.info(f"[Block测试] 任务 {task_key} 下载失败({error_type_str})，direct 模式不等待中心 APK") + else: + self.monitor.set_failure_bucket(task_key, "failed") + self.redis.sadd("task:failed", task_key) + logger.info(f"[Block测试] 任务 {task_key} 因非下载错误({error_type_str})直接标记为永久失败") + self._dispatch_alert(task_key, worker_id, error_info, message, retry_count, download_errors=download_errors) + return + + auto_physical_reroute = self._should_auto_reroute_to_physical(task_details, error_info) + + self.redis.hset("task:status", task_key, json.dumps(task_status)) + + related_magic_label_collected = ( + bool(app_magic_label) + and bool(package_name) + and self.analytics.repo.has_successful_related_magic_label(app_magic_label, package_name) + ) + if related_magic_label_collected: + qualified_status = dict(task_status) + qualified_status['status'] = 'qualified' + qualified_status['end_time'] = time.time() + qualified_status['collection_status_reason'] = '关联应用已采集' + self.redis.hset("task:status", task_key, json.dumps(qualified_status)) + self.redis.srem("task:failed", task_key) + self.redis.sadd("task:completed", task_key) + self.monitor.set_failure_bucket(task_key, "light_restricted") + logger.warning( + f"[轻度受限] 任务 {task_key} 测试失败，但 app_magic_label={app_magic_label} 的关联应用已采集成功，" + f"当前任务按轻度受限记录" + ) + return + + if auto_physical_reroute: + self.monitor.set_failure_bucket(task_key, "retry") + rerouted_task_details = self._build_physical_reroute_task_details(task_details) + rerouted_task_details['auto_physical_fallback_applied'] = True + rerouted_task_details['auto_physical_fallback_from_worker'] = worker_id + rerouted_task_details['auto_physical_fallback_error_type'] = error_type_str + self.redis.hset("task:details", task_key, json.dumps(rerouted_task_details)) + + reroute_status = dict(task_status) + reroute_status['status'] = 'pending' + reroute_status['retry_count'] = retry_count + self.redis.hset("task:status", task_key, json.dumps(reroute_status)) + self.redis.srem("task:completed", task_key) + self.redis.srem("task:failed", task_key) + self._queue_retry_task(task_key, rerouted_task_details, prioritize=self._is_hot_new_app_task(rerouted_task_details)) + logger.warning( + f"[真机改派] 任务 {task_key} 在 Worker {worker_id} 上报 {error_type_str}，" + f"已改为仅 physical 执行并重新放回待分发" + ) + elif is_app_crash: + self.monitor.set_failure_bucket(task_key, "failed") + self.redis.sadd("task:failed", task_key) + logger.error(f"[永久失败] 任务 {task_key} 应用闪退，直接标记为永久失败") + elif is_no_retry_error and not is_terminal_local_fallback_failed: + # 非下载类的不可重试错误（APP_ERROR / BUSINESS_ERROR）→ 直接永久失败 + self.monitor.set_failure_bucket(task_key, "failed") + self.redis.sadd("task:failed", task_key) + logger.error(f"[永久失败] 任务 {task_key} 因不可重试错误({error_type_str})，直接标记为永久失败") + else: + # 可重试的下载错误 或 终态下载错误（REGION_RESTRICTED等 + local失败） + # 中心 APK 缓存模式下等待 US/MinIO 下载；direct 模式下由 Worker 自行下载。 + if self._uses_central_apk_cache(): + if error_info.category in (ErrorCategory.DOWNLOAD_ERROR,): + self._mark_download_error_awaiting_apk(package_name) + elif error_info.category == ErrorCategory.INFRA_ERROR and error_info.code == InfraError.DOWNLOAD_FAILED: + self._mark_download_error_awaiting_apk(package_name) + + if is_terminal_local_fallback_failed and self._uses_central_apk_cache(): + logger.warning( + f"[等待中心APK] 任务 {task_key} Google Play 终态错误({error_type_str})且本地兜底失败，" + f"标记为等待中心 APK 下载服务（可能通过不同地区或 APKPure 成功下载）" + ) + + self.monitor.set_failure_bucket(task_key, "retry") + self.redis.srem("task:failed", task_key) + retry_task_details = self._mark_retry_avoid_worker(task_key, task_details, worker_id) + retry_status = dict(task_status) + retry_status["status"] = "pending_retry" + self.redis.hset("task:status", task_key, json.dumps(retry_status)) + if self._is_hot_new_app_task(retry_task_details): + self._queue_retry_task(task_key, retry_task_details, prioritize=True) + logger.info( + f"[热门重试] 新增热门应用任务 {task_key} 第{new_retry_count}次失败，" + f"已优先重新入队并避开 Worker {worker_id}" + ) + logger.info(f"[等待分析] 任务 {task_key} 已记录为可重试（第{new_retry_count}次失败），等待应用画像决定是否重新入队") + self._dispatch_alert(task_key, worker_id, error_info, message, retry_count, download_errors=download_errors) + + def clear_all(self): + """清空所有任务数据（谨慎使用）""" + keys_to_delete = [ + *self.TASK_QUEUE_CLEANUP_NAMES, + "task:status", + "task:details", + "worker:tasks", + "task:completed", + "task:failed", + "workers:info", + "workers:idle", + "workers:busy", + "workers:meta", + ] + for key in keys_to_delete: + self.redis.delete(key) + for key in self.redis.scan_iter(match="alert:*"): + self.redis.delete(key) + for key in self.redis.scan_iter(match="cooldown:*"): + self.redis.delete(key) + for key in self.redis.scan_iter(match="worker:recovery_observe:*"): + self.redis.delete(key) + for key in self.redis.scan_iter(match="worker:recovery_pending:*"): + self.redis.delete(key) + logger.info("已清空所有任务数据") + self._apk_registry = None + self._minio_storage = None + + # ── APK prefetch integration ─────────────────────────────────── + + def _get_apk_registry(self): + if self._apk_registry is None: + from apk_cloud.registry import ApkRegistry + self._apk_registry = ApkRegistry() + return self._apk_registry + + def _get_minio_storage(self): + """获取 MinioStorage 实例，如果 MINIO_ENABLED=False 则返回 None""" + if self._apk_download_mode() != "minio": + return None + if self._minio_storage is None: + from config import MINIO_ENABLED + if not MINIO_ENABLED: + return None + from apk_cloud.storage import MinioStorage + self._minio_storage = MinioStorage.create() # 使用工厂方法 + return self._minio_storage + + def _resolve_apk_for_task(self, package_name: str, last_updated: str) -> dict: + if not self._uses_central_apk_cache(): + return {} + registry = self._get_apk_registry() + if not registry.is_fresh_enough(package_name, last_updated): + return {} + smb_dir = registry.get_smb_apk_dir(package_name) + if not smb_dir: + return {} + entry = registry.get(package_name) or {} + apk_files = entry.get("apk_files", []) + return { + "local_apk_dir": smb_dir, + "local_apk_files": apk_files, + "apk_download_date": entry.get("download_date", ""), + } + + def _enrich_task_with_apk(self, task_details: dict) -> dict: + package_name = str(task_details.get("package_name") or "").strip() + last_updated = str(task_details.get("last_updated") or "").strip() + if not package_name: + return task_details + apk_info = self._resolve_apk_for_task(package_name, last_updated) + if apk_info: + task_details = dict(task_details) + task_details["available_sources"] = ["local", "google_play", "apkpure"] + task_details["local_apk_dir"] = apk_info["local_apk_dir"] + task_details["local_apk_files"] = apk_info["local_apk_files"] + return task_details + + def _mark_download_error_awaiting_apk(self, package_name: str, app_name: str = "", last_updated: str = "", force: bool = False) -> None: + if not package_name: + return + if not self._uses_central_apk_cache(): + logger.info( + "[APK] 当前为 direct 模式，Worker 将自行下载，跳过中心 APK 等待队列: %s", + package_name, + ) + return + registry = self._get_apk_registry() + if not force and registry.is_fresh_enough(package_name, last_updated): + logger.info("[APK] 本地已有新鲜APK，无需等待中心 APK 缓存刷新: %s", package_name) + return + registry.mark_pending_download(package_name) + logger.info("[APK] 下载失败，标记为等待APK: %s", package_name) + try: + storage = self._get_minio_storage() + if storage is None: + logger.info("[APK] Minio 未启用，跳过优先下载队列推送: %s", package_name) + return + storage.prepend_to_download_queue([{ + "package_name": package_name, + "app_name": app_name or package_name, + "last_updated": last_updated or "", + "force": force, + }]) + logger.info("[APK] 下载失败，任务已加入下载队列开头: %s force=%s", package_name, force) + except Exception as exc: + logger.warning("[APK] 下载失败，推送优先下载队列失败: %s", exc) + + def _fail_block_tasks_waiting_for_apk(self, package_name: str) -> None: + task_keys = self._find_task_keys_by_package(package_name) + for task_key in task_keys: + if not task_key.endswith("_block"): + continue + status_json = self.redis.hget("task:status", task_key) + if not status_json: + continue + status_data = json.loads(status_json) + if status_data.get("status") != "pending_retry": + continue + status_data["status"] = "failed" + status_data["last_fail_message"] = "中心 APK 下载失败，block任务永久失败" + self.redis.hset("task:status", task_key, json.dumps(status_data)) + self.redis.sadd("task:failed", task_key) + logger.info(f"[Block测试] 中心 APK 下载失败，block任务 {task_key} 标记为永久失败") + + def _log_apk_install_failure(self, task_key: str, package_name: str, app_name: str, message: str) -> None: + import csv, io + from config import REPORT_DIR + import os as _os + _os.makedirs(REPORT_DIR, exist_ok=True) + report_path = _os.path.join(REPORT_DIR, "apk_install_failures.csv") + existed = _os.path.isfile(report_path) + with open(report_path, "a", newline="", encoding="utf-8") as f: + writer = csv.writer(f) + if not existed: + writer.writerow(["时间", "任务键", "包名", "应用名", "失败原因"]) + writer.writerow([ + time.strftime("%Y-%m-%d %H:%M:%S"), + task_key, + package_name, + app_name, + message, + ]) + + def force_reapk(self, package_name: str) -> int: + registry = self._get_apk_registry() + pkg_dir = registry.storage_dir.rstrip("/\\") + "/" + package_name + import shutil + if os.path.isdir(pkg_dir): + shutil.rmtree(pkg_dir, ignore_errors=True) + logger.info("[force_reapk] 已清除本地APK目录: %s", pkg_dir) + try: + with registry._connect() as conn: + conn.execute("DELETE FROM apk_registry WHERE package_name = ?", (package_name,)) + except Exception: + pass + registry.clear_pending_download(package_name) + + task_keys = self._find_task_keys_by_package(package_name) + task_info = None + for tk in task_keys: + d = self.redis.hget("task:details", tk) + if d: + task_info = json.loads(d) + break + app_name = "" + last_updated = "" + if task_info: + app_name = str(task_info.get("app_name") or "").strip() + last_updated = str(task_info.get("last_updated") or "").strip() + self._mark_download_error_awaiting_apk(package_name, app_name=app_name, last_updated=last_updated, force=True) + + for task_key in task_keys: + self.redis.hset("task:status", task_key, json.dumps({ + "status": "pending_retry", + "end_time": time.time(), + "retry_count": 0, + "last_fail_message": "force_reapk: 等待重新下载", + })) + logger.info("[force_reapk] 已重置 %s (%d 个任务)，等待中心 APK 缓存刷新", package_name, len(task_keys)) + return len(task_keys) + + def _re_enqueue_apk_tasks_for_package(self, package_name: str) -> int: + task_keys = self._find_task_keys_by_package(package_name) + re_enqueued = 0 + for task_key in task_keys: + status_json = self.redis.hget("task:status", task_key) + if not status_json: + continue + try: + status = json.loads(status_json) + except (json.JSONDecodeError, TypeError): + continue + if status.get("status") not in ("pending_retry", "failed", "pending"): + continue + task_details_json = self.redis.hget("task:details", task_key) + task_details = json.loads(task_details_json) if task_details_json else {} + self._remove_task_from_queue(task_key) + self.redis.srem("task:completed", task_key) + self.redis.srem("task:failed", task_key) + self.redis.hset("task:status", task_key, json.dumps({ + "status": "pending", + "end_time": time.time(), + "retry_count": int(status.get("retry_count") or 0), + })) + self._queue_retry_task(task_key, task_details) + re_enqueued += 1 + logger.info("[APK] APK就绪，重新入队: %s (%s)", task_key, package_name) + registry = self._get_apk_registry() + registry.clear_pending_download(package_name) + return re_enqueued + + def _try_re_enqueue_apk_ready_tasks(self) -> int: + registry = self._get_apk_registry() + ready_packages = registry.get_pending_that_are_available() + if not ready_packages: + return 0 + + re_enqueued = 0 + for package_name in ready_packages: + task_keys = self._find_task_keys_by_package(package_name) + for task_key in task_keys: + status_json = self.redis.hget("task:status", task_key) + if not status_json: + continue + try: + status = json.loads(status_json) + except (json.JSONDecodeError, TypeError): + continue + if status.get("status") not in ("pending_retry", "failed", "pending"): + continue + + task_details_json = self.redis.hget("task:details", task_key) + task_details = json.loads(task_details_json) if task_details_json else {} + + self._remove_task_from_queue(task_key) + self.redis.srem("task:completed", task_key) + self.redis.srem("task:failed", task_key) + self.redis.hset("task:status", task_key, json.dumps({ + "status": "pending", + "end_time": time.time(), + "retry_count": int(status.get("retry_count") or 0), + })) + self._queue_retry_task(task_key, task_details) + re_enqueued += 1 + logger.info("[APK] APK就绪，重新入队: %s (%s)", task_key, package_name) + + registry.clear_pending_download(package_name) + return re_enqueued + + def _find_task_keys_by_package(self, package_name: str) -> List[str]: + found = [] + all_details = self.redis.hgetall("task:details") + for task_key, details_json in (all_details or {}).items(): + try: + details = json.loads(details_json) + except (json.JSONDecodeError, TypeError): + continue + if str(details.get("package_name") or "").strip() == package_name: + found.append(task_key) + return found + + def sync_download_queue(self) -> int: + """同步下载队列（仅当 Minio 启用时）""" + try: + storage = self._get_minio_storage() + if storage is None: + # Minio 未启用，跳过队列同步 + return 0 + except Exception as exc: + logger.warning("[APK] MinIO不可达，跳过下载队列同步: %s", exc) + return 0 + + registry = self._get_apk_registry() + + current_queue = storage.pull_download_queue() + filtered = [] + removed_completed = 0 + removed_failed = 0 + for task in current_queue: + pkg = str(task.get("package_name") or "").strip() + last_updated = str(task.get("last_updated") or "").strip() + if not pkg: + continue + if registry.is_fresh_enough(pkg, last_updated): + removed_completed += 1 + continue + if last_updated and registry.has_failed_with_same_version(pkg, last_updated): + removed_failed += 1 + continue + filtered.append(task) + + if removed_completed or removed_failed: + logger.info("[APK] 队列清理: 移除已完成 %d, 移除已失败 %d, 剩余 %d", + removed_completed, removed_failed, len(filtered)) + + pending_pkgs = registry.get_all_pending_downloads() + existing_pkgs = {t.get("package_name", "") for t in filtered} + new_tasks = [] + for pkg in pending_pkgs: + if pkg in existing_pkgs: + continue + task_info = self._find_task_keys_by_package(pkg) + app_name = "" + last_updated = "" + for task_key in task_info: + details_json = self.redis.hget("task:details", task_key) + if details_json: + details = json.loads(details_json) + app_name = str(details.get("app_name") or app_name or "").strip() or app_name + last_updated = str(details.get("last_updated") or last_updated or "").strip() or last_updated + if app_name and last_updated: + break + if last_updated and registry.has_failed_with_same_version(pkg, last_updated): + continue + new_tasks.insert(0, { + "package_name": pkg, + "app_name": app_name or pkg, + "last_updated": last_updated or "", + }) + + queue = new_tasks + filtered + max_queue = 1000 + if len(queue) > max_queue: + logger.info("[APK] 下载队列超过上限 %d，截断尾部 %d 条", max_queue, len(queue) - max_queue) + queue = queue[:max_queue] + try: + storage.push_download_queue(queue) + except Exception as exc: + logger.warning("[APK] 推送下载队列失败: %s", exc) + return 0 + + if new_tasks: + logger.info("[APK] 队列新增 %d 个优先任务, 队列总数 %d", len(new_tasks), len(queue)) + return len(new_tasks) + + def _get_stale_download_tasks(self, registry) -> List[Dict[str, str]]: + if not hasattr(self, '_last_stale_check'): + self._last_stale_check = 0.0 + if time.time() - self._last_stale_check < APK_DOWNLOAD_QUEUE_INTERVAL: + return [] + self._last_stale_check = time.time() + + try: + stale_records = registry.get_stale_records() + except Exception as exc: + logger.warning("[APK] 查询过期APK记录失败: %s", exc) + return [] + + tasks = [] + for record in stale_records: + package_name = str(record.get("package_name") or "").strip() + if not package_name: + continue + tasks.append({ + "package_name": package_name, + "app_name": str(record.get("app_name") or "").strip(), + "last_updated": str(record.get("last_updated") or "").strip(), + }) + if tasks: + logger.info("[APK] 检测到 %d 个过期本地APK，将重新下载", len(tasks)) + return tasks + + def poll_apk_results(self) -> int: + """轮询 APK 下载结果（仅当 Minio 启用时）""" + try: + storage = self._get_minio_storage() + if storage is None: + # Minio 未启用，跳过结果轮询 + return 0 + except Exception as exc: + logger.warning("[APK] MinIO不可达，跳过结果轮询: %s", exc) + return 0 + + try: + results = storage.list_download_results() + except Exception as exc: + logger.warning("[APK] 轮询下载结果失败: %s", exc) + return 0 + + if not results: + return self._try_re_enqueue_apk_ready_tasks() + + registry = self._get_apk_registry() + imported = 0 + for package_name, result in results.items(): + status = result.get("status", "") + + try: + self._record_result_to_db(package_name, result) + except Exception as exc: + logger.warning("[APK] 记录下载结果到DB失败 %s: %s", package_name, exc) + + if status != "ok": + if status == "failed": + storage.delete_download_result(package_name) + registry.clear_pending_download(package_name) + self._fail_block_tasks_waiting_for_apk(package_name) + continue + try: + local_paths = storage.download_apk(package_name, result, registry.storage_dir) + if not registry.verify_apk_versions(package_name): + logger.warning("[APK] 版本码不一致，丢弃: %s", package_name) + import shutil + shutil.rmtree(registry.storage_dir + "/" + package_name, ignore_errors=True) + continue + registry.import_from_download_result(package_name, result, local_dir=registry.storage_dir) + storage.delete_download_result(package_name) + imported += 1 + re_enqueued = self._re_enqueue_apk_tasks_for_package(package_name) + logger.info("[APK] 下载完成: %s (%d files), 重新入队 %d 个任务", + package_name, len(local_paths), re_enqueued) + version_dir = f"{result.get('download_date', '')}_{result.get('version_code', '')}" + if version_dir.strip("_"): + storage.delete_apk_version(package_name, version_dir) + logger.info("[APK] 已清理MinIO源文件: %s/%s", package_name, version_dir) + storage.cleanup_old_versions(package_name, keep=0) + except Exception as exc: + logger.warning("[APK] 处理结果失败 %s: %s", package_name, exc) + + if imported > 0: + logger.info("[APK] 本轮导入%d个新APK，刷新待分发任务...", imported) + self.refresh_tasks_from_app_summary() + self._enforce_apk_storage_limits() + + re_enqueued = self._try_re_enqueue_apk_ready_tasks() + if re_enqueued: + logger.info("[APK] 重新入队 %d 个等待APK的任务", re_enqueued) + return imported + re_enqueued + + def _record_result_to_db(self, package_name: str, result: dict) -> None: + registry = self._get_apk_registry() + package_name = str(package_name or "").strip() + if not package_name: + return + + task_keys = self._find_task_keys_by_package(package_name) + if not task_keys: + registry.record_download_from_result(package_name, "", result) + logger.info("[APK] 记录下载结果(无对应任务): %s status=%s", + package_name, result.get("status")) + return + + for task_key in task_keys: + task_details_json = self.redis.hget("task:details", task_key) + if not task_details_json: + continue + try: + task_details = json.loads(task_details_json) + except (json.JSONDecodeError, TypeError): + continue + country_code = str(task_details.get("country_code") or "").strip() + registry.record_download_from_result(package_name, country_code, result) + logger.info("[APK] 记录下载结果: %s/%s status=%s", + package_name, country_code, result.get("status")) + break + + def _manage_stale_apks(self) -> int: + """定期清理过期本地 APK 文件，并推送到下载队列重新下载。 + + 返回清理的过期 APK 数量。 + """ + registry = self._get_apk_registry() + try: + stale_records = registry.get_stale_records() + except Exception as exc: + logger.warning("[APK] 查询过期 APK 失败: %s", exc) + return 0 + + if not stale_records: + return 0 + + cleaned = 0 + for record in stale_records: + package_name = str(record.get("package_name") or "").strip() + if not package_name: + continue + try: + registry.cleanup_stale_local(package_name) + cleaned += 1 + logger.info("[APK] 已清理过期本地 APK: %s", package_name) + except Exception as exc: + logger.warning("[APK] 清理过期 APK 失败 %s: %s", package_name, exc) + + if cleaned > 0: + logger.info("[APK] 过期 APK 清理: %d 个包已删除，将在下次队列同步时重新下载", cleaned) + return cleaned + + def _enforce_apk_storage_limits(self) -> None: + """强制执行 APK 存储限制（仅当 Minio 启用时）""" + try: + storage = self._get_minio_storage() + if storage is None: + # Minio 未启用，跳过存储限制检查 + return + except Exception: + return + try: + minio_result = storage.enforce_storage_limit() + if minio_result.get("deleted_versions"): + logger.info( + "[APK] MinIO 超限清理: %.1fGB → %.1fGB (删除 %d 个版本)", + minio_result["before_bytes"] / (1024 ** 3), + minio_result["after_bytes"] / (1024 ** 3), + minio_result["deleted_versions"], + ) + except Exception as exc: + logger.warning("[APK] MinIO 存储限制检查失败: %s", exc) diff --git a/remote_worker_controller.py b/remote_worker_controller.py new file mode 100644 index 0000000..7c15fc0 --- /dev/null +++ b/remote_worker_controller.py @@ -0,0 +1,1543 @@ +import base64 +import codecs +import time +import traceback +from typing import Any, Callable, Dict, List, Optional, Union + +import paramiko + +from config import ( + EMULATOR_ADB_IP_OFFSET, + LOCAL_IMPORT_DIR, + MUMU_BRIDGE_IP_OFFSET, + MUMU_NETWORK_GATEWAYS, + MUMU_NETWORK_SCRIPT_PATH, + MUMU_RECOVER_AHK_EXE, + MUMU_RECOVER_SCRIPT_PATH, + MUMU_MANAGER_PATH, + MUMU_RESTART_SETTLE_SECONDS, + MUMU_VM_INDEX, + NET_BRIDGE_CARD, + RUNBATCH_FORCE_KILL_IMAGES, + RUNBATCH_LEGACY_END_WORKER_IPS, + RUN_PIPELINE_DEFAULT_OPTIONS, + SSH_ACTION_TIMEOUT_SECONDS, + SSH_CONNECT_TIMEOUT_SECONDS, + SSH_DEFAULT_PASSWORD, + SHARE_SMB_USER, + SHARE_SMB_PASSWORD, + extract_worker_ip, + get_repo_url, + get_clean_backup_path, + get_share_smb_target, +) +from log_manager import logger + +RemoteCommand = Union[str, Dict[str, str]] +MUMU_NETWORK_GATEWAY_RULES = ";".join(f"{prefix}={gateway}" for prefix, gateway in MUMU_NETWORK_GATEWAYS.items()) + +class RemoteWorkerController: + def __init__( + self, + worker_inventory: List[Dict[str, Any]], + *, + ssh_connect_timeout: int = SSH_CONNECT_TIMEOUT_SECONDS, + ssh_action_timeout: int = SSH_ACTION_TIMEOUT_SECONDS, + default_password: str = SSH_DEFAULT_PASSWORD, + ): + self.worker_inventory = {item["worker_id"]: dict(item) for item in worker_inventory} + self.ssh_connect_timeout = ssh_connect_timeout + self.ssh_action_timeout = ssh_action_timeout + self.default_password = default_password + + def get_worker(self, worker_id: str) -> Dict[str, Any]: + worker = self.worker_inventory.get(worker_id) + if not worker: + raise KeyError(f"unknown worker_id: {worker_id}") + return worker + + def run_action( + self, + worker_id: str, + action: str, + timeout: Optional[int] = None, + options: Optional[Dict[str, Any]] = None, + progress_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None, + ) -> Dict[str, Any]: + worker = self.get_worker(worker_id) + if action == "clone": + return self._run_clone(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + if action == "pull": + return self._run_pull(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + if action == "setup": + return self._run_setup(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + if action == "pull_pcap_files": + return self._run_pull_pcap_files(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + if action == "stop_worker": + return self._run_stop_worker(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + if action == "status": + return self._run_status(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + if action == "reboot": + return self._run_reboot(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + if action == "recover_mumu": + return self._run_recover_mumu(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + if action == "restart_mumu": + return self._run_restart_mumu(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + if action == "start_worker": + return self._run_start_worker( + worker_id, + worker, + timeout=timeout, + progress_callback=progress_callback, + ) + if action == "run_pipeline": + return self._run_pipeline(worker_id, worker, timeout=timeout, options=options, progress_callback=progress_callback) + if action == "execute_command": + return self._run_execute_command(worker_id, worker, timeout=timeout, options=options, progress_callback=progress_callback) + if action == "fix_adb_connection": + return self._run_fix_adb_connection(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + if action == "configure_mumu_network": + return self._run_configure_mumu_network(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + if action == "recover_mumu_full": + return self._run_recover_mumu_full(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + raise ValueError(f"unsupported action: {action}") + + def _run_clone( + self, + worker_id: str, + worker: Dict[str, Any], + *, + timeout: Optional[int] = None, + progress_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None, + ) -> Dict[str, Any]: + return self._run_ssh_command( + worker_id, + worker, + "clone", + self._clone_command(worker), + timeout=timeout, + progress_callback=progress_callback, + ) + + def _run_execute_command( + self, + worker_id: str, + worker: Dict[str, Any], + *, + timeout: Optional[int] = None, + options: Optional[Dict[str, Any]] = None, + progress_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None, + ) -> Dict[str, Any]: + command = (options or {}).get("command", "") + if not command: + return { + "ok": False, + "worker_id": worker_id, + "action": "execute_command", + "message": "no command provided", + } + return self._run_ssh_command( + worker_id, + worker, + "execute_command", + command, + timeout=timeout, + progress_callback=progress_callback, + ) + + def _run_fix_adb_connection( + self, + worker_id: str, + worker: Dict[str, Any], + *, + timeout: Optional[int] = None, + progress_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None, + ) -> Dict[str, Any]: + host = str(worker.get("ssh_host") or "").strip() + parts = host.split(".") + if len(parts) != 4: + return { + "ok": False, + "worker_id": worker_id, + "action": "fix_adb_connection", + "message": f"无法解析 Worker IP: {host}", + } + try: + last_octet = int(parts[3]) + EMULATOR_ADB_IP_OFFSET + except (ValueError, IndexError): + return { + "ok": False, + "worker_id": worker_id, + "action": "fix_adb_connection", + "message": f"无法计算模拟器 IP: {host}", + } + emulator_ip = f"{parts[0]}.{parts[1]}.{parts[2]}.{last_octet}" + command = f"adb disconnect && adb connect {emulator_ip}:5555" + return self._run_ssh_command( + worker_id, + worker, + "fix_adb_connection", + command, + timeout=timeout, + progress_callback=progress_callback, + ) + + def _run_configure_mumu_network( + self, + worker_id: str, + worker: Dict[str, Any], + *, + timeout: Optional[int] = None, + progress_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None, + ) -> Dict[str, Any]: + manager = MUMU_MANAGER_PATH + vm_index = MUMU_VM_INDEX + default_adapter = NET_BRIDGE_CARD + script_path = MUMU_NETWORK_SCRIPT_PATH + command = " ".join( + [ + self._cmd_quote(script_path), + self._cmd_quote(manager), + str(vm_index), + self._cmd_quote(default_adapter), + str(MUMU_BRIDGE_IP_OFFSET), + self._cmd_quote(MUMU_NETWORK_GATEWAY_RULES), + ] + ) + inner = f"if not exist {self._cmd_quote(script_path)} (echo [ERROR] script not found: {script_path} 1>&2 & exit /b 1) & call {command}" + return self._run_ssh_command( + worker_id, worker, "configure_mumu_network", + f"cmd.exe /d /q /s /c {self._cmd_quote(inner)}", + timeout=timeout, progress_callback=progress_callback, + ) + + def _run_pull( + self, + worker_id: str, + worker: Dict[str, Any], + *, + timeout: Optional[int] = None, + progress_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None, + ) -> Dict[str, Any]: + return self._run_ssh_command( + worker_id, + worker, + "pull", + self._pull_command(worker), + timeout=timeout, + progress_callback=progress_callback, + ) + + def _run_setup( + self, + worker_id: str, + worker: Dict[str, Any], + *, + timeout: Optional[int] = None, + progress_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None, + ) -> Dict[str, Any]: + return self._run_ssh_command( + worker_id, + worker, + "setup", + self._setup_command(worker), + timeout=timeout, + progress_callback=progress_callback, + ) + + def _run_stop_worker( + self, + worker_id: str, + worker: Dict[str, Any], + *, + timeout: Optional[int] = None, + progress_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None, + ) -> Dict[str, Any]: + return self._run_ssh_command( + worker_id, + worker, + "stop_worker", + self._stop_worker_command(worker), + timeout=timeout, + progress_callback=progress_callback, + ) + + def _run_pull_pcap_files( + self, + worker_id: str, + worker: Dict[str, Any], + *, + timeout: Optional[int] = None, + progress_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None, + ) -> Dict[str, Any]: + pull_result = self._run_pull_pcap_files_once(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + steps: List[Dict[str, Any]] = [self._step_result("pull_pcap_files", pull_result)] + if not pull_result.get("ok"): + return { + "ok": False, + "worker_id": worker_id, + "action": "pull_pcap_files", + "message": pull_result.get("message", "pull_pcap_files failed"), + "failed_step": "pull_pcap_files", + "steps": steps, + "stdout": pull_result.get("stdout", ""), + "stderr": pull_result.get("stderr", ""), + } + + return { + "ok": True, + "worker_id": worker_id, + "action": "pull_pcap_files", + "message": pull_result.get("message", "pull_pcap_files completed"), + "steps": steps, + "stdout": pull_result.get("stdout", ""), + "stderr": pull_result.get("stderr", ""), + } + + def _run_pull_pcap_files_once( + self, + worker_id: str, + worker: Dict[str, Any], + *, + timeout: Optional[int] = None, + progress_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None, + ) -> Dict[str, Any]: + return self._run_ssh_command( + worker_id, + worker, + "pull_pcap_files", + self._pull_pcap_files_command(worker), + timeout=timeout, + progress_callback=progress_callback, + ) + + def _run_status( + self, + worker_id: str, + worker: Dict[str, Any], + *, + timeout: Optional[int] = None, + progress_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None, + ) -> Dict[str, Any]: + return self._run_ssh_command( + worker_id, + worker, + "status", + self._status_command(worker), + timeout=timeout, + progress_callback=progress_callback, + ) + + def _run_reboot( + self, + worker_id: str, + worker: Dict[str, Any], + *, + timeout: Optional[int] = None, + progress_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None, + ) -> Dict[str, Any]: + return self._run_ssh_command( + worker_id, + worker, + "reboot", + self._reboot_command(), + timeout=timeout, + progress_callback=progress_callback, + ) + + def _run_recover_mumu( + self, + worker_id: str, + worker: Dict[str, Any], + *, + timeout: Optional[int] = None, + progress_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None, + ) -> Dict[str, Any]: + command = self._recover_mumu_command(worker) + primary_timeout = (timeout or self.ssh_action_timeout) // 2 + result = self._run_ssh_command( + worker_id, worker, "recover_mumu", + command, timeout=primary_timeout, + progress_callback=progress_callback, + ) + if result.get("ok"): + return result + logger.warning( + "[WorkerAction] recover_mumu 首次失败 worker=%s, 自动重试1次", + worker_id, + ) + self._emit_progress( + progress_callback, worker_id, "recover_mumu", "retry", + "MuMu 镜像恢复首次失败，自动重试 1 次", status="running", + ) + result = self._run_ssh_command( + worker_id, worker, "recover_mumu", + command, timeout=primary_timeout, + progress_callback=progress_callback, + ) + if not result.get("ok"): + logger.warning( + "[WorkerAction] recover_mumu 重试依然失败 worker=%s message=%s", + worker_id, result.get("message", ""), + ) + return result + + def _run_recover_mumu_full( + self, + worker_id: str, + worker: Dict[str, Any], + *, + timeout: Optional[int] = None, + progress_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None, + ) -> Dict[str, Any]: + steps: List[Dict[str, Any]] = [] + + recover_result = self._run_recover_mumu(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + steps.append(self._step_result("recover_mumu", recover_result)) + if not recover_result.get("ok"): + return self._pipeline_failure(worker_id, {}, steps, "recover_mumu", recover_result) + + restart_result = self._run_restart_mumu(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + steps.append(self._step_result("restart_mumu", restart_result)) + if not restart_result.get("ok"): + return self._pipeline_failure(worker_id, {}, steps, "restart_mumu", restart_result) + + # wait for MuMu to boot before configuring network + self._wait_for_mumu_restart(worker_id, progress_callback=progress_callback) + + network_result = self._run_configure_mumu_network(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + steps.append(self._step_result("configure_mumu_network", network_result)) + + adb_result = self._run_fix_adb_connection(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + steps.append(self._step_result("fix_adb_connection", adb_result)) + + return { + "ok": True, + "worker_id": worker_id, + "action": "recover_mumu_full", + "message": "完整恢复完成（镜像恢复 → 重启 → 网络配置 → ADB连接）", + "steps": steps, + "stdout": recover_result.get("stdout", ""), + "stderr": recover_result.get("stderr", ""), + } + + def _run_restart_mumu( + self, + worker_id: str, + worker: Dict[str, Any], + *, + timeout: Optional[int] = None, + progress_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None, + ) -> Dict[str, Any]: + return self._run_ssh_command( + worker_id, + worker, + "restart_mumu", + self._restart_mumu_task_command(worker), + timeout=timeout, + progress_callback=progress_callback, + ) + + def _run_start_worker( + self, + worker_id: str, + worker: Dict[str, Any], + *, + timeout: Optional[int] = None, + progress_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None, + ) -> Dict[str, Any]: + return self._run_ssh_command( + worker_id, + worker, + "start_worker", + self._start_worker_command(worker), + timeout=timeout, + progress_callback=progress_callback, + ) + + def _run_pipeline( + self, + worker_id: str, + worker: Dict[str, Any], + *, + timeout: Optional[int] = None, + options: Optional[Dict[str, Any]] = None, + progress_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None, + ) -> Dict[str, Any]: + spec = dict(RUN_PIPELINE_DEFAULT_OPTIONS) + if options: + spec.update(options) + steps: List[Dict[str, Any]] = [] + + if not any( + spec.get(key) + for key in ( + "clone_if_missing", + "stop_worker", + "git_pull", + "setup", + "pull_pcap_files", + "recover_mumu", + "restart_mumu", + "start_worker", + ) + ): + return { + "ok": True, + "worker_id": worker_id, + "action": "run_pipeline", + "message": "no pipeline steps selected", + "steps": steps, + "selected_options": spec, + } + + if spec.get("clone_if_missing"): + result = self._run_ssh_command( + worker_id, + worker, + "clone_if_missing", + self._clone_command(worker), + timeout=timeout, + progress_callback=progress_callback, + ) + steps.append(self._step_result("clone_if_missing", result)) + if not result.get("ok"): + return self._pipeline_failure(worker_id, spec, steps, "clone_if_missing", result) + + if spec.get("stop_worker"): + result = self._run_stop_worker(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + steps.append(self._step_result("stop_worker", result)) + if not result.get("ok"): + return self._pipeline_failure(worker_id, spec, steps, "stop_worker", result) + + if spec.get("git_pull"): + result = self._run_pull(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + steps.append(self._step_result("git_pull", result)) + if not result.get("ok"): + return self._pipeline_failure(worker_id, spec, steps, "git_pull", result) + + if spec.get("setup"): + result = self._run_setup(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + steps.append(self._step_result("setup", result)) + if not result.get("ok"): + return self._pipeline_failure(worker_id, spec, steps, "setup", result) + + pcap_pulled = False + restart_ready = False + if spec.get("pull_pcap_files") and spec.get("recover_mumu"): + if not restart_ready: + result = self._run_restart_mumu(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + steps.append(self._step_result("restart_mumu_before_pull_pcap_files", result)) + if not result.get("ok"): + return self._pipeline_failure(worker_id, spec, steps, "restart_mumu_before_pull_pcap_files", result) + self._wait_for_mumu_restart(worker_id, progress_callback=progress_callback) + restart_ready = True + result = self._run_pull_pcap_files_once(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + steps.append(self._step_result("pull_pcap_files", result)) + if not result.get("ok"): + return self._pipeline_failure(worker_id, spec, steps, "pull_pcap_files", result) + pcap_pulled = True + + if spec.get("recover_mumu"): + result = self._run_recover_mumu(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + steps.append(self._step_result("recover_mumu", result)) + if not result.get("ok"): + return self._pipeline_failure(worker_id, spec, steps, "recover_mumu", result) + restart_ready = False + + if spec.get("restart_mumu"): + result = self._run_restart_mumu(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + steps.append(self._step_result("restart_mumu", result)) + if not result.get("ok"): + return self._pipeline_failure(worker_id, spec, steps, "restart_mumu", result) + self._run_configure_mumu_network(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + self._run_fix_adb_connection(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + if spec.get("start_worker") or (spec.get("pull_pcap_files") and not pcap_pulled): + self._wait_for_mumu_restart(worker_id, progress_callback=progress_callback) + restart_ready = True + + if spec.get("start_worker"): + if spec.get("pull_pcap_files") and not pcap_pulled: + if not restart_ready: + result = self._run_restart_mumu(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + steps.append(self._step_result("restart_mumu_before_pull_pcap_files", result)) + if not result.get("ok"): + return self._pipeline_failure(worker_id, spec, steps, "restart_mumu_before_pull_pcap_files", result) + self._wait_for_mumu_restart(worker_id, progress_callback=progress_callback) + restart_ready = True + result = self._run_pull_pcap_files_once(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + steps.append(self._step_result("pull_pcap_files", result)) + if not result.get("ok"): + return self._pipeline_failure(worker_id, spec, steps, "pull_pcap_files", result) + pcap_pulled = True + result = self._run_start_worker( + worker_id, + worker, + timeout=timeout, + progress_callback=progress_callback, + ) + steps.append(self._step_result("start_worker", result)) + if not result.get("ok"): + return self._pipeline_failure(worker_id, spec, steps, "start_worker", result) + + if spec.get("pull_pcap_files") and not pcap_pulled: + if not restart_ready: + result = self._run_restart_mumu(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + steps.append(self._step_result("restart_mumu_before_pull_pcap_files", result)) + if not result.get("ok"): + return self._pipeline_failure(worker_id, spec, steps, "restart_mumu_before_pull_pcap_files", result) + self._wait_for_mumu_restart(worker_id, progress_callback=progress_callback) + result = self._run_pull_pcap_files_once(worker_id, worker, timeout=timeout, progress_callback=progress_callback) + steps.append(self._step_result("pull_pcap_files", result)) + if not result.get("ok"): + return self._pipeline_failure(worker_id, spec, steps, "pull_pcap_files", result) + + return { + "ok": True, + "worker_id": worker_id, + "action": "run_pipeline", + "message": "pipeline completed", + "steps": steps, + "selected_options": spec, + } + + @staticmethod + def _step_result(step: str, result: Dict[str, Any]) -> Dict[str, Any]: + item = { + "step": step, + "ok": result.get("ok"), + "message": result.get("message", ""), + } + command = str(result.get("command", "") or "").strip() + if command: + item["command"] = command + # 保留 stdout/stderr，确保失败步骤的详细报错信息不会丢失 + for key in ("stdout", "stderr"): + value = str(result.get(key, "") or "").strip() + if value: + item[key] = value + return item + + @staticmethod + def _pipeline_failure( + worker_id: str, + spec: Dict[str, Any], + steps: List[Dict[str, Any]], + failed_step: str, + result: Dict[str, Any], + ) -> Dict[str, Any]: + return { + "ok": False, + "worker_id": worker_id, + "action": "run_pipeline", + "message": result.get("message", f"{failed_step} failed"), + "failed_step": failed_step, + "steps": steps, + "selected_options": spec, + "stdout": result.get("stdout", ""), + "stderr": result.get("stderr", ""), + } + + def _clone_command(self, worker: Dict[str, Any]) -> str: + repo_dir = self._repo_dir(worker) + repo_url = get_repo_url(worker["ssh_target"]) + return ( + f'if exist "{repo_dir}\\.git" ' + f'(cd /d "{repo_dir}" && git config lfs.fetchexclude "" && git lfs pull) ' + f'else (git clone "{repo_url}" "{repo_dir}" && cd /d "{repo_dir}" && git config lfs.fetchexclude "" && git lfs pull)' + ) + + def _pull_command(self, worker: Dict[str, Any]) -> str: + repo_dir = self._repo_dir(worker) + is_physical = str(worker.get("device_type") or "").strip().lower() == "physical" + if is_physical: + pull_cmd = "git stash && git pull && (git stash pop || ver>nul)" + else: + pull_cmd = "git reset --hard HEAD && git pull" + return ( + f'if exist "{repo_dir}\\.git" ' + f'(cd /d "{repo_dir}" && {pull_cmd} && git config lfs.fetchexclude "" && git lfs pull) ' + f'else (echo repository is not initialized 1>&2 & exit /b 1)' + ) + + def _setup_command(self, worker: Dict[str, Any]) -> str: + repo_dir = self._repo_dir(worker) + bootstrap_python = self._bootstrap_python(worker) + # 纯 CLI 操作，无需 GUI/用户会话，直接通过 SSH 执行 + cmd_line = ( + f'cd /d "{repo_dir}" && ' + f'"{bootstrap_python}" -m pip config set global.index-url https://mirrors.aliyun.com/pypi/simple && ' + f'"{bootstrap_python}" -m pip config set install.trusted-host mirrors.aliyun.com && ' + f'"{bootstrap_python}" setup.py' + ) + return self._repo_required_command(worker, cmd_line) + + def _stop_worker_command(self, worker: Dict[str, Any]) -> RemoteCommand: + if self._use_legacy_runbatch_end(worker): + return self._legacy_runbatch_end_command() + return self._runbatch_taskkill_command() + + def _status_command(self, worker: Dict[str, Any]) -> RemoteCommand: + repo_dir = self._repo_dir(worker) + venv_python = self._venv_python(worker) + script = "\n".join( + [ + f"$repoDir = {self._powershell_quote(repo_dir)}", + f"$venvPython = {self._powershell_quote(venv_python)}", + ( + "$workerRunning = @(" + "Get-CimInstance Win32_Process -ErrorAction SilentlyContinue | " + "Where-Object { " + "$_.Name -in @('python.exe', 'pythonw.exe') -and " + "$_.CommandLine -like '*batch_run.py*' " + "}" + ").Count -gt 0" + ), + "Write-Output ('repo_dir=' + $repoDir)", + "Write-Output ('repo_exists=' + [int](Test-Path -LiteralPath (Join-Path $repoDir '.git')))", + "Write-Output ('venv_exists=' + [int](Test-Path -LiteralPath $venvPython))", + "Write-Output ('worker_running=' + [int]$workerRunning)", + ] + ) + return self._powershell_stdin_command(script) + + def _reboot_command(self) -> str: + return '"C:\\Windows\\System32\\shutdown.exe" /r /f /t 10 /c "Restart initiated by remote controller"' + + def _recover_mumu_command(self, worker: Dict[str, Any]) -> str: + manager = MUMU_MANAGER_PATH + vm_index = MUMU_VM_INDEX + source = get_clean_backup_path(worker["ssh_target"]) + local_import_dir = LOCAL_IMPORT_DIR + script_path = MUMU_RECOVER_SCRIPT_PATH + script_call = " ".join( + [ + self._cmd_quote(script_path), + self._cmd_quote(source), + self._cmd_quote(local_import_dir), + self._cmd_quote(manager), + str(vm_index), + self._cmd_quote(NET_BRIDGE_CARD), + self._cmd_quote(MUMU_RECOVER_AHK_EXE), + str(MUMU_BRIDGE_IP_OFFSET), + self._cmd_quote(MUMU_NETWORK_GATEWAY_RULES), + ] + ) + inner_command = ( + f"if not exist {self._cmd_quote(script_path)} " + f"(echo [ERROR] Recover script not found: {script_path} 1>&2 & exit /b 1) & call {script_call}" + ) + return f"cmd.exe /d /q /s /c {self._cmd_quote(inner_command)}" + + def _restart_mumu_command(self, worker: Dict[str, Any]) -> str: + manager = MUMU_MANAGER_PATH + vm_index = MUMU_VM_INDEX + del worker + return f'"{manager}" control -v {vm_index} restart' + + def _restart_mumu_task_command(self, worker: Dict[str, Any]) -> RemoteCommand: + return self._scheduled_task_run_command( + "RestartMuMu", + self._restart_mumu_command(worker), + direct_command=True, + ) + + def _wait_for_mumu_restart( + self, + worker_id: str, + *, + progress_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None, + ) -> None: + message = f"等待 MuMu 重启稳定 {MUMU_RESTART_SETTLE_SECONDS} 秒" + self._emit_progress(progress_callback, worker_id, "run_pipeline", "wait_mumu_restart", message, status="running") + time.sleep(MUMU_RESTART_SETTLE_SECONDS) + self._emit_progress(progress_callback, worker_id, "run_pipeline", "wait_mumu_restart", "MuMu 重启等待完成", status="success") + + def _start_worker_command(self, worker: Dict[str, Any]) -> RemoteCommand: + repo_dir = self._repo_dir(worker) + batch_script = f"{self._repo_dir(worker)}\\batch_run.py" + venv_python = self._venv_python(worker) + inner_command = f"{venv_python} {batch_script}" + return self._scheduled_task_run_command( + "RunBatch", + inner_command, + keep_open=True, + stop_existing=True, + unlimited_execution_time=True, + working_directory=repo_dir, + ) + + def _pull_pcap_files_command(self, worker: Dict[str, Any]) -> str: + utils_dir = f"{self._repo_dir(worker)}\\utils_android" + # ADB 拉取 pcap 文件，纯 CLI 操作，无需计划任务 + inner_command = f'cd /d "{utils_dir}" && python pull_pcapdroid_files.py' + return self._repo_required_command(worker, inner_command) + + def _repo_required_command(self, worker: Dict[str, Any], command: str) -> str: + repo_dir = self._repo_dir(worker) + return f'if not exist "{repo_dir}\\.git" (echo repository is not initialized 1>&2 & exit /b 1) && {command}' + + def _repo_dir(self, worker: Dict[str, Any]) -> str: + return str(worker["repo_dir"]).replace("/", "\\") + + @staticmethod + def _legacy_runbatch_end_command() -> str: + return ( + 'cmd.exe /d /v:on /s /c "set STOPPED=0' + ' & schtasks /end /tn ""RunBatch"" >nul 2>&1 && set STOPPED=1' + ' & if !STOPPED!==1 (echo worker stopped) else (echo worker is not running)"' + ) + + def _runbatch_taskkill_command(self) -> str: + command_script = self._runbatch_taskkill_command_script() + return f"cmd.exe /d /v:on /q /s /c {self._cmd_quote(command_script)}" + + def _runbatch_taskkill_command_script(self) -> str: + command_parts = [ + "set STOPPED=0", + 'schtasks /end /tn "RunBatch" >nul 2>&1 && set STOPPED=1', + ] + for image_name in RUNBATCH_FORCE_KILL_IMAGES: + command_parts.append(f"taskkill /F /T /IM {image_name} >nul 2>&1 && set STOPPED=1") + command_parts.extend( + [ + 'if !STOPPED!==1 (echo worker stopped) else (echo worker is not running)', + "start /min taskkill /F /T /IM cmd.exe", + ] + ) + return " & ".join(command_parts) + + @staticmethod + def _use_legacy_runbatch_end(worker: Dict[str, Any]) -> bool: + for key in ("worker_id", "ssh_host", "ssh_target"): + worker_ip = extract_worker_ip(worker.get(key, "")) + if worker_ip in RUNBATCH_LEGACY_END_WORKER_IPS: + return True + return False + + @staticmethod + def _cmd_escape(value: str) -> str: + return str(value).replace('"', '""') + + @classmethod + def _cmd_quote(cls, value: str) -> str: + return f'"{cls._cmd_escape(value)}"' + + @classmethod + def _cmd_task_command(cls, inner_command: str) -> str: + del cls + return f'C:\\Windows\\System32\\cmd.exe /d /q /s /c "{inner_command}"' + + @staticmethod + def _cmd_sleep(seconds: int) -> str: + return f"ping -n {max(1, int(seconds)) + 1} 127.0.0.1 >nul" + + @staticmethod + def _powershell_quote(value: str) -> str: + return "'" + str(value).replace("'", "''") + "'" + + @staticmethod + def _cmd_switch_with_command(cmd_switch: str, inner_command: str) -> str: + return f"{cmd_switch} {inner_command}" + + @staticmethod + def _powershell_encoded_command(script: str) -> str: + payload = base64.b64encode(str(script).encode("utf-16le")).decode("ascii") + return f"powershell -NoProfile -ExecutionPolicy Bypass -EncodedCommand {payload}" + + @staticmethod + def _powershell_stdin_command(script: str) -> Dict[str, str]: + return { + "command": "powershell -NoProfile -ExecutionPolicy Bypass -Command -", + "stdin": str(script), + } + + @staticmethod + def _task_signal_name(task_name: str) -> str: + sanitized = "".join(ch if ch.isalnum() else "_" for ch in str(task_name)) + return f"AUTOOL_{sanitized}_DONE" + + def _scheduled_task_runner_script( + self, + task_name: str, + inner_command: str, + *, + wait: bool, + working_directory: Optional[str], + ) -> str: + artifact_dir = "%TEMP%\\autool-dispatcher-remote-worker" + stdout_path = f"{artifact_dir}\\{task_name}.stdout.log" + stderr_path = f"{artifact_dir}\\{task_name}.stderr.log" + exit_path = f"{artifact_dir}\\{task_name}.exitcode" + lines = [ + "@echo off", + "setlocal", + f'if not exist "{artifact_dir}" mkdir "{artifact_dir}"', + ] + if working_directory: + lines.append(f'cd /d "{working_directory}"') + if wait: + lines.extend( + [ + f'del /f /q "{stdout_path}" "{stderr_path}" "{exit_path}" >nul 2>&1', + f'({inner_command}) 1>"{stdout_path}" 2>"{stderr_path}"', + 'set "RC=%ERRORLEVEL%"', + f'> "{exit_path}" echo %RC%', + f'waitfor /SI {self._task_signal_name(task_name)} >nul 2>&1', + "exit /b %RC%", + ] + ) + else: + lines.extend( + [ + inner_command, + "exit /b %ERRORLEVEL%", + ] + ) + return "\r\n".join(lines) + "\r\n" + + def _scheduled_task_command_via_cmd( + self, + task_name: str, + inner_command: str, + *, + wait: bool, + keep_open: bool, + stop_existing: bool, + working_directory: Optional[str], + ) -> RemoteCommand: + artifact_dir = "%TEMP%\\autool-dispatcher-remote-worker" + runner_path = f"{artifact_dir}\\{task_name}.cmd" + stdout_path = f"{artifact_dir}\\{task_name}.stdout.log" + stderr_path = f"{artifact_dir}\\{task_name}.stderr.log" + exit_path = f"{artifact_dir}\\{task_name}.exitcode" + signal_name = self._task_signal_name(task_name) + cmd_switch = "/k" if keep_open else "/c" + runner_script = self._scheduled_task_runner_script( + task_name, + inner_command, + wait=wait, + working_directory=working_directory, + ) + task_command = f'C:\\Windows\\System32\\cmd.exe {cmd_switch} {runner_path}' + command_parts = [ + 'if not exist "%TEMP%\\autool-dispatcher-remote-worker" mkdir "%TEMP%\\autool-dispatcher-remote-worker"', + f'type con > {self._cmd_quote(runner_path)}', + f'schtasks /create /tn {self._cmd_quote(task_name)} /tr {self._cmd_quote(task_command)} /sc ONCE /st 00:00 /rl HIGHEST /it /f >nul', + "if errorlevel 1 exit /b %ERRORLEVEL%", + ] + if stop_existing: + command_parts.append(f'schtasks /end /tn {self._cmd_quote(task_name)} >nul 2>&1') + if wait: + command_parts.extend( + [ + f'del /f /q "{stdout_path}" "{stderr_path}" "{exit_path}" >nul 2>&1', + f'schtasks /run /tn {self._cmd_quote(task_name)} >nul', + "if errorlevel 1 exit /b %ERRORLEVEL%", + f'waitfor {signal_name} /t {int(self.ssh_action_timeout)} >nul', + "if errorlevel 1 exit /b %ERRORLEVEL%", + f'set /p TASK_RC=<{self._cmd_quote(exit_path)}', + f'if exist "{stdout_path}" type "{stdout_path}"', + f'if exist "{stderr_path}" type "{stderr_path}" 1>&2', + "exit /b %TASK_RC%", + ] + ) + else: + command_parts.extend( + [ + f'schtasks /run /tn {self._cmd_quote(task_name)} >nul', + "if errorlevel 1 exit /b %ERRORLEVEL%", + f'echo {task_name} started', + ] + ) + command_script = " & ".join(command_parts) + return { + "command": f"cmd.exe /d /q /s /c {self._cmd_quote(command_script)}", + "stdin": runner_script, + } + + def _scheduled_task_direct_run_command_via_cmd( + self, + task_name: str, + task_command: str, + *, + stop_existing: bool = False, + ) -> str: + command_parts = [ + ( + f'schtasks /create /tn {self._cmd_quote(task_name)} ' + f'/tr {self._cmd_quote(task_command)} /sc ONCE /st 00:00 /rl HIGHEST /it /f >nul' + ), + "if errorlevel 1 exit /b %ERRORLEVEL%", + ] + if stop_existing: + command_parts.append(f'schtasks /end /tn {self._cmd_quote(task_name)} >nul 2>&1') + command_parts.extend( + [ + f'schtasks /run /tn {self._cmd_quote(task_name)} >nul', + "if errorlevel 1 exit /b %ERRORLEVEL%", + f"echo {task_name} started", + ] + ) + command_script = " & ".join(command_parts) + return f"cmd.exe /d /q /s /c {self._cmd_quote(command_script)}" + + def _scheduled_task_direct_wait_command_via_cmd( + self, + task_name: str, + inner_command: str, + ) -> str: + artifact_dir = "%TEMP%\\autool-dispatcher-remote-worker" + stdout_path = f"{artifact_dir}\\{task_name}.stdout.log" + stderr_path = f"{artifact_dir}\\{task_name}.stderr.log" + exit_path = f"{artifact_dir}\\{task_name}.exitcode" + loop_limit = max(1, int(self.ssh_action_timeout // 2)) + task_command = self._cmd_task_command( + ( + f'({inner_command}) 1>"{stdout_path}" 2>"{stderr_path}"' + f' & echo %ERRORLEVEL%>"{exit_path}"' + ) + ) + command_parts = [ + f'if not exist "{artifact_dir}" mkdir "{artifact_dir}"', + f'del /f /q "{stdout_path}" "{stderr_path}" "{exit_path}" >nul 2>&1', + ( + f'schtasks /create /tn {self._cmd_quote(task_name)} ' + f'/tr {self._cmd_quote(task_command)} /sc ONCE /st 00:00 /rl HIGHEST /it /f >nul' + ), + "if errorlevel 1 exit /b %ERRORLEVEL%", + f'schtasks /run /tn {self._cmd_quote(task_name)} >nul', + "if errorlevel 1 exit /b %ERRORLEVEL%", + ( + f'for /l %i in (1,1,{loop_limit}) do @if exist "{exit_path}" ' + f'(goto {task_name}_done) else @ping -n 3 127.0.0.1 >nul' + ), + "exit /b 124", + f":{task_name}_done", + f'set /p TASK_RC=<{self._cmd_quote(exit_path)}', + f'if exist "{stdout_path}" type "{stdout_path}"', + f'if exist "{stderr_path}" type "{stderr_path}" 1>&2', + "exit /b %TASK_RC%", + ] + command_script = " & ".join(command_parts) + return f"cmd.exe /d /q /s /c {self._cmd_quote(command_script)}" + + def _scheduled_task_command_via_powershell( + self, + task_name: str, + inner_command: str, + *, + wait: bool, + keep_open: bool, + hidden: bool, + stop_existing: bool, + unlimited_execution_time: bool, + working_directory: Optional[str], + ) -> RemoteCommand: + stdout_path = f"%TEMP%\\autool-dispatcher-remote-worker\\{task_name}.stdout.log" + stderr_path = f"%TEMP%\\autool-dispatcher-remote-worker\\{task_name}.stderr.log" + task_inner_command = inner_command + if wait: + task_inner_command = f'({inner_command}) 1>"{stdout_path}" 2>"{stderr_path}"' + cmd_switch = "/k" if keep_open else "/c" + if hidden: + hidden_script = " ".join( + [ + f"$cmdCommand = {self._powershell_quote(task_inner_command)};", + f"& cmd.exe {cmd_switch} $cmdCommand;", + "exit $LASTEXITCODE;", + ] + ) + task_execute = "powershell.exe" + task_arguments = ( + "-NoProfile -WindowStyle Hidden -ExecutionPolicy Bypass " + f"-EncodedCommand {base64.b64encode(hidden_script.encode('utf-16le')).decode('ascii')}" + ) + else: + task_execute = "C:\\Windows\\System32\\cmd.exe" + task_arguments = self._cmd_switch_with_command(cmd_switch, task_inner_command) + script_parts = [ + f"$taskName = {self._powershell_quote(task_name)};", + f"$taskExecute = {self._powershell_quote(task_execute)};", + f"$taskArguments = {self._powershell_quote(task_arguments)};", + f"$stopExisting = {'$true' if stop_existing else '$false'};", + f"$unlimitedExecutionTime = {'$true' if unlimited_execution_time else '$false'};", + f"$workingDirectory = {self._powershell_quote(working_directory or '')};", + "$needsCreate = $true;", + "try {", + " $task = Get-ScheduledTask -TaskName $taskName -ErrorAction Stop;", + " $action = @($task.Actions)[0];", + " $settings = $task.Settings;", + " $actionMatches = $action.Execute -eq $taskExecute -and $action.Arguments -eq $taskArguments;", + " $currentWorkingDirectory = '';", + " if ($null -ne $action.WorkingDirectory) { $currentWorkingDirectory = $action.WorkingDirectory };", + " $workingDirectoryMatches = $currentWorkingDirectory -eq $workingDirectory;", + " $executionTimeMatches = (-not $unlimitedExecutionTime) -or $settings.ExecutionTimeLimit -eq 'PT0S';", + " $multipleInstancesMatches = (-not $stopExisting) -or $settings.MultipleInstances -eq 'StopExisting';", + " if ($actionMatches -and $workingDirectoryMatches -and $executionTimeMatches -and $multipleInstancesMatches) { $needsCreate = $false };", + "} catch {};", + "if ($needsCreate) {", + " if ($workingDirectory) {", + " $action = New-ScheduledTaskAction -Execute $taskExecute -Argument $taskArguments -WorkingDirectory $workingDirectory;", + " } else {", + " $action = New-ScheduledTaskAction -Execute $taskExecute -Argument $taskArguments;", + " };", + " $trigger = New-ScheduledTaskTrigger -Once -At ([datetime]'2099-01-01T00:00:00');", + " $currentUser = [System.Security.Principal.WindowsIdentity]::GetCurrent().Name;", + " $principal = New-ScheduledTaskPrincipal -UserId $currentUser -LogonType Interactive -RunLevel Highest;", + " if ($unlimitedExecutionTime -and $stopExisting) {", + " $taskSettings = New-ScheduledTaskSettingsSet -ExecutionTimeLimit ([TimeSpan]::Zero) -MultipleInstances StopExisting;", + " } elseif ($unlimitedExecutionTime) {", + " $taskSettings = New-ScheduledTaskSettingsSet -ExecutionTimeLimit ([TimeSpan]::Zero);", + " } elseif ($stopExisting) {", + " $taskSettings = New-ScheduledTaskSettingsSet -MultipleInstances StopExisting;", + " } else {", + " $taskSettings = $null;", + " };", + " if ($null -ne $taskSettings) {", + " Register-ScheduledTask -TaskName $taskName -Action $action -Trigger $trigger -Principal $principal -Settings $taskSettings -Force -ErrorAction Stop | Out-Null;", + " } else {", + " Register-ScheduledTask -TaskName $taskName -Action $action -Trigger $trigger -Principal $principal -Force -ErrorAction Stop | Out-Null;", + " };", + "};", + ] + if stop_existing and task_name == "RunBatch": + script_parts.append("& schtasks /end /tn $taskName 2>$null | Out-Null;") + if wait: + script_parts.extend( + [ + ( + "$artifactDir = " + "[System.IO.Path]::GetDirectoryName(" + "[System.Environment]::ExpandEnvironmentVariables(" + f"{self._powershell_quote(stdout_path)}));" + ), + "New-Item -ItemType Directory -Path $artifactDir -Force | Out-Null;", + ( + "$stdoutPath = " + f"[System.Environment]::ExpandEnvironmentVariables({self._powershell_quote(stdout_path)});" + ), + ( + "$stderrPath = " + f"[System.Environment]::ExpandEnvironmentVariables({self._powershell_quote(stderr_path)});" + ), + "Remove-Item -LiteralPath $stdoutPath, $stderrPath -Force -ErrorAction SilentlyContinue;", + "$baseline = [datetime]::MinValue;", + "try { $baseline = (Get-ScheduledTaskInfo -TaskName $taskName -ErrorAction Stop).LastRunTime } catch {};", + ] + ) + script_parts.extend( + [ + "& schtasks /run /tn $taskName | Out-Null;", + "if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE };", + ] + ) + if wait: + script_parts.extend( + [ + "while ($true) {", + " $task = Get-ScheduledTask -TaskName $taskName -ErrorAction Stop;", + " $info = Get-ScheduledTaskInfo -TaskName $taskName -ErrorAction Stop;", + " $hasRun = $info.LastRunTime -gt $baseline;", + " if ($hasRun -and $task.State -ne 'Running') { break };", + " Start-Sleep -Milliseconds 500;", + "};", + "if (Test-Path -LiteralPath $stdoutPath) {", + " [Console]::Out.Write([System.IO.File]::ReadAllText($stdoutPath));", + "};", + "if (Test-Path -LiteralPath $stderrPath) {", + " [Console]::Error.Write([System.IO.File]::ReadAllText($stderrPath));", + "};", + "exit [int]$info.LastTaskResult;", + ] + ) + else: + script_parts.append(f"Write-Output {self._powershell_quote(task_name + ' started')};") + script = " ".join(script_parts) + return self._powershell_stdin_command(script) + + def _scheduled_task_run_command( + self, + task_name: str, + inner_command: str, + *, + keep_open: bool = False, + hidden: bool = False, + stop_existing: bool = False, + unlimited_execution_time: bool = False, + working_directory: Optional[str] = None, + direct_command: bool = False, + ) -> RemoteCommand: + return self._scheduled_task_command( + task_name, + inner_command, + wait=False, + keep_open=keep_open, + hidden=hidden, + stop_existing=stop_existing, + unlimited_execution_time=unlimited_execution_time, + working_directory=working_directory, + direct_command=direct_command, + ) + + def _scheduled_task_wait_command( + self, + task_name: str, + inner_command: str, + *, + hidden: bool = True, + direct_command: bool = False, + ) -> RemoteCommand: + return self._scheduled_task_command( + task_name, + inner_command, + wait=True, + keep_open=False, + hidden=hidden, + stop_existing=False, + unlimited_execution_time=False, + working_directory=None, + direct_command=direct_command, + ) + + def _scheduled_task_stop_command( + self, + task_name: str, + success_message: str, + idle_message: str, + ) -> RemoteCommand: + command_script = ( + f'schtasks /query /tn {self._cmd_quote(task_name)} >nul 2>&1' + f' || (echo {idle_message} & exit /b 0)' + f' & schtasks /end /tn {self._cmd_quote(task_name)} >nul 2>&1' + f' || (echo {idle_message} & exit /b 0)' + f' & echo {success_message}' + ) + return f"cmd.exe /d /q /s /c {self._cmd_quote(command_script)}" + + def _scheduled_task_command( + self, + task_name: str, + inner_command: str, + *, + wait: bool, + keep_open: bool, + hidden: bool, + stop_existing: bool, + unlimited_execution_time: bool, + working_directory: Optional[str], + direct_command: bool = False, + ) -> RemoteCommand: + if not unlimited_execution_time and not direct_command: + del hidden + return self._scheduled_task_command_via_cmd( + task_name, + inner_command, + wait=wait, + keep_open=keep_open, + stop_existing=stop_existing, + working_directory=working_directory, + ) + return self._scheduled_task_command_via_powershell( + task_name, + inner_command, + wait=wait, + keep_open=keep_open, + hidden=hidden, + stop_existing=stop_existing, + unlimited_execution_time=unlimited_execution_time, + working_directory=working_directory, + ) + + def _pull_pcap_script_path(self, worker: Dict[str, Any]) -> str: + return f"{self._repo_dir(worker)}\\utils_android\\pull_pcapdroid_files.py" + + def _venv_python(self, worker: Dict[str, Any]) -> str: + return f"{self._repo_dir(worker)}\\venv\\Scripts\\python.exe" + + def _bootstrap_python(self, worker: Dict[str, Any]) -> str: + return str(worker["python_exe"]).replace("/", "\\") + + def _run_ssh_command( + self, + worker_id: str, + worker: Dict[str, Any], + action: str, + command: RemoteCommand, + *, + timeout: Optional[int] = None, + progress_callback: Optional[Callable[[str, Dict[str, Any]], None]] = None, + ) -> Dict[str, Any]: + ssh = paramiko.SSHClient() + ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + target = f"{worker['ssh_host']}:{int(worker['ssh_port'])}" + if isinstance(command, dict): + command_text = str(command.get("command") or "") + stdin_text = str(command.get("stdin") or "") + else: + command_text = str(command) + stdin_text = "" + + if SHARE_SMB_USER and SHARE_SMB_PASSWORD: + share_target = get_share_smb_target(worker["ssh_target"]) + repo_url = get_repo_url(worker["ssh_target"]) + share_markers = [] + if share_target: + share_markers.extend([share_target.lower(), share_target.replace("\\\\", "//", 1).lower()]) + if repo_url.startswith("//"): + repo_host = repo_url[2:].split("/", 1)[0].strip() + if repo_host: + share_markers.extend([f"//{repo_host}".lower(), f"\\\\{repo_host}".lower()]) + if not share_target: + share_target = f"\\\\{repo_host}" + combined_command = "\n".join(part for part in (command_text, stdin_text) if part).lower() + requires_share_auth = action in {"clone", "pull", "recover_mumu"} or any( + marker in combined_command for marker in share_markers + ) + if share_target and requires_share_auth: + net_use_cmd = f'net use {share_target} "{SHARE_SMB_PASSWORD}" /user:"{SHARE_SMB_USER}" >nul 2>&1' + command_text = f"{net_use_cmd} & {command_text}" + + logger.info( + "[WorkerAction] start worker=%s action=%s target=%s timeout=%s", + worker_id, + action, + target, + timeout or self.ssh_action_timeout, + ) + try: + self._emit_progress(progress_callback, worker_id, action, "ssh_connect", "建立 SSH 连接", status="running") + try: + ssh.connect( + hostname=worker["ssh_host"], + port=int(worker["ssh_port"]), + username=worker["ssh_user"], + password=worker.get("ssh_password") or self.default_password, + timeout=self.ssh_connect_timeout, + auth_timeout=self.ssh_connect_timeout, + banner_timeout=self.ssh_connect_timeout, + look_for_keys=True, + allow_agent=True, + ) + except Exception as exc: + detail = traceback.format_exc().strip() + message = f"SSH 连接失败: {exc}" + self._emit_progress(progress_callback, worker_id, action, "ssh_connect", message, status="failed") + logger.error( + "[WorkerAction] ssh_connect failed worker=%s action=%s target=%s err=%s", + worker_id, + action, + target, + self._preview_text(detail), + ) + return { + "ok": False, + "worker_id": worker_id, + "action": action, + "stage": "ssh_connect", + "message": message, + "stderr": detail, + "command": command_text, + } + self._emit_progress(progress_callback, worker_id, action, "ssh_connect", "SSH 已连接", status="success") + self._emit_progress(progress_callback, worker_id, action, "ssh_exec", "发送远端命令", status="running") + try: + stdin, stdout, stderr = ssh.exec_command(command_text, timeout=timeout or self.ssh_action_timeout) + if stdin_text: + stdin.write(stdin_text) + if not stdin_text.endswith("\n"): + stdin.write("\n") + stdin.flush() + stdin.channel.shutdown_write() + except Exception as exc: + detail = traceback.format_exc().strip() + message = f"远端命令启动失败: {exc}" + self._emit_progress(progress_callback, worker_id, action, "ssh_exec", message, status="failed") + logger.error( + "[WorkerAction] ssh_exec failed worker=%s action=%s target=%s err=%s", + worker_id, + action, + target, + self._preview_text(detail), + ) + return { + "ok": False, + "worker_id": worker_id, + "action": action, + "stage": "ssh_exec", + "message": message, + "stderr": detail, + "command": command_text, + } + stdin.close() + self._emit_progress(progress_callback, worker_id, action, "ssh_exec", "命令已启动，等待远端反馈", status="success") + channel = stdout.channel + deadline = time.time() + float(timeout or self.ssh_action_timeout) + stdout_chunks = bytearray() + stderr_chunks = bytearray() + while True: + while channel.recv_ready(): + stdout_chunks.extend(channel.recv(4096)) + while channel.recv_stderr_ready(): + stderr_chunks.extend(channel.recv_stderr(4096)) + if channel.exit_status_ready(): + break + if time.time() > deadline: + stdout_text = self._decode_remote_output(stdout_chunks).strip() + stderr_text = self._decode_remote_output(stderr_chunks).strip() + message = f"{action} 执行超时" + self._emit_progress(progress_callback, worker_id, action, "wait_remote_exit", message, status="failed") + logger.error( + "[WorkerAction] timeout worker=%s action=%s target=%s stdout=%s stderr=%s", + worker_id, + action, + target, + self._preview_text(stdout_text), + self._preview_text(stderr_text), + ) + return { + "ok": False, + "worker_id": worker_id, + "action": action, + "stage": "wait_remote_exit", + "message": f"{action} timed out", + "stdout": stdout_text, + "stderr": stderr_text, + "command": command_text, + } + time.sleep(0.1) + + while channel.recv_ready(): + stdout_chunks.extend(channel.recv(4096)) + while channel.recv_stderr_ready(): + stderr_chunks.extend(channel.recv_stderr(4096)) + exit_status = channel.recv_exit_status() + stdout_text = self._decode_remote_output(stdout_chunks).strip() + stderr_text = self._decode_remote_output(stderr_chunks).strip() + result = self._format_result(worker_id, worker, action, exit_status, stdout_text, stderr_text, command_text) + clean_stdout = result.get("stdout", "") + clean_stderr = result.get("stderr", "") + result_message = self._build_result_progress_message(result, exit_status, clean_stdout, clean_stderr) + self._emit_progress( + progress_callback, + worker_id, + action, + "remote_result", + result_message, + status="success" if result.get("ok") else "failed", + ) + log_method = logger.info if result.get("ok") else logger.error + log_method( + "[WorkerAction] finish worker=%s action=%s target=%s ok=%s failed_step=%s stage=%s msg=%s stdout=%s stderr=%s", + worker_id, + action, + target, + result.get("ok"), + result.get("failed_step", ""), + result.get("stage", ""), + self._preview_text(result.get("message", "")), + self._preview_text(clean_stdout), + self._preview_text(clean_stderr), + ) + return result + finally: + ssh.close() + + @staticmethod + def _preview_text(value: Any, limit: int = 240) -> str: + text = str(value or "").replace("\r", " ").replace("\n", " | ").strip() + if len(text) <= limit: + return text + return text[:limit] + "..." + + @staticmethod + def _decode_remote_output(payload: Any) -> str: + data = bytes(payload or b"") + if not data: + return "" + if data.startswith((codecs.BOM_UTF16_LE, codecs.BOM_UTF16_BE)): + for encoding in ("utf-16", "utf-16le", "utf-16be"): + try: + return data.decode(encoding) + except UnicodeDecodeError: + continue + if data.count(b"\x00") * 5 >= len(data): + for encoding in ("utf-16le", "utf-16be"): + try: + return data.decode(encoding) + except UnicodeDecodeError: + continue + for encoding in ("utf-8-sig", "utf-8", "gb18030", "cp936", "cp1252"): + try: + return data.decode(encoding) + except UnicodeDecodeError: + continue + return data.decode("utf-8", errors="replace") + + @classmethod + def _build_result_progress_message( + cls, + result: Dict[str, Any], + exit_status: int, + stdout_text: str, + stderr_text: str, + ) -> str: + parts = [f"exit={exit_status}"] + message = cls._preview_text(result.get("message", ""), limit=120) + stdout_preview = cls._preview_text(stdout_text, limit=120) + stderr_preview = cls._preview_text(stderr_text, limit=120) + if message: + parts.append(f"msg={message}") + if stdout_preview: + parts.append(f"stdout={stdout_preview}") + if stderr_preview: + parts.append(f"stderr={stderr_preview}") + if len(parts) == 1: + parts.append("msg=empty") + return " | ".join(parts) + + @staticmethod + def _emit_progress( + progress_callback: Optional[Callable[[str, Dict[str, Any]], None]], + worker_id: str, + action: str, + step: str, + message: str, + **extra: Any, + ) -> None: + if not progress_callback: + return + payload = {"progress": True, "action": action, "step": step, "message": message, **extra} + progress_callback(worker_id, payload) + + def _format_result( + self, + worker_id: str, + worker: Dict[str, Any], + action: str, + exit_status: int, + stdout: str, + stderr: str, + command: str, + ) -> Dict[str, Any]: + del worker + clean_stdout = str(stdout or "").strip() + clean_stderr = str(stderr or "").strip() + default_message = f"{action} completed" if exit_status == 0 else f"{action} failed (exit={exit_status})" + if exit_status == 0: + # 成功时优先显示 stdout + message = clean_stdout or clean_stderr or default_message + else: + # 失败时优先显示 stderr（真正的报错），其次是 stdout（可能包含合并后的报错），最后是默认消息 + message = clean_stderr or clean_stdout or default_message + return { + "ok": exit_status == 0, + "worker_id": worker_id, + "action": action, + "message": message, + "command": command, + "stdout": clean_stdout, + "stderr": clean_stderr, + "raw_stdout": stdout, + "raw_stderr": stderr, + } diff --git a/requirements.txt b/requirements.txt new file mode 100755 index 0000000..72c9571 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,22 @@ +bcrypt==5.0.0 +blinker==1.9.0 +certifi==2026.2.25 +cffi==2.0.0 +charset-normalizer==3.4.7 +click==8.3.2 +cryptography==46.0.7 +flask==3.1.3 +idna==3.11 +invoke==3.0.3 +itsdangerous==2.2.0 +jinja2==3.1.6 +markupsafe==3.0.3 +paramiko==4.0.0 +pycparser==3.0 +PyYAML==6.0.3 +pynacl==1.6.2 +redis==5.0.1 +requests==2.33.1 +tqdm==4.67.3 +urllib3==2.6.3 +werkzeug==3.1.8 diff --git a/result_codes.py b/result_codes.py new file mode 100644 index 0000000..1bc9fe3 --- /dev/null +++ b/result_codes.py @@ -0,0 +1,331 @@ +""" +统一采集结果码模块 + +错误分类原则：按责任方划分 +- SUCCESS: 成功 +- INFRA_ERROR: 基础设施错误 (环境/设备问题) +- APP_ERROR: 应用错误 (应用本身的问题) +- BUSINESS_ERROR: 业务限制 (应用功能限制) +""" +from enum import IntEnum +from dataclasses import dataclass +from typing import Any, Dict, Optional, Tuple + + +class ResultCode(IntEnum): + """采集结果退出码""" + SUCCESS = 0 + ERROR_USER = 1 + ADB_ERROR = 2 + NETWORK_ERROR = 3 + SIMULATOR_ERROR = 4 + EXPLORATION_STUCK = 5 + APP_CRASH = 6 + ERROR_GENERAL = -1 + + +class ErrorCategory(IntEnum): + """错误大类 - 按责任方划分""" + SUCCESS = 0 + INFRA_ERROR = 1 + APP_ERROR = 2 + BUSINESS_ERROR = 3 + DOWNLOAD_ERROR = 4 + + +class InfraError(IntEnum): + """基础设施错误原因""" + ADB_ERROR = 1 + NETWORK_ERROR = 2 + EMULATOR_CRASH = 3 + EMULATOR_RECOVERY_FAILED = 4 + PCAP_START_FAILED = 5 + DOWNLOAD_FAILED = 6 + AIRTEST_INIT_FAILED = 7 + EMULATOR_START_FAILED = 8 + WORKER_TIMEOUT = 9 + + +class DownloadError(IntEnum): + """下载错误原因 - 细分""" + OTHER = 0 + REGION_RESTRICTED = 1 + ACCOUNT_BANNED = 2 + APP_NOT_FOUND = 3 + INSTALL_FAILED = 4 + DOWNLOAD_TIMEOUT = 5 + NETWORK_ERROR = 6 + SOURCE_UNAVAILABLE = 7 + ALL_SOURCES_FAILED = 8 + COUNTRY_LOCKED = 9 # legacy aggregated terminal code + INCOMPATIBLE = 10 + + +class AppError(IntEnum): + """应用错误原因""" + CRASH = 1 + LAUNCH_ERROR = 2 + DISCONTINUED = 3 + INCOMPATIBLE = 4 + NEED_UPDATE = 5 + ROOT_MODE_UNSUPPORTED = 6 + + +class BusinessError(IntEnum): + """业务限制原因""" + OTHER = 0 + LOGIN_FAILED = 1 + REGISTER_FAILED = 2 + PAYMENT_REQUIRED = 3 + REGION_RESTRICTED = 4 + INVITE_ONLY = 5 + VIRTUAL_PHONE_INVALID = 6 + VIRTUAL_IDENTITY_INVALID = 7 + VALIDATION_FAILED = 8 + PHYSICAL_ID_REQUIRED = 9 + NO_SCENARIO = 10 + NORMAL_NO_ACTION = 11 + + +INFRA_ERROR_DESC = { + InfraError.ADB_ERROR: "ADB 断联", + InfraError.NETWORK_ERROR: "网络异常", + InfraError.EMULATOR_CRASH: "模拟器崩溃", + InfraError.EMULATOR_RECOVERY_FAILED: "模拟器恢复失败", + InfraError.PCAP_START_FAILED: "PCAP 启动失败", + InfraError.DOWNLOAD_FAILED: "下载失败", + InfraError.AIRTEST_INIT_FAILED: "Airtest 初始化失败", + InfraError.EMULATOR_START_FAILED: "模拟器启动失败", + InfraError.WORKER_TIMEOUT: "Worker 心跳超时", +} + +DOWNLOAD_ERROR_DESC = { + DownloadError.OTHER: "下载失败", + DownloadError.REGION_RESTRICTED: "Google Play 单国家锁区", + DownloadError.ACCOUNT_BANNED: "账号被封禁", + DownloadError.APP_NOT_FOUND: "Google Play 单国家未找到应用", + DownloadError.INCOMPATIBLE: "Google Play 当前设备不兼容", + DownloadError.INSTALL_FAILED: "安装失败", + DownloadError.DOWNLOAD_TIMEOUT: "下载超时", + DownloadError.NETWORK_ERROR: "下载网络异常", + DownloadError.SOURCE_UNAVAILABLE: "本地导入源不可用", + DownloadError.ALL_SOURCES_FAILED: "所有下载源失败", + DownloadError.COUNTRY_LOCKED: "历史聚合下载错误（Google Play 终态且本地兜底失败）", +} + +APP_ERROR_DESC = { + AppError.CRASH: "应用闪退", + AppError.LAUNCH_ERROR: "启动异常", + AppError.DISCONTINUED: "应用停服", + AppError.INCOMPATIBLE: "应用不兼容", + AppError.NEED_UPDATE: "应用持续跳转谷歌商店", + AppError.ROOT_MODE_UNSUPPORTED: "Root 环境下无法使用", +} + +BUSINESS_ERROR_DESC = { + BusinessError.OTHER: "其他原因", + BusinessError.LOGIN_FAILED: "登录失败", + BusinessError.REGISTER_FAILED: "注册失败", + BusinessError.PAYMENT_REQUIRED: "需付费", + BusinessError.REGION_RESTRICTED: "地区限制", + BusinessError.INVITE_ONLY: "需邀请码", + BusinessError.VIRTUAL_PHONE_INVALID: "虚拟手机号无效", + BusinessError.VIRTUAL_IDENTITY_INVALID: "虚拟身份无效", + BusinessError.VALIDATION_FAILED: "校验失败", + BusinessError.PHYSICAL_ID_REQUIRED: "需实体证件", + BusinessError.NO_SCENARIO: "无登录注册场景", + BusinessError.NORMAL_NO_ACTION: "测试正常无新动作", +} + + +STUCK_TO_ERROR: Dict[int, Tuple[ErrorCategory, int, str]] = { + 0: (ErrorCategory.BUSINESS_ERROR, BusinessError.OTHER, "其他原因"), + 1: (ErrorCategory.BUSINESS_ERROR, BusinessError.LOGIN_FAILED, "登录注册需人工辅助"), + 2: (ErrorCategory.APP_ERROR, AppError.LAUNCH_ERROR, "启动异常"), + 3: (ErrorCategory.SUCCESS, 0, "测试正常无新动作"), + 4: (ErrorCategory.BUSINESS_ERROR, BusinessError.REGION_RESTRICTED, "地区限制"), + 5: (ErrorCategory.BUSINESS_ERROR, BusinessError.PAYMENT_REQUIRED, "需付费"), + 6: (ErrorCategory.BUSINESS_ERROR, BusinessError.VIRTUAL_PHONE_INVALID, "虚拟手机号无效"), + 7: (ErrorCategory.BUSINESS_ERROR, BusinessError.VIRTUAL_IDENTITY_INVALID, "虚拟身份无效"), + 8: (ErrorCategory.BUSINESS_ERROR, BusinessError.VALIDATION_FAILED, "注册校验失败"), + 9: (ErrorCategory.BUSINESS_ERROR, BusinessError.INVITE_ONLY, "非开放注册"), + 10: (ErrorCategory.APP_ERROR, AppError.DISCONTINUED, "应用停服"), + 11: (ErrorCategory.BUSINESS_ERROR, BusinessError.PHYSICAL_ID_REQUIRED, "需实体证件"), + 12: (ErrorCategory.APP_ERROR, AppError.ROOT_MODE_UNSUPPORTED, "Root 环境下无法使用"), +} + + +NO_RETRY_ERRORS = ( + {(ErrorCategory.APP_ERROR, error) for error in AppError} + | { + (ErrorCategory.BUSINESS_ERROR, error) + for error in BusinessError + if error not in {BusinessError.NO_SCENARIO, BusinessError.NORMAL_NO_ACTION} + } +) +DOWNLOAD_TERMINAL_GOOGLE_PLAY_ERRORS = { + int(DownloadError.REGION_RESTRICTED), + int(DownloadError.APP_NOT_FOUND), + int(DownloadError.INCOMPATIBLE), +} +def _normalize_error_category(category: Any) -> Optional[ErrorCategory]: + if isinstance(category, ErrorCategory): + return category + if isinstance(category, str): + try: + return ErrorCategory[category] + except KeyError: + return None + try: + return ErrorCategory(int(category)) + except (TypeError, ValueError): + return None + + +def _normalize_error_code(code: Any) -> Optional[int]: + try: + return int(code) + except (TypeError, ValueError): + return None + + +def _normalize_download_errors(download_errors: Any) -> Dict[str, Dict[str, Any]]: + return download_errors if isinstance(download_errors, dict) else {} + + +def _is_terminal_download_combo(error_code: Any, download_errors: Any = None) -> bool: + normalized_code = _normalize_error_code(error_code) + if normalized_code == int(DownloadError.COUNTRY_LOCKED): + return True + + errors = _normalize_download_errors(download_errors) + if not errors: + return False + + google_play_codes = [] + local_codes = [] + for source, detail in errors.items(): + if not isinstance(detail, dict): + return False + detail_code = _normalize_error_code(detail.get("code")) + if detail_code is None: + return False + source_name = str(source or "").strip().lower() + if source_name.startswith("google_play:"): + google_play_codes.append(detail_code) + elif source_name == "local": + local_codes.append(detail_code) + + if not google_play_codes or not local_codes: + return False + if any(code not in DOWNLOAD_TERMINAL_GOOGLE_PLAY_ERRORS for code in google_play_codes): + return False + return True + + +def is_no_retry_error(category: Any, code: Any, download_errors: Any = None) -> bool: + normalized_category = _normalize_error_category(category) + normalized_code = _normalize_error_code(code) + if normalized_category is None or normalized_code is None: + return False + if (normalized_category, normalized_code) in NO_RETRY_ERRORS: + return True + if normalized_category == ErrorCategory.DOWNLOAD_ERROR: + return _is_terminal_download_combo(normalized_code, download_errors=download_errors) + return False + + +@dataclass +class ErrorInfo: + """标准化错误信息""" + category: ErrorCategory + code: int + reason: str + details: Optional[str] = None + + def to_report_dict(self) -> dict: + """转换为上报格式""" + result = { + "category": self.category.name, + "code": self.code, + "reason": self.reason, + } + if self.details: + result["details"] = self.details + return result + + @classmethod + def success(cls) -> "ErrorInfo": + return cls(ErrorCategory.SUCCESS, 0, "成功") + + @classmethod + def from_stuck_reason(cls, stuck_code: int, details: str = None) -> "ErrorInfo": + """从 stuck_reason_code 创建 ErrorInfo""" + category, code, reason = STUCK_TO_ERROR.get( + stuck_code, + (ErrorCategory.BUSINESS_ERROR, BusinessError.OTHER, "其他原因") + ) + return cls(category, code, reason, details) + + @classmethod + def infra(cls, error: InfraError, details: str = None) -> "ErrorInfo": + """创建基础设施错误""" + return cls( + ErrorCategory.INFRA_ERROR, + error, + INFRA_ERROR_DESC[error], + details + ) + + @classmethod + def app(cls, error: AppError, details: str = None) -> "ErrorInfo": + """创建应用错误""" + return cls( + ErrorCategory.APP_ERROR, + error, + APP_ERROR_DESC[error], + details + ) + + @classmethod + def business(cls, error: BusinessError, details: str = None) -> "ErrorInfo": + """创建业务限制错误""" + return cls( + ErrorCategory.BUSINESS_ERROR, + error, + BUSINESS_ERROR_DESC[error], + details + ) + + @classmethod + def download(cls, error: "DownloadError", details: str = None) -> "ErrorInfo": + """创建下载错误""" + return cls( + ErrorCategory.DOWNLOAD_ERROR, + error, + DOWNLOAD_ERROR_DESC.get(error, "下载失败"), + details + ) + + @classmethod + def from_report_dict(cls, data: dict) -> "ErrorInfo": + """从上报字典解析""" + if not data: + return cls(ErrorCategory.SUCCESS, 0, "成功") + + category_str = data.get('category', 'SUCCESS') + try: + category = ErrorCategory[category_str] + except KeyError: + category = ErrorCategory.SUCCESS + + code = data.get('code', 0) + reason = data.get('reason', '') + details = data.get('details') + + return cls(category, code, reason, details) + + def is_no_retry(self, download_errors: Any = None) -> bool: + """判断是否为不可重试错误""" + return is_no_retry_error(self.category, self.code, download_errors=download_errors) + diff --git a/schema_adapter.py b/schema_adapter.py new file mode 100644 index 0000000..70a098d --- /dev/null +++ b/schema_adapter.py @@ -0,0 +1,208 @@ +#!/usr/bin/env python3 +"""新数据库 schema 轻量访问适配器。""" + +import json +import time +from typing import Dict, List, Any, Optional + + +class SchemaAdapter: + """只访问 app_catalog / collection_task，不做旧表兼容。""" + + def __init__(self, db_connection, use_new_schema: bool = True): + self.conn = db_connection + self.use_new = use_new_schema + + def _table_exists(self, table_name: str) -> bool: + """检查表是否存在""" + result = self.conn.execute( + "SELECT name FROM sqlite_master WHERE type = 'table' AND name=?", + (table_name,) + ).fetchone() + return result is not None + + def get_collection_row(self, package_name: str) -> Optional[Dict]: + """ + 获取应用的完整信息（元数据 + 最新执行状态） + + 直接读取 app_catalog，并按 collection_task 最新记录补充状态。 + """ + if not (self.use_new and self._table_exists('app_catalog') and self._table_exists('collection_task')): + return None + catalog = self.conn.execute( + "SELECT * FROM app_catalog WHERE package_name = ?", + (package_name,) + ).fetchone() + if not catalog: + return None + task = self.conn.execute(""" + SELECT *, + CASE WHEN total_traffic_bytes > 0 + THEN self_traffic_bytes * 100.0 / total_traffic_bytes + ELSE 0 END AS self_ratio + FROM collection_task + WHERE package_name = ? + ORDER BY + datetime(COALESCE(NULLIF(completed_at, ''), NULLIF(started_at, ''), created_at)) DESC, + attempt DESC, + batch_tag DESC, + run_kind DESC + LIMIT 1 + """, (package_name,)).fetchone() + result = dict(catalog) + if task: + task_dict = dict(task) + result.update({ + "latest_status": task_dict.get("execution_status"), + "error_category": task_dict.get("error_category"), + "error_code": task_dict.get("error_code"), + "num_nodes": task_dict.get("num_nodes"), + "self_ratio": task_dict.get("self_ratio"), + "total_traffic_bytes": task_dict.get("total_traffic_bytes"), + "latest_test_time": task_dict.get("completed_at"), + }) + return result + + def upsert_catalog_entry(self, entry: Dict): + """ + 插入/更新应用目录条目（元数据） + + UPSERT INTO app_catalog + """ + now = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime()) + + if self.use_new and self._table_exists('app_catalog'): + self.conn.execute(""" + INSERT INTO app_catalog ( + package_name, app_name, app_magic_label, downloads, source_order, + last_updated, country_code, device_type, task_payload_json, + last_update_interval_days, batch_tags, is_active, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(package_name) DO UPDATE SET + app_name=excluded.app_name, + app_magic_label=COALESCE(NULLIF(excluded.app_magic_label, ''), app_catalog.app_magic_label), + downloads=excluded.downloads, + source_order=excluded.source_order, + last_updated=COALESCE(NULLIF(excluded.last_updated, ''), app_catalog.last_updated), + country_code=excluded.country_code, + device_type=excluded.device_type, + task_payload_json=excluded.task_payload_json, + last_update_interval_days=excluded.last_update_interval_days, + is_active=excluded.is_active, + updated_at=excluded.updated_at + """, ( + entry['package_name'], entry.get('app_name'), entry.get('app_magic_label'), + entry.get('downloads'), entry.get('source_order', 0), + entry.get('last_updated', ''), entry.get('country_code', ''), + entry.get('device_type', ''), json.dumps(entry.get('task_payload', {}), ensure_ascii=False), + entry.get('last_update_interval_days', 0), + json.dumps(entry.get('batch_tags', [])), + entry.get('is_active', 1), + now + )) + + def insert_collection_task(self, task: Dict): + """ + 插入任务执行记录 + + INSERT INTO collection_task + """ + now_iso = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime()) + + if self.use_new and self._table_exists('collection_task'): + self.conn.execute(""" + INSERT INTO collection_task ( + package_name, batch_tag, run_kind, attempt, + execution_status, error_category, error_code, error_reason, + worker_id, completed_at, duration_seconds, + total_traffic_bytes, self_traffic_bytes, server_traffic_bytes, + num_nodes, droidbot_steps, gui_agent_steps, + created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(package_name, batch_tag, run_kind, attempt) DO UPDATE SET + execution_status=excluded.execution_status, + error_category=excluded.error_category, + error_code=excluded.error_code, + error_reason=excluded.error_reason, + worker_id=excluded.worker_id, + completed_at=excluded.completed_at, + duration_seconds=excluded.duration_seconds, + total_traffic_bytes=excluded.total_traffic_bytes, + self_traffic_bytes=excluded.self_traffic_bytes, + server_traffic_bytes=excluded.server_traffic_bytes, + num_nodes=excluded.num_nodes + """, ( + task['package_name'], task.get('batch_tag', 'unknown'), + task.get('run_kind', 'ranking'), task.get('attempt', 1), + task.get('execution_status'), task.get('error_category'), + task.get('error_code'), task.get('error_reason'), + task.get('worker_id'), now_iso, task.get('duration_seconds'), + task.get('total_traffic_bytes', 0), task.get('self_traffic_bytes', 0), + task.get('server_traffic_bytes', 0), task.get('num_nodes', 0), + task.get('droidbot_steps', 0), task.get('gui_agent_steps', 0), + now_iso + )) + + def list_pending_tasks(self, limit: int = 1000) -> List[Dict]: + """ + 列出待处理任务 + + 返回 active catalog 中没有成功最新记录的任务。 + """ + if self.use_new and self._table_exists('app_catalog') and self._table_exists('collection_task'): + rows = self.conn.execute(""" + SELECT ac.package_name, ac.app_name, ac.downloads, ac.source_order + FROM app_catalog ac + WHERE COALESCE(ac.is_active, 1) = 1 + AND NOT EXISTS ( + SELECT 1 + FROM collection_task ct + WHERE ct.package_name = ac.package_name + AND ct.execution_status = 'success' + ) + ORDER BY + CASE WHEN ac.source_order IS NULL THEN 1 ELSE 0 END ASC, + ac.source_order ASC, + ac.package_name ASC + LIMIT ? + """, (limit,)).fetchall() + return [dict(r) for r in rows] + + return [] + + +# 使用示例 +if __name__ == '__main__': + import sqlite3 + import os + + USE_NEW_SCHEMA = os.getenv('USE_NEW_SCHEMA', 'true').lower() == 'true' + + conn = sqlite3.connect('runtime/main/monitoring.sqlite3') + conn.row_factory = sqlite3.Row + + adapter = SchemaAdapter(conn, use_new_schema=USE_NEW_SCHEMA) + + # 测试查询 + row = adapter.get_collection_row('com.example.test') + print(f"查询结果: {row}") + + # 测试写入 + adapter.upsert_catalog_entry({ + 'package_name': 'com.test.app', + 'app_name': 'Test App', + 'downloads': 10000, + 'is_active': 1, + }) + + adapter.insert_collection_task({ + 'package_name': 'com.test.app', + 'batch_tag': '2026-6-15', + 'execution_status': 'success', + 'total_traffic_bytes': 5000000, + }) + + conn.commit() + conn.close() + + print("✅ 适配器测试通过") diff --git a/schema_new.sql b/schema_new.sql new file mode 100644 index 0000000..7995df8 --- /dev/null +++ b/schema_new.sql @@ -0,0 +1,195 @@ +-- ============================================================================ +-- 新数据库 Schema 定义 (v4 重构) +-- ============================================================================ +-- 设计理念： +-- 1. 保留完整的原始数据（worker返回的所有字段） +-- 2. collection_task 既是任务列表也是执行记录，用 attempt 追踪重试 +-- 3. 不做预计算，traffic 文件路径存为 JSON 数组，按需重算流量 +-- 4. 运行时代码直接读取新表并在查询层按需计算分析结果 +-- 5. 不存经过判断的受限状态，只存 worker 原始 success/failed +-- +-- 运行库只创建基础表和索引，不创建兼容视图/投影 +-- ============================================================================ + +PRAGMA journal_mode = WAL; +PRAGMA foreign_keys = OFF; + +-- ============================================================================ +-- 表 1: collection_task — 采集任务表（核心） +-- 既是任务列表也是执行记录。复合主键 (package_name, batch_tag, run_kind, attempt) +-- ============================================================================ +CREATE TABLE IF NOT EXISTS collection_task ( + -- 复合主键 + package_name TEXT NOT NULL, + batch_tag TEXT NOT NULL, -- 批次标签，如 '2026-6-4' + run_kind TEXT NOT NULL DEFAULT 'ranking' + CHECK(run_kind IN ('ranking', 'block', 'model', 'manual')), + attempt INTEGER NOT NULL DEFAULT 1,-- 采集轮次（重试自增） + + -- 任务基本信息 + task_key TEXT, -- 任务唯一 key + app_name TEXT, -- 应用显示名称 + app_magic_label TEXT, -- 应用魔法标签（关联同一应用的不同变体） + + -- 任务类型标识 + is_new_app INTEGER DEFAULT 0, -- ranking任务: 新增应用(1) / 更新应用(0) + + -- 任务状态 + task_status TEXT DEFAULT 'pending', -- 'pending' / 'running' / 'completed' + execution_status TEXT, -- worker原始状态: 'success'/'failed'/'stop' + + -- Worker信息 + worker_id TEXT, + + -- 时间信息 + created_at TEXT DEFAULT (datetime('now', 'localtime')), + started_at TEXT, -- 任务开始时间 + completed_at TEXT, -- 任务完成时间 + + -- 时长统计（worker返回） + duration_seconds REAL, + download_duration_seconds REAL, + execution_duration_seconds REAL, + analysis_duration_seconds REAL, + + -- 错误信息（worker返回的原始错误） + error_category TEXT, -- 'INFRA_ERROR'/'DOWNLOAD_ERROR'/'APP_ERROR'/'BUSINESS_ERROR' + error_code INTEGER, -- 错误代码 + error_reason TEXT, -- 错误原因 + error_details TEXT, -- 详细错误信息 + crashed_source TEXT, -- 导致崩溃的下载源 + + -- 执行统计（worker返回的原始字段） + droidbot_steps INTEGER DEFAULT 0, + gui_agent_steps INTEGER DEFAULT 0, + total_steps INTEGER DEFAULT 0, + num_nodes INTEGER DEFAULT 0, -- UI节点数 + num_reached_activities INTEGER DEFAULT 0, -- 到达的activity数量 + app_num_total_activities INTEGER DEFAULT 0, -- 应用总activity数量 + + -- 流量统计（worker返回的原始数据） + total_traffic_bytes INTEGER DEFAULT 0, + self_traffic_bytes INTEGER DEFAULT 0, + server_traffic_bytes INTEGER DEFAULT 0, + unrecognized_traffic_bytes INTEGER DEFAULT 0, + model_flow_count INTEGER DEFAULT 0, + model_traffic_bytes INTEGER DEFAULT 0, + + -- Metrics（worker返回的额外指标） + login_count INTEGER DEFAULT 0, + register_count INTEGER DEFAULT 0, + stuck_reason_code INTEGER, + guiagent_message TEXT, + scenario_triggered INTEGER DEFAULT 0, + + -- 下载信息 + download_source TEXT, -- 'google_play' / 'local' + is_retry INTEGER DEFAULT 0, + + -- 执行追踪 + exit_code INTEGER, + trace_json TEXT, -- 执行轨迹 JSON + + -- 流量文件路径列表（JSON数组） + traffic_file_paths TEXT, -- ['\\\\lfs.../pkg/...', ...] + + PRIMARY KEY (package_name, batch_tag, run_kind, attempt) +); + +CREATE INDEX IF NOT EXISTS idx_collection_task_status + ON collection_task(task_status, run_kind, batch_tag); +CREATE INDEX IF NOT EXISTS idx_collection_task_execution + ON collection_task(execution_status, run_kind); +CREATE INDEX IF NOT EXISTS idx_collection_task_pkg + ON collection_task(package_name, run_kind); +CREATE INDEX IF NOT EXISTS idx_collection_task_time + ON collection_task(created_at); +CREATE INDEX IF NOT EXISTS idx_collection_task_worker + ON collection_task(worker_id, completed_at); +CREATE INDEX IF NOT EXISTS idx_collection_task_magic + ON collection_task(app_magic_label); +CREATE INDEX IF NOT EXISTS idx_collection_task_batch + ON collection_task(batch_tag, run_kind); + +-- ============================================================================ +-- 表 2: app_catalog — 应用目录元数据 +-- 应用的分类、标签、优先级等管理信息（与采集结果分离） +-- ============================================================================ +CREATE TABLE IF NOT EXISTS app_catalog ( + package_name TEXT PRIMARY KEY, + app_name TEXT, + + -- 分类管理 + batch_tags TEXT, -- JSON数组：完整的批次标签列表 + app_magic_label TEXT, -- 应用魔法标签 + last_updated TEXT, -- 应用商店最后更新时间 + country_code TEXT, -- 采集国家/地区 + device_type TEXT, -- emulator / physical / any + task_payload_json TEXT, -- 下发给 worker 的完整任务 payload + last_update_interval_days INTEGER DEFAULT 0, -- 版本更新时间间隔 + + -- 优先级管理 + task_queue TEXT DEFAULT 'default', -- 'default' / 'priority' + task_priority INTEGER DEFAULT 50, + source_order INTEGER, -- 在源列表中的顺序 + + -- 状态标记 + is_active INTEGER DEFAULT 1, -- 是否激活 + is_blocked INTEGER DEFAULT 0, -- 是否被屏蔽 + + -- 应用信息 + category TEXT, -- 应用分类 + downloads INTEGER, -- 下载量 + + created_at TEXT DEFAULT (datetime('now', 'localtime')), + updated_at TEXT DEFAULT (datetime('now', 'localtime')) +); + +CREATE INDEX IF NOT EXISTS idx_app_catalog_priority + ON app_catalog(task_priority DESC); +CREATE INDEX IF NOT EXISTS idx_app_catalog_magic + ON app_catalog(app_magic_label); +CREATE INDEX IF NOT EXISTS idx_app_catalog_active + ON app_catalog(is_active, source_order); + +-- ============================================================================ +-- 表 3: apk_registry — APK文件注册表 +-- 结构与原系统完全一致（由 apk_cloud/registry.py 维护），此处仅用于迁移数据落地 +-- ============================================================================ +CREATE TABLE IF NOT EXISTS apk_registry ( + package_name TEXT PRIMARY KEY, + download_date TEXT NOT NULL, + download_time REAL NOT NULL, + local_dir TEXT, + smb_dir TEXT, + version_name TEXT, + source TEXT, + apk_files_json TEXT, + updated_at REAL NOT NULL +); + +-- ============================================================================ +-- 表 4: worker_activity_summary — Worker状态时间分布（每小时汇总） +-- ============================================================================ +CREATE TABLE IF NOT EXISTS worker_activity_summary ( + worker_id TEXT NOT NULL, + stat_date TEXT NOT NULL, -- YYYY-MM-DD + stat_hour INTEGER NOT NULL, -- 0-23（-1 表示全天汇总） + + -- 状态时长统计（秒） + idle_duration_seconds INTEGER DEFAULT 0, + busy_duration_seconds INTEGER DEFAULT 0, + offline_duration_seconds INTEGER DEFAULT 0, + + -- 任务统计 + task_count INTEGER DEFAULT 0, + success_count INTEGER DEFAULT 0, + failed_count INTEGER DEFAULT 0, + + PRIMARY KEY (worker_id, stat_date, stat_hour) +); + +CREATE INDEX IF NOT EXISTS idx_worker_activity_date + ON worker_activity_summary(stat_date DESC); + +-- 运行时代码直接查询真实表，不创建投影或视图。 diff --git a/scripts/add_workers.py b/scripts/add_workers.py new file mode 100644 index 0000000..f509704 --- /dev/null +++ b/scripts/add_workers.py @@ -0,0 +1,499 @@ +#!/usr/bin/env python3 +# -*- encoding=utf8 -*- +""" +Worker 批量添加脚本 + +扫描指定 IP 范围内的设备，通过 SSH 获取 MAC 地址， +自动生成 Worker 条目并合并写入 config/config.json 的 WORKER_INVENTORY。 + +使用示例: + # 扫描 192.168.1.10 到 192.168.1.50 + python scripts/add_workers.py --range 192.168.1.10-50 + + # 扫描多个范围 + python scripts/add_workers.py --range 192.168.1.10-50 192.168.2.50-80 + + # 指定标签 + python scripts/add_workers.py --range 192.168.1.10-50 --tags group_b + + # 仅扫描，不写入 + python scripts/add_workers.py --range 192.168.1.10-50 --dry-run + + # 跳过连通性检测 + python scripts/add_workers.py --range 192.168.1.10-50 --skip-ping +""" + +import argparse +import json +import os +import re +import socket +import sys +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import Dict, List, Optional, Tuple + +# 项目根目录 +BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, BASE_DIR) + +from config import ( + CONFIG_PATH, + SSH_DEFAULT_USER, + SSH_DEFAULT_PASSWORD, + SSH_DEFAULT_PORT, + WORKER_INVENTORY_CONFIG_KEY, + WORKER_INVENTORY_PATH, +) + +# 默认 Worker 配置 +DEFAULT_REPO_DIR = "D:/autool" +DEFAULT_PYTHON_EXE = "python" + +# 并发设置 +MAX_PING_WORKERS = 50 +MAX_SSH_WORKERS = 10 +PING_TIMEOUT = 2 # 秒 +SSH_TIMEOUT = 10 # 秒 + + +def parse_ip_range(range_str: str) -> List[str]: + """ + 解析 IP 范围字符串，支持以下格式： + - 单个 IP: 192.168.1.10 + - 末段范围: 192.168.1.10-50 + - CIDR: 192.168.1.0/24 + """ + range_str = range_str.strip() + + # CIDR 格式: 192.168.1.0/24 + cidr_match = re.fullmatch(r"(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})/(\d{1,2})", range_str) + if cidr_match: + base_ip = cidr_match.group(1) + prefix_len = int(cidr_match.group(2)) + if prefix_len < 24 or prefix_len > 32: + raise ValueError(f"仅支持 /24 ~ /32 的 CIDR 范围: {range_str}") + parts = base_ip.split(".") + base = int(parts[3]) + host_bits = 32 - prefix_len + count = 2 ** host_bits + # 排除网络地址和广播地址 + start = max(base, 1) + end = min(base + count - 1, 254) + return [f"{parts[0]}.{parts[1]}.{parts[2]}.{i}" for i in range(start, end + 1)] + + # 末段范围格式: 192.168.1.10-50 + range_match = re.fullmatch(r"(\d{1,3}\.\d{1,3}\.\d{1,3})\.(\d{1,3})-(\d{1,3})", range_str) + if range_match: + prefix = range_match.group(1) + start = int(range_match.group(2)) + end = int(range_match.group(3)) + if start > end: + raise ValueError(f"起始地址大于结束地址: {range_str}") + if end > 254: + raise ValueError(f"IP 末段超出范围 (最大 254): {range_str}") + return [f"{prefix}.{i}" for i in range(start, end + 1)] + + # 单个 IP: 192.168.1.10 + ip_match = re.fullmatch(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", range_str) + if ip_match: + return [range_str] + + raise ValueError(f"无法解析 IP 范围格式: {range_str}") + + +def check_ssh_port(ip: str, port: int = SSH_DEFAULT_PORT, timeout: float = PING_TIMEOUT) -> bool: + """尝试通过 TCP 连接 SSH 端口检查主机是否在线""" + try: + with socket.create_connection((ip, port), timeout=timeout): + return True + except (socket.timeout, ConnectionRefusedError, OSError): + return False + + +def scan_reachable_hosts(ip_list: List[str], port: int = SSH_DEFAULT_PORT, skip_ping: bool = False) -> List[str]: + """并发端口扫描，返回 SSH 端口可达的主机列表""" + if skip_ping: + print(f" 跳过连通性检测，将直接尝试 SSH 登录 {len(ip_list)} 台设备") + return ip_list + + print(f" 正在扫描 {len(ip_list)} 个 IP 地址的 SSH 端口...") + reachable = [] + + with ThreadPoolExecutor(max_workers=MAX_PING_WORKERS) as executor: + futures = {executor.submit(check_ssh_port, ip, port): ip for ip in ip_list} + for future in as_completed(futures): + ip = futures[future] + try: + if future.result(): + reachable.append(ip) + print(f" ✓ {ip} 在线 (端口可达)") + else: + print(f" ✗ {ip} 不可达 (端口关闭或超时)") + except Exception as e: + print(f" ✗ {ip} 检测异常: {e}") + + # 按 IP 排序 + reachable.sort(key=lambda x: tuple(int(p) for p in x.split("."))) + print(f" 扫描完成，{len(reachable)} 台设备 SSH 端口可连接\n") + return reachable + + +def get_mac_via_ssh( + ip: str, + user: str = SSH_DEFAULT_USER, + password: str = SSH_DEFAULT_PASSWORD, + port: int = SSH_DEFAULT_PORT, +) -> Optional[str]: + """ + 通过 SSH 连接 Windows 设备并获取物理网卡 MAC 地址。 + 使用 getmac /fo csv /nh 获取 MAC 列表，选取以太网适配器的 MAC。 + """ + try: + import paramiko + except ImportError: + print(" [错误] 需要 paramiko 库，请运行: pip install paramiko") + sys.exit(1) + + ssh = paramiko.SSHClient() + ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + + try: + ssh.connect( + hostname=ip, + port=port, + username=user, + password=password, + timeout=SSH_TIMEOUT, + auth_timeout=SSH_TIMEOUT, + banner_timeout=SSH_TIMEOUT, + look_for_keys=True, + allow_agent=True, + ) + + # 获取 MAC 地址列表 + _, stdout, _ = ssh.exec_command("getmac /fo csv /nh", timeout=SSH_TIMEOUT) + output = stdout.read().decode("utf-8", errors="ignore").strip() + + if not output: + return None + + # 解析 getmac 输出，格式: "MAC地址","传输名称","..." + # 优先选取以太网/有线网卡的 MAC，跳过虚拟网卡和断开的连接 + best_mac = None + for line in output.splitlines(): + line = line.strip() + if not line: + continue + # 去掉引号分割 + parts = [p.strip().strip('"') for p in line.split(",")] + if len(parts) < 2: + continue + mac = parts[0].strip() + transport = parts[1].strip() if len(parts) > 1 else "" + + # 跳过无效 MAC + if not re.match(r"([0-9A-Fa-f]{2}[:-]){5}[0-9A-Fa-f]{2}", mac): + continue + + # 跳过已断开的连接 + if "已断开" in transport or "Disconnected" in transport.lower() or "Media disconnected" in transport.lower(): + continue + + # 统一 MAC 格式为 XX:XX:XX:XX:XX:XX（大写、冒号分隔） + mac = mac.upper().replace("-", ":") + + # 优先选取硬件以太网适配器 + transport_lower = transport.lower() + if any(kw in transport_lower for kw in ["ethernet", "以太网", "realtek", "intel"]): + best_mac = mac + break + + # 记录第一个有效的 MAC 作为备选 + if best_mac is None: + best_mac = mac + + return best_mac + except Exception as e: + print(f" ✗ {ip} SSH 连接失败: {e}") + return None + finally: + ssh.close() + + +def collect_worker_info( + ip_list: List[str], + tags: List[str], + user: str = SSH_DEFAULT_USER, + password: str = SSH_DEFAULT_PASSWORD, + port: int = SSH_DEFAULT_PORT, +) -> List[Dict]: + """批量通过 SSH 收集设备信息，生成 worker 条目""" + print(f" 正在通过 SSH 获取 {len(ip_list)} 台设备的 MAC 地址...") + workers = [] + + def _collect_one(ip: str) -> Optional[Dict]: + mac = get_mac_via_ssh(ip, user=user, password=password, port=port) + if mac: + worker_id = f"{ip}_{mac}" + print(f" ✓ {ip} → MAC: {mac} → worker_id: {worker_id}") + return { + "worker_id": worker_id, + "ssh_target": ip, + "repo_dir": DEFAULT_REPO_DIR, + "python_exe": DEFAULT_PYTHON_EXE, + "tags": list(tags), + } + else: + print(f" ✗ {ip} 无法获取 MAC 地址") + return None + + with ThreadPoolExecutor(max_workers=MAX_SSH_WORKERS) as executor: + futures = {executor.submit(_collect_one, ip): ip for ip in ip_list} + for future in as_completed(futures): + result = future.result() + if result: + workers.append(result) + + # 按 IP 排序 + workers.sort(key=lambda w: tuple(int(p) for p in w["ssh_target"].split("."))) + print(f" 成功获取 {len(workers)} 台设备信息\n") + return workers + + +def load_existing_inventory(path: str) -> Tuple[List[Dict], set, set]: + """加载现有 Worker 列表，返回列表、已存在 worker_id 和已存在 IP 集合""" + if not os.path.exists(path): + return [], set(), set() + + with open(path, "r", encoding="utf-8") as f: + payload = json.load(f) + + data = payload.get(WORKER_INVENTORY_CONFIG_KEY, []) if isinstance(payload, dict) else payload + + if not isinstance(data, list): + print(f" [警告] {path} 格式异常（非数组），将创建新文件") + return [], set(), set() + + existing_ids = set() + # 同时收集已存在的 ssh_target（IP）用于去重 + existing_ips = set() + for item in data: + if isinstance(item, dict): + wid = item.get("worker_id", "") + if wid: + existing_ids.add(wid) + ip = item.get("ssh_target", "") + if ip: + existing_ips.add(ip) + + return data, existing_ids, existing_ips + + +def save_inventory(path: str, workers: List[Dict]) -> None: + """写入统一配置；如果传入旧数组文件路径，则按旧格式写回。""" + payload = {} + if os.path.exists(path): + with open(path, "r", encoding="utf-8") as f: + try: + payload = json.load(f) + except json.JSONDecodeError: + payload = {} + + if isinstance(payload, dict): + payload[WORKER_INVENTORY_CONFIG_KEY] = workers + with open(path, "w", encoding="utf-8") as f: + json.dump(payload, f, indent=2, ensure_ascii=False) + f.write("\n") + return + + with open(path, "w", encoding="utf-8") as f: + json.dump(workers, f, indent=2, ensure_ascii=False) + f.write("\n") + + +def merge_and_save( + existing: List[Dict], + existing_ids: set, + existing_ips: set, + new_workers: List[Dict], + path: str, + dry_run: bool = False, +) -> int: + """合并新旧 Worker 列表并写入文件""" + to_add = [] + skipped_dup_id = 0 + skipped_dup_ip = 0 + + for worker in new_workers: + wid = worker["worker_id"] + ip = worker["ssh_target"] + if wid in existing_ids: + print(f" [跳过] worker_id 已存在: {wid}") + skipped_dup_id += 1 + continue + if ip in existing_ips: + print(f" [跳过] IP 已存在于其他 worker: {ip}") + skipped_dup_ip += 1 + continue + to_add.append(worker) + + if not to_add: + print(" 没有新的 Worker 需要添加") + return 0 + + print(f"\n 将添加 {len(to_add)} 台新 Worker:") + for w in to_add: + tag_str = f" tags={w['tags']}" if w["tags"] else "" + print(f" + {w['worker_id']}{tag_str}") + + if skipped_dup_id > 0: + print(f" （已跳过 {skipped_dup_id} 台 worker_id 重复的设备）") + if skipped_dup_ip > 0: + print(f" （已跳过 {skipped_dup_ip} 台 IP 重复的设备）") + + if dry_run: + print("\n [Dry Run] 未写入文件") + return len(to_add) + + # 合并并写入 + merged = existing + to_add + # 备份原文件 + if os.path.exists(path): + backup_path = path + ".bak" + with open(path, "r", encoding="utf-8") as f: + backup_content = f.read() + with open(backup_path, "w", encoding="utf-8") as f: + f.write(backup_content) + print(f"\n 已备份原文件到: {backup_path}") + + save_inventory(path, merged) + + print(f" 已写入 {path}，共 {len(merged)} 台 Worker") + return len(to_add) + + +def main(): + parser = argparse.ArgumentParser( + description="Worker 批量添加工具 — 扫描 IP 范围并自动添加到 config/config.json", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +示例: + python scripts/add_workers.py --range 192.168.1.10-50 + python scripts/add_workers.py --range 192.168.1.10-50 192.168.2.50-80 + python scripts/add_workers.py --range 192.168.1.10-50 --tags group_b + python scripts/add_workers.py --range 192.168.1.10-50 --dry-run + python scripts/add_workers.py --range 192.168.1.10-50 --skip-ping + """, + ) + parser.add_argument( + "--range", + nargs="+", + required=True, + dest="ip_ranges", + help="IP 范围，支持格式: 192.168.1.10 / 192.168.1.10-50 / 192.168.1.0/24", + ) + parser.add_argument( + "--tags", + nargs="*", + default=[], + help="为新添加的 Worker 设置标签", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="仅扫描和预览，不写入文件", + ) + parser.add_argument( + "--skip-ping", + action="store_true", + help="跳过连通性检测，直接尝试 SSH 登录", + ) + parser.add_argument( + "--inventory", + default=WORKER_INVENTORY_PATH, + help=f"配置文件或旧 worker_inventory.json 路径 (默认: {CONFIG_PATH})", + ) + parser.add_argument( + "--user", + default=SSH_DEFAULT_USER, + help=f"SSH 用户名 (默认: {SSH_DEFAULT_USER})", + ) + parser.add_argument( + "--password", + default=SSH_DEFAULT_PASSWORD, + help="SSH 密码", + ) + parser.add_argument( + "--port", + type=int, + default=SSH_DEFAULT_PORT, + help=f"SSH 端口 (默认: {SSH_DEFAULT_PORT})", + ) + + args = parser.parse_args() + + print("=" * 60) + print("Worker 批量添加工具") + print("=" * 60) + + # 1. 解析 IP 范围 + all_ips = [] + for r in args.ip_ranges: + try: + ips = parse_ip_range(r) + print(f" 范围 {r} → {len(ips)} 个 IP") + all_ips.extend(ips) + except ValueError as e: + print(f" [错误] {e}") + sys.exit(1) + + # 去重并排序 + all_ips = sorted(set(all_ips), key=lambda x: tuple(int(p) for p in x.split("."))) + print(f"\n 共计 {len(all_ips)} 个唯一 IP 地址\n") + + if not all_ips: + print(" 没有有效的 IP 地址") + return + + # 2. 端口扫描 + print("[1/4] 连通性检测") + reachable = scan_reachable_hosts(all_ips, port=args.port, skip_ping=args.skip_ping) + + if not reachable: + print(" 没有在线的设备") + return + + # 3. SSH 获取 MAC 地址 + print("[2/4] 获取设备信息") + new_workers = collect_worker_info( + reachable, + tags=args.tags, + user=args.user, + password=args.password, + port=args.port, + ) + + if not new_workers: + print(" 未能获取任何设备信息") + return + + # 4. 加载现有 inventory + print("[3/4] 加载现有 Worker 列表") + existing, existing_ids, existing_ips = load_existing_inventory(args.inventory) + print(f" 当前共有 {len(existing)} 台 Worker\n") + + # 5. 合并写入 + print("[4/4] 合并并写入") + added = merge_and_save( + existing, existing_ids, existing_ips, new_workers, args.inventory, dry_run=args.dry_run, + ) + + print("\n" + "=" * 60) + print(f"完成！新增 {added} 台 Worker") + if args.dry_run: + print("（Dry Run 模式，未实际写入）") + print("=" * 60) + + +if __name__ == "__main__": + main() diff --git a/scripts/analyze_daily_quality.py b/scripts/analyze_daily_quality.py new file mode 100644 index 0000000..4a744a5 --- /dev/null +++ b/scripts/analyze_daily_quality.py @@ -0,0 +1,728 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +""" +Analyze a single day's app results and export one CSV report. + +Rules: +- success + light_restricted => qualified +- severe_restricted => failed +- failure breakdown shares follow the dashboard's monitoring view: + overall_share_percent = failure_duration / total_duration + error_share_percent = failure_duration / failed_duration +""" + +from __future__ import annotations + +import argparse +import csv +import os +import sys +import tempfile +from collections import Counter +from datetime import datetime +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional, Tuple + + +SCRIPT_DIR = Path(__file__).resolve().parent +PROJECT_ROOT = SCRIPT_DIR.parent +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +from analytics import ( # noqa: E402 + AnalyticsService, + _build_legacy_report_row, + _classify_restriction_status, + _iter_report_paths, + _load_task_csv_packages, + _parse_local_datetime, + _prefer_legacy_report_candidate, + _read_dict_rows, +) +from config import ( # noqa: E402 + ANALYTICS_TPDPI_APP_LIST, + ANALYTICS_TPDPI_URL_LIB, + ANALYTICS_TRAFFIC_ROOT, + FAILED_TASKS_CSV, + REPORT_DIR, + RETRY_TASKS_CSV, + SUCCESS_TASKS_CSV, + TASK_CSV_PATH, +) +from monitoring import ( # noqa: E402 + SHANGHAI_TZ, + _classify_failure, + _humanize_failure_domain, + _humanize_failure_reason, + _humanize_failure_subtype, +) + + +CSV_COLUMNS = [ + "record_type", + "date", + "generated_at", + "name", + "label", + "value", + "count", + "share_percent", + "overall_share_percent", + "error_share_percent", + "domain_share_percent", + "duration_seconds", + "avg_duration_seconds", + "qualified", + "app_name", + "package_name", + "report_type", + "latest_status", + "restriction_status", + "retryability", + "latest_worker_id", + "latest_test_time", + "latest_failure_type", + "latest_task_detail", + "failure_domain", + "failure_domain_label", + "failure_subtype", + "failure_subtype_label", + "failure_reason", + "num_nodes", + "unique_domain_count", + "unique_second_level_domain_count", + "total_traffic_bytes", + "self_traffic_bytes", + "server_traffic_bytes", + "unrecognized_traffic_bytes", + "self_ratio", + "recognition_ratio", + "artifact_status", +] + + +def _date_str_from_ts(ts: float) -> str: + return datetime.fromtimestamp(ts, SHANGHAI_TZ).strftime("%Y-%m-%d") + + +def _resolve_allowed_packages(task_csv_path: str) -> Optional[set]: + packages = _load_task_csv_packages(task_csv_path) + return {item for item in packages if item} or None + + +def _load_daily_latest_rows( + *, + date_str: str, + success_csv_path: str, + failed_csv_path: str, + retry_csv_path: str, + allowed_packages: Optional[set] = None, +) -> Dict[str, Dict[str, Any]]: + latest_rows: Dict[str, Dict[str, Any]] = {} + report_specs = [ + ("success", success_csv_path, True), + ("failed", failed_csv_path, False), + ("retry", retry_csv_path, False), + ] + for report_type, report_path, allow_glob in report_specs: + sequence = 0 + for current_path in _iter_report_paths(report_path, allow_glob=allow_glob): + for row in _read_dict_rows(current_path): + package_name = str(row.get("包名") or row.get("package_name") or "").strip() + if not package_name: + continue + if allowed_packages is not None and package_name not in allowed_packages: + continue + row_time = _parse_local_datetime(row.get("时间") or row.get("test_time") or "") + if not row_time or _date_str_from_ts(row_time) != date_str: + continue + sequence += 1 + sort_key = (row_time, sequence) + current_row = latest_rows.get(package_name) + if not _prefer_legacy_report_candidate(current_row, report_type, sort_key): + continue + latest_rows[package_name] = { + "_sort_key": sort_key, + "_report_type": report_type, + "package_name": package_name, + "report_type": report_type, + **_build_legacy_report_row(row, report_type, row_time), + } + for row in latest_rows.values(): + row.pop("_sort_key", None) + row.pop("_report_type", None) + return latest_rows + + +def _build_failure_breakdown(items: List[Dict[str, Any]]) -> Dict[str, Any]: + total_duration = sum(float(item.get("duration_seconds") or 0.0) for item in items) + failed_items = [item for item in items if not item["qualified"]] + failed_duration = sum(float(item.get("duration_seconds") or 0.0) for item in failed_items) + + domain_buckets: Dict[str, Dict[str, Any]] = {} + for item in failed_items: + failure_domain = item["failure_domain"] + failure_subtype = item["failure_subtype"] + duration_seconds = float(item.get("duration_seconds") or 0.0) + reason_label = item["failure_reason"] + message_label = item["latest_task_detail"] or reason_label + sample = { + "app_name": item["app_name"], + "package_name": item["package_name"], + "latest_status": item["latest_status"], + "restriction_status": item["restriction_status"], + "retryability": item["retryability"], + "duration_seconds": round(duration_seconds, 2), + "self_ratio": round(float(item.get("self_ratio") or 0.0), 2), + "latest_failure_type": item["latest_failure_type"], + "latest_task_detail": item["latest_task_detail"], + } + + domain_entry = domain_buckets.setdefault( + failure_domain, + { + "domain": failure_domain, + "label": _humanize_failure_domain(failure_domain), + "duration_seconds": 0.0, + "task_count": 0, + "_subtypes": {}, + "_reasons": {}, + "_messages": {}, + "_apps": [], + }, + ) + domain_entry["duration_seconds"] += duration_seconds + domain_entry["task_count"] += 1 + domain_entry["_apps"].append(sample) + + domain_reason = domain_entry["_reasons"].setdefault( + reason_label, + {"label": reason_label, "duration_seconds": 0.0, "task_count": 0}, + ) + domain_reason["duration_seconds"] += duration_seconds + domain_reason["task_count"] += 1 + + domain_message = domain_entry["_messages"].setdefault( + message_label, + {"label": message_label, "duration_seconds": 0.0, "task_count": 0}, + ) + domain_message["duration_seconds"] += duration_seconds + domain_message["task_count"] += 1 + + subtype_entry = domain_entry["_subtypes"].setdefault( + failure_subtype, + { + "domain": failure_domain, + "subtype": failure_subtype, + "label": _humanize_failure_subtype( + failure_domain, + failure_subtype, + item["latest_failure_type"], + item["latest_task_detail"], + ), + "duration_seconds": 0.0, + "task_count": 0, + "_reasons": {}, + "_messages": {}, + "_apps": [], + }, + ) + subtype_entry["duration_seconds"] += duration_seconds + subtype_entry["task_count"] += 1 + subtype_entry["_apps"].append(sample) + + subtype_reason = subtype_entry["_reasons"].setdefault( + reason_label, + {"label": reason_label, "duration_seconds": 0.0, "task_count": 0}, + ) + subtype_reason["duration_seconds"] += duration_seconds + subtype_reason["task_count"] += 1 + + subtype_message = subtype_entry["_messages"].setdefault( + message_label, + {"label": message_label, "duration_seconds": 0.0, "task_count": 0}, + ) + subtype_message["duration_seconds"] += duration_seconds + subtype_message["task_count"] += 1 + + failure_domain_breakdown: List[Dict[str, Any]] = [] + failure_subtype_breakdown: List[Dict[str, Any]] = [] + for domain_entry in sorted(domain_buckets.values(), key=lambda item: item["duration_seconds"], reverse=True): + domain_duration = float(domain_entry["duration_seconds"]) + subtypes: List[Dict[str, Any]] = [] + for subtype_entry in sorted(domain_entry["_subtypes"].values(), key=lambda item: item["duration_seconds"], reverse=True): + subtype_duration = float(subtype_entry["duration_seconds"]) + reasons = [ + { + "label": current["label"], + "duration_seconds": round(float(current["duration_seconds"]), 2), + "task_count": int(current["task_count"]), + "overall_share_percent": round((float(current["duration_seconds"]) / total_duration) * 100, 2) + if total_duration > 0 + else 0.0, + "error_share_percent": round((float(current["duration_seconds"]) / failed_duration) * 100, 2) + if failed_duration > 0 + else 0.0, + "domain_share_percent": round((float(current["duration_seconds"]) / domain_duration) * 100, 2) + if domain_duration > 0 + else 0.0, + "subtype_share_percent": round((float(current["duration_seconds"]) / subtype_duration) * 100, 2) + if subtype_duration > 0 + else 0.0, + } + for current in sorted(subtype_entry["_reasons"].values(), key=lambda item: item["duration_seconds"], reverse=True) + ] + messages = [ + { + "label": current["label"], + "message": current["label"], + "duration_seconds": round(float(current["duration_seconds"]), 2), + "task_count": int(current["task_count"]), + "overall_share_percent": round((float(current["duration_seconds"]) / total_duration) * 100, 2) + if total_duration > 0 + else 0.0, + "error_share_percent": round((float(current["duration_seconds"]) / failed_duration) * 100, 2) + if failed_duration > 0 + else 0.0, + "domain_share_percent": round((float(current["duration_seconds"]) / domain_duration) * 100, 2) + if domain_duration > 0 + else 0.0, + "subtype_share_percent": round((float(current["duration_seconds"]) / subtype_duration) * 100, 2) + if subtype_duration > 0 + else 0.0, + } + for current in sorted(subtype_entry["_messages"].values(), key=lambda item: item["duration_seconds"], reverse=True) + ] + subtype_item = { + "domain": subtype_entry["domain"], + "subtype": subtype_entry["subtype"], + "label": subtype_entry["label"], + "duration_seconds": round(subtype_duration, 2), + "task_count": int(subtype_entry["task_count"]), + "avg_duration_seconds": round(subtype_duration / subtype_entry["task_count"], 2) + if subtype_entry["task_count"] + else 0.0, + "overall_share_percent": round((subtype_duration / total_duration) * 100, 2) + if total_duration > 0 + else 0.0, + "error_share_percent": round((subtype_duration / failed_duration) * 100, 2) + if failed_duration > 0 + else 0.0, + "domain_share_percent": round((subtype_duration / domain_duration) * 100, 2) + if domain_duration > 0 + else 0.0, + "reasons": reasons, + "messages": messages, + "apps": sorted( + subtype_entry["_apps"], + key=lambda current: (float(current["duration_seconds"]), current["package_name"]), + reverse=True, + ), + } + subtypes.append(subtype_item) + failure_subtype_breakdown.append(subtype_item) + + failure_domain_breakdown.append( + { + "domain": domain_entry["domain"], + "label": domain_entry["label"], + "duration_seconds": round(domain_duration, 2), + "task_count": int(domain_entry["task_count"]), + "overall_share_percent": round((domain_duration / total_duration) * 100, 2) if total_duration > 0 else 0.0, + "error_share_percent": round((domain_duration / failed_duration) * 100, 2) if failed_duration > 0 else 0.0, + "share_percent": round((domain_duration / failed_duration) * 100, 2) if failed_duration > 0 else 0.0, + "subtypes": subtypes, + "reasons": [ + { + "label": current["label"], + "duration_seconds": round(float(current["duration_seconds"]), 2), + "task_count": int(current["task_count"]), + "overall_share_percent": round((float(current["duration_seconds"]) / total_duration) * 100, 2) + if total_duration > 0 + else 0.0, + "error_share_percent": round((float(current["duration_seconds"]) / failed_duration) * 100, 2) + if failed_duration > 0 + else 0.0, + "domain_share_percent": round((float(current["duration_seconds"]) / domain_duration) * 100, 2) + if domain_duration > 0 + else 0.0, + } + for current in sorted(domain_entry["_reasons"].values(), key=lambda item: item["duration_seconds"], reverse=True) + ], + "messages": [ + { + "label": current["label"], + "message": current["label"], + "duration_seconds": round(float(current["duration_seconds"]), 2), + "task_count": int(current["task_count"]), + "overall_share_percent": round((float(current["duration_seconds"]) / total_duration) * 100, 2) + if total_duration > 0 + else 0.0, + "error_share_percent": round((float(current["duration_seconds"]) / failed_duration) * 100, 2) + if failed_duration > 0 + else 0.0, + "domain_share_percent": round((float(current["duration_seconds"]) / domain_duration) * 100, 2) + if domain_duration > 0 + else 0.0, + } + for current in sorted(domain_entry["_messages"].values(), key=lambda item: item["duration_seconds"], reverse=True) + ], + "apps": sorted( + domain_entry["_apps"], + key=lambda current: (float(current["duration_seconds"]), current["package_name"]), + reverse=True, + ), + } + ) + + failure_reasons = Counter(item["failure_reason"] for item in failed_items if item["failure_reason"]) + return { + "failure_domain_breakdown": failure_domain_breakdown, + "failure_subtype_breakdown": failure_subtype_breakdown, + "failure_reasons": [ + {"label": label, "count": count} + for label, count in failure_reasons.most_common() + ], + "failed_task_count": len(failed_items), + "failed_duration_seconds": round(failed_duration, 2), + "total_duration_seconds": round(total_duration, 2), + } + + +def _report_overview(items: List[Dict[str, Any]]) -> Dict[str, Any]: + total = len(items) + success_count = sum(1 for item in items if item["restriction_status"] == "success") + light_count = sum(1 for item in items if item["restriction_status"] == "light_restricted") + severe_count = sum(1 for item in items if item["restriction_status"] == "severe_restricted") + retryable_count = sum( + 1 + for item in items + if item["restriction_status"] == "severe_restricted" and item["retryability"] == "retryable" + ) + non_retryable_count = sum( + 1 + for item in items + if item["restriction_status"] == "severe_restricted" and item["retryability"] == "non_retryable" + ) + qualified_count = success_count + light_count + failed_count = severe_count + total_duration = sum(float(item.get("duration_seconds") or 0.0) for item in items) + qualified_duration = sum(float(item.get("duration_seconds") or 0.0) for item in items if item["qualified"]) + failed_duration = sum(float(item.get("duration_seconds") or 0.0) for item in items if not item["qualified"]) + return { + "total_app_count": total, + "success_count": success_count, + "light_restricted_count": light_count, + "failed_count": failed_count, + "qualified_count": qualified_count, + "qualified_rate_percent": round((qualified_count / total) * 100, 2) if total > 0 else 0.0, + "severe_retryable_count": retryable_count, + "severe_non_retryable_count": non_retryable_count, + "total_duration_seconds": round(total_duration, 2), + "qualified_duration_seconds": round(qualified_duration, 2), + "failed_duration_seconds": round(failed_duration, 2), + } + + +def build_daily_quality_report( + *, + date_str: str, + task_csv_path: str, + success_csv_path: str, + failed_csv_path: str, + retry_csv_path: str, + traffic_root: str, + app_list_path: str, + url_lib_path: str, +) -> Dict[str, Any]: + allowed_packages = _resolve_allowed_packages(task_csv_path) + daily_rows = _load_daily_latest_rows( + date_str=date_str, + success_csv_path=success_csv_path, + failed_csv_path=failed_csv_path, + retry_csv_path=retry_csv_path, + allowed_packages=allowed_packages, + ) + if not daily_rows: + return { + "date": date_str, + "generated_at": datetime.now(SHANGHAI_TZ).strftime("%Y-%m-%d %H:%M:%S"), + "source": { + "task_csv_path": task_csv_path, + "success_csv_path": success_csv_path, + "failed_csv_path": failed_csv_path, + "retry_csv_path": retry_csv_path, + "traffic_root": traffic_root, + "app_list_path": app_list_path, + "url_lib_path": url_lib_path, + }, + "overview": _report_overview([]), + "failure_analysis": _build_failure_breakdown([]), + "apps": [], + } + + with tempfile.TemporaryDirectory(prefix="daily_quality_") as tmp_dir: + temp_db_path = os.path.join(tmp_dir, "daily_quality.sqlite3") + service = AnalyticsService( + db_path=temp_db_path, + task_csv_path=task_csv_path, + success_tasks_csv_path=success_csv_path, + failed_tasks_csv_path=failed_csv_path, + retry_tasks_csv_path=retry_csv_path, + traffic_root=traffic_root, + app_list_path=app_list_path, + url_lib_path=url_lib_path, + start_worker=False, + ) + items: List[Dict[str, Any]] = [] + for package_name, row in sorted(daily_rows.items(), key=lambda item: item[0]): + summary = service.rebuild_package_now(package_name, latest_task_override=row) + restriction_status, retryability = _classify_restriction_status( + summary.get("latest_status", ""), + summary.get("self_ratio", 0.0), + summary.get("latest_failure_type", ""), + ) + qualified = restriction_status in {"success", "light_restricted"} + failure_domain = "" + failure_subtype = "" + if not qualified: + failure_domain, failure_subtype = _classify_failure( + summary.get("latest_failure_type"), + summary.get("latest_task_detail"), + ) + items.append( + { + "app_name": summary.get("app_name") or package_name, + "package_name": package_name, + "report_type": row.get("report_type", ""), + "latest_status": summary.get("latest_status", ""), + "restriction_status": restriction_status, + "retryability": retryability, + "qualified": qualified, + "latest_worker_id": summary.get("latest_worker_id", ""), + "latest_test_time": summary.get("latest_test_time", 0.0), + "latest_failure_type": summary.get("latest_failure_type", ""), + "latest_task_detail": summary.get("latest_task_detail", ""), + "duration_seconds": round(float(summary.get("duration_seconds") or 0.0), 2), + "num_nodes": int(summary.get("num_nodes") or 0), + "unique_domain_count": int(summary.get("unique_domain_count") or 0), + "unique_second_level_domain_count": int(summary.get("unique_second_level_domain_count") or 0), + "total_traffic_bytes": int(summary.get("total_traffic_bytes") or 0), + "self_traffic_bytes": int(summary.get("self_traffic_bytes") or 0), + "server_traffic_bytes": int(summary.get("server_traffic_bytes") or 0), + "unrecognized_traffic_bytes": int(summary.get("unrecognized_traffic_bytes") or 0), + "self_ratio": round(float(summary.get("self_ratio") or 0.0), 2), + "recognition_ratio": round(float(summary.get("recognition_ratio") or 0.0), 2), + "artifact_status": summary.get("artifact_status", "missing"), + "failure_domain": failure_domain, + "failure_domain_label": _humanize_failure_domain(failure_domain) if failure_domain else "", + "failure_subtype": failure_subtype, + "failure_subtype_label": _humanize_failure_subtype( + failure_domain, + failure_subtype, + summary.get("latest_failure_type"), + summary.get("latest_task_detail"), + ) + if failure_domain + else "", + "failure_reason": _humanize_failure_reason( + summary.get("latest_failure_type"), + summary.get("latest_task_detail"), + ) + if not qualified + else "", + } + ) + + items.sort( + key=lambda item: ( + 0 if not item["qualified"] else 1, + -float(item["duration_seconds"]), + item["package_name"], + ) + ) + return { + "date": date_str, + "generated_at": datetime.now(SHANGHAI_TZ).strftime("%Y-%m-%d %H:%M:%S"), + "source": { + "task_csv_path": task_csv_path, + "success_csv_path": success_csv_path, + "failed_csv_path": failed_csv_path, + "retry_csv_path": retry_csv_path, + "traffic_root": traffic_root, + "app_list_path": app_list_path, + "url_lib_path": url_lib_path, + "package_filter_enabled": bool(allowed_packages), + }, + "overview": _report_overview(items), + "failure_analysis": _build_failure_breakdown(items), + "apps": items, + } + + +def build_daily_quality_csv_rows(report: Dict[str, Any]) -> List[Dict[str, Any]]: + date_str = str(report.get("date") or "") + generated_at = str(report.get("generated_at") or "") + overview = report.get("overview") or {} + failure_analysis = report.get("failure_analysis") or {} + rows: List[Dict[str, Any]] = [] + + overview_specs = [ + ("total_app_count", "应用总数"), + ("success_count", "成功"), + ("light_restricted_count", "轻度受限"), + ("failed_count", "失败"), + ("qualified_count", "合格"), + ("qualified_rate_percent", "合格率"), + ("severe_retryable_count", "严重受限可重试"), + ("severe_non_retryable_count", "严重受限不可重试"), + ("total_duration_seconds", "总耗时"), + ("qualified_duration_seconds", "合格耗时"), + ("failed_duration_seconds", "失败耗时"), + ] + for key, label in overview_specs: + rows.append( + { + "record_type": "overview", + "date": date_str, + "generated_at": generated_at, + "name": key, + "label": label, + "value": overview.get(key, ""), + } + ) + + for item in failure_analysis.get("failure_domain_breakdown") or []: + rows.append( + { + "record_type": "failure_domain", + "date": date_str, + "generated_at": generated_at, + "name": item.get("domain", ""), + "label": item.get("label", ""), + "count": item.get("task_count", 0), + "share_percent": item.get("share_percent", 0.0), + "overall_share_percent": item.get("overall_share_percent", 0.0), + "error_share_percent": item.get("error_share_percent", 0.0), + "duration_seconds": item.get("duration_seconds", 0.0), + } + ) + + for item in failure_analysis.get("failure_subtype_breakdown") or []: + rows.append( + { + "record_type": "failure_subtype", + "date": date_str, + "generated_at": generated_at, + "name": item.get("subtype", ""), + "label": item.get("label", ""), + "count": item.get("task_count", 0), + "overall_share_percent": item.get("overall_share_percent", 0.0), + "error_share_percent": item.get("error_share_percent", 0.0), + "domain_share_percent": item.get("domain_share_percent", 0.0), + "duration_seconds": item.get("duration_seconds", 0.0), + "avg_duration_seconds": item.get("avg_duration_seconds", 0.0), + "failure_domain": item.get("domain", ""), + } + ) + + for item in report.get("apps") or []: + rows.append( + { + "record_type": "app", + "date": date_str, + "generated_at": generated_at, + "qualified": "yes" if item.get("qualified") else "no", + "app_name": item.get("app_name", ""), + "package_name": item.get("package_name", ""), + "report_type": item.get("report_type", ""), + "latest_status": item.get("latest_status", ""), + "restriction_status": item.get("restriction_status", ""), + "retryability": item.get("retryability", ""), + "latest_worker_id": item.get("latest_worker_id", ""), + "latest_test_time": item.get("latest_test_time", ""), + "latest_failure_type": item.get("latest_failure_type", ""), + "latest_task_detail": item.get("latest_task_detail", ""), + "failure_domain": item.get("failure_domain", ""), + "failure_domain_label": item.get("failure_domain_label", ""), + "failure_subtype": item.get("failure_subtype", ""), + "failure_subtype_label": item.get("failure_subtype_label", ""), + "failure_reason": item.get("failure_reason", ""), + "duration_seconds": item.get("duration_seconds", 0.0), + "num_nodes": item.get("num_nodes", 0), + "unique_domain_count": item.get("unique_domain_count", 0), + "unique_second_level_domain_count": item.get("unique_second_level_domain_count", 0), + "total_traffic_bytes": item.get("total_traffic_bytes", 0), + "self_traffic_bytes": item.get("self_traffic_bytes", 0), + "server_traffic_bytes": item.get("server_traffic_bytes", 0), + "unrecognized_traffic_bytes": item.get("unrecognized_traffic_bytes", 0), + "self_ratio": item.get("self_ratio", 0.0), + "recognition_ratio": item.get("recognition_ratio", 0.0), + "artifact_status": item.get("artifact_status", ""), + } + ) + return rows + + +def write_daily_quality_csv(report: Dict[str, Any], output_path: str) -> None: + rows = build_daily_quality_csv_rows(report) + output_dir = os.path.dirname(output_path) + if output_dir: + os.makedirs(output_dir, exist_ok=True) + with open(output_path, "w", newline="", encoding="utf-8-sig") as handle: + writer = csv.DictWriter(handle, fieldnames=CSV_COLUMNS) + writer.writeheader() + for row in rows: + normalized = {column: row.get(column, "") for column in CSV_COLUMNS} + writer.writerow(normalized) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Generate a single-day quality analysis CSV report.") + parser.add_argument("--date", required=True, help="Target date in YYYY-MM-DD format.") + parser.add_argument( + "--output", + default="", + help="Output CSV path. Default: runtime//reports/daily_quality_.csv", + ) + parser.add_argument("--task-csv", default=TASK_CSV_PATH, help="Package list CSV path.") + parser.add_argument("--success-csv", default=SUCCESS_TASKS_CSV, help="success_tasks.csv path or glob.") + parser.add_argument("--failed-csv", default=FAILED_TASKS_CSV, help="failed_tasks.csv path.") + parser.add_argument("--retry-csv", default=RETRY_TASKS_CSV, help="retry_tasks.csv path.") + parser.add_argument("--traffic-root", default=ANALYTICS_TRAFFIC_ROOT, help="traffic_data root path.") + parser.add_argument("--app-list", default=ANALYTICS_TPDPI_APP_LIST, help="TPDPI app list CSV path.") + parser.add_argument("--url-lib", default=ANALYTICS_TPDPI_URL_LIB, help="TPDPI url lib CSV path.") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + try: + datetime.strptime(args.date, "%Y-%m-%d") + except ValueError: + print(f"Invalid --date: {args.date}", file=sys.stderr) + return 2 + + output_path = args.output.strip() + if not output_path: + output_path = os.path.join(REPORT_DIR, f"daily_quality_{args.date.replace('-', '')}.csv") + + report = build_daily_quality_report( + date_str=args.date, + task_csv_path=args.task_csv, + success_csv_path=args.success_csv, + failed_csv_path=args.failed_csv, + retry_csv_path=args.retry_csv, + traffic_root=args.traffic_root, + app_list_path=args.app_list, + url_lib_path=args.url_lib, + ) + + write_daily_quality_csv(report, output_path) + print(output_path) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/analyze_non_retryable_popularity.py b/scripts/analyze_non_retryable_popularity.py new file mode 100644 index 0000000..5947e74 --- /dev/null +++ b/scripts/analyze_non_retryable_popularity.py @@ -0,0 +1,306 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +import argparse +import csv +import json +import shutil +import sqlite3 +import sys +import tempfile +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional, Tuple + + +SCRIPT_DIR = Path(__file__).resolve().parent +PROJECT_ROOT = SCRIPT_DIR.parent +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +from config import MONITORING_DB_PATH, TASK_CSV_PATH # noqa: E402 +from analytics import AnalyticsRepository # noqa: E402 +from result_codes import ( # noqa: E402 + APP_ERROR_DESC, + BUSINESS_ERROR_DESC, + DOWNLOAD_ERROR_DESC, + INFRA_ERROR_DESC, + AppError, + BusinessError, + DownloadError, + InfraError, +) + + +def _parse_downloads(value: Any) -> Optional[int]: + text = str(value or "").strip().upper() + if not text: + return None + normalized = text.replace(",", "").replace(" ", "") + if normalized.endswith("+"): + normalized = normalized[:-1] + multiplier = 1 + if normalized.endswith("K"): + multiplier = 1_000 + normalized = normalized[:-1] + elif normalized.endswith("M"): + multiplier = 1_000_000 + normalized = normalized[:-1] + elif normalized.endswith("B"): + multiplier = 1_000_000_000 + normalized = normalized[:-1] + try: + return int(float(normalized) * multiplier) + except (TypeError, ValueError): + return None + + +def _humanize_error_type(error_type: str) -> str: + normalized = str(error_type or "").strip() + if not normalized: + return "未标记错误" + category_name, _, code_text = normalized.partition("/") + try: + code = int(code_text) + except (TypeError, ValueError): + return normalized + try: + if category_name == "INFRA_ERROR": + return INFRA_ERROR_DESC.get(InfraError(code), normalized) + if category_name == "APP_ERROR": + return APP_ERROR_DESC.get(AppError(code), normalized) + if category_name == "BUSINESS_ERROR": + return BUSINESS_ERROR_DESC.get(BusinessError(code), normalized) + if category_name == "DOWNLOAD_ERROR": + return DOWNLOAD_ERROR_DESC.get(DownloadError(code), normalized) + except ValueError: + return normalized + return normalized + + +def _load_downloads_by_package(csv_path: str) -> Dict[str, Optional[int]]: + path = Path(csv_path) + if not path.exists(): + return {} + result: Dict[str, Optional[int]] = {} + with path.open("r", encoding="utf-8-sig", newline="") as handle: + for row in csv.DictReader(handle): + package_name = str(row.get("package_name") or row.get("包名") or "").strip() + if not package_name or package_name in result: + continue + result[package_name] = _parse_downloads( + row.get("downloads") + or row.get("download") + or row.get("下载量") + ) + return result + + +def _copy_sqlite_snapshot(db_path: str) -> Tuple[str, tempfile.TemporaryDirectory[str]]: + source = Path(db_path) + temp_dir: tempfile.TemporaryDirectory[str] = tempfile.TemporaryDirectory(prefix="dpi-db-snapshot-") + snapshot_path = Path(temp_dir.name) / source.name + shutil.copy2(source, snapshot_path) + for suffix in ("-wal", "-shm"): + sidecar = Path(f"{db_path}{suffix}") + if sidecar.exists(): + shutil.copy2(sidecar, Path(f"{snapshot_path}{suffix}")) + return str(snapshot_path), temp_dir + + +def _open_readable_connection(db_path: str) -> Tuple[sqlite3.Connection, Optional[tempfile.TemporaryDirectory[str]], bool]: + try: + connection = sqlite3.connect(db_path) + connection.row_factory = sqlite3.Row + connection.execute("SELECT name FROM sqlite_master WHERE type = 'table' LIMIT 1").fetchone() + return connection, None, False + except sqlite3.Error: + snapshot_path, temp_dir = _copy_sqlite_snapshot(db_path) + connection = sqlite3.connect(snapshot_path) + connection.row_factory = sqlite3.Row + connection.execute("SELECT name FROM sqlite_master WHERE type = 'table' LIMIT 1").fetchone() + return connection, temp_dir, True + + +def _extract_downloads_from_payload(task_payload_json: Optional[str]) -> Optional[int]: + if not task_payload_json: + return None + try: + payload = json.loads(task_payload_json) + except (TypeError, ValueError): + return None + if not isinstance(payload, dict): + return None + original_row = payload.get("original_row") + if isinstance(original_row, dict): + for key in ("downloads", "download", "下载量"): + if key in original_row: + return _parse_downloads(original_row.get(key)) + for key in ("downloads", "download", "下载量"): + if key in payload: + return _parse_downloads(payload.get(key)) + return None + + +def build_reason_popularity_report( + db_path: str = MONITORING_DB_PATH, + csv_path: str = TASK_CSV_PATH, + *, + threshold: int = 100_000, +) -> Dict[str, Any]: + connection, temp_dir, used_snapshot = _open_readable_connection(db_path) + try: + repo = AnalyticsRepository.__new__(AnalyticsRepository) + summaries = repo._list_catalog_summaries(connection) + rows = [ + item for item in summaries + if item.get("restriction_status") == "severe_restricted" + and item.get("retryability") == "non_retryable" + ] + rows.sort(key=lambda item: (str(item.get("latest_failure_type") or ""), str(item.get("package_name") or ""))) + finally: + connection.close() + if temp_dir is not None: + temp_dir.cleanup() + + downloads_by_package = _load_downloads_by_package(csv_path) + + reason_buckets: Dict[str, Dict[str, Any]] = {} + missing_download_packages: List[str] = [] + matched_downloads_apps = 0 + for row in rows: + package_name = str(row.get("package_name") or "").strip() + downloads = downloads_by_package.get(package_name) + if downloads is None: + downloads = _extract_downloads_from_payload(row.get("task_payload_json")) + if downloads is None and row.get("downloads") is not None: + downloads = int(row.get("downloads") or 0) + if downloads is None: + missing_download_packages.append(package_name) + continue + matched_downloads_apps += 1 + error_type = str(row.get("latest_failure_type") or "").strip() + bucket = reason_buckets.setdefault( + error_type, + { + "error_type": error_type, + "reason_label": _humanize_error_type(error_type), + "hot_count": 0, + "non_hot_count": 0, + "total_count": 0, + }, + ) + if downloads > threshold: + bucket["hot_count"] += 1 + else: + bucket["non_hot_count"] += 1 + bucket["total_count"] += 1 + + report_rows = sorted( + reason_buckets.values(), + key=lambda item: (-int(item["total_count"]), -int(item["hot_count"]), str(item["error_type"])), + ) + return { + "db_path": db_path, + "csv_path": csv_path, + "threshold": int(threshold), + "used_snapshot": used_snapshot, + "total_non_retryable_apps": len(rows), + "matched_downloads_apps": matched_downloads_apps, + "missing_downloads_apps": len(missing_download_packages), + "missing_download_packages": missing_download_packages, + "rows": report_rows, + } + + +def _iter_table_rows(report_rows: Iterable[Dict[str, Any]]) -> Iterable[List[str]]: + for item in report_rows: + yield [ + item["reason_label"], + item["error_type"] or "-", + str(item["hot_count"]), + str(item["non_hot_count"]), + str(item["total_count"]), + ] + + +def _format_table(report_rows: List[Dict[str, Any]]) -> str: + headers = ["失败原因", "错误码", "热门", "非热门", "总数"] + rows = list(_iter_table_rows(report_rows)) + if not rows: + return "没有符合条件的应用。" + widths = [len(header) for header in headers] + for row in rows: + for index, cell in enumerate(row): + widths[index] = max(widths[index], len(cell)) + + def _render(row: List[str]) -> str: + return " | ".join(cell.ljust(widths[index]) for index, cell in enumerate(row)) + + divider = "-+-".join("-" * width for width in widths) + lines = [_render(headers), divider] + lines.extend(_render(row) for row in rows) + return "\n".join(lines) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="统计严重受限不可重试应用中，各失败原因对应的热门/非热门数量。" + ) + parser.add_argument("--db-path", default=MONITORING_DB_PATH, help="monitoring sqlite 路径。") + parser.add_argument("--csv-path", default=TASK_CSV_PATH, help="包含 downloads 列的应用清单 CSV。") + parser.add_argument( + "--threshold", + type=int, + default=100_000, + help="热门阈值。downloads > threshold 视为热门，默认 100000。", + ) + parser.add_argument( + "--format", + choices=("table", "json"), + default="table", + help="输出格式，默认 table。", + ) + parser.add_argument( + "--show-missing-packages", + action="store_true", + help="额外打印缺失下载量的包名。", + ) + return parser + + +def main() -> int: + parser = build_parser() + args = parser.parse_args() + report = build_reason_popularity_report( + db_path=args.db_path, + csv_path=args.csv_path, + threshold=args.threshold, + ) + + if args.format == "json": + print(json.dumps(report, ensure_ascii=False, indent=2)) + return 0 + + print(f"数据库: {report['db_path']}") + print(f"应用清单: {report['csv_path']}") + print(f"热门阈值: downloads > {report['threshold']}") + print("统计范围: restriction_status=severe_restricted AND retryability=non_retryable") + print(f"严重受限不可重试应用总数: {report['total_non_retryable_apps']}") + print(f"已匹配下载量应用数: {report['matched_downloads_apps']}") + print(f"下载量缺失应用数: {report['missing_downloads_apps']}") + if report["used_snapshot"]: + print("读取方式: 已自动复制 live sqlite 快照后再统计") + print() + print(_format_table(report["rows"])) + + if args.show_missing_packages and report["missing_download_packages"]: + print() + print("下载量缺失包名:") + for package_name in report["missing_download_packages"]: + print(package_name) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/analyze_tag_collection_stats.py b/scripts/analyze_tag_collection_stats.py new file mode 100644 index 0000000..aa3decc --- /dev/null +++ b/scripts/analyze_tag_collection_stats.py @@ -0,0 +1,422 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +""" +统计指定 tag 下 new_app / app_update 的采集质量分布和耗时分布。 +支持通过 --tag2 传入额外 tag，合并两个 tag 的应用（去重）再统计。 + +输出: +- 总体统计 (total/success/failed_terminal/pending) +- failed_terminal / pending 失败细分 +- task_execution 耗时分布 (下载/采集/异常耗时, 按 failure_bucket 分) +- worker_state_event 状态时间分布 +""" + +from __future__ import annotations + +import argparse +import sqlite3 +import sys +from pathlib import Path +from typing import Any, Dict, List, Tuple + +SCRIPT_DIR = Path(__file__).resolve().parent +PROJECT_ROOT = SCRIPT_DIR.parent +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +from config import MONITORING_DB_PATH # noqa: E402 +from analytics import AnalyticsRepository # noqa: E402 +from result_codes import ( # noqa: E402 + APP_ERROR_DESC, + BUSINESS_ERROR_DESC, + DOWNLOAD_ERROR_DESC, + INFRA_ERROR_DESC, + AppError, + BusinessError, + DownloadError, + InfraError, +) + +# ── helpers ──────────────────────────────────────────────── + + +def _matches_tag(summary: Dict[str, Any], tag: str, *, strict: bool) -> bool: + tags = [str(item or "").strip() for item in (summary.get("incremental_batch_tags") or [])] + if not tag or not tags: + return False + task_type = str(summary.get("collection_task_type") or "new_app").strip() or "new_app" + if strict and task_type == "new_app": + return tags[0] == tag + return tag in tags + + +def _in_scope(summary: Dict[str, Any], tag: str, tag2: str | None) -> bool: + if not summary.get("catalog_active"): + return False + if _matches_tag(summary, tag, strict=True): + return True + return bool(tag2 and _matches_tag(summary, tag2, strict=False)) + + +def _humanize_failure_type(error_type: str) -> str: + normalized = str(error_type or "").strip() + if not normalized: + return "未标记" + category_name, _, code_text = normalized.partition("/") + try: + code = int(code_text) + except (TypeError, ValueError): + return normalized + try: + if category_name == "INFRA_ERROR": + return INFRA_ERROR_DESC.get(InfraError(code), normalized) + if category_name == "APP_ERROR": + return APP_ERROR_DESC.get(AppError(code), normalized) + if category_name == "BUSINESS_ERROR": + return BUSINESS_ERROR_DESC.get(BusinessError(code), normalized) + if category_name == "DOWNLOAD_ERROR": + return DOWNLOAD_ERROR_DESC.get(DownloadError(code), normalized) + except ValueError: + return normalized + return normalized + + +def _secs_to_human(seconds: float) -> str: + if seconds < 60: + return f"{seconds:.0f}s" + if seconds < 3600: + return f"{seconds / 60:.1f}m" + return f"{seconds / 3600:.2f}h" + + +def _timing_category(row: sqlite3.Row) -> str: + execution_status = str(row["execution_status"] or "").strip() + error_category = str(row["error_category"] or "").strip() + if execution_status == "success": + return "正常下载、采集" + if error_category == "APP_ERROR": + return "商店跳转更新、启动异常、探索停滞等" + if error_category == "BUSINESS_ERROR": + return "登录/注册失败、地区/付费限制" + if error_category == "INFRA_ERROR": + return "模拟器启动失败、ADB断联、抓包异常" + if error_category == "DOWNLOAD_ERROR": + return "下载超时、网络异常" + return "其他" + + +def _empty_total_time() -> Dict[str, Any]: + return {"tot_dl": 0, "tot_col": 0, "tot_total": 0, "tot_cnt": 0} + + +def _run(db_path: str, tag: str, tag2: str | None = None) -> Dict[str, Any]: + repo = AnalyticsRepository(db_path=db_path) + with repo._connect() as conn: + summaries = [ + item for item in repo._list_catalog_summaries(conn) + if _in_scope(item, tag, tag2) + ] + package_names = [str(item.get("package_name") or "").strip() for item in summaries if str(item.get("package_name") or "").strip()] + package_task_type = { + str(item.get("package_name") or "").strip(): str(item.get("collection_task_type") or "new_app").strip() or "new_app" + for item in summaries + } + + task_rows = [] + worker_rows = [] + if package_names: + placeholders = ",".join("?" for _ in package_names) + task_rows = conn.execute( + f""" + SELECT package_name, execution_status, error_category, + download_duration_seconds, execution_duration_seconds, duration_seconds + FROM collection_task + WHERE package_name IN ({placeholders}) + """, + package_names, + ).fetchall() + + worker_ids = sorted( + { + str(item.get("latest_worker_id") or "").strip() + for item in summaries + if str(item.get("latest_worker_id") or "").strip() + } + ) + if worker_ids: + worker_placeholders = ",".join("?" for _ in worker_ids) + worker_rows = conn.execute( + f""" + SELECT state, COUNT(*) AS event_count + FROM worker_state_event + WHERE worker_id IN ({worker_placeholders}) + GROUP BY state + ORDER BY event_count DESC + """, + worker_ids, + ).fetchall() + + summary_buckets: Dict[str, Dict[str, Any]] = {} + failure_buckets: Dict[Tuple[str, str], Dict[str, Any]] = {} + pending_buckets: Dict[Tuple[str, str], Dict[str, Any]] = {} + for item in summaries: + task_type = str(item.get("collection_task_type") or "new_app").strip() or "new_app" + bucket = summary_buckets.setdefault( + task_type, + {"task_type": task_type, "total": 0, "success": 0, "failed": 0, "pending": 0}, + ) + bucket["total"] += 1 + collection_status = str(item.get("collection_status") or "").strip() + latest_status = str(item.get("latest_status") or "").strip() + failure_type = str(item.get("latest_failure_type") or "").strip() or "UNKNOWN" + if collection_status == "qualified": + bucket["success"] += 1 + elif collection_status == "failed_terminal": + bucket["failed"] += 1 + failure_buckets.setdefault( + (task_type, failure_type), + {"task_type": task_type, "failure_type": failure_type, "cnt": 0}, + )["cnt"] += 1 + elif collection_status == "pending" and latest_status not in {"", "success"}: + bucket["pending"] += 1 + pending_buckets.setdefault( + (task_type, failure_type), + {"task_type": task_type, "failure_type": failure_type, "cnt": 0}, + )["cnt"] += 1 + + timing_buckets: Dict[Tuple[str, str], Dict[str, Any]] = {} + total_time = _empty_total_time() + for row in task_rows: + task_type = package_task_type.get(str(row["package_name"] or "").strip(), "new_app") + category = _timing_category(row) + bucket = timing_buckets.setdefault( + (task_type, category), + { + "task_type": task_type, + "category": category, + "cnt": 0, + "sum_dl_s": 0.0, + "sum_col_s": 0.0, + "sum_total_s": 0.0, + }, + ) + dl_seconds = float(row["download_duration_seconds"] or 0.0) + col_seconds = float(row["execution_duration_seconds"] or 0.0) + total_seconds = float(row["duration_seconds"] or 0.0) + bucket["cnt"] += 1 + bucket["sum_dl_s"] += dl_seconds + bucket["sum_col_s"] += col_seconds + bucket["sum_total_s"] += total_seconds + total_time["tot_cnt"] += 1 + total_time["tot_dl"] += dl_seconds + total_time["tot_col"] += col_seconds + total_time["tot_total"] += total_seconds + + timing_rows = [] + category_order = { + "正常下载、采集": 0, + "商店跳转更新、启动异常、探索停滞等": 1, + "登录/注册失败、地区/付费限制": 2, + "模拟器启动失败、ADB断联、抓包异常": 3, + "下载超时、网络异常": 4, + "其他": 5, + } + for bucket in timing_buckets.values(): + cnt = int(bucket["cnt"] or 0) + timing_rows.append( + { + "task_type": bucket["task_type"], + "category": bucket["category"], + "cnt": cnt, + "avg_dl_s": int(round(float(bucket["sum_dl_s"] or 0.0) / cnt)) if cnt else 0, + "avg_col_s": int(round(float(bucket["sum_col_s"] or 0.0) / cnt)) if cnt else 0, + "avg_total_s": int(round(float(bucket["sum_total_s"] or 0.0) / cnt)) if cnt else 0, + "sum_dl_s": int(round(float(bucket["sum_dl_s"] or 0.0))), + "sum_col_s": int(round(float(bucket["sum_col_s"] or 0.0))), + "sum_total_s": int(round(float(bucket["sum_total_s"] or 0.0))), + } + ) + timing_rows.sort(key=lambda row: (row["task_type"], category_order.get(row["category"], 99))) + + worker_total = sum(int(row["event_count"] or 0) for row in worker_rows) + worker_items = [] + for row in worker_rows: + event_count = int(row["event_count"] or 0) + worker_items.append( + { + "state": row["state"], + "event_count": event_count, + "event_pct": round(event_count * 100.0 / worker_total, 1) if worker_total else 0.0, + } + ) + + return { + "summary": sorted(summary_buckets.values(), key=lambda row: row["task_type"]), + "failure_breakdown": sorted(failure_buckets.values(), key=lambda row: (row["task_type"], -int(row["cnt"]))), + "pending_breakdown": sorted(pending_buckets.values(), key=lambda row: (row["task_type"], -int(row["cnt"]))), + "timing": timing_rows, + "worker": worker_items, + "total_time": { + "tot_dl": int(round(float(total_time["tot_dl"] or 0.0))), + "tot_col": int(round(float(total_time["tot_col"] or 0.0))), + "tot_total": int(round(float(total_time["tot_total"] or 0.0))), + "tot_cnt": int(total_time["tot_cnt"] or 0), + }, + } + + +# ── printers ─────────────────────────────────────────────── + + +def _print_section(title: str) -> None: + print() + print(f"── {title} ──") + + +def _print_task_type_stats(label: str, rows: List[Dict[str, Any]]) -> None: + if not rows: + print(f" ({label}) 无数据") + return + for r in rows: + total = int(r["total"]) + success = int(r["success"]) + failed = int(r["failed"]) + pending = int(r["pending"]) + success_pct = f"{success / total * 100:.1f}%" if total else "0%" + failed_pct = f"{failed / total * 100:.1f}%" if total else "0%" + print( + f" [{r['task_type']}] " + f"总 {total} | " + f"成功 {success} ({success_pct}) | " + f"失败 {failed} ({failed_pct}) | " + f"待重采 {pending}" + ) + + +def _print_failure_breakdown(status_label: str, rows: List[Dict[str, Any]]) -> None: + if not rows: + print(f" (无数据)") + return + by_type: Dict[str, List[Dict[str, Any]]] = {} + for r in rows: + by_type.setdefault(r["task_type"], []).append(r) + for task_type in sorted(by_type): + items = by_type[task_type] + total = sum(int(r["cnt"]) for r in items) + print(f" [{task_type}] {status_label} 共 {total} 个：") + for r in items: + ft = r["failure_type"] + cnt = int(r["cnt"]) + pct = f"{cnt / total * 100:.1f}%" + print(f" {_humanize_failure_type(ft):<30s} ({ft}) : {cnt} ({pct})") + + +def _print_timing_distribution(timing_rows: List[Dict[str, Any]], tot: Dict[str, Any]) -> None: + if not timing_rows: + print(" (无 task_execution 数据)") + return + tot_dl = int(tot.get("tot_dl") or 0) + tot_col = int(tot.get("tot_col") or 0) + tot_total = int(tot.get("tot_total") or 0) + tot_cnt = int(tot.get("tot_cnt") or 0) + + print() + print(f" 总执行次数: {tot_cnt} | 下载耗时: {_secs_to_human(tot_dl)} | 采集耗时: {_secs_to_human(tot_col)} | 总耗时: {_secs_to_human(tot_total)}") + print() + + print(f" {'task_type':10s} | {'错误大类':32s} | {'次数':>5s} | {'avg_dl':>7s} | {'avg_col':>7s} | {'avg总':>7s} | {'总耗时':>9s} | {'%':>6s}") + print(f" {'':-<10s}-+-{'':-<32s}-+-{'':->5s}-+-{'':->7s}-+-{'':->7s}-+-{'':->7s}-+-{'':->9s}-+-{'':->6s}") + + by_type: Dict[str, List[Dict[str, Any]]] = {} + for r in timing_rows: + by_type.setdefault(r["task_type"], []).append(r) + + for task_type in sorted(by_type): + items = by_type[task_type] + first = True + for r in items: + category = r["category"] + cnt = int(r["cnt"]) + avg_dl = r["avg_dl_s"] + avg_col = r["avg_col_s"] + avg_tot = r["avg_total_s"] + sum_tot = int(r["sum_total_s"] or 0) + sum_pct = f"{sum_tot / tot_total * 100:.1f}%" if tot_total else "0%" + label = task_type if first else "" + print( + f" {label:10s} | {category:32s} | {cnt:5d} | " + f"{str(avg_dl) + 's':>7s} | {str(avg_col) + 's':>7s} | {str(avg_tot) + 's':>7s} | " + f"{_secs_to_human(sum_tot):>9s} | {sum_pct:>6s}" + ) + first = False + + +def _print_worker_state_distribution(worker_rows: List[Dict[str, Any]]) -> None: + if not worker_rows: + print(" (无 worker_state_event 数据)") + return + print() + print(f" {'state':30s} | {'event_count':>10s} | {'占比':>7s}") + print(f" {'':-<30s}-+-{'':->10s}-+-{'':->7s}") + for r in worker_rows: + state = r["state"] + cnt = int(r["event_count"]) + pct = f"{r['event_pct']:.1f}%" + print(f" {state:30s} | {cnt:10d} | {pct:>7s}") + + +# ── args / main ──────────────────────────────────────────── + + +def main() -> int: + parser = argparse.ArgumentParser( + description="按 incremental_batch_tag 统计采集质量分布和耗时分布" + ) + parser.add_argument("--tag", required=True, help="目标 incremental_batch_tag (严格匹配 new_app)") + parser.add_argument("--tag2", default=None, help="额外 incremental_batch_tag (new_app 仅需包含此 tag)") + parser.add_argument( + "--db-path", + default=MONITORING_DB_PATH, + help="monitoring.sqlite3 路径", + ) + args = parser.parse_args() + + tag = str(args.tag).strip() + if not tag: + print("错误: --tag 不能为空") + return 1 + + tag2 = str(args.tag2).strip() if args.tag2 else None + if tag2 == "": + tag2 = None + + data = _run(args.db_path, tag, tag2) + + print(f"Tag = \"{tag}\"") + if tag2: + print(f"Tag2 = \"{tag2}\"") + + _print_section("总体统计") + _print_task_type_stats("sum", data["summary"]) + + _print_section("failed_terminal 失败细分") + _print_failure_breakdown("failed_terminal", data["failure_breakdown"]) + + if data["pending_breakdown"]: + _print_section("pending 待重采细分 (按上次失败类型)") + _print_failure_breakdown("pending", data["pending_breakdown"]) + + _print_section("task_execution 耗时分布 (错误大类 × task_type)") + _print_timing_distribution(data["timing"], data["total_time"]) + + if data["worker"]: + _print_section("worker_state_event 状态分布 (latest_worker_id 匹配)") + _print_worker_state_distribution(data["worker"]) + + print() + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/backfill_model_traffic.py b/scripts/backfill_model_traffic.py new file mode 100644 index 0000000..253e3a2 --- /dev/null +++ b/scripts/backfill_model_traffic.py @@ -0,0 +1,144 @@ +#!/usr/bin/env python3 +import argparse +import os +import sys +from typing import List + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) + +from analytics import AnalyticsRepository, AnalyticsService, _is_model_traffic_line, _parse_traffic_size +from config import ANALYTICS_TRAFFIC_ROOT, MODEL_TRAFFIC_THRESHOLD, MONITORING_DB_PATH + + +def _split_packages(raw: str) -> List[str]: + return [item.strip() for item in str(raw or "").split(",") if item.strip()] + + +def _load_all_packages(service: AnalyticsService, tag: str = "", limit: int = 0) -> List[str]: + normalized_tag = str(tag or "").strip() + with service.repo._connect() as connection: + summaries = service.repo._list_catalog_summaries(connection) + packages = [] + for item in sorted(summaries, key=lambda value: str(value.get("package_name") or "")): + if normalized_tag and normalized_tag not in (item.get("incremental_batch_tags") or []): + continue + package_name = str(item.get("package_name") or "").strip() + if package_name: + packages.append(package_name) + return packages[:limit] if limit > 0 else packages + + +def _update_model_stats(repo: AnalyticsRepository, package_name: str, flow_count: int, traffic_bytes: int) -> bool: + with repo._write_lock, repo._connect() as connection: + latest = connection.execute( + """ + SELECT package_name, batch_tag, run_kind, attempt + FROM collection_task + WHERE package_name = ? + ORDER BY + datetime(COALESCE(NULLIF(completed_at, ''), NULLIF(started_at, ''), created_at)) DESC, + attempt DESC, + batch_tag DESC, + run_kind DESC + LIMIT 1 + """, + (package_name,), + ).fetchone() + if not latest: + return False + connection.execute( + """ + UPDATE collection_task + SET model_flow_count = ?, + model_traffic_bytes = ? + WHERE package_name = ? + AND batch_tag = ? + AND run_kind = ? + AND attempt = ? + """, + ( + flow_count, + traffic_bytes, + latest["package_name"], + latest["batch_tag"], + latest["run_kind"], + latest["attempt"], + ), + ) + return True + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Backfill model traffic stats for all or selected packages." + ) + parser.add_argument("--db-path", default=MONITORING_DB_PATH) + parser.add_argument("--traffic-root", default=ANALYTICS_TRAFFIC_ROOT) + parser.add_argument("--packages", default="", help="Comma-separated package names. Defaults to all packages.") + parser.add_argument("--tag", default="", help="Only process packages whose incremental_batch_tag contains this tag.") + parser.add_argument("--limit", type=int, default=0) + parser.add_argument("--dry-run", action="store_true") + parser.add_argument("--auto-promote", action="store_true", + help="Auto-set model_eligible=1 and collection_status=pending if model_flow_count > 0.") + args = parser.parse_args() + + scanned = 0 + matched = 0 + backfilled = 0 + promoted = 0 + missing_task = 0 + service = AnalyticsService( + db_path=args.db_path, + traffic_root=args.traffic_root, + start_worker=False, + ) + try: + packages = _split_packages(args.packages) or _load_all_packages(service, args.tag, args.limit) + + for package_name in packages: + scanned += 1 + traffic_files = service._find_traffic_files(package_name) + if not traffic_files: + continue + model_flow_count = 0 + model_traffic_bytes = 0 + for file_path in traffic_files: + try: + with open(file_path, "r", encoding="utf-8", errors="ignore") as handle: + for raw_line in handle: + line = raw_line.strip() + if not line: + continue + fields = [item.strip() for item in line.split(",")] + if not fields or fields[0].strip() != package_name: + continue + if _is_model_traffic_line(fields): + model_flow_count += 1 + model_traffic_bytes += _parse_traffic_size(fields[6]) + except OSError: + continue + if model_flow_count == 0: + continue + matched += 1 + eligible = 1 if 0 < model_flow_count <= MODEL_TRAFFIC_THRESHOLD else 0 + if args.dry_run: + print(f"DRY-RUN {package_name}: model_flow={model_flow_count} model_bytes={model_traffic_bytes} eligible={eligible}") + continue + if not _update_model_stats(service.repo, package_name, model_flow_count, model_traffic_bytes): + missing_task += 1 + continue + backfilled += 1 + if args.auto_promote and eligible: + promoted += service.set_packages_pending([package_name], reason="model_traffic_backfill") + print(f"backfilled {package_name}: model_flow={model_flow_count} model_bytes={model_traffic_bytes} eligible={eligible}") + finally: + service.close() + + print(f"done scanned={scanned} matched={matched} backfilled={backfilled} promoted={promoted} missing_task={missing_task}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/download.py b/scripts/download.py new file mode 100644 index 0000000..fdcff7f --- /dev/null +++ b/scripts/download.py @@ -0,0 +1,220 @@ +import os +import time +import requests +import csv +from requests.adapters import HTTPAdapter +from urllib3.util.retry import Retry +from tqdm import tqdm +import concurrent.futures + +# --- 配置项 --- +csv_file_path = '/home/tplink/code/autool-dispatcher/saved_csv/download_error_2026-4-21.csv' +DOWNLOAD_FOLDER = r"/srv/samba/disk2/dpi/mumu_apk" +TEMPLATE_APKPURE_DOWNLOAD = "https://d.apkpure.net/b/XAPK/{package_name}?version=latest" +MAX_WORKERS = 10 +# 最大线程数 +RETRY_DELAY_SECONDS = 5 # 每次重试的等待时间 + +# 禁用由 verify=False 引起的 InsecureRequestWarning 警告 +import urllib3 + +urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) + +# --- 创建下载目录 --- +if not os.path.exists(DOWNLOAD_FOLDER): + print(f"创建下载目录: {DOWNLOAD_FOLDER}") + try: + os.makedirs(DOWNLOAD_FOLDER) + except OSError as e: + print(f"错误: 无法创建下载目录 '{DOWNLOAD_FOLDER}'。请检查权限或路径是否正确。") + exit() + + +def requests_session_with_retries(retries=3, backoff_factor=0.3, status_forcelist=(500, 502, 504), session=None): + session = session or requests.Session() + retry = Retry(total=retries, read=retries, connect=retries, backoff_factor=backoff_factor, + status_forcelist=status_forcelist) + adapter = HTTPAdapter(max_retries=retry) + session.mount('http://', adapter) + session.mount('https://', adapter) + return session + + +def get_file_size(url): + """尝试获取文件总大小，如果HEAD请求失败，则尝试GET请求。""" + try: + response = requests.head(url, verify=False, timeout=10) + response.raise_for_status() + content_length = response.headers.get('Content-Length') + if content_length: + return int(content_length) + except requests.RequestException: + pass + + try: + response = requests.get(url, stream=True, verify=False, timeout=10) + response.raise_for_status() + content_length = response.headers.get('Content-Length') + if content_length: + response.close() + return int(content_length) + except requests.RequestException as e: + tqdm.write(f"警告: 无法获取文件大小。原因: {e}") + return 0 + + +def download_file(package_name): + """ + 从指定URL下载文件，支持断点续传和进度显示。 + 此函数现在是线程安全的，使用tqdm.write进行输出。 + """ + download_url = TEMPLATE_APKPURE_DOWNLOAD.format(package_name=package_name) + print(download_url) + package_dir = os.path.join(DOWNLOAD_FOLDER, package_name) + os.makedirs(package_dir, exist_ok=True) + final_filename = os.path.join(package_dir, f"{package_name}.xapk") + partial_filename = final_filename + ".part" + + if os.path.exists(final_filename): + tqdm.write(f"文件 '{os.path.basename(final_filename)}' 已完整下载，跳过。") + return package_name, True + + total_size = get_file_size(download_url) + + resumable_size = 0 + headers = { + 'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/108.0.0.0 Safari/537.36', + 'Referer': f'https://apkpure.net/1/{package_name}' + } + if os.path.exists(partial_filename): + resumable_size = os.path.getsize(partial_filename) + headers['Range'] = f'bytes={resumable_size}-' + + try: + tqdm.write(f"\n--- 开始下载: {package_name} ---") + session = requests_session_with_retries() + with session.get(download_url, headers=headers, stream=True, timeout=60, verify=False) as response: + response.raise_for_status() + + mode = 'ab' if resumable_size > 0 else 'wb' + + with open(partial_filename, mode) as f, tqdm( + total=total_size, initial=resumable_size, + unit='B', unit_scale=True, unit_divisor=1024, + desc=f"{package_name}", ascii=True, leave=False, miniters=1 + ) as pbar: + for chunk in response.iter_content(chunk_size=1024 * 1024): + if not chunk: + break + f.write(chunk) + pbar.update(len(chunk)) + + if os.path.getsize(partial_filename) != total_size and total_size > 0: + raise Exception("文件下载不完整，大小不匹配。") + + os.rename(partial_filename, final_filename) + tqdm.write(f"下载完成，文件已保存为: {os.path.basename(final_filename)}") + return package_name, True + + except Exception as e: + tqdm.write(f"\n请求异常: {package_name} - {e}") + return package_name, False + + +def run_downloads_in_parallel(packages_to_download): + """ + 使用多线程并发下载文件列表。 + """ + if not packages_to_download: + print("没有需要下载的包。") + return [] + + print(f"\n--- 开始多线程下载 {len(packages_to_download)} 个包 (最大线程数: {MAX_WORKERS}) ---") + + failed_packages = [] + + with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor: + future_to_package = {executor.submit(download_file, pkg): pkg for pkg in packages_to_download} + + main_pbar = tqdm(total=len(packages_to_download), desc="总下载进度") + + for future in concurrent.futures.as_completed(future_to_package): + package_name, success = future.result() + if not success: + failed_packages.append(package_name) + main_pbar.update(1) + + main_pbar.close() + + return failed_packages + + +def retry_failed_downloads(failed_packages): + """ + 反复重试下载失败的包，直到所有包都下载成功。 + """ + if not failed_packages: + print("\n所有包都已成功下载！") + return + + while failed_packages: + print(f"\n--- 以下包下载失败，开始重试: {len(failed_packages)} 个 ---") + for pkg in failed_packages: + print(f"- {pkg}") + + print(f"等待 {RETRY_DELAY_SECONDS} 秒后开始下一次重试...") + time.sleep(RETRY_DELAY_SECONDS) + + failed_packages = run_downloads_in_parallel(failed_packages) + + print("\n--- 所有包已成功下载！程序退出。---") + + + +# --- 主程序入口 --- +if __name__ == "__main__": + packages_from_csv = [] + if not os.path.exists(csv_file_path): + print(f"错误: 找不到文件 '{csv_file_path}'。请确保文件存在于脚本的同一目录下。") + exit() + + with open(csv_file_path, 'r', encoding='utf-8') as f: + reader = csv.DictReader(f) + for row in reader: + pkg = row.get('package_name', '').strip() + if pkg: + packages_from_csv.append(pkg) + + success_packages = set() + success_csv_path = 'success_tasks.csv' + if os.path.exists(success_csv_path): + with open(success_csv_path, 'r', encoding='utf-8') as f: + reader = csv.DictReader(f) + for row in reader: + pkg = row.get('包名', '').strip() + if pkg: + success_packages.add(pkg) + print(f"已成功处理的包: {len(success_packages)} 个") + + packages_from_csv = [pkg for pkg in packages_from_csv if pkg not in success_packages] + + existing_files = set() + for name in os.listdir(DOWNLOAD_FOLDER): + subdir = os.path.join(DOWNLOAD_FOLDER, name) + if os.path.isdir(subdir): + for f in os.listdir(subdir): + if f.endswith('.xapk'): + existing_files.add(f.replace('.xapk', '')) + print(len(packages_from_csv)) + print(existing_files) + print(existing_files.intersection(set(packages_from_csv))) + packages_to_download = [pkg for pkg in packages_from_csv if pkg not in existing_files] + packages_to_download = set(packages_to_download) + if not packages_to_download: + print("所有包都已存在，无需下载。") + else: + # 首次下载尝试 + failed_downloads = run_downloads_in_parallel(packages_to_download) + + # 进入重试循环，直到所有包都成功 + retry_failed_downloads(failed_downloads) \ No newline at end of file diff --git a/scripts/import_catalog_last_updated.py b/scripts/import_catalog_last_updated.py new file mode 100644 index 0000000..8ed2685 --- /dev/null +++ b/scripts/import_catalog_last_updated.py @@ -0,0 +1,179 @@ +#!/usr/bin/env python3 + +import argparse +import json +import os +import sqlite3 +import sys +import time +from typing import Any, Dict, List, Tuple + + +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if PROJECT_ROOT not in sys.path: + sys.path.insert(0, PROJECT_ROOT) + +from analytics import _last_updated_ts, _normalize_last_updated, _read_dict_rows +from config import MONITORING_DB_PATH + + +def _package_name_from_row(row: Dict[str, Any]) -> str: + return str(row.get("package_name") or row.get("包名") or row.get("package") or "").strip() + + +def _last_updated_from_row(row: Dict[str, Any]) -> str: + return _normalize_last_updated( + row.get("last_updated") + or row.get("最后更新") + or row.get("更新时间") + or "" + ) + + +def _safe_loads(value: Any) -> Dict[str, Any]: + if not value: + return {} + try: + payload = json.loads(value) + except (TypeError, ValueError): + return {} + return payload if isinstance(payload, dict) else {} + + +def _payload_with_last_updated(raw_payload: Any, last_updated: str) -> str: + payload = _safe_loads(raw_payload) + if not payload: + return str(raw_payload or "") + payload["last_updated"] = last_updated + original_row = payload.get("original_row") + if isinstance(original_row, dict): + original_row = dict(original_row) + original_row["last_updated"] = last_updated + original_row["最后更新"] = last_updated + original_row["更新时间"] = last_updated + payload["original_row"] = original_row + return json.dumps(payload, ensure_ascii=False) + + +def _load_csv_updates(csv_paths: List[str]) -> Tuple[Dict[str, str], int, int]: + updates: Dict[str, str] = {} + listed_rows = 0 + rows_with_date = 0 + for csv_path in csv_paths: + for row in _read_dict_rows(csv_path): + listed_rows += 1 + package_name = _package_name_from_row(row) + last_updated = _last_updated_from_row(row) + if not package_name or not last_updated: + continue + rows_with_date += 1 + existing = updates.get(package_name, "") + if not existing or _last_updated_ts(last_updated) > _last_updated_ts(existing): + updates[package_name] = last_updated + return updates, listed_rows, rows_with_date + + +def import_catalog_last_updated(args) -> int: + csv_paths = [str(item or "").strip() for item in args.csv_paths if str(item or "").strip()] + if not csv_paths: + raise ValueError("at least one csv path is required") + apply_changes = bool(args.apply) + update_payload = not bool(args.skip_payload) + incoming_updates, listed_rows, rows_with_date = _load_csv_updates(csv_paths) + now_iso = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + changed = 0 + missing_packages = [] + older_or_same = 0 + with sqlite3.connect(str(args.db_path or "").strip()) as connection: + connection.row_factory = sqlite3.Row + table_row = connection.execute( + """ + SELECT name + FROM sqlite_master + WHERE type = 'table' + AND name = 'app_catalog' + """ + ).fetchone() + if not table_row: + raise SystemExit(f"app_catalog not found in {args.db_path}") + for package_name, incoming_last_updated in sorted(incoming_updates.items()): + row = connection.execute( + """ + SELECT package_name, last_updated, task_payload_json + FROM app_catalog + WHERE package_name = ? + """, + (package_name,), + ).fetchone() + if not row: + missing_packages.append(package_name) + continue + existing_last_updated = _normalize_last_updated(row["last_updated"] or "") + if existing_last_updated and _last_updated_ts(existing_last_updated) >= _last_updated_ts(incoming_last_updated): + older_or_same += 1 + continue + changed += 1 + if apply_changes: + if update_payload: + connection.execute( + """ + UPDATE app_catalog + SET last_updated = ?, + task_payload_json = ?, + updated_at = ? + WHERE package_name = ? + """, + ( + incoming_last_updated, + _payload_with_last_updated(row["task_payload_json"], incoming_last_updated), + now_iso, + package_name, + ), + ) + else: + connection.execute( + """ + UPDATE app_catalog + SET last_updated = ?, + updated_at = ? + WHERE package_name = ? + """, + (incoming_last_updated, now_iso, package_name), + ) + missing_preview = ",".join(missing_packages[:10]) + print( + "dry_run={dry_run} db={db} csv_files={csv_files} listed_rows={listed_rows} " + "rows_with_date={rows_with_date} unique_packages_with_date={unique_packages} " + "updated={changed} older_or_same={older_or_same} missing={missing}{missing_suffix}".format( + dry_run=0 if apply_changes else 1, + db=str(args.db_path or "").strip(), + csv_files=len(csv_paths), + listed_rows=listed_rows, + rows_with_date=rows_with_date, + unique_packages=len(incoming_updates), + changed=changed, + older_or_same=older_or_same, + missing=len(missing_packages), + missing_suffix=f" missing_preview={missing_preview}" if missing_preview else "", + ) + ) + return 0 + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Import last_updated from one or more CSV files without changing catalog collection state." + ) + parser.add_argument("csv_paths", nargs="+", help="CSV files containing package_name/last_updated columns.") + parser.add_argument("--db-path", default=MONITORING_DB_PATH, help="Target monitoring sqlite path.") + parser.add_argument("--apply", action="store_true", help="Apply changes. Without this, only prints counts.") + parser.add_argument("--skip-payload", action="store_true", help="Do not update task_payload_json.") + return parser + + +def main() -> int: + return import_catalog_last_updated(build_parser().parse_args()) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/legacy_seed_analytics.py b/scripts/legacy_seed_analytics.py new file mode 100644 index 0000000..3fed9e0 --- /dev/null +++ b/scripts/legacy_seed_analytics.py @@ -0,0 +1,260 @@ +#!/usr/bin/env python3 + +import argparse +import csv +import os +import sqlite3 +import sys +from typing import List, Optional + + +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if PROJECT_ROOT not in sys.path: + sys.path.insert(0, PROJECT_ROOT) + +from analytics import AnalyticsRepository, _classify_restriction_status, _make_task_key +from monitoring import MonitoringRepository + + +def _table_columns(connection: sqlite3.Connection, table_name: str) -> List[str]: + return [row[1] for row in connection.execute(f"PRAGMA table_info({table_name})").fetchall()] + + +def import_monitoring(old_db_path: str, target_db_path: str): + MonitoringRepository(db_path=target_db_path) + AnalyticsRepository(db_path=target_db_path) + + with sqlite3.connect(old_db_path) as source, sqlite3.connect(target_db_path) as target: + source.row_factory = sqlite3.Row + target.row_factory = sqlite3.Row + source_columns = set(_table_columns(source, "task_execution")) + target_columns = set(_table_columns(target, "task_execution")) + common_columns = sorted((source_columns & target_columns) - {"execution_id"}) + select_columns = ["execution_id", *common_columns] + placeholders = ", ".join(["?"] * len(select_columns)) + insert_columns = ", ".join(select_columns) + update_columns = ", ".join([f"{column}=excluded.{column}" for column in common_columns]) + + rows = source.execute(f"SELECT {insert_columns} FROM task_execution").fetchall() + for row in rows: + target.execute( + f""" + INSERT INTO task_execution ({insert_columns}) + VALUES ({placeholders}) + ON CONFLICT(execution_id) DO UPDATE SET {update_columns} + """, + [row[column] for column in select_columns], + ) + target.commit() + print(f"Imported {len(rows)} task_execution rows into {target_db_path}") + + +def _parse_bool(value: str) -> bool: + return str(value or "").strip().lower() in {"1", "true", "yes"} + + +def _split_names(raw_value: str) -> List[str]: + return [item.strip() for item in str(raw_value or "").split(",") if item.strip()] + + +def _resolve_self_ratio(self_ratio: object, self_traffic_bytes: object = 0, total_traffic_bytes: object = 0) -> float: + try: + normalized_ratio = float(self_ratio or 0.0) + except (TypeError, ValueError): + normalized_ratio = 0.0 + if normalized_ratio > 0: + return normalized_ratio + try: + normalized_self_bytes = float(self_traffic_bytes or 0.0) + except (TypeError, ValueError): + normalized_self_bytes = 0.0 + try: + normalized_total_bytes = float(total_traffic_bytes or 0.0) + except (TypeError, ValueError): + normalized_total_bytes = 0.0 + if normalized_self_bytes > 0 and normalized_total_bytes > 0: + return round((normalized_self_bytes / normalized_total_bytes) * 100, 2) + return 0.0 + + +def import_app_summary(summary_csv_path: str, target_db_path: str): + repo = AnalyticsRepository(db_path=target_db_path) + imported = 0 + with open(summary_csv_path, "r", encoding="utf-8-sig", newline="") as handle: + reader = csv.DictReader(handle) + for row in reader: + package_name = str(row.get("package_name") or "").strip() + if not package_name: + continue + unique_domains = _split_names(row.get("unique_domain_names", "")) + second_level_domains = _split_names(row.get("unique_second_level_domains", "")) + summary = { + "app_name": str(row.get("app_name") or "").strip(), + "latest_task_key": "", + "latest_status": "success" if _parse_bool(row.get("test_success", "")) else "failed", + "latest_test_time": 0.0, + "latest_worker_id": str(row.get("test_host") or "").strip(), + "latest_task_detail": str(row.get("task_detail") or "").strip(), + "latest_failure_type": str(row.get("failure_type") or "").strip(), + "unique_domain_count": len(unique_domains), + "unique_domain_names": unique_domains, + "unique_second_level_domain_count": len(second_level_domains), + "unique_second_level_domains": second_level_domains, + "droidbot_steps": int(float(row.get("droidbot_steps") or 0)), + "gui_agent_steps": int(float(row.get("gui_agent_steps") or 0)), + "duration_seconds": float(row.get("duration_seconds") or 0.0), + "num_nodes": int(float(row.get("num_nodes") or 0)), + "num_reached_activities": int(float(row.get("num_reached_activities") or 0)), + "app_num_total_activities": int(float(row.get("app_num_total_activities") or 0)), + "total_traffic_bytes": 0, + "self_traffic_bytes": 0, + "server_traffic_bytes": 0, + "unrecognized_traffic_bytes": 0, + "self_ratio": float(row.get("self_ratio") or 0.0), + "recognition_ratio": 0.0, + "artifact_status": "partial", + } + restriction_status, retryability = _classify_restriction_status( + summary["latest_status"], + _resolve_self_ratio(summary["self_ratio"], summary["self_traffic_bytes"], summary["total_traffic_bytes"]), + summary["latest_failure_type"], + ) + summary["restriction_status"] = restriction_status + summary["retryability"] = retryability + repo.replace_package_snapshot(package_name, summary, [], [], []) + imported += 1 + print(f"Imported {imported} app summary rows into {target_db_path}") + + +def _get_latest_task_fallback(connection: sqlite3.Connection, package_name: str) -> Optional[sqlite3.Row]: + task_columns = set(_table_columns(connection, "task_execution")) + if not task_columns: + return None + selected_columns = [ + "status" if "status" in task_columns else "'' AS status", + "error_type" if "error_type" in task_columns else "'' AS error_type", + "result_detail" if "result_detail" in task_columns else "'' AS result_detail", + "error_message" if "error_message" in task_columns else "'' AS error_message", + "worker_id" if "worker_id" in task_columns else "'' AS worker_id", + "task_key" if "task_key" in task_columns else "'' AS task_key", + "COALESCE(task_ended_at, task_started_at, last_updated_at, 0) AS latest_time", + ] + return connection.execute( + f""" + SELECT {', '.join(selected_columns)} + FROM task_execution + WHERE package_name = ? + ORDER BY COALESCE(task_ended_at, task_started_at, last_updated_at, 0) DESC, execution_id DESC + LIMIT 1 + """, + (package_name,), + ).fetchone() + + +def repair_app_summary(target_db_path: str, only_missing: bool = False): + repo = AnalyticsRepository(db_path=target_db_path) + repaired = 0 + skipped = 0 + with repo._connect() as connection: + rows = repo._list_catalog_summaries(connection) + for row in rows: + package_name = str(row["package_name"] or "").strip() + if not package_name: + skipped += 1 + continue + if only_missing and row["restriction_status"] and row["retryability"] and row["latest_status"]: + skipped += 1 + continue + + latest_status = str(row["latest_status"] or "").strip() + latest_failure_type = str(row["latest_failure_type"] or "").strip() + fallback = None + if not latest_status or not latest_failure_type: + fallback_row = _get_latest_task_fallback(connection, package_name) + fallback = dict(fallback_row) if fallback_row else None + if fallback: + latest_status = latest_status or str(fallback["status"] or "").strip() + latest_failure_type = latest_failure_type or str(fallback["error_type"] or "").strip() + if not latest_status and not latest_failure_type: + skipped += 1 + continue + + app_name = str(row.get("app_name") or package_name).strip() or package_name + repo.replace_package_snapshot( + package_name, + { + "app_name": app_name, + "app_magic_label": row.get("app_magic_label", ""), + "latest_task_key": str((fallback or {}).get("task_key") or row.get("latest_task_key") or _make_task_key(app_name, package_name)).strip(), + "latest_status": latest_status, + "latest_worker_id": str((fallback or {}).get("worker_id") or row.get("latest_worker_id") or "").strip(), + "latest_failure_type": latest_failure_type, + "latest_task_detail": str( + (fallback or {}).get("result_detail") + or (fallback or {}).get("error_message") + or row.get("latest_task_detail") + or "" + ).strip(), + "latest_test_time": float((fallback or {}).get("latest_time") or row.get("latest_test_time") or 0.0), + "collection_status_reason": latest_failure_type or latest_status, + "collection_task_type": row.get("collection_task_type") or "new_app", + "duration_seconds": row.get("duration_seconds", 0), + "droidbot_steps": row.get("droidbot_steps", 0), + "gui_agent_steps": row.get("gui_agent_steps", 0), + "num_nodes": row.get("num_nodes", 0), + "num_reached_activities": row.get("num_reached_activities", 0), + "app_num_total_activities": row.get("app_num_total_activities", 0), + "total_traffic_bytes": row.get("total_traffic_bytes", 0), + "self_traffic_bytes": row.get("self_traffic_bytes", 0), + "server_traffic_bytes": row.get("server_traffic_bytes", 0), + "unrecognized_traffic_bytes": row.get("unrecognized_traffic_bytes", 0), + "model_flow_count": row.get("model_flow_count", 0), + "model_traffic_bytes": row.get("model_traffic_bytes", 0), + "task_payload": row.get("task_payload") or {}, + "downloads": row.get("downloads"), + "source_order": row.get("source_order"), + }, + [], + [], + [], + ) + repaired += 1 + suffix = " (only missing rows)" if only_missing else "" + print(f"Repaired {repaired} app_catalog/collection_task rows{suffix} in {target_db_path}; skipped {skipped}") + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Seed analytics tables from legacy files") + subparsers = parser.add_subparsers(dest="command", required=True) + + monitoring_parser = subparsers.add_parser("import-monitoring", help="Import task_execution rows from a legacy monitoring sqlite db") + monitoring_parser.add_argument("--from", dest="from_path", required=True, help="Legacy monitoring sqlite path") + monitoring_parser.add_argument("--to", dest="to_path", required=True, help="Target runtime sqlite path") + + summary_parser = subparsers.add_parser("import-app-summary", help="Import legacy app_domain_summary.csv into app_catalog and collection_task") + summary_parser.add_argument("--from", dest="from_path", required=True, help="Legacy app_domain_summary.csv path") + summary_parser.add_argument("--to", dest="to_path", required=True, help="Target runtime sqlite path") + + repair_parser = subparsers.add_parser("repair-app-summary", help="Backfill missing latest collection_task rows from task_execution") + repair_parser.add_argument("--to", dest="to_path", required=True, help="Target runtime sqlite path") + repair_parser.add_argument( + "--only-missing", + action="store_true", + help="Only repair rows where restriction_status/retryability are still empty", + ) + return parser + + +def main(): + parser = build_parser() + args = parser.parse_args() + if args.command == "import-monitoring": + import_monitoring(args.from_path, args.to_path) + elif args.command == "import-app-summary": + import_app_summary(args.from_path, args.to_path) + elif args.command == "repair-app-summary": + repair_app_summary(args.to_path, only_missing=bool(args.only_missing)) + + +if __name__ == "__main__": + main() diff --git a/scripts/manage_app_catalog.py b/scripts/manage_app_catalog.py new file mode 100644 index 0000000..5484f03 --- /dev/null +++ b/scripts/manage_app_catalog.py @@ -0,0 +1,1630 @@ +#!/usr/bin/env python3 + +import argparse +import csv +import glob +import json +import os +import sqlite3 +import sys +import time +from typing import Any, List, Optional + +import redis as redis_lib + + +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if PROJECT_ROOT not in sys.path: + sys.path.insert(0, PROJECT_ROOT) + +from analytics import ( + AnalyticsRepository, + AnalyticsService, + DEFAULT_INCREMENTAL_TOP_N, + LIGHT_RESTRICTED_MIN_NUM_NODES, + _humanize_error_type, + _incremental_batch_tags_to_json, + _last_updated_ts, + _normalize_app_magic_label, + _normalize_country_codes, + _normalize_incremental_batch_tags, + _normalize_last_updated, + _normalize_task_device_type, + _resolve_downloads_for_package, + build_task_payload_from_row, +) +from config import ANALYTICS_TRAFFIC_ROOT, MONITORING_DB_PATH, REDIS_HOST, REDIS_PORT, REDIS_DB, TASK_CSV_PATH, channel_name + + +def _build_service(db_path: str) -> AnalyticsService: + return AnalyticsService(db_path=db_path, start_worker=False) + + +def _load_active_catalog_packages(service: AnalyticsService) -> List[str]: + with service.repo._connect() as connection: + rows = connection.execute( + """ + SELECT package_name + FROM app_catalog + WHERE COALESCE(is_active, 1) = 1 + ORDER BY package_name ASC + """ + ).fetchall() + return [str(row["package_name"]).strip() for row in rows if row["package_name"]] + + +def _load_catalog_snapshot(service: AnalyticsService) -> dict: + with service.repo._connect() as connection: + summaries = service.repo._list_catalog_summaries(connection) + snapshot = {} + for row in summaries: + package_name = str(row["package_name"] or "").strip() + if not package_name: + continue + snapshot[package_name] = { + "catalog_active": 1 if row["catalog_active"] else 0, + "app_magic_label": str(row["app_magic_label"] or "").strip(), + "last_updated": _normalize_last_updated(row["last_updated"] or ""), + "collection_status": str(row["collection_status"] or "").strip(), + "latest_status": str(row["latest_status"] or "").strip(), + "restriction_status": str(row["restriction_status"] or "").strip(), + "collection_task_type": str(row["collection_task_type"] or "new_app").strip() or "new_app", + } + return snapshot + + +def _read_csv_rows(csv_path: str) -> List[dict]: + with open(csv_path, "r", newline="", encoding="utf-8-sig") as handle: + reader = csv.DictReader(handle) + rows = [] + for row in reader: + normalized = {} + for key, value in (row or {}).items(): + normalized[str(key or "").strip()] = str(value or "").strip() + rows.append(normalized) + return rows + + +def _load_csv_package_names(csv_path: str) -> List[str]: + if not csv_path or not os.path.exists(csv_path): + return [] + packages: List[str] = [] + seen = set() + for row in _read_csv_rows(csv_path): + package_name = str( + row.get("package_name") + or row.get("包名") + or row.get("package") + or "" + ).strip() + if not package_name or package_name in seen: + continue + seen.add(package_name) + packages.append(package_name) + return packages + + +def _csv_value_matches_stored_original(key: str, raw_value: Any, stored_value: Any) -> bool: + normalized_key = str(key or "").strip() + raw_text = str(raw_value or "").strip() + stored_text = str(stored_value or "").strip() + if normalized_key in {"last_updated", "最后更新", "更新时间"}: + raw_date = _normalize_last_updated(raw_text) + stored_date = _normalize_last_updated(stored_text) + return raw_date == stored_date if raw_date or stored_date else raw_text == stored_text + if normalized_key in {"device_type", "设备类型"}: + return _normalize_task_device_type(raw_text) == _normalize_task_device_type(stored_text) + return raw_text == stored_text + + +def _validate_catalog_payload_storage(service: AnalyticsService, incoming_payloads: List[tuple]) -> List[str]: + issues: List[str] = [] + with service.repo._connect() as connection: + for _, row, payload in incoming_payloads: + package_name = str(payload.get("package_name") or "").strip() + if not package_name: + continue + catalog_row = connection.execute( + """ + SELECT last_updated, country_code, device_type, task_payload_json + FROM app_catalog + WHERE package_name = ? + """, + (package_name,), + ).fetchone() + if not catalog_row or not str(catalog_row["task_payload_json"] or "").strip(): + issues.append(package_name) + continue + incoming_last_updated = _normalize_last_updated(payload.get("last_updated", "")) + if incoming_last_updated and _normalize_last_updated(catalog_row["last_updated"] or "") != incoming_last_updated: + issues.append(package_name) + continue + if str(catalog_row["country_code"] or "").strip() != str(payload.get("country_code") or "").strip(): + issues.append(package_name) + continue + if _normalize_task_device_type(catalog_row["device_type"] or "") != _normalize_task_device_type(payload.get("device_type") or ""): + issues.append(package_name) + continue + stored_payload = {} + try: + stored_payload = json.loads(catalog_row["task_payload_json"] or "{}") + except (TypeError, ValueError): + issues.append(package_name) + continue + original_row = stored_payload.get("original_row") + if not isinstance(original_row, dict): + issues.append(package_name) + continue + for key, value in row.items(): + normalized_key = str(key or "").strip() + if not normalized_key: + continue + if normalized_key not in original_row: + issues.append(package_name) + break + if not _csv_value_matches_stored_original(normalized_key, value, original_row.get(normalized_key, "")): + issues.append(package_name) + break + return sorted(set(issues)) + + +def _count_pending_collection_task_rows(service: AnalyticsService, package_names: List[str]) -> int: + normalized = [str(item or "").strip() for item in package_names if str(item or "").strip()] + if not normalized: + return 0 + placeholders = ",".join("?" for _ in normalized) + with service.repo._connect() as connection: + row = connection.execute( + f""" + SELECT COUNT(*) AS total + FROM collection_task + WHERE package_name IN ({placeholders}) + AND task_status = 'pending' + """, + normalized, + ).fetchone() + return int((row["total"] or 0) if row else 0) + + +def _normalize_manual_test_status(raw_status: str) -> str: + normalized = str(raw_status or "").strip() + if normalized in {"测试正常"}: + return "success" + if normalized in {"轻度受限"}: + return "light_restricted" + if normalized in {"严重受限"}: + return "severe_restricted" + raise ValueError("test_status must be 测试正常, 轻度受限, or 严重受限") + + +def _find_traffic_files_under_root(traffic_root: str, package_name: str) -> List[str]: + normalized_root = str(traffic_root or "").strip() + normalized_package_name = str(package_name or "").strip() + if not normalized_root or not normalized_package_name or not os.path.isdir(normalized_root): + return [] + pattern = os.path.join(normalized_root, "**", normalized_package_name, "traffic_count*.txt") + return sorted(path for path in glob.glob(pattern, recursive=True) if os.path.isfile(path)) + + +def add_app(args) -> int: + service = _build_service(args.db_path) + try: + created = service.add_catalog_app( + app_name=args.app_name, + package_name=args.package_name, + country_code=args.country_code, + device_type=args.device_type, + app_magic_label=getattr(args, "app_magic_label", ""), + last_updated=getattr(args, "last_updated", ""), + ) + print(f"added={created} package={args.package_name}") + return 0 + finally: + service.close() + + +def sync_from_csv(args) -> int: + csv_path = str(args.csv_path or "").strip() or TASK_CSV_PATH + batch_tags = str( + getattr(args, "batch_tags", "") + or getattr(args, "incremental_batch_tag", "") + or "" + ).strip() + dry_run = bool(getattr(args, "dry_run", False)) + service = _build_service(args.db_path) + try: + rows = _read_csv_rows(csv_path) + catalog_snapshot = _load_catalog_snapshot(service) + active_before = { + package_name + for package_name, item in catalog_snapshot.items() + if int(item.get("catalog_active") or 0) == 1 + } + history_packages = set(catalog_snapshot.keys()) + incoming_payloads = [] + seen_packages = set() + for index, row in enumerate(rows): + payload = build_task_payload_from_row(row) + package_name = str(payload.get("package_name") or "").strip() + if not package_name or package_name in seen_packages: + continue + seen_packages.add(package_name) + incoming_payloads.append((index, row, payload)) + incoming_packages = {str(payload["package_name"]).strip() for _, _, payload in incoming_payloads} + history_new_packages = sorted(incoming_packages - history_packages) + added_packages = history_new_packages + added_count = len(added_packages) + reduced_count = len(active_before - incoming_packages) + reactivated_packages = sorted((incoming_packages & history_packages) - active_before) + app_magic_label_updates = 0 + last_updated_fills = 0 + version_updates = 0 + older_last_updated_rows = 0 + unchanged_existing = 0 + pending_task_packages = set(added_packages) + for _, _, payload in incoming_payloads: + package_name = str(payload.get("package_name") or "").strip() + existing = catalog_snapshot.get(package_name) + if not existing: + continue + changed = False + incoming_magic_label = _normalize_app_magic_label(payload.get("app_magic_label", "")) + existing_magic_label = _normalize_app_magic_label(existing.get("app_magic_label", "")) + if incoming_magic_label and incoming_magic_label != existing_magic_label: + app_magic_label_updates += 1 + changed = True + incoming_last_updated = _normalize_last_updated(payload.get("last_updated", "")) + existing_last_updated = _normalize_last_updated(existing.get("last_updated", "")) + if incoming_last_updated and not existing_last_updated: + last_updated_fills += 1 + changed = True + pending_task_packages.add(package_name) + elif incoming_last_updated and existing_last_updated: + incoming_ts = _last_updated_ts(incoming_last_updated) + existing_ts = _last_updated_ts(existing_last_updated) + if incoming_ts > existing_ts: + version_updates += 1 + changed = True + pending_task_packages.add(package_name) + elif incoming_ts < existing_ts: + older_last_updated_rows += 1 + if existing.get("collection_status") in {"pending", "failed_terminal"}: + pending_task_packages.add(package_name) + if not changed: + unchanged_existing += 1 + if dry_run: + print( + "dry_run=1 csv={csv} listed={listed} upserted=0 added={added} reduced={reduced} " + "deactivated={deactivated} history_new={history_new} reactivated={reactivated} " + "app_magic_label_updates={app_magic_label_updates} last_updated_fills={last_updated_fills} " + "version_updates={version_updates} older_last_updated_rows={older_last_updated_rows} " + "unchanged_existing={unchanged_existing} would_create_pending_tasks={pending_tasks}{batch_suffix}".format( + csv=csv_path, + listed=len(incoming_packages), + added=added_count, + reduced=reduced_count, + deactivated=reduced_count, + history_new=len(history_new_packages), + reactivated=len(reactivated_packages), + app_magic_label_updates=app_magic_label_updates, + last_updated_fills=last_updated_fills, + version_updates=version_updates, + older_last_updated_rows=older_last_updated_rows, + unchanged_existing=unchanged_existing, + pending_tasks=len(pending_task_packages), + batch_suffix=( + f" batch_tags={batch_tags} incremental_marked={len(incoming_packages)}" + if batch_tags + else "" + ), + ) + ) + return 0 + pending_tasks_before = _count_pending_collection_task_rows(service, sorted(pending_task_packages)) + result = service.replace_catalog_from_rows(rows) + incremental_marked = 0 + if batch_tags: + service.clear_incremental_batch_tags(batch_tags) + incremental_marked = service.set_incremental_batch_tag_for_packages(sorted(incoming_packages), batch_tags) + pending_tasks_after = _count_pending_collection_task_rows(service, sorted(pending_task_packages)) + pending_tasks_created = max(0, pending_tasks_after - pending_tasks_before) + payload_storage_issues = _validate_catalog_payload_storage(service, incoming_payloads) + print( + "csv={csv} listed={listed} upserted={upserted} added={added} reduced={reduced} " + "deactivated={deactivated} history_new={history_new} reactivated={reactivated} " + "pending_tasks_created={pending_tasks_created} payload_storage_issues={payload_issues}{batch_suffix}".format( + csv=csv_path, + listed=int(result.get("listed") or 0), + upserted=int(result.get("upserted") or 0), + added=added_count, + reduced=reduced_count, + deactivated=int(result.get("deactivated") or 0), + history_new=len(history_new_packages), + reactivated=len(reactivated_packages), + pending_tasks_created=pending_tasks_created, + payload_issues=len(payload_storage_issues), + batch_suffix=( + f" batch_tags={batch_tags} incremental_marked={incremental_marked}" + if batch_tags + else "" + ), + ) + ) + try: + r = redis_lib.Redis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True) + r.publish(channel_name("catalog:updated"), json.dumps({"added": added_count})) + except Exception: + pass + return 0 + finally: + service.close() + + +def enqueue_high_priority(args) -> int: + csv_path = str(args.csv_path or "").strip() + batch_tags = str(getattr(args, "batch_tags", "") or "").strip() + force_pending = bool(getattr(args, "force_pending", False)) + if not csv_path: + raise ValueError("csv_path is required") + service = _build_service(args.db_path) + try: + rows = _read_csv_rows(csv_path) + catalog_snapshot = _load_catalog_snapshot(service) if not force_pending else {} + updated = 0 + skipped = 0 + seen_packages = set() + for row in rows: + payload = build_task_payload_from_row(row) + package_name = str(payload.get("package_name") or "").strip() + if not package_name or package_name in seen_packages: + continue + seen_packages.add(package_name) + if not force_pending: + existing = catalog_snapshot.get(package_name) + if existing and existing.get("collection_status", "") not in ("", "pending", "failed_terminal"): + skipped += 1 + continue + app_name = str(payload.get("app_name") or "").strip() + service.repo.upsert_high_priority_app(package_name, app_name, payload) + updated += 1 + incremental_marked = 0 + if batch_tags: + service.clear_incremental_batch_tags(batch_tags) + incremental_marked = service.set_incremental_batch_tag_for_packages(sorted(seen_packages), batch_tags) + batch_suffix = f" batch_tags={batch_tags} incremental_marked={incremental_marked}" if batch_tags else "" + force_suffix = " force_pending=1" if force_pending else "" + skip_suffix = f" skipped={skipped}" if skipped else "" + print(f"enqueue-high-priority csv={csv_path} updated={updated}{skip_suffix}{force_suffix}{batch_suffix}") + try: + r = redis_lib.Redis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True) + r.publish(channel_name("catalog:updated"), json.dumps({"added": updated, "priority": "high"})) + except Exception: + pass + return 0 + finally: + service.close() + + +def mark_added_from_csv_diff(args) -> int: + old_csv_path = str(getattr(args, "old_csv_path", "") or "").strip() + new_csv_path = str(args.new_csv_path or "").strip() + compare_with_db = bool(getattr(args, "compare_with_db", False)) + batch_tags = str(getattr(args, "batch_tags", "") or "").strip() + if not new_csv_path: + raise ValueError("new_csv_path is required") + if not batch_tags: + raise ValueError("batch_tags is required") + if not compare_with_db and not old_csv_path: + raise ValueError("old_csv_path is required when --compare-with-db is not set") + + new_packages = set(_load_csv_package_names(new_csv_path)) + service = _build_service(args.db_path) + try: + if compare_with_db: + # 与数据库历史采集全量应用进行比较，找出新增包名 + db_packages = set(_load_active_catalog_packages(service)) + # 同时加载非活跃的包名，确保覆盖全量历史数据 + with service.repo._connect() as connection: + all_rows = connection.execute( + "SELECT package_name FROM app_catalog" + ).fetchall() + all_db_packages = {str(row["package_name"]).strip() for row in all_rows if row["package_name"]} + old_packages = all_db_packages + source_label = "db(all_history)" + else: + old_packages = set(_load_csv_package_names(old_csv_path)) + source_label = f"csv({old_csv_path})" + + added_packages = sorted(new_packages - old_packages) + service.clear_incremental_batch_tags(batch_tags) + marked = service.set_incremental_batch_tag_for_packages(added_packages, batch_tags) + finally: + service.close() + print( + "baseline={baseline} new_csv={new_csv} new_total={new_total} added={csv_added} batch_tags={batch_tag} marked={marked}".format( + baseline=source_label, + new_csv=new_csv_path, + new_total=len(new_packages), + csv_added=len(added_packages), + batch_tag=batch_tags, + marked=marked, + ) + ) + return 0 + + +def show_overview(args) -> int: + batch_tags = str(getattr(args, "batch_tags", "") or "").strip() + top_n = int(getattr(args, "top_n", DEFAULT_INCREMENTAL_TOP_N) or DEFAULT_INCREMENTAL_TOP_N) + service = _build_service(args.db_path) + try: + overview = service.get_overview( + batch_tags=batch_tags, + top_n=top_n, + ) + finally: + service.close() + print( + "scope={scope} batch_tag={batch_tag} top_n={top_n} batch_total={batch_total} missing_rank={missing_rank} total={total} qualified={qualified} abnormal={abnormal} pending={pending} analyzed={analyzed}".format( + scope=overview.get("scope") or ("incremental_topn" if batch_tags else "topn_catalog"), + batch_tag=batch_tags or "-", + top_n=int(overview.get("top_n") or top_n), + batch_total=int(overview.get("base_scope_app_count") or 0), + missing_rank=int(overview.get("base_scope_missing_source_order_count") or 0), + total=int(overview.get("app_count") or 0), + qualified=int(overview.get("qualified_count") or 0), + abnormal=int(overview.get("failed_terminal_count") or 0), + pending=int(overview.get("pending_count") or 0), + analyzed=int(overview.get("analyzed_app_count") or 0), + ) + ) + return 0 + + +def set_batch_tag_from_csv(args) -> int: + csv_path = str(args.csv_path).strip() + tag = str(args.tag).strip() + if not csv_path or not os.path.exists(csv_path): + print(f"错误: CSV 文件不存在: {csv_path}") + return 1 + if not tag: + print("错误: --tag 不能为空") + return 1 + service = _build_service(args.db_path) + try: + catalog_snapshot = _load_catalog_snapshot(service) + rows = _read_csv_rows(csv_path) + to_tag: List[str] = [] + skipped = 0 + new_apps = 0 + updated_apps = 0 + reclassified = 0 + seen = set() + for row in rows: + package_name = str(row.get("package_name") or row.get("包名") or "").strip() + if not package_name or package_name in seen: + continue + seen.add(package_name) + existing = catalog_snapshot.get(package_name) + if not existing: + to_tag.append(package_name) + new_apps += 1 + continue + incoming_last_updated = _normalize_last_updated(row.get("last_updated", "")) + existing_last_updated = _normalize_last_updated(existing.get("last_updated", "")) + incoming_ts = _last_updated_ts(incoming_last_updated) if incoming_last_updated else 0.0 + existing_ts = _last_updated_ts(existing_last_updated) if existing_last_updated else 0.0 + is_version_update = incoming_ts > existing_ts > 0.0 + is_fill_update = bool(incoming_last_updated) and not existing_last_updated + is_version_or_fill = is_version_update or is_fill_update + incoming_magic_label = _normalize_app_magic_label(row.get("app_magic_label", "")) + existing_magic_label = _normalize_app_magic_label(existing.get("app_magic_label", "")) + is_magic_label_changed = bool(incoming_magic_label) and incoming_magic_label != existing_magic_label + if is_version_or_fill: + to_tag.append(package_name) + updated_apps += 1 + elif is_magic_label_changed: + to_tag.append(package_name) + reclassified += 1 + else: + skipped += 1 + if not to_tag: + print(f"tag={tag} csv_total={len(seen)} new={new_apps} updated={updated_apps} reclassified={reclassified} skipped={skipped} → 无需打tag的应用") + return 0 + marked = service.set_incremental_batch_tag_for_packages(to_tag, tag) + print( + f"tag={tag} csv_total={len(seen)} new={new_apps} updated={updated_apps} " + f"reclassified={reclassified} skipped={skipped} marked={marked}" + ) + finally: + service.close() + return 0 + + +def clear_batch_tag(args) -> int: + tag = str(args.tag).strip() + if not tag: + print("错误: --tag 不能为空") + return 1 + service = _build_service(args.db_path) + try: + with service.repo._write_lock, service.repo._connect() as conn: + conn.row_factory = sqlite3.Row + rows = conn.execute( + """ + SELECT package_name, batch_tags + FROM app_catalog + WHERE COALESCE(batch_tags, '') LIKE '%"' + || ? + || '"%' + """ + , + (tag,), + ).fetchall() + cleared = 0 + skipped = 0 + now_iso = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + for row in rows: + summary = service.repo.get_collection_row(row["package_name"]) or {} + tags = _normalize_incremental_batch_tags(row["batch_tags"]) + task_type = str(summary.get("collection_task_type") or "").strip() + is_new_app = task_type == "new_app" + has_multiple_tags = len(tags) > 1 + if is_new_app and has_multiple_tags: + # 仅对 old batch 就已标记为 new_app 且被错误追加了第二个 tag 的包做清除 + new_tags = [t for t in tags if t != tag] + conn.execute( + """ + UPDATE app_catalog + SET batch_tags = ?, + updated_at = ? + WHERE package_name = ? + """, + (_incremental_batch_tags_to_json(new_tags), now_iso, row["package_name"]), + ) + cleared += 1 + else: + skipped += 1 + print(f"tag={tag} cleared={cleared} skipped={skipped}") + finally: + service.close() + return 0 + + +def deactivate_stale_tagged(args) -> int: + """将指定 tag 下 new_app 且多 tag（被错误追加到旧批次）的应用 deactivate""" + tag = str(args.tag).strip() + if not tag: + print("错误: --tag 不能为空") + return 1 + service = _build_service(args.db_path) + try: + with service.repo._write_lock, service.repo._connect() as conn: + conn.row_factory = sqlite3.Row + rows = conn.execute( + """ + SELECT package_name, batch_tags + FROM app_catalog + WHERE COALESCE(is_active, 1) = 1 + AND COALESCE(batch_tags, '') LIKE '%"' + || ? + || '"%' + """ + , + (tag,), + ).fetchall() + to_deactivate = [] + for r in rows: + summary = service.repo.get_collection_row(r["package_name"]) or {} + tags = _normalize_incremental_batch_tags(r["batch_tags"]) + task_type = str(summary.get("collection_task_type") or "").strip() + if task_type == "new_app" and len(tags) > 1: + to_deactivate.append(r["package_name"]) + if not to_deactivate: + print(f"tag={tag} 无需 deactivate 的应用") + return 0 + now_iso = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + placeholders = ",".join("?" for _ in to_deactivate) + cursor = conn.execute( + f""" + UPDATE app_catalog + SET is_active = 0, updated_at = ? + WHERE package_name IN ({placeholders}) + """, + [now_iso, *to_deactivate], + ) + deactivated = int(cursor.rowcount or 0) + print(f"tag={tag} deactivated={deactivated}") + finally: + service.close() + return 0 + + +def set_pending(args) -> int: + packages: List[str] = [str(item or "").strip() for item in (args.packages or []) if str(item or "").strip()] + tag_filter = str(getattr(args, "tag", "") or "").strip() + tag_exact = bool(getattr(args, "tag_exact", False)) + error_type_prefix = str(getattr(args, "error_type_prefix", "") or "").strip() + catalog_active = getattr(args, "catalog_active", None) + collection_task_type = str(getattr(args, "collection_task_type", "") or "").strip() + collection_status_filter = str(getattr(args, "status", "") or "").strip() + reason = str(args.reason or "manual_pending").strip() or "manual_pending" + dry_run = bool(getattr(args, "dry_run", False)) + + has_filters = any([tag_filter, error_type_prefix, catalog_active is not None, collection_task_type, collection_status_filter]) + + service = _build_service(args.db_path) + try: + summaries = _list_catalog_summaries(service) + selected = summaries + if packages: + package_set = set(packages) + selected = [item for item in selected if item["package_name"] in package_set] + if tag_filter: + if tag_exact: + selected = [ + item for item in selected + if _normalize_incremental_batch_tags(item.get("incremental_batch_tags")) == [tag_filter] + ] + else: + selected = [ + item for item in selected + if tag_filter in (item.get("incremental_batch_tags") or []) + ] + if error_type_prefix: + selected = [ + item for item in selected + if str(item.get("latest_failure_type") or "").startswith(error_type_prefix) + ] + if catalog_active is not None: + selected = [item for item in selected if int(bool(item.get("catalog_active"))) == int(catalog_active)] + if collection_task_type: + if collection_task_type in ("new_app", "new_app_or_null"): + selected = [item for item in selected if (item.get("collection_task_type") or "new_app") == "new_app"] + else: + selected = [item for item in selected if item.get("collection_task_type") == collection_task_type] + if collection_status_filter: + selected = [item for item in selected if item.get("collection_status") == collection_status_filter] + + if packages and not has_filters: + if dry_run: + print(f"would_update={len(selected)} status=pending [DRY RUN]") + return 0 + updated = service.set_packages_pending([item["package_name"] for item in selected], reason=reason) + print(f"updated={updated} status=pending") + return 0 + + if not packages and not has_filters: + print("No filter criteria or packages provided.") + return 1 + + if dry_run: + print(f"would_update={len(selected)} status=pending [DRY RUN]") + return 0 + updated = service.set_packages_pending([item["package_name"] for item in selected], reason=reason) + print(f"updated={updated} status=pending") + return 0 + finally: + service.close() + + +def _extract_country_codes_for_export(task_payload_json: str, fallback_country_code: str) -> List[str]: + payload = {} + if task_payload_json: + try: + payload = json.loads(task_payload_json) + except (TypeError, ValueError): + payload = {} + if not isinstance(payload, dict): + payload = {} + raw_country_codes = payload.get("country_codes") or payload.get("country_code") or fallback_country_code + return _normalize_country_codes(raw_country_codes) + + +def _build_google_play_urls(package_name: str, country_codes: List[str]) -> str: + normalized_package_name = str(package_name or "").strip() + if not normalized_package_name: + return "" + normalized_country_codes = _normalize_country_codes(country_codes) + if not normalized_country_codes: + return f"https://play.google.com/store/apps/details?id={normalized_package_name}" + return ",".join( + f"https://play.google.com/store/apps/details?id={normalized_package_name}&gl={country_code}" + for country_code in normalized_country_codes + ) + + +def _normalize_optional_top_n(value: Any) -> Optional[int]: + if value is None: + return None + normalized_text = str(value).strip() + if not normalized_text: + return None + try: + normalized_top_n = int(normalized_text) + except (TypeError, ValueError) as exc: + raise ValueError("top_n must be a positive integer") from exc + if normalized_top_n <= 0: + raise ValueError("top_n must be a positive integer") + return normalized_top_n + + +def _list_catalog_summaries(service: AnalyticsService) -> List[dict]: + with service.repo._connect() as connection: + return service.repo._list_catalog_summaries(connection) + + +def _filter_summaries_by_scope(summaries: List[dict], *, tag: str = "", top_n: Any = None) -> List[dict]: + normalized_tag = str(tag or "").strip() + normalized_top_n = _normalize_optional_top_n(top_n) + filtered = list(summaries) + if normalized_tag: + filtered = [item for item in filtered if normalized_tag in (item.get("incremental_batch_tags") or [])] + if normalized_top_n is not None: + filtered = [ + item for item in filtered + if item.get("source_order") is not None and int(item.get("source_order") or 0) < normalized_top_n + ] + return filtered + + +def _scope_parts(tag: str = "", top_n: Any = None) -> List[str]: + parts: List[str] = [] + if str(tag or "").strip(): + parts.append(f"tag={str(tag).strip()}") + normalized_top_n = _normalize_optional_top_n(top_n) + if normalized_top_n is not None: + parts.append(f"top_n={normalized_top_n}") + return parts + + +def _set_device_type_for_packages(service: AnalyticsService, summaries: List[dict], device_type: str) -> int: + normalized_device_type = _normalize_task_device_type(device_type) + if not normalized_device_type: + return 0 + now_iso = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + updated = 0 + with service.repo._write_lock, service.repo._connect() as connection: + for summary in summaries: + package_name = str(summary.get("package_name") or "").strip() + if not package_name: + continue + payload = dict(summary.get("task_payload") or {}) + payload["device_type"] = normalized_device_type + original_row = payload.get("original_row") + if isinstance(original_row, dict): + original_row_copy = dict(original_row) + original_row_copy["device_type"] = normalized_device_type + payload["original_row"] = original_row_copy + connection.execute( + """ + UPDATE app_catalog + SET device_type = ?, + task_payload_json = ?, + updated_at = ? + WHERE package_name = ? + """, + (normalized_device_type, json.dumps(payload, ensure_ascii=False), now_iso, package_name), + ) + updated += 1 + return updated + + +def set_pending_by_error_type(args) -> int: + error_types: List[str] = [str(item or "").strip() for item in args.error_types if str(item or "").strip()] + error_type_prefix = str(getattr(args, "error_type_prefix", "") or "").strip() + all_non_retryable = bool(getattr(args, "all_non_retryable", False)) + non_retryable_only = bool(getattr(args, "non_retryable_only", False)) + pending_only = bool(getattr(args, "pending_only", False)) + tag_filter = str(getattr(args, "tag", "") or "").strip() + if not error_types and not error_type_prefix and not all_non_retryable: + raise ValueError("at least one exact error type, error type prefix, or --all-non-retryable is required") + if non_retryable_only and pending_only: + raise ValueError("--non-retryable-only and --pending-only cannot be used together") + if pending_only and all_non_retryable: + raise ValueError("--pending-only cannot be used with --all-non-retryable") + reason = str(args.reason or "manual_pending_by_error_type").strip() or "manual_pending_by_error_type" + requested_device_type = str(getattr(args, "device_type", "") or "").strip() + normalized_device_type = _normalize_task_device_type(requested_device_type) if requested_device_type else "" + if requested_device_type and not normalized_device_type: + raise ValueError("device_type must be emulator, physical, or any") + + scope_name = "failed_terminal_non_retryable" + if pending_only: + scope_name = "pending_only" + scope_parts: List[str] = [] + if error_types: + scope_parts.append(f"error_types={','.join(error_types)}") + if error_type_prefix: + scope_parts.append(f"error_type_prefix={error_type_prefix}") + if all_non_retryable: + scope_parts.append("all_non_retryable=1") + if non_retryable_only: + scope_parts.append("non_retryable_only=1") + if pending_only: + scope_parts.append("pending_only=1") + if normalized_device_type: + scope_parts.append(f"device_type={normalized_device_type}") + scope_parts.extend(_scope_parts(tag_filter, getattr(args, "top_n", None))) + + service = _build_service(args.db_path) + try: + summaries = _filter_summaries_by_scope( + _list_catalog_summaries(service), + tag=tag_filter, + top_n=getattr(args, "top_n", None), + ) + if pending_only: + summaries = [item for item in summaries if item.get("collection_status") == "pending"] + else: + summaries = [ + item for item in summaries + if item.get("collection_status") == "failed_terminal" + and item.get("restriction_status") == "severe_restricted" + and item.get("retryability") == "non_retryable" + ] + if error_types: + summaries = [item for item in summaries if item.get("latest_failure_type") in error_types] + if error_type_prefix: + summaries = [ + item for item in summaries + if str(item.get("latest_failure_type") or "").startswith(error_type_prefix) + ] + if normalized_device_type: + _set_device_type_for_packages(service, summaries, normalized_device_type) + updated = service.set_packages_pending([item["package_name"] for item in summaries], reason=reason) + finally: + service.close() + + scope_text = " ".join(scope_parts) if scope_parts else "all_non_retryable=0" + print( + f"updated={updated} status=pending " + f"scope={scope_name} " + f"{scope_text}" + ) + return 0 + + +def set_pending_success_zero_self_ratio(args) -> int: + reason = ( + str(args.reason or "manual_pending_success_zero_self_ratio").strip() + or "manual_pending_success_zero_self_ratio" + ) + tag_filter = str(getattr(args, "tag", "") or "").strip() + scope_parts = _scope_parts(tag_filter, getattr(args, "top_n", None)) + service = _build_service(args.db_path) + try: + summaries = _filter_summaries_by_scope( + _list_catalog_summaries(service), + tag=tag_filter, + top_n=getattr(args, "top_n", None), + ) + selected = [ + item for item in summaries + if item.get("collection_status") == "qualified" + and item.get("latest_status") == "success" + and item.get("restriction_status") in {"success", "light_restricted"} + and float(item.get("self_ratio") or 0.0) <= float(args.max_self_ratio) + and int(item.get("num_nodes") or 0) < LIGHT_RESTRICTED_MIN_NUM_NODES + ] + updated = service.set_packages_pending([item["package_name"] for item in selected], reason=reason) + finally: + service.close() + scope_text = f" {' '.join(scope_parts)}" if scope_parts else "" + print( + "updated=" + f"{updated} status=pending latest_status=success max_self_ratio={float(args.max_self_ratio)} " + f"max_num_nodes={LIGHT_RESTRICTED_MIN_NUM_NODES - 1}{scope_text}" + ) + return 0 + + +def set_pending_retryable(args) -> int: + reason = str(args.reason or "manual_pending_retryable").strip() or "manual_pending_retryable" + tag_filter = str(getattr(args, "tag", "") or "").strip() + scope_parts = _scope_parts(tag_filter, getattr(args, "top_n", None)) + current_status = str(args.current_status or "").strip() + if current_status: + scope_parts.append(f"current_status={current_status}") + service = _build_service(args.db_path) + try: + summaries = _filter_summaries_by_scope( + _list_catalog_summaries(service), + tag=tag_filter, + top_n=getattr(args, "top_n", None), + ) + selected = [ + item for item in summaries + if item.get("retryability") == "retryable" + and item.get("latest_status") != "success" + ] + if current_status: + selected = [item for item in selected if item.get("collection_status") == current_status] + updated = service.set_packages_pending([item["package_name"] for item in selected], reason=reason) + finally: + service.close() + scope_text = f" {' '.join(scope_parts)}" if scope_parts else "" + print(f"updated={updated} status=pending retryability=retryable{scope_text}") + return 0 + + +def set_existing_non_retryable_light_restricted(args) -> int: + reason = ( + str(args.reason or "manual_existing_non_retryable_light_restricted").strip() + or "manual_existing_non_retryable_light_restricted" + ) + min_num_nodes = int(args.min_num_nodes) + tag_filter = str(getattr(args, "tag", "") or "").strip() + scope_parts = _scope_parts(tag_filter, getattr(args, "top_n", None)) + service = _build_service(args.db_path) + try: + summaries = _filter_summaries_by_scope( + _list_catalog_summaries(service), + tag=tag_filter, + top_n=getattr(args, "top_n", None), + ) + selected = [ + item for item in summaries + if item.get("retryability") == "non_retryable" + and item.get("restriction_status") == "severe_restricted" + and item.get("collection_status") == "failed_terminal" + and int(item.get("num_nodes") or 0) >= min_num_nodes + ] + # 新 schema 不落库保存派生状态；满足该条件的包重新入 pending 以便重采。 + updated = service.set_packages_pending([item["package_name"] for item in selected], reason=reason) + finally: + service.close() + scope_text = f" {' '.join(scope_parts)}" if scope_parts else "" + print( + f"updated={updated} status=pending " + f"reason={reason} min_num_nodes={min_num_nodes}{scope_text}" + ) + return 0 + + +def export_apps_by_error_type(args) -> int: + error_types: List[str] = [str(item or "").strip() for item in args.error_types if str(item or "").strip()] + error_type_prefix = str(getattr(args, "error_type_prefix", "") or "").strip() + output_path = str(getattr(args, "output", "") or "").strip() + non_retryable_only = bool(getattr(args, "non_retryable_only", False)) + tag_filter = str(getattr(args, "tag", "") or "").strip() + if not error_types and not error_type_prefix and not non_retryable_only: + raise ValueError("at least one exact error type, error type prefix, or --non-retryable-only is required") + if not output_path: + raise ValueError("output path is required") + + output_dir = os.path.dirname(output_path) + if output_dir: + os.makedirs(output_dir, exist_ok=True) + + scope_parts = _scope_parts(tag_filter, getattr(args, "top_n", None)) + + service = _build_service(args.db_path) + try: + rows = _filter_summaries_by_scope( + _list_catalog_summaries(service), + tag=tag_filter, + top_n=getattr(args, "top_n", None), + ) + finally: + service.close() + + exported_rows = [] + for row in rows: + if non_retryable_only and not ( + row.get("restriction_status") == "severe_restricted" + and row.get("retryability") == "non_retryable" + ): + continue + package_name = str(row["package_name"] or "").strip() + matched_error_type = str(row["latest_failure_type"] or "").strip() + if error_types and matched_error_type not in error_types: + continue + if error_type_prefix and not matched_error_type.startswith(error_type_prefix): + continue + downloads = _resolve_downloads_for_package( + package_name, + row["task_payload_json"], + downloads=row["downloads"], + ) + country_codes = _extract_country_codes_for_export( + row["task_payload_json"], + row["country_code"] or "", + ) + exported_rows.append( + { + "app_name": row["app_name"] or "", + "package_name": package_name, + "error_type": matched_error_type, + "error_type_description": _humanize_error_type(matched_error_type), + "reason_message": row["latest_task_detail"] or "", + "downloads": "" if downloads is None else int(downloads), + "google_play_urls": _build_google_play_urls(package_name, country_codes), + } + ) + + with open(output_path, "w", newline="", encoding="utf-8-sig") as handle: + writer = csv.DictWriter( + handle, + fieldnames=[ + "app_name", + "package_name", + "error_type", + "error_type_description", + "reason_message", + "downloads", + "google_play_urls", + ], + ) + writer.writeheader() + writer.writerows(exported_rows) + + exact_text = ",".join(error_types) if error_types else "-" + prefix_text = error_type_prefix or "-" + print( + f"exported=1 rows={len(exported_rows)} output={output_path} " + f"error_types={exact_text} error_type_prefix={prefix_text} " + f"non_retryable_only={int(non_retryable_only)}" + f"{(' ' + ' '.join(scope_parts)) if scope_parts else ''}" + ) + return 0 + + +def rebuild_from_csv(args) -> int: + csv_path = str(args.csv_path or "").strip() or TASK_CSV_PATH + packages = _load_csv_package_names(csv_path) + if not packages: + print(f"rebuilt=0 partial=0 failed=0 csv={csv_path}") + return 0 + + service = _build_service(args.db_path) + rebuilt = 0 + partial = 0 + failed = 0 + failed_packages: List[str] = [] + try: + for package_name in packages: + try: + summary = service.rebuild_package_now(package_name) + rebuilt += 1 + if summary.get("artifact_status") != "complete": + partial += 1 + except Exception: + failed += 1 + failed_packages.append(package_name) + suffix = f" failed_packages={','.join(failed_packages[:10])}" if failed_packages else "" + print(f"rebuilt={rebuilt} partial={partial} failed={failed} csv={csv_path}{suffix}") + return 0 if failed == 0 else 1 + finally: + service.close() + + +def rebuild_status_from_csv(args) -> int: + csv_path = str(args.csv_path or "").strip() + traffic_root = str(args.traffic_root or "").strip() or ANALYTICS_TRAFFIC_ROOT + rows = _read_csv_rows(csv_path) + normalized_rows = [row for row in rows if str(row.get("package_name") or "").strip()] + if not normalized_rows: + print(f"updated=0 mode=manual-status csv={csv_path} traffic_root={traffic_root}") + return 0 + + service = AnalyticsService( + db_path=args.db_path, + traffic_root=traffic_root, + start_worker=False, + ) + repo = AnalyticsRepository(db_path=args.db_path) + original_find_traffic_files = service._find_traffic_files + service._find_traffic_files = lambda package_name: _find_traffic_files_under_root(traffic_root, package_name) + + updated = 0 + failed = 0 + failed_packages: List[str] = [] + try: + for row in normalized_rows: + package_name = str(row.get("package_name") or "").strip() + manual_status = _normalize_manual_test_status(row.get("test_status", "")) + exception_reason = str(row.get("exception_reason") or "").strip() + if manual_status == "severe_restricted" and not exception_reason: + failed += 1 + failed_packages.append(package_name) + continue + + existing_row = repo.get_collection_row(package_name) or {} + app_name = str(existing_row.get("app_name") or package_name).strip() or package_name + app_magic_label = str(row.get("app_magic_label") or existing_row.get("app_magic_label") or "").strip() + last_updated = str(row.get("last_updated") or existing_row.get("last_updated") or "").strip() + latest_task_override = { + "task_key": f"{app_name}_{package_name}", + "app_name": app_name, + "package_name": package_name, + "worker_id": "manual_csv", + "latest_time": time.time(), + } + if manual_status == "success": + latest_task_override.update( + { + "status": "success", + "error_type": "", + "result_detail": "", + "download_errors": {}, + } + ) + elif manual_status == "light_restricted": + latest_task_override.update( + { + "status": "failed", + "error_type": "", + "result_detail": "", + "download_errors": {}, + } + ) + else: + latest_task_override.update( + { + "status": "failed", + "error_type": exception_reason, + "result_detail": exception_reason, + "download_errors": {}, + } + ) + + try: + service.rebuild_package_now(package_name, latest_task_override=latest_task_override) + now_iso = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + with repo._write_lock, repo._connect() as connection: + connection.execute( + """ + UPDATE app_catalog + SET app_magic_label = COALESCE(NULLIF(?, ''), app_magic_label), + last_updated = COALESCE(NULLIF(?, ''), last_updated), + updated_at = ? + WHERE package_name = ? + """, + (app_magic_label, last_updated, now_iso, package_name), + ) + updated += 1 + except Exception: + failed += 1 + failed_packages.append(package_name) + finally: + service._find_traffic_files = original_find_traffic_files + service.close() + + suffix = f" failed_packages={','.join(failed_packages[:10])}" if failed_packages else "" + print( + f"updated={updated} failed={failed} mode=manual-status " + f"csv={csv_path} traffic_root={traffic_root}{suffix}" + ) + return 0 if failed == 0 else 1 + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="Manage app_catalog and collection_task entries.") + parser.add_argument("--db-path", default=MONITORING_DB_PATH, help="Target monitoring sqlite path.") + subparsers = parser.add_subparsers(dest="command", required=True) + + add_parser = subparsers.add_parser("add-app", help="Add one app to app_catalog and create a pending collection_task.") + add_parser.add_argument("--app-name", required=True, help="App name.") + add_parser.add_argument("--package-name", required=True, help="Package name.") + add_parser.add_argument("--app-magic-label", default="", help="Related app magic label.") + add_parser.add_argument("--last-updated", default="", help="App last updated date, e.g. 2026/04/07.") + add_parser.add_argument("--country-code", default="", help="Country code string, e.g. US,JP.") + add_parser.add_argument("--device-type", default="", help="Device type: emulator, physical, or any.") + add_parser.set_defaults(func=add_app) + + sync_parser = subparsers.add_parser( + "sync-from-csv", + help="Sync current catalog from csv, keep missing apps in DB but hide them from queue and dashboard.", + ) + sync_parser.add_argument( + "--csv-path", + default=TASK_CSV_PATH, + help="CSV path containing app_name/package_name rows.", + ) + sync_parser.add_argument( + "--incremental-batch-tag", + default="", + help="Optional batch tag to write into app_catalog.batch_tags for packages in this CSV sync.", + ) + sync_parser.add_argument( + "--dry-run", + action="store_true", + help="Preview sync impact without changing app_catalog or batch tags.", + ) + sync_parser.set_defaults(func=sync_from_csv) + + mark_added_parser = subparsers.add_parser( + "mark-added-from-csv-diff", + help="Compare old/new CSV or new CSV vs DB history, and mark added packages with an incremental batch tag.", + ) + mark_added_parser.add_argument("--old-csv-path", default="", help="Baseline CSV path (not required when --compare-with-db is set).") + mark_added_parser.add_argument("--new-csv-path", required=True, help="Latest CSV path.") + mark_added_parser.add_argument( + "--compare-with-db", + action="store_true", + help="Use all historical packages in DB as baseline instead of old CSV.", + ) + mark_added_parser.add_argument( + "--incremental-batch-tag", + required=True, + help="Batch tag to write onto packages that exist in new CSV but not in baseline.", + ) + mark_added_parser.set_defaults(func=mark_added_from_csv_diff) + + show_overview_parser = subparsers.add_parser( + "show-overview", + help="Show analytics overview for the selected incremental batch within TopN of the full catalog.", + ) + show_overview_parser.add_argument( + "--incremental-batch-tag", + default="", + help="Optional incremental batch tag. Empty means use the current active catalog.", + ) + show_overview_parser.add_argument( + "--top-n", + type=int, + default=DEFAULT_INCREMENTAL_TOP_N, + help=f"Only count apps whose full-catalog rank is within TopN. Default: {DEFAULT_INCREMENTAL_TOP_N}.", + ) + show_overview_parser.set_defaults(func=show_overview) + + pending_parser = subparsers.add_parser("set-pending", help="Mark existing apps as pending for recollection.") + pending_parser.add_argument("packages", nargs="*", help="Package names to mark pending (optional, use with filter flags).") + pending_parser.add_argument("--tag", default="", help="Only affect apps with this batch_tags.") + pending_parser.add_argument("--error-type-prefix", default="", help="Prefix match for latest_failure_type, e.g. DOWNLOAD_ERROR.") + pending_parser.add_argument("--catalog-active", type=int, choices=[0, 1], default=None, help="Filter by catalog_active (0 or 1).") + pending_parser.add_argument("--collection-task-type", default="", help="Filter by collection_task_type, e.g. new_app or app_update.") + pending_parser.add_argument("--reason", default="manual_pending", help="collection_status_reason to write.") + pending_parser.add_argument("--dry-run", action="store_true", help="Preview what would be updated without making changes.") + pending_parser.add_argument("--tag-exact", action="store_true", help="When set, use exact JSON array match ([\"tag\"]) instead of LIKE for --tag.") + pending_parser.add_argument( + "--status", + default="", + choices=["qualified", "failed_terminal", "pending"], + help="Filter by current collection_status.", + ) + pending_parser.set_defaults(func=set_pending) + + rebuild_from_csv_parser = subparsers.add_parser( + "rebuild-from-csv", + help="Rebuild analytics by package list, or apply manual status updates from package_name/test_status/exception_reason csv.", + ) + rebuild_from_csv_parser.add_argument( + "--csv-path", + default=TASK_CSV_PATH, + help="CSV path containing package_name rows.", + ) + rebuild_from_csv_parser.set_defaults(func=rebuild_from_csv) + + rebuild_status_from_csv_parser = subparsers.add_parser( + "rebuild-status-from-csv", + help="Read package_name/test_status/exception_reason from csv, rebuild traffic metrics from a folder, and apply manual statuses.", + ) + rebuild_status_from_csv_parser.add_argument( + "--csv-path", + required=True, + help="CSV path containing package_name,test_status,exception_reason rows.", + ) + rebuild_status_from_csv_parser.add_argument( + "--traffic-root", + default=ANALYTICS_TRAFFIC_ROOT, + help="Folder to recursively search for /traffic_count*.txt files.", + ) + rebuild_status_from_csv_parser.set_defaults(func=rebuild_status_from_csv) + + pending_by_error_parser = subparsers.add_parser( + "set-pending-by-error-type", + help="Mark failed terminal non_retryable apps as pending by latest_failure_type.", + ) + pending_by_error_parser.add_argument("error_types", nargs="*", help="Exact latest_failure_type values to match.") + pending_by_error_parser.add_argument( + "--error-type-prefix", + default="", + help="Prefix match for latest_failure_type, e.g. DOWNLOAD_ERROR/.", + ) + pending_by_error_parser.add_argument( + "--all-non-retryable", + action="store_true", + help="Apply to all current failed terminal non_retryable apps.", + ) + pending_by_error_parser.add_argument( + "--non-retryable-only", + action="store_true", + help="Only modify current failed_terminal non_retryable apps.", + ) + pending_by_error_parser.add_argument( + "--pending-only", + action="store_true", + help="Only modify current pending apps.", + ) + pending_by_error_parser.add_argument( + "--reason", + default="manual_pending_by_error_type", + help="collection_status_reason to write.", + ) + pending_by_error_parser.add_argument( + "--device-type", + default="", + help="Optional retry device type override: emulator, physical, or any.", + ) + pending_by_error_parser.add_argument( + "--tag", + default="", + help="Only affect apps with this batch_tags.", + ) + pending_by_error_parser.add_argument( + "--top-n", + type=int, + default=None, + help="Only affect rows whose source_order is within TopN, using source_order < N.", + ) + pending_by_error_parser.set_defaults(func=set_pending_by_error_type) + + pending_success_zero_self_ratio_parser = subparsers.add_parser( + "set-pending-success-zero-self-ratio", + help=( + "Mark qualified success apps with self_ratio <= threshold and " + f"num_nodes < {LIGHT_RESTRICTED_MIN_NUM_NODES} as pending." + ), + ) + pending_success_zero_self_ratio_parser.add_argument( + "--max-self-ratio", + type=float, + default=0.0, + help="Maximum self_ratio to match. Default: 0.0.", + ) + pending_success_zero_self_ratio_parser.add_argument( + "--reason", + default="manual_pending_success_zero_self_ratio", + help="collection_status_reason to write.", + ) + pending_success_zero_self_ratio_parser.add_argument( + "--tag", + default="", + help="Only affect apps with this batch_tags.", + ) + pending_success_zero_self_ratio_parser.add_argument( + "--top-n", + type=int, + default=None, + help="Only affect rows whose source_order is within TopN, using source_order < N.", + ) + pending_success_zero_self_ratio_parser.set_defaults(func=set_pending_success_zero_self_ratio) + + pending_retryable_parser = subparsers.add_parser( + "set-pending-retryable", + help="Mark retryable non-success apps back to pending, useful after removing retry caps.", + ) + pending_retryable_parser.add_argument( + "--current-status", + default="", + help="Only reset rows currently in this collection_status, e.g. failed_terminal.", + ) + pending_retryable_parser.add_argument( + "--reason", + default="manual_pending_retryable", + help="collection_status_reason to write.", + ) + pending_retryable_parser.add_argument( + "--tag", + default="", + help="Only affect apps with this batch_tags.", + ) + pending_retryable_parser.add_argument( + "--top-n", + type=int, + default=None, + help="Only affect rows whose source_order is within TopN, using source_order < N.", + ) + pending_retryable_parser.set_defaults(func=set_pending_retryable) + + existing_non_retryable_parser = subparsers.add_parser( + "set-existing-non-retryable-light-restricted", + help="Requeue matching failed non_retryable apps by creating pending collection_task rows.", + ) + existing_non_retryable_parser.add_argument( + "--min-num-nodes", + type=int, + default=LIGHT_RESTRICTED_MIN_NUM_NODES, + help=f"Minimum num_nodes to match. Default: {LIGHT_RESTRICTED_MIN_NUM_NODES}.", + ) + existing_non_retryable_parser.add_argument( + "--reason", + default="manual_existing_non_retryable_light_restricted", + help="collection_status_reason to write.", + ) + existing_non_retryable_parser.add_argument( + "--tag", + default="", + help="Only affect apps with this batch_tags.", + ) + existing_non_retryable_parser.add_argument( + "--top-n", + type=int, + default=None, + help="Only affect rows whose source_order is within TopN, using source_order < N.", + ) + existing_non_retryable_parser.set_defaults(func=set_existing_non_retryable_light_restricted) + + export_by_error_parser = subparsers.add_parser( + "export-by-error-type", + help="Export matching apps by latest_failure_type to csv.", + ) + export_by_error_parser.add_argument("error_types", nargs="*", help="Exact latest_failure_type values to match.") + export_by_error_parser.add_argument( + "--error-type-prefix", + default="", + help="Prefix match for latest_failure_type, e.g. DOWNLOAD_ERROR/.", + ) + export_by_error_parser.add_argument( + "--non-retryable-only", + action="store_true", + help="Only export rows currently judged as severe_restricted non_retryable; can be used alone to export all non-retryable apps.", + ) + export_by_error_parser.add_argument( + "--tag", + default="", + help="Only export apps with this batch_tags.", + ) + export_by_error_parser.add_argument( + "--top-n", + type=int, + default=None, + help="Only export rows whose source_order is within TopN, using source_order < N.", + ) + export_by_error_parser.add_argument("--output", required=True, help="Target csv output path.") + export_by_error_parser.set_defaults(func=export_apps_by_error_type) + + set_batch_tag_parser = subparsers.add_parser( + "set-batch-tag", + help="Add tag only to new/updated apps in a CSV — skips unchanged existing to avoid cross-batch pollution.", + ) + set_batch_tag_parser.add_argument("--csv-path", required=True, help="CSV path containing package_name column.") + set_batch_tag_parser.add_argument("--tag", required=True, help="batch_tags to set on matching apps.") + set_batch_tag_parser.set_defaults(func=set_batch_tag_from_csv) + + clear_batch_tag_parser = subparsers.add_parser( + "clear-batch-tag", + help="Remove tag only from new_app apps with multiple tags (cross-batch pollution).", + ) + clear_batch_tag_parser.add_argument("--tag", required=True, help="batch_tags to remove.") + clear_batch_tag_parser.set_defaults(func=clear_batch_tag) + + deactivate_stale_parser = subparsers.add_parser( + "deactivate-stale-tagged", + help="Deactivate new_app apps that were cross-batch tagged (multi-tag pollution).", + ) + deactivate_stale_parser.add_argument("--tag", required=True, help="batch_tags to inspect.") + deactivate_stale_parser.set_defaults(func=deactivate_stale_tagged) + + enqueue_high_parser = subparsers.add_parser( + "enqueue-high-priority", + help="Mark apps from CSV as high priority and enqueue for immediate collection.", + ) + enqueue_high_parser.add_argument("--csv-path", required=True, help="CSV path containing package_name column.") + enqueue_high_parser.add_argument( + "--incremental-batch-tag", + default="", + help="Optional batch tag to mark packages enqueued by this command.", + ) + enqueue_high_parser.add_argument( + "--force-pending", + action="store_true", + help="Force all apps in CSV to pending, including ones already successfully collected.", + ) + enqueue_high_parser.set_defaults(func=enqueue_high_priority) + + load_block_parser = subparsers.add_parser( + "load-block-tasks", + help="Load qualified apps as block test tasks (from CSV or directly from DB).", + ) + load_block_parser.add_argument( + "--csv-path", + default="", + help="Optional: CSV file with package_name column. If not provided, loads all qualified apps from DB.", + ) + load_block_parser.add_argument( + "--all", + dest="load_all", + action="store_true", + help="加载数据库中所有应用（忽略 collection_status），与 --csv-path 配合时跳过 qualified 校验。", + ) + load_block_parser.set_defaults(func=load_block_tasks) + + force_reapk_parser = subparsers.add_parser( + "force-reapk", + help="清除指定应用的本地APK缓存并触发US重新下载。", + ) + force_reapk_parser.add_argument("package_name", help="要重新下载的应用包名") + force_reapk_parser.set_defaults(func=force_reapk) + + return parser + + +def force_reapk(args) -> int: + pkg = str(args.package_name or "").strip() + if not pkg: + print("package_name is required") + return 1 + try: + from redis_task_distribute import RedisTaskDispatcher + dispatcher = RedisTaskDispatcher() + n = dispatcher.force_reapk(pkg) + print(f"[force_reapk] 已重置 {pkg} ({n} 个任务)，清除本地缓存并等待中心 APK 缓存刷新") + return 0 + except Exception as e: + print(f"[force_reapk] 失败: {e}") + return 1 + + +def load_block_tasks(args) -> int: + csv_path = str(args.csv_path or "").strip() + load_all = bool(getattr(args, "load_all", False)) + service = _build_service(args.db_path) + try: + if csv_path: + rows = _read_csv_rows(csv_path) + if load_all: + db_apps = service.repo.list_all_catalog_apps() + else: + db_apps = service.repo.list_qualified_apps() + db_lookup = {item["package_name"]: item for item in db_apps} + packages = [] + skipped = 0 + seen = set() + for row in rows: + payload = build_task_payload_from_row(row) + package_name = str(payload.get("package_name") or "").strip() + if not package_name or package_name in seen: + continue + seen.add(package_name) + if package_name not in db_lookup: + skipped += 1 + print(f"[跳过] {package_name}: 未在数据库中") + continue + db_entry = db_lookup[package_name] + app_name = str(db_entry.get("app_name") or payload.get("app_name") or "").strip() + packages.append({ + "package_name": package_name, + "app_name": app_name, + "task_payload": db_entry.get("task_payload") or payload, + }) + if skipped: + print(f"[CSV校验] 共跳过 {skipped} 个未在数据库中的包名") + else: + if load_all: + packages = service.repo.list_all_catalog_apps() + else: + packages = service.repo.list_qualified_apps() + + if not packages: + label = "数据库" if load_all else "qualified" + print(f"[Block测试] 没有符合条件的{label}应用") + return 0 + + print(f"[Block测试] 准备加载 {len(packages)} 个任务") + try: + from redis_task_distribute import RedisTaskDispatcher + dispatcher = RedisTaskDispatcher() + count = dispatcher.load_block_tasks(packages) + if count > 0: + r = redis_lib.Redis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True) + r.publish(channel_name("catalog:updated"), json.dumps({"added": count, "priority": "high_block"})) + print(f"[Block测试] 成功加载 {count} 个 block 任务到高优队列") + except Exception as e: + print(f"[Block测试] 加载失败: {e}") + return 1 + return 0 + finally: + service.close() + + +def main() -> int: + parser = build_parser() + args = parser.parse_args() + return int(args.func(args)) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/mumu_restore_v2.ahk b/scripts/mumu_restore_v2.ahk new file mode 100644 index 0000000..5b3be80 --- /dev/null +++ b/scripts/mumu_restore_v2.ahk @@ -0,0 +1,144 @@ +#Requires AutoHotkey v2.0 +;@Ahk2Exe-ConsoleApp +#SingleInstance Force +; 在脚本开头尝试附加到启动它的命令行控制台 (cmd 或 PowerShell) +; -1 表示附加到父进程的控制台 +;try DllCall("AttachConsole", "Int", -1) +SendMode("Input") +SetWorkingDir(A_ScriptDir) +FileEncoding("UTF-8-RAW") +CoordMode("Mouse", "Screen") + +DetectHiddenWindows(true) +DetectHiddenText(true) + +; ============================================================ +; 命令行参数: +; A_Args[1] - 备份文件路径 (默认 D:\mumu_backups\taskagent.mumudata) +; A_Args[2] - VM 索引 (可选, 默认 2) +; +; 用法: +; mumu_restore_v2.exe "D:\mumu_backups\taskagent.mumudata" 2 +; +; 返回值: +; 0 = 成功, 1 = 失败 +; ============================================================ + +backupFile := A_Args.Length >= 1 ? A_Args[1] : "D:\mumu_backups\taskagent.mumudata" +vmIndex := A_Args.Length >= 2 ? A_Args[2] : 2 +resultFile := A_Args.Length >= 3 ? A_Args[3] : "D:\mumu_backups\result.txt" +if !backupFile { + FileAppend("[ERROR] Missing backup file argument`n", "*") + ExitApp(1) +} + +if !FileExist(backupFile) { + FileAppend("[ERROR] Backup file not found: " backupFile "`n", "*") + ExitApp(1) +} + +; ---------- MuMu 管理器信息 (用 Window Spy 抓取后填入) ---------- +MuMuNxMainTitle := "ahk_class Qt5156QWindowIcon" ; 如 "ahk_class Qt5QWindowIcon" +MuMuNxMainPath := "C:\Program Files\Netease\MuMu\nx_main\MuMuNxMain.exe" + +; ---------- 超时设置 ---------- +IMPORT_TIMEOUT_SECONDS := 900 ; 导入最大等待时间 (15 分钟) +WINDOW_WAIT_SECONDS := 30 ; 等待窗口出现的最大时间 + +; ============================================================ +; 辅助函数 +; ============================================================ + +WriteResult(content) { + global resultFile + try { + ; 建议在 v2 中先判断文件是否存在，因为如果文件被占用导致删除失败，v2 会抛出严重错误 + if FileExist(resultFile) { + FileDelete(resultFile) + } + ; 参数需要用括号，文本编码等字符串需要用双引号 + FileAppend(content, resultFile, "UTF-8-RAW") + } catch { + ; 字符串和变量之间通过空格拼接，向控制台输出的星号 "*" 也需要加双引号 + FileAppend("FAIL: Cannot write result file: " resultFile "`n", "*") + } +} + +; ============================================================ +; 主流程 +; ============================================================ + +; 1. 启动 MuMu 管理器 +; TODO: 如果管理器已在运行, 先关闭再重新打开, 保证窗口状态干净 +ProcessClose("MuMuNxMain.exe") +Sleep(2000) + +Run('"' MuMuNxMainPath '"') + +; 2. 等待管理器窗口出现 +if !WinWait(MuMuNxMainTitle, , WINDOW_WAIT_SECONDS) + WriteResult("FAIL: MuMu manager window did not appear within " WINDOW_WAIT_SECONDS "s") + +WinActivate(MuMuNxMainTitle) +Sleep(2000) + +; 3. 导入备份文件 +; TODO: 根据实际录制结果替换下方的坐标 / 操作步骤 +; 典型流程: +; a) 点击 "导入" 按钮 (或 菜单 -> 导入) +; b) 在弹出的文件选择对话框中输入备份文件路径并确认 +; c) 等待导入进度条完成 +; 最大化窗口 移动到左上角 +WinMaximize +WinSetAlwaysOnTop(1, MuMuNxMainTitle) +; --- 步骤 a: 点击导入按钮 --- +Click(810, 420) +Sleep(1000) +Click(860, 608) +Sleep(1000) +CoordMode("Mouse", "Screen") +Click(600, 338) +; --- 步骤 b: 文件选择对话框 (通常是 Windows 通用对话框) --- +; 方式1: 如果支持直接输入路径 (推荐, 比逐级浏览更稳定) +Sleep(1000) +dialogWin := "选择备份文件 ahk_class #32770 ahk_exe MuMuNxMain.exe" +; 3. 激活该窗口并等待其完全激活 + WinActivate(dialogWin) + WinWaitActive(dialogWin) + ;WinSetAlwaysOnTop(1, dialogWin) + ; --- 进阶技巧：直接把文件路径填进去，而不是模拟键盘敲字 --- + ; 在标准的文件对话框中，输入文件名的文本框控件名称通常是 "Edit1" + + ; 直接设置文本（比用 Send 发送按键更稳定，不受输入法干扰） + ControlSetText(backupFile, "Edit1", dialogWin) + + ; 短暂延迟，确保系统反应过来 + Sleep(200) + + ; 4. 点击“打开”按钮 (默认的打开按钮控件通常是 "Button1") + ControlClick("Button1", dialogWin) + Send("{Enter}") +; --- 步骤 c: 等待导入完成 --- +; 方式C - 检测按钮状态 (如 "完成" 按钮出现): +WinActivate(MuMuNxMainTitle) +WinSetAlwaysOnTop(1, MuMuNxMainTitle) +CoordMode("Mouse", "Client") + importDeadline := A_TickCount + IMPORT_TIMEOUT_SECONDS * 1000 + Loop { + WinActivate(MuMuNxMainTitle) + color := PixelGetColor(568,427, "RGB") + try FileAppend("[color] " color "`n", "*") + if color = 0x0F7B0F + break + if A_TickCount > importDeadline + WriteResult("FAIL: Import timed out") + Sleep(2000) + } + +; 4. 导入完成后, 点击忽略, 关闭管理器 +Click(600, 440) +Sleep(200) + +WinClose(MuMuNxMainTitle) + +WriteResult("OK") diff --git a/scripts/rebuild_zero_traffic_analytics.py b/scripts/rebuild_zero_traffic_analytics.py new file mode 100644 index 0000000..7aaa6fa --- /dev/null +++ b/scripts/rebuild_zero_traffic_analytics.py @@ -0,0 +1,109 @@ +#!/usr/bin/env python3 +import argparse +import os +import sys +from typing import List + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) + +from analytics import AnalyticsService +from config import ANALYTICS_TRAFFIC_ROOT, MONITORING_DB_PATH + + +def _split_packages(raw: str) -> List[str]: + return [item.strip() for item in str(raw or "").split(",") if item.strip()] + + +def _has_traffic_domains(traffic_files: List[str], package_name: str) -> bool: + """Check whether any traffic file contains at least one valid traffic entry for the package.""" + for file_path in traffic_files: + try: + with open(file_path, "r", encoding="utf-8", errors="ignore") as handle: + for raw_line in handle: + line = raw_line.strip() + if not line: + continue + fields = [item.strip() for item in line.split(",")] + if len(fields) < 7: + continue + if fields[0].strip() == package_name: + return True + except OSError: + continue + return False + + +def _load_zero_traffic_packages(service: AnalyticsService, tag: str = "", limit: int = 0) -> List[str]: + normalized_tag = str(tag or "").strip() + with service.repo._connect() as connection: + summaries = service.repo._list_catalog_summaries(connection) + rows = [] + for item in summaries: + if not item.get("catalog_active"): + continue + if normalized_tag and normalized_tag not in (item.get("incremental_batch_tags") or []): + continue + if int(item.get("total_traffic_bytes") or 0) != 0: + continue + if str(item.get("artifact_status") or "") not in {"missing", "partial"}: + continue + rows.append(item) + rows.sort(key=lambda item: (float(item.get("updated_at") or 0.0), item.get("package_name") or ""), reverse=True) + packages = [str(item.get("package_name") or "").strip() for item in rows if str(item.get("package_name") or "").strip()] + return packages[:limit] if limit > 0 else packages + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Rebuild analytics snapshots for zero-traffic packages that already have traffic_count files." + ) + parser.add_argument("--db-path", default=MONITORING_DB_PATH) + parser.add_argument("--traffic-root", default=ANALYTICS_TRAFFIC_ROOT) + parser.add_argument("--packages", default="", help="Comma-separated package names. Defaults to zero-traffic packages.") + parser.add_argument("--tag", default="", help="Only rebuild packages whose incremental_batch_tag contains this tag.") + parser.add_argument("--limit", type=int, default=0) + parser.add_argument("--dry-run", action="store_true") + args = parser.parse_args() + + service = AnalyticsService( + db_path=args.db_path, + traffic_root=args.traffic_root, + start_worker=False, + ) + try: + packages = _split_packages(args.packages) or _load_zero_traffic_packages(service, args.tag, args.limit) + + scanned = 0 + matched = 0 + rebuilt = 0 + skipped = 0 + for package_name in packages: + scanned += 1 + traffic_files = service._find_traffic_files(package_name) + if not traffic_files: + skipped += 1 + continue + if not _has_traffic_domains(traffic_files, package_name): + skipped += 1 + continue + matched += 1 + if args.dry_run: + print(f"DRY-RUN {package_name}: {len(traffic_files)} traffic files") + continue + summary = service.rebuild_package_now(package_name) + rebuilt += 1 + print( + f"rebuilt {package_name}: files={len(traffic_files)} " + f"artifact={summary.get('artifact_status')} bytes={summary.get('total_traffic_bytes')}" + ) + finally: + service.close() + + print(f"done scanned={scanned} matched={matched} rebuilt={rebuilt} skipped={skipped}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/recover_mumu.cmd b/scripts/recover_mumu.cmd new file mode 100644 index 0000000..bc5af0d --- /dev/null +++ b/scripts/recover_mumu.cmd @@ -0,0 +1,238 @@ +@echo off +setlocal EnableExtensions EnableDelayedExpansion + +set "BACKUP_FULL_PATH=%~1" +set "LOCAL_IMPORT_DIR=%~2" +set "MUMU_MANAGER_PATH=%~3" +set "VM_ID=%~4" +set "DEFAULT_NET_BRIDGE_CARD=%~5" +set "AHK_EXE=%~6" +set "BRIDGE_IP_OFFSET=%~7" +set "GATEWAY_RULES=%~8" + +if not defined BACKUP_FULL_PATH call :fail "BACKUP_FULL_PATH argument is required" +if not defined LOCAL_IMPORT_DIR set "LOCAL_IMPORT_DIR=D:\mumu_backups" +if not defined MUMU_MANAGER_PATH set "MUMU_MANAGER_PATH=C:\Program Files\Netease\MuMu\nx_main\MuMuManager.exe" +if not defined VM_ID set "VM_ID=2" +if not defined DEFAULT_NET_BRIDGE_CARD set "DEFAULT_NET_BRIDGE_CARD=Realtek PCIe GbE Family Controller" +if not defined AHK_EXE call :fail "AHK_EXE argument is required" +if not defined BRIDGE_IP_OFFSET set "BRIDGE_IP_OFFSET=100" +if not defined GATEWAY_RULES call :fail "GATEWAY_RULES argument is required, format: 192.168.1=192.168.1.1;192.168.2=192.168.2.1" + +set "SUBNET_MASK=255.255.255.0" +set "DNS1=8.8.8.8" +set "DNS2=1.1.1.1" +set "POLL_MAX=300" + +for %%I in ("%BACKUP_FULL_PATH%") do ( + set "BACKUP_DIR=%%~dpI" + set "BACKUP_FILENAME=%%~nxI" +) +if "!BACKUP_DIR:~-1!"=="\" set "BACKUP_DIR=!BACKUP_DIR:~0,-1!" + +set "IMPORT_FILE=%LOCAL_IMPORT_DIR%\%BACKUP_FILENAME%" +set "RESULT_FILE=%LOCAL_IMPORT_DIR%\result.txt" +set "AHK_BAT=%LOCAL_IMPORT_DIR%\_run_ahk.bat" + +echo ================================================== +echo Recover MuMu image started +echo ================================================== +echo [INFO] Backup source: %BACKUP_FULL_PATH% +echo [INFO] Local import dir: %LOCAL_IMPORT_DIR% +echo [INFO] MuMu manager: %MUMU_MANAGER_PATH% +echo [INFO] VM index: %VM_ID% +echo [INFO] AHK executable: %AHK_EXE% + +if not exist "%MUMU_MANAGER_PATH%" call :fail "MuMuManager.exe not found: %MUMU_MANAGER_PATH%" +if not exist "%BACKUP_FULL_PATH%" call :fail "Remote backup not found: %BACKUP_FULL_PATH%" +if not exist "%LOCAL_IMPORT_DIR%" ( + mkdir "%LOCAL_IMPORT_DIR%" + if errorlevel 1 call :fail "Failed to create local import dir: %LOCAL_IMPORT_DIR%" +) + +call :sync_backup +if errorlevel 1 exit /b 1 + +call :run_restore +if errorlevel 1 exit /b 1 + +call :configure_network +if errorlevel 1 exit /b 1 + +echo ================================================== +echo MuMu restore and network config completed successfully. +echo ================================================== +exit /b 0 + +:sync_backup +echo [Phase 1] Check and synchronize backup image +set "NEED_COPY_BACKUP=1" +set "REMOTE_BACKUP_TIME=" +set "LOCAL_BACKUP_TIME=" + +for %%I in ("%BACKUP_FULL_PATH%") do set "REMOTE_BACKUP_TIME=%%~tI" +echo [INFO] Remote backup time = !REMOTE_BACKUP_TIME! + +if exist "%IMPORT_FILE%" ( + for %%I in ("%IMPORT_FILE%") do set "LOCAL_BACKUP_TIME=%%~tI" + echo [INFO] Local backup time = !LOCAL_BACKUP_TIME! + if /i "!REMOTE_BACKUP_TIME!"=="!LOCAL_BACKUP_TIME!" ( + echo [INFO] Backup timestamps match. Skip copying backup file. + set "NEED_COPY_BACKUP=0" + ) +) else ( + echo [INFO] Local backup file not found. A fresh copy is required. +) + +if "!NEED_COPY_BACKUP!"=="1" ( + echo [INFO] Copying backup file from %BACKUP_DIR% to %LOCAL_IMPORT_DIR% ... + robocopy "%BACKUP_DIR%" "%LOCAL_IMPORT_DIR%" "%BACKUP_FILENAME%" /R:2 /W:2 /COPY:DAT /DCOPY:T /NFL /NDL /NJH /NJS /NP >nul + if !errorlevel! gtr 7 call :fail "Failed to copy backup file to %IMPORT_FILE%" +) +exit /b 0 + +:run_restore +echo [Phase 2] Run MuMu restore through RunAHK scheduled task +echo [INFO] Shutting down MuMu emulator... +"%MUMU_MANAGER_PATH%" control -v %VM_ID% shutdown +call :sleep 15 + +echo [INFO] Killing MuMu-related processes if they are running... +call :kill_mumu_processes +call :sleep 5 + +del /f /q "%RESULT_FILE%" >nul 2>&1 +( + echo @echo off + echo "%AHK_EXE%" "%IMPORT_FILE%" %VM_ID% "%RESULT_FILE%" +) > "%AHK_BAT%" +if errorlevel 1 call :fail "Failed to write AHK runner: %AHK_BAT%" + +schtasks /delete /tn "RunAHK" /f >nul 2>&1 +schtasks /create /tn "RunAHK" /tr "%AHK_BAT%" /sc ONCE /st 00:00 /rl HIGHEST /it /f >nul +if errorlevel 1 call :fail "Failed to create RunAHK task" + +schtasks /run /tn "RunAHK" >nul +if errorlevel 1 call :fail "Failed to start RunAHK task" + +echo [INFO] Waiting for restore result file: %RESULT_FILE% +for /l %%I in (1,1,%POLL_MAX%) do ( + if exist "%RESULT_FILE%" goto :check_restore_result + call :sleep 2 +) +call :fail "AHK restore timed out" + +:check_restore_result +call :sleep 3 +set "RESTORE_RESULT=" +set /p RESTORE_RESULT=<"%RESULT_FILE%" +echo [INFO] AHK result: !RESTORE_RESULT! +echo !RESTORE_RESULT! | findstr /i /c:"OK" >nul +if errorlevel 1 call :fail "AHK restore failed: !RESTORE_RESULT!" +exit /b 0 + +:configure_network +echo [Phase 3] Configure MuMu network +call :detect_network_adapter +if errorlevel 1 exit /b 1 + +call :detect_bridge_network +if errorlevel 1 exit /b 1 + +call :set_mumu_setting net_bridge_card "!NetAdapterName!" +if errorlevel 1 exit /b 1 +call :set_mumu_setting net_bridge_ip_mode static +if errorlevel 1 exit /b 1 +call :set_mumu_setting net_bridge_ip_addr "!BRIDGE_IP!" +if errorlevel 1 exit /b 1 +call :set_mumu_setting net_bridge_gateway "!GATEWAY!" +if errorlevel 1 exit /b 1 +call :set_mumu_setting net_bridge_subnet_mask %SUBNET_MASK% +if errorlevel 1 exit /b 1 +call :set_mumu_setting net_bridge_dns1 %DNS1% +if errorlevel 1 exit /b 1 +call :set_mumu_setting net_bridge_dns2 %DNS2% +if errorlevel 1 exit /b 1 + +echo [INFO] Network configuration completed successfully. +exit /b 0 + +:kill_mumu_processes +for %%P in (MuMuPlayer.exe MuMuNxDevice.exe MuMuNxMain.exe MuMuManager.exe) do ( + taskkill /F /IM %%P >nul 2>&1 +) +exit /b 0 + +:detect_network_adapter +set "NetAdapterName=" +for /f "tokens=* skip=1" %%a in ('wmic nic where "NetConnectionStatus=2" get Name /value 2^>nul') do ( + if "!NetAdapterName!"=="" ( + for /f "tokens=2 delims==" %%b in ("%%a") do set "NetAdapterName=%%b" + ) +) +if not "!NetAdapterName!"=="" ( + for /f "tokens=*" %%a in ("!NetAdapterName!") do set "NetAdapterName=%%a" +) +if "!NetAdapterName!"=="" set "NetAdapterName=%DEFAULT_NET_BRIDGE_CARD%" +echo [INFO] net_bridge_card = !NetAdapterName! +exit /b 0 + +:detect_bridge_network +set "LOCAL_IP=" +set "LOCAL_PREFIX=" +for /f "tokens=2 delims=:" %%a in ('ipconfig ^| findstr /r /c:"[0-9][0-9]*\.[0-9][0-9]*\.[0-9][0-9]*\.[0-9][0-9]*"') do ( + if not defined LOCAL_IP ( + set "CANDIDATE_IP=%%a" + set "CANDIDATE_IP=!CANDIDATE_IP: =!" + for /f "tokens=1-4 delims=." %%b in ("!CANDIDATE_IP!") do ( + call :find_gateway "%%b.%%c.%%d" + if defined GATEWAY ( + set "LOCAL_IP=!CANDIDATE_IP!" + set "LOCAL_PREFIX=%%b.%%c.%%d" + ) + ) + ) +) + +if not defined LOCAL_IP call :fail "Failed to detect local IPv4 address matching GATEWAY_RULES: %GATEWAY_RULES%" +echo [INFO] Local IP: !LOCAL_IP! + +set "BRIDGE_IP=" +for /f "tokens=1-4 delims=." %%a in ("!LOCAL_IP!") do ( + set /a LAST_OCTET=%%d + %BRIDGE_IP_OFFSET% + if !LAST_OCTET! gtr 254 call :fail "Local IP last octet is too large for bridge IP derivation" + set "BRIDGE_IP=%%a.%%b.%%c.!LAST_OCTET!" +) + +echo [INFO] net_bridge_ip_addr = !BRIDGE_IP! +echo [INFO] net_bridge_gateway = !GATEWAY! +exit /b 0 + +:find_gateway +set "GATEWAY=" +set "TARGET_PREFIX=%~1" +for %%R in ("%GATEWAY_RULES:;=" "%") do ( + for /f "tokens=1,* delims==" %%K in ("%%~R") do ( + if "%%K"=="!TARGET_PREFIX!" set "GATEWAY=%%L" + ) +) +exit /b 0 + +:set_mumu_setting +set "SETTING_KEY=%~1" +set "SETTING_VALUE=%~2" +echo [INFO] !SETTING_KEY! = !SETTING_VALUE! +"%MUMU_MANAGER_PATH%" setting -v %VM_ID% -k "!SETTING_KEY!" -val "!SETTING_VALUE!" +if errorlevel 1 call :fail "Failed to set !SETTING_KEY!" +exit /b 0 + +:sleep +set /a SLEEP_COUNT=%~1 + 1 +ping -n !SLEEP_COUNT! 127.0.0.1 >nul +exit /b 0 + +:fail +echo. +echo [ERROR] %~1 1>&2 +echo Recover MuMu aborted. 1>&2 +exit 1 diff --git a/scripts/reset_catalog_update_state.py b/scripts/reset_catalog_update_state.py new file mode 100644 index 0000000..71ef07c --- /dev/null +++ b/scripts/reset_catalog_update_state.py @@ -0,0 +1,189 @@ +#!/usr/bin/env python3 + +import argparse +import json +import os +import sqlite3 +import sys +import time +from typing import Any, Dict + + +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if PROJECT_ROOT not in sys.path: + sys.path.insert(0, PROJECT_ROOT) + +from config import MONITORING_DB_PATH + + +def _safe_loads(value: Any) -> Dict[str, Any]: + if not value: + return {} + try: + payload = json.loads(value) + except (TypeError, ValueError): + return {} + return payload if isinstance(payload, dict) else {} + + +def _clear_payload_last_updated(raw_payload: Any) -> str: + payload = _safe_loads(raw_payload) + if not payload: + return str(raw_payload or "") + payload["last_updated"] = "" + original_row = payload.get("original_row") + if isinstance(original_row, dict): + original_row = dict(original_row) + original_row["last_updated"] = "" + original_row["最后更新"] = "" + original_row["更新时间"] = "" + payload["original_row"] = original_row + return json.dumps(payload, ensure_ascii=False) + + +def reset_catalog_update_state(args) -> int: + db_path = str(args.db_path or "").strip() + apply_changes = bool(args.apply) + clear_payload = bool(args.clear_payload) + now_iso = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + with sqlite3.connect(db_path) as connection: + connection.row_factory = sqlite3.Row + table_row = connection.execute( + """ + SELECT name + FROM sqlite_master + WHERE type = 'table' + AND name = 'app_catalog' + """ + ).fetchone() + if not table_row: + raise SystemExit(f"app_catalog not found in {db_path}") + + counts = connection.execute( + """ + SELECT + COUNT(*) AS total, + COALESCE(SUM(CASE WHEN COALESCE(last_updated, '') != '' THEN 1 ELSE 0 END), 0) AS last_updated_rows, + COALESCE(SUM(CASE WHEN COALESCE(last_update_interval_days, 0) != 0 THEN 1 ELSE 0 END), 0) AS interval_rows + FROM app_catalog + """ + ).fetchone() + task_counts = connection.execute( + """ + SELECT + COALESCE(SUM(CASE WHEN COALESCE(is_new_app, 1) = 0 THEN 1 ELSE 0 END), 0) AS app_update_rows, + COALESCE( + SUM( + CASE + WHEN task_status = 'pending' + AND COALESCE(is_new_app, 1) = 0 + AND COALESCE(error_reason, '') LIKE 'catalog_version_update:%' + THEN 1 + ELSE 0 + END + ), + 0 + ) AS version_pending_rows + FROM collection_task + """ + ).fetchone() + + if not apply_changes: + print( + "dry_run=1 db={db} total={total} clear_last_updated={last_updated_rows} " + "reset_interval_rows={interval_rows} app_update_tasks={app_update_rows} " + "delete_version_pending={version_pending_rows}".format( + db=db_path, + total=int(counts["total"] or 0), + last_updated_rows=int(counts["last_updated_rows"] or 0), + interval_rows=int(counts["interval_rows"] or 0), + app_update_rows=int(task_counts["app_update_rows"] or 0), + version_pending_rows=int(task_counts["version_pending_rows"] or 0), + ) + ) + return 0 + + cursor = connection.execute( + """ + DELETE FROM collection_task + WHERE task_status = 'pending' + AND COALESCE(is_new_app, 1) = 0 + AND COALESCE(error_reason, '') LIKE 'catalog_version_update:%' + """ + ) + deleted_version_pending = int(cursor.rowcount or 0) + + cursor = connection.execute( + """ + UPDATE app_catalog + SET last_update_interval_days = 0, + updated_at = ? + WHERE COALESCE(last_update_interval_days, 0) != 0 + """, + (now_iso,), + ) + reset_interval_rows = int(cursor.rowcount or 0) + + if clear_payload: + payload_rows = connection.execute( + """ + SELECT package_name, task_payload_json + FROM app_catalog + WHERE COALESCE(task_payload_json, '') != '' + """ + ).fetchall() + for row in payload_rows: + connection.execute( + """ + UPDATE app_catalog + SET task_payload_json = ?, + updated_at = ? + WHERE package_name = ? + """, + (_clear_payload_last_updated(row["task_payload_json"]), now_iso, row["package_name"]), + ) + + cursor = connection.execute( + """ + UPDATE app_catalog + SET last_updated = '', + updated_at = ? + WHERE COALESCE(last_updated, '') != '' + """, + (now_iso,), + ) + cleared_last_updated = int(cursor.rowcount or 0) + print( + "dry_run=0 db={db} cleared_last_updated={cleared_last_updated} " + "reset_interval_rows={reset_interval_rows} deleted_version_pending={deleted_version_pending} " + "clear_payload={clear_payload}".format( + db=db_path, + cleared_last_updated=cleared_last_updated, + reset_interval_rows=reset_interval_rows, + deleted_version_pending=deleted_version_pending, + clear_payload=int(clear_payload), + ) + ) + return 0 + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Clear app_catalog last_updated and delete pending catalog_version_update collection_task rows." + ) + parser.add_argument("--db-path", default=MONITORING_DB_PATH, help="Target monitoring sqlite path.") + parser.add_argument("--apply", action="store_true", help="Apply changes. Without this, only prints counts.") + parser.add_argument( + "--clear-payload", + action="store_true", + help="Also clear last_updated inside task_payload_json/original_row.", + ) + return parser + + +def main() -> int: + return reset_catalog_update_state(build_parser().parse_args()) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/static/vendor/bootstrap/bootstrap.bundle.min.js b/static/vendor/bootstrap/bootstrap.bundle.min.js new file mode 100644 index 0000000..13539e2 --- /dev/null +++ b/static/vendor/bootstrap/bootstrap.bundle.min.js @@ -0,0 +1 @@ +window.bootstrap=window.bootstrap||{}; diff --git a/static/vendor/bootstrap/bootstrap.min.css b/static/vendor/bootstrap/bootstrap.min.css new file mode 100644 index 0000000..14219d3 --- /dev/null +++ b/static/vendor/bootstrap/bootstrap.min.css @@ -0,0 +1,83 @@ +:root{ + --bs-body-font-family:"Segoe UI","PingFang SC","Microsoft YaHei",sans-serif; + --bs-body-font-size:1rem; + --bs-body-line-height:1.5; + --bs-border-color:rgba(32,25,18,.12); + --bs-border-radius:1rem; + --bs-border-radius-lg:1.5rem; + --bs-card-spacer-y:1rem; + --bs-card-spacer-x:1rem; + --bs-primary:#9a4d29; + --bs-secondary:#74685c; + --bs-success:#25674a; + --bs-warning:#a6791d; + --bs-danger:#9d2f2f; +} +*,:before,:after{box-sizing:border-box} +.container-fluid{width:100%;padding-right:1rem;padding-left:1rem;margin-right:auto;margin-left:auto} +.row{display:flex;flex-wrap:wrap;margin-right:-.75rem;margin-left:-.75rem} +.row>*{padding-right:.75rem;padding-left:.75rem;flex-shrink:0;width:100%;max-width:100%} +.col-12{flex:0 0 auto;width:100%} +@media(min-width:992px){ + .col-lg-5{flex:0 0 auto;width:41.66666667%} + .col-lg-7{flex:0 0 auto;width:58.33333333%} +} +.g-3{row-gap:1rem} +.g-4{row-gap:1.5rem} +.mb-0{margin-bottom:0!important} +.mb-2{margin-bottom:.5rem!important} +.mb-3{margin-bottom:1rem!important} +.mb-4{margin-bottom:1.5rem!important} +.mt-2{margin-top:.5rem!important} +.mt-3{margin-top:1rem!important} +.me-2{margin-right:.5rem!important} +.p-0{padding:0!important} +.w-100{width:100%!important} +.d-flex{display:flex!important} +.d-inline-flex{display:inline-flex!important} +.d-grid{display:grid!important} +.align-items-center{align-items:center!important} +.align-items-start{align-items:flex-start!important} +.justify-content-between{justify-content:space-between!important} +.justify-content-end{justify-content:flex-end!important} +.justify-content-center{justify-content:center!important} +.flex-wrap{flex-wrap:wrap!important} +.gap-2{gap:.5rem!important} +.gap-3{gap:1rem!important} +.small,.form-text{font-size:.875rem} +.text-muted{color:var(--bs-secondary)!important} +.fw-semibold{font-weight:600!important} +.fw-bold{font-weight:700!important} +.badge{display:inline-block;padding:.35em .65em;font-size:.75em;font-weight:700;line-height:1;border-radius:999px} +.bg-light{background:rgba(255,255,255,.7)!important} +.bg-body-secondary{background:rgba(32,25,18,.06)!important} +.rounded-4{border-radius:1.5rem!important} +.shadow-sm{box-shadow:0 .35rem 1.2rem rgba(32,25,18,.08)!important} +.card{position:relative;display:flex;flex-direction:column;min-width:0;word-wrap:break-word;background:rgba(255,251,246,.92);border:1px solid var(--bs-border-color);border-radius:var(--bs-border-radius-lg)} +.card-header{padding:1rem 1.25rem;border-bottom:1px solid var(--bs-border-color);background:transparent} +.card-body{flex:1 1 auto;padding:1.25rem} +.form-select,.form-control,.btn{font:inherit;border:1px solid var(--bs-border-color);border-radius:999px;padding:.625rem .95rem;background:rgba(255,255,255,.82);color:inherit} +.form-select{appearance:none} +.btn{cursor:pointer;display:inline-flex;align-items:center;justify-content:center;text-decoration:none} +.btn-primary{background:rgba(154,77,41,.14);border-color:rgba(154,77,41,.22)} +.btn-outline-secondary{background:rgba(32,25,18,.04)} +.btn-outline-danger{background:rgba(157,47,47,.12);border-color:rgba(157,47,47,.22);color:var(--bs-danger)} +.accordion{display:grid;gap:.85rem} +.accordion-item{border:1px solid var(--bs-border-color);border-radius:1rem;background:rgba(255,255,255,.62);overflow:hidden} +.accordion-button{width:100%;border:0;background:transparent;padding:1rem 1.15rem;text-align:left;display:flex;align-items:center;justify-content:space-between;gap:1rem;font:inherit;cursor:pointer} +.accordion-button.collapsed{background:transparent} +.accordion-body{padding:0 1.15rem 1.15rem} +.collapse{display:none} +.collapse.show{display:block} +.nav{display:flex;flex-wrap:wrap;padding-left:0;margin-bottom:0;list-style:none} +.nav-tabs{border-bottom:1px solid var(--bs-border-color);gap:.5rem} +.nav-link{display:block;padding:.6rem .9rem;border:1px solid transparent;border-radius:999px;cursor:pointer;background:transparent} +.nav-link.active{background:rgba(154,77,41,.12);border-color:rgba(154,77,41,.18)} +.tab-content>.tab-pane{display:none} +.tab-content>.active{display:block} +.table-responsive{width:100%;overflow:auto} +.table{width:100%;margin-bottom:0;border-collapse:collapse} +.table th,.table td{padding:.75rem .7rem;border-bottom:1px solid var(--bs-border-color);vertical-align:top;text-align:left} +.table-hover tbody tr:hover{background:rgba(32,25,18,.04)} +.progress{display:flex;height:.7rem;background:rgba(32,25,18,.08);border-radius:999px;overflow:hidden} +.progress-bar{display:flex;justify-content:flex-end;overflow:hidden;color:#fff;text-align:center;white-space:nowrap;background:linear-gradient(90deg,rgba(201,121,71,.96),rgba(154,77,41,.82));transition:width .2s ease} diff --git a/test_fake_worker.py b/test_fake_worker.py new file mode 100644 index 0000000..46667fd --- /dev/null +++ b/test_fake_worker.py @@ -0,0 +1,298 @@ +#!/usr/bin/env python3 +""" +Fake Worker - 模拟真实 Worker 上报采集结果 + +功能： +1. 模拟成功采集 +2. 模拟失败采集（各种错误类型） +3. 模拟重试场景 +4. 测试 analytics.replace_package_snapshot 的完整流程 +""" + +import time +import json +import sys +import os + +# 添加项目路径 +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from analytics import AnalyticsService + + +class FakeWorker: + """模拟 Worker 上报采集结果""" + + def __init__(self, worker_id="FAKE-WORKER-001"): + self.worker_id = worker_id + self.analytics = AnalyticsService(db_path="runtime/main/monitoring.sqlite3") + + def report_success(self, package_name: str, batch_tag: str = "test-batch"): + """模拟成功采集上报""" + print(f"\n[{self.worker_id}] 上报成功采集: {package_name}") + + summary = { + "package_name": package_name, + "app_name": f"Test App {package_name}", + "app_magic_label": f"magic_{package_name}", + "latest_task_key": f"{batch_tag}_{package_name}_ranking_1", + "latest_status": "success", + "restriction_status": "success", + "retryability": "retryable", + "latest_test_time": time.time(), + "latest_worker_id": self.worker_id, + "latest_task_detail": "采集成功", + "latest_failure_type": "", + "collection_status": "qualified", + "collection_status_reason": "success", + # 流量数据 + "total_traffic_bytes": 5000000, + "self_traffic_bytes": 3000000, + "server_traffic_bytes": 1500000, + "unrecognized_traffic_bytes": 500000, + "self_ratio": 60.0, + "recognition_ratio": 90.0, + # 执行数据 + "droidbot_steps": 50, + "gui_agent_steps": 20, + "duration_seconds": 180.5, + "num_nodes": 25, + "num_reached_activities": 8, + "app_num_total_activities": 15, + # 模型数据 + "model_flow_count": 10, + "model_traffic_bytes": 2000000, + "model_eligible": 1, + # 元数据 + "downloads": 1000000, + "source_order": 1, + "collection_task_type": "new_app", + } + + domain_rows = [ + { + "domain": "api.example.com", + "domain_type": "api", + "traffic_bytes": 2000000, + "flow_count": 50, + "traffic_ratio": 40.0, + "domain_traffic_ratio": 66.7, + "organization": "Example Inc", + "matched_pattern": "*.example.com", + "match_state": "matched", + }, + { + "domain": "cdn.test.com", + "domain_type": "cdn", + "traffic_bytes": 1000000, + "flow_count": 30, + "traffic_ratio": 20.0, + "domain_traffic_ratio": 33.3, + }, + ] + + component_rows = [ + { + "component_name": package_name, + "is_self": True, + "traffic_bytes": 3000000, + "share_percent": 60.0, + }, + { + "component_name": "com.google.android.gms", + "is_self": False, + "traffic_bytes": 1500000, + "share_percent": 30.0, + }, + ] + + source_rows = [] + + self.analytics.repo.replace_package_snapshot( + package_name, summary, domain_rows, component_rows, source_rows + ) + + print(f" ✅ 上报完成") + return summary + + def report_failure(self, package_name: str, error_type: str = "DOWNLOAD_ERROR/404"): + """模拟失败采集上报""" + print(f"\n[{self.worker_id}] 上报失败采集: {package_name} ({error_type})") + + summary = { + "package_name": package_name, + "app_name": f"Test App {package_name}", + "latest_task_key": f"test_{package_name}_ranking_1", + "latest_status": "failed", + "restriction_status": "severe_restricted", + "retryability": "non_retryable" if "404" in error_type else "retryable", + "latest_test_time": time.time(), + "latest_worker_id": self.worker_id, + "latest_task_detail": f"下载失败: {error_type}", + "latest_failure_type": error_type, + "collection_status": "restricted", + "collection_status_reason": error_type, + # 失败任务无流量数据 + "total_traffic_bytes": 0, + "self_traffic_bytes": 0, + "server_traffic_bytes": 0, + "num_nodes": 0, + "downloads": 500000, + "source_order": 10, + } + + self.analytics.repo.replace_package_snapshot( + package_name, summary, [], [], [] + ) + + print(f" ✅ 上报完成") + return summary + + def report_light_restricted(self, package_name: str): + """模拟轻度受限采集（成功但质量不足）""" + print(f"\n[{self.worker_id}] 上报轻度受限采集: {package_name}") + + summary = { + "package_name": package_name, + "app_name": f"Test App {package_name}", + "latest_task_key": f"test_{package_name}_ranking_1", + "latest_status": "success", + "restriction_status": "light_restricted", + "retryability": "retryable", + "latest_test_time": time.time(), + "latest_worker_id": self.worker_id, + "latest_task_detail": "采集成功但质量不足", + "collection_status": "restricted", + "collection_status_reason": "low_quality", + # 质量不足：nodes<5 或 self_ratio<30 + "total_traffic_bytes": 1000000, + "self_traffic_bytes": 200000, # self_ratio=20% < 30% + "server_traffic_bytes": 500000, + "unrecognized_traffic_bytes": 300000, + "self_ratio": 20.0, + "recognition_ratio": 70.0, + "num_nodes": 3, # < 5 + "downloads": 800000, + "source_order": 5, + } + + self.analytics.repo.replace_package_snapshot( + package_name, summary, [], [], [] + ) + + print(f" ✅ 上报完成") + return summary + + +def test_basic_flow(): + """测试基本流程""" + print("=" * 60) + print("测试1：基本采集流程") + print("=" * 60) + + worker = FakeWorker() + + # 1. 成功采集 + worker.report_success("com.test.app1", "batch-2026-06-15") + + # 2. 失败采集 + worker.report_failure("com.test.app2", "DOWNLOAD_ERROR/404") + + # 3. 轻度受限 + worker.report_light_restricted("com.test.app3") + + print("\n✅ 基本流程测试完成") + + +def test_retry_scenario(): + """测试重试场景""" + print("\n" + "=" * 60) + print("测试2：重试场景（同一应用多次上报）") + print("=" * 60) + + worker = FakeWorker() + package = "com.test.retry.app" + + # 第1次：失败 + print("\n--- 第1次尝试（失败）---") + worker.report_failure(package, "APP_ERROR/1") + + time.sleep(1) + + # 第2次：成功 + print("\n--- 第2次尝试（成功）---") + worker.report_success(package, "batch-retry") + + print("\n✅ 重试场景测试完成") + + +def test_query_after_report(): + """测试上报后查询""" + print("\n" + "=" * 60) + print("测试3：上报后查询验证") + print("=" * 60) + + worker = FakeWorker() + package = "com.test.query.app" + + # 上报 + worker.report_success(package, "batch-query-test") + + # 查询验证 + print(f"\n查询应用详情: {package}") + analytics = AnalyticsService(db_path="runtime/main/monitoring.sqlite3") + result = analytics.repo.get_collection_row(package) + + if result: + print(f" ✅ 查询成功:") + print(f" app_name: {result.get('app_name')}") + print(f" latest_status: {result.get('latest_status')}") + print(f" restriction_status: {result.get('restriction_status')}") + print(f" collection_status: {result.get('collection_status')}") + print(f" total_traffic_bytes: {result.get('total_traffic_bytes')}") + print(f" num_nodes: {result.get('num_nodes')}") + else: + print(f" ❌ 查询失败") + + print("\n✅ 查询验证测试完成") + + +def test_batch_report(): + """测试批量上报""" + print("\n" + "=" * 60) + print("测试4：批量上报（模拟多个Worker）") + print("=" * 60) + + workers = [ + FakeWorker(f"WORKER-{i:03d}") + for i in range(1, 4) + ] + + for i, worker in enumerate(workers, 1): + print(f"\n--- Worker {i} 上报 ---") + worker.report_success(f"com.test.batch.app{i}", "batch-multi") + + print("\n✅ 批量上报测试完成") + + +if __name__ == "__main__": + print("\n" + "🚀" * 30) + print("Fake Worker 测试开始") + print("🚀" * 30) + + try: + # 运行所有测试 + test_basic_flow() + test_retry_scenario() + test_query_after_report() + test_batch_report() + + print("\n" + "🎉" * 30) + print("所有测试通过！") + print("🎉" * 30) + + except Exception as e: + print(f"\n❌ 测试失败: {e}") + import traceback + traceback.print_exc() + sys.exit(1) diff --git a/test_minio_config.py b/test_minio_config.py new file mode 100644 index 0000000..7b26734 --- /dev/null +++ b/test_minio_config.py @@ -0,0 +1,48 @@ +#!/usr/bin/env python3 +"""测试 Minio 配置是否生效""" + +import sys +import os + +# 确保从当前目录加载 +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from config import MINIO_ENABLED, APK_DOWNLOAD_MODE + +print("=" * 60) +print("Minio 配置测试") +print("=" * 60) + +print(f"\nMINIO_ENABLED: {MINIO_ENABLED}") +print(f"APK_DOWNLOAD_MODE: {APK_DOWNLOAD_MODE}") + +if MINIO_ENABLED: + print("\n❌ 配置未生效：Minio 仍然启用") + print(" 请检查 config/prod.json 中的 MINIO_ENABLED 设置") + sys.exit(1) +else: + print("\n✓ 配置已生效：Minio 已禁用") + +# 测试 MinioStorage.create() 工厂方法 +print("\n测试 MinioStorage.create() 工厂方法...") +try: + from apk_cloud.storage import MinioStorage + + storage = MinioStorage.create() + + if storage is None: + print("✓ MinioStorage.create() 返回 None（符合预期）") + else: + print(f"❌ MinioStorage.create() 返回了实例：{storage}") + sys.exit(1) + +except ImportError as e: + print(f"⚠️ 无法导入 MinioStorage（可能缺少 minio 依赖）：{e}") + print(" 这是正常的，因为 Minio 未启用时不需要安装 minio 库") +except Exception as e: + print(f"❌ 测试失败：{e}") + sys.exit(1) + +print("\n" + "=" * 60) +print("✓ 所有测试通过") +print("=" * 60) diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_analytics_new_schema.py b/tests/test_analytics_new_schema.py new file mode 100644 index 0000000..45b3157 --- /dev/null +++ b/tests/test_analytics_new_schema.py @@ -0,0 +1,241 @@ +import json + +from analytics import AnalyticsService, _make_task_key +from redis_task_distribute import RedisTaskDispatcher + + +class FakeRedis: + def __init__(self): + self.lists = {} + self.hashes = {} + self.sets = {} + + def delete(self, key): + self.lists.pop(key, None) + self.hashes.pop(key, None) + self.sets.pop(key, None) + + def lpush(self, key, value): + self.lists.setdefault(key, []).insert(0, value) + + def rpush(self, key, value): + self.lists.setdefault(key, []).append(value) + + def rpop(self, key): + values = self.lists.setdefault(key, []) + return values.pop() if values else None + + def lrange(self, key, start, end): + values = self.lists.get(key, []) + stop = None if end == -1 else end + 1 + return values[start:stop] + + def llen(self, key): + return len(self.lists.get(key, [])) + + def lpos(self, key, value): + try: + return self.lists.get(key, []).index(value) + except ValueError: + return None + + def lrem(self, key, count, value): + values = self.lists.get(key, []) + original_len = len(values) + self.lists[key] = [item for item in values if item != value] + return original_len - len(self.lists[key]) + + def hset(self, key, field, value): + self.hashes.setdefault(key, {})[field] = value + + def hget(self, key, field): + return self.hashes.get(key, {}).get(field) + + def hgetall(self, key): + return dict(self.hashes.get(key, {})) + + def hvals(self, key): + return list(self.hashes.get(key, {}).values()) + + def hlen(self, key): + return len(self.hashes.get(key, {})) + + def hdel(self, key, field): + self.hashes.get(key, {}).pop(field, None) + + def sadd(self, key, value): + self.sets.setdefault(key, set()).add(value) + + def srem(self, key, value): + self.sets.setdefault(key, set()).discard(value) + + def scard(self, key): + return len(self.sets.get(key, set())) + + def smembers(self, key): + return set(self.sets.get(key, set())) + + def scan_iter(self, match=None): + return iter(()) + + +class FakeMonitor: + def set_worker_state(self, *args, **kwargs): + return None + + def set_failure_bucket(self, *args, **kwargs): + return None + + +class FakeWorker: + def __init__(self, analytics, dispatcher, worker_id="fake-worker"): + self.analytics = analytics + self.dispatcher = dispatcher + self.worker_id = worker_id + + def report( + self, + *, + app_name, + package_name, + status, + failure_type="", + num_nodes=12, + self_traffic_bytes=600, + server_traffic_bytes=300, + total_traffic_bytes=1000, + model_flow_count=0, + ): + summary = { + "app_name": app_name, + "app_magic_label": f"magic:{package_name}", + "latest_task_key": _make_task_key(app_name, package_name), + "latest_status": status, + "latest_worker_id": self.worker_id, + "latest_failure_type": failure_type, + "latest_task_detail": failure_type, + "collection_status_reason": failure_type or "success", + "collection_task_type": "new_app", + "duration_seconds": 42, + "droidbot_steps": 30, + "gui_agent_steps": 12, + "num_nodes": num_nodes, + "total_traffic_bytes": total_traffic_bytes, + "self_traffic_bytes": self_traffic_bytes, + "server_traffic_bytes": server_traffic_bytes, + "unrecognized_traffic_bytes": max(0, total_traffic_bytes - self_traffic_bytes - server_traffic_bytes), + "model_flow_count": model_flow_count, + "model_traffic_bytes": model_flow_count * 100, + } + self.analytics.repo.replace_package_snapshot(package_name, summary, [], [], []) + persisted = self.analytics.repo.get_collection_row(package_name) + self.dispatcher._handle_analytics_snapshot(package_name, persisted) + return persisted + + +def make_dispatcher(redis_conn, analytics): + dispatcher = RedisTaskDispatcher.__new__(RedisTaskDispatcher) + dispatcher.redis = redis_conn + dispatcher.analytics = analytics + dispatcher.worker_inventory = {} + dispatcher.managed_worker_ids = set() + dispatcher.only_managed_workers_can_dispatch = False + dispatcher.worker_online_timeout = 300 + dispatcher.task_routing_rules = {"package_name": {}, "task_key": {}} + dispatcher.monitor = FakeMonitor() + dispatcher.notifier = None + dispatcher._apk_registry = None + dispatcher._minio_storage = None + return dispatcher + + +def task_status(redis_conn, task_key): + payload = redis_conn.hget("task:status", task_key) + return json.loads(payload) if payload else {} + + +def test_repository_initializes_new_tables_without_views(tmp_path): + analytics = AnalyticsService(db_path=str(tmp_path / "analytics.sqlite3"), start_worker=False) + + with analytics.repo._connect() as connection: + names = { + row["name"] + for row in connection.execute("SELECT name FROM sqlite_master WHERE type IN ('table', 'view')") + } + + assert "app_catalog" in names + assert "collection_task" in names + assert "app_collect_summary" not in names + assert "v_collection_latest" not in names + + analytics.replace_catalog_from_rows( + [{"app_name": "Tagged App", "package_name": "com.example.tagged"}] + ) + assert analytics.set_incremental_batch_tag_for_packages(["com.example.tagged"], "batch-a") == 1 + assert "batch-a" in analytics.repo.get_collection_row("com.example.tagged")["incremental_batch_tags"] + assert analytics.clear_incremental_batch_tags("batch-a") == 1 + assert "batch-a" not in analytics.repo.get_collection_row("com.example.tagged")["incremental_batch_tags"] + + +def test_fake_worker_reports_drive_new_schema_and_redis_states(tmp_path): + analytics = AnalyticsService(db_path=str(tmp_path / "analytics.sqlite3"), start_worker=False) + analytics.replace_catalog_from_rows( + [ + {"app_name": "Good App", "package_name": "com.example.good", "downloads": "100K"}, + {"app_name": "Gone App", "package_name": "com.example.gone", "downloads": "1K"}, + {"app_name": "Retry App", "package_name": "com.example.retry", "downloads": "2M"}, + {"app_name": "Model App", "package_name": "com.example.model", "downloads": "3M"}, + ] + ) + redis_conn = FakeRedis() + dispatcher = make_dispatcher(redis_conn, analytics) + + assert dispatcher.load_tasks_from_app_summary() == 4 + + worker = FakeWorker(analytics, dispatcher) + good = worker.report(app_name="Good App", package_name="com.example.good", status="success") + gone = worker.report( + app_name="Gone App", + package_name="com.example.gone", + status="failed", + failure_type="APP_ERROR/3", + num_nodes=0, + total_traffic_bytes=0, + self_traffic_bytes=0, + server_traffic_bytes=0, + ) + retry = worker.report( + app_name="Retry App", + package_name="com.example.retry", + status="failed", + failure_type="DOWNLOAD_ERROR/5", + num_nodes=0, + total_traffic_bytes=0, + self_traffic_bytes=0, + server_traffic_bytes=0, + ) + model = worker.report( + app_name="Model App", + package_name="com.example.model", + status="success", + model_flow_count=2, + ) + + assert good["collection_status"] == "qualified" + assert gone["collection_status"] == "failed_terminal" + assert retry["collection_status"] == "pending" + assert model["model_eligible"] == 1 + + good_key = _make_task_key("Good App", "com.example.good") + gone_key = _make_task_key("Gone App", "com.example.gone") + retry_key = _make_task_key("Retry App", "com.example.retry") + + assert task_status(redis_conn, good_key)["status"] == "qualified" + assert good_key in redis_conn.smembers("task:completed") + assert task_status(redis_conn, gone_key)["status"] == "failed" + assert gone_key in redis_conn.smembers("task:failed") + assert task_status(redis_conn, retry_key)["status"] == "pending" + assert retry_key in redis_conn.lrange("task:queue:default", 0, -1) + + assert dispatcher.refresh_tasks_from_app_summary() == 1 + assert _make_task_key("Model App", "com.example.model") + "_model" in redis_conn.lrange("task:queue:low", 0, -1) diff --git a/tests/test_block_task.py b/tests/test_block_task.py new file mode 100644 index 0000000..7d8595f --- /dev/null +++ b/tests/test_block_task.py @@ -0,0 +1,1156 @@ +# -*- encoding=utf8 -*- +""" +Block 测试任务功能单元测试 + +覆盖: +1. Block 任务加载（task_key / task_details / 高优先队列） +2. Worker 收到的 payload 包含 blocked 参数 +3. Block 任务重试逻辑（下载错误重试，其他错误直接失败） +4. Analytics snapshot 跳过 block 任务 +5. Block 测试报告生成 +""" + +import json +import sys +import os +import tempfile +import csv + +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, PROJECT_ROOT) + +import pytest + +from analytics import _make_task_key + + +class FakeRedis: + """模拟 Redis 操作""" + + def __init__(self): + self.lists: dict = {} + self.hashes: dict = {} + self.sets: dict = {} + self.published: list = [] + + def delete(self, key): + self.lists.pop(key, None) + self.hashes.pop(key, None) + self.sets.pop(key, None) + + def lpush(self, key, value): + self.lists.setdefault(key, []).insert(0, value) + + def rpush(self, key, value): + self.lists.setdefault(key, []).append(value) + + def rpop(self, key): + values = self.lists.setdefault(key, []) + return values.pop() if values else None + + def lrange(self, key, start, end): + values = self.lists.get(key, []) + stop = None if end == -1 else end + 1 + return values[start:stop] + + def llen(self, key): + return len(self.lists.get(key, [])) + + def lpos(self, key, value): + try: + return self.lists.get(key, []).index(value) + except ValueError: + return None + + def lrem(self, key, count, value): + values = self.lists.get(key, []) + original_len = len(values) + self.lists[key] = [item for item in values if item != value] + return original_len - len(self.lists[key]) + + def hset(self, key, *args, **kwargs): + if len(args) == 2: + field, value = args + self.hashes.setdefault(key, {})[field] = value + elif len(args) == 1 and isinstance(args[0], dict): + self.hashes.setdefault(key, {}).update(args[0]) + + def hget(self, key, field): + return self.hashes.get(key, {}).get(field) + + def hgetall(self, key): + return dict(self.hashes.get(key, {})) + + def hvals(self, key): + return list(self.hashes.get(key, {}).values()) + + def hlen(self, key): + return len(self.hashes.get(key, {})) + + def hexists(self, key, field): + return field in self.hashes.get(key, {}) + + def hdel(self, key, field): + self.hashes.get(key, {}).pop(field, None) + + def sadd(self, key, value): + self.sets.setdefault(key, set()).add(value) + + def srem(self, key, value): + self.sets.setdefault(key, set()).discard(value) + + def scard(self, key): + return len(self.sets.get(key, set())) + + def smembers(self, key): + return set(self.sets.get(key, set())) + + def sismember(self, key, value): + return value in self.sets.get(key, set()) + + def exists(self, key): + return int(key in self.hashes or key in self.lists or key in self.sets) + + def publish(self, channel, message): + self.published.append((channel, message)) + + def keys(self, pattern="*"): + result = [] + for name in list(self.lists) + list(self.hashes) + list(self.sets): + if pattern == "*" or name.startswith(pattern.replace("*", "")): + result.append(name) + return result + + def scan_iter(self, match="*"): + for key in self.keys(match): + yield key + + def setex(self, key, ttl, value): + pass + + def pipeline(self, transaction=True): + return _FakePipeline(self) + + @staticmethod + def lock(lock_key, timeout=5, blocking_timeout=1): + class _FakeLock: + def __enter__(self2): + pass + + def __exit__(self2, *args): + pass + + return _FakeLock() + + +class _FakePipeline: + """模拟 Redis pipeline""" + + def __init__(self, redis_conn): + self._redis = redis_conn + + def __enter__(self): + return self + + def __exit__(self, *args): + pass + + def sadd(self, key, value): + self._redis.sadd(key, value) + return self + + def srem(self, key, value): + self._redis.srem(key, value) + return self + + def hset(self, key, field, value): + self._redis.hset(key, field, value) + return self + + def hdel(self, key, field): + self._redis.hdel(key, field) + return self + + def delete(self, key): + self._redis.delete(key) + return self + + def expire(self, key, ttl): + return self + + def execute(self): + pass + + +class StubAnalytics: + """模拟 Analytics 服务""" + + def __init__(self, pending=None, model=None, qualified=None): + self.pending = pending or [] + self.model = model or [] + self.qualified = qualified or [] + self.enqueued_incremental = [] + + def list_pending_collection_tasks(self): + return list(self.pending) + + def list_model_eligible_apps(self): + return list(self.model) + + def enqueue_incremental(self, package_name, **kwargs): + self.enqueued_incremental.append((package_name, kwargs)) + + +class StubMonitor: + """模拟 Monitoring 服务""" + + def __init__(self): + self.failure_buckets: dict = {} + self.assigned_tasks: list = [] + self.reports: list = [] + + def set_failure_bucket(self, task_key, bucket): + self.failure_buckets[task_key] = bucket + + def handle_task_assigned(self, worker_id, task_key, **kwargs): + self.assigned_tasks.append((worker_id, task_key)) + + def handle_worker_report(self, **kwargs): + self.reports.append(kwargs) + + def set_worker_state(self, worker_id, state, **kwargs): + pass + + def register_worker(self, worker_id, **kwargs): + pass + + def mark_worker_offline(self, worker_id, reason): + pass + + def close_controller_session(self, **kwargs): + pass + + def start_controller_session(self, **kwargs): + pass + + def heartbeat_controller_session(self): + pass + + +class StubNotifier: + def send_weCom_alert(self, *args, **kwargs): + pass + + def send_weChat_alert(self, *args, **kwargs): + pass + + +class StubAnalyticsRepo: + def get_collection_row(self, package_name): + return None + + def get_latest_task_execution(self, package_name): + return None + + def has_successful_related_magic_label(self, app_magic_label, package_name): + return False + + def list_qualified_apps(self): + return list(self.qualified_apps) if hasattr(self, 'qualified_apps') else [] + + +class StubAnalyticsWithRepo: + def __init__(self, qualified_apps=None): + self.repo = StubAnalyticsRepo() + if qualified_apps: + self.repo.qualified_apps = qualified_apps + self.enqueued_incremental = [] + + def list_pending_collection_tasks(self): + return [] + + def list_model_eligible_apps(self): + return [] + + def list_qualified_apps(self): + return self.repo.list_qualified_apps() + + def enqueue_incremental(self, **kwargs): + self.enqueued_incremental.append(kwargs) + + def close(self): + pass + + +def make_dispatcher(redis_conn=None, analytics=None, monitor=None, notifier=None): + """构建一个最小化的 RedisTaskDispatcher 用于测试""" + from redis_task_distribute import RedisTaskDispatcher + + dispatcher = RedisTaskDispatcher.__new__(RedisTaskDispatcher) + dispatcher.redis = redis_conn or FakeRedis() + dispatcher.analytics = analytics or StubAnalytics() + dispatcher.monitor = monitor or StubMonitor() + dispatcher.notifier = notifier or StubNotifier() + dispatcher.worker_inventory = {} + dispatcher.managed_worker_ids = set() + dispatcher.only_managed_workers_can_dispatch = False + dispatcher.worker_online_timeout = 300 + dispatcher.task_routing_rules = {"package_name": {}, "task_key": {}} + dispatcher._apk_registry = None + dispatcher._minio_storage = None + + import types + def _mark_noop(self, package_name, app_name="", last_updated=""): + return None + dispatcher._mark_download_error_awaiting_apk = types.MethodType(_mark_noop, dispatcher) + + def _dispatch_alert_noop(self, *args, **kwargs): + pass + dispatcher._dispatch_alert = types.MethodType(_dispatch_alert_noop, dispatcher) + + def _registry_is_fresh_enough_noop(package_name, last_updated): + return False + class _StubRegistry: + is_fresh_enough = staticmethod(_registry_is_fresh_enough_noop) + dispatcher._get_apk_registry = lambda: _StubRegistry + return dispatcher + + +def qualified_entry(app_name, package_name, country_code="US", device_type="emulator"): + """构建一个合格的 app 条目""" + return { + "package_name": package_name, + "app_name": app_name, + "task_payload": { + "app_name": app_name, + "package_name": package_name, + "country_code": country_code, + "device_type": device_type, + "original_row": {}, + }, + } + + +# ============================================================= +# 1. Block 任务加载测试 +# ============================================================= + + +class TestBlockTaskLoading: + """测试 block 任务的加载行为""" + + def test_block_task_has_correct_task_key_suffix(self): + """block 任务 key 以 _block 结尾""" + redis_conn = FakeRedis() + dispatcher = make_dispatcher(redis_conn=redis_conn) + + entries = [qualified_entry("TestApp", "com.test.block")] + count = dispatcher.load_block_tasks(entries) + assert count == 1 + + expected_key = "TestApp_com.test.block_block" + assert redis_conn.hexists("task:details", expected_key), ( + f"expected task_key '{expected_key}' in task:details" + ) + + def test_block_tasks_pushed_to_high_priority_queue(self): + """block 任务入高优队列""" + redis_conn = FakeRedis() + dispatcher = make_dispatcher(redis_conn=redis_conn) + + entries = [ + qualified_entry("App1", "com.app1"), + qualified_entry("App2", "com.app2"), + ] + count = dispatcher.load_block_tasks(entries) + assert count == 2 + + high_queue = redis_conn.lrange("task:queue:high", 0, -1) + assert len(high_queue) == 2 + for task_key in high_queue: + assert task_key.endswith("_block") + + def test_block_task_details_contain_is_block_task_flag(self): + """task_details 包含 is_block_task=True""" + redis_conn = FakeRedis() + dispatcher = make_dispatcher(redis_conn=redis_conn) + + entries = [qualified_entry("BlockApp", "com.block.test")] + dispatcher.load_block_tasks(entries) + + task_key = "BlockApp_com.block.test_block" + details_json = redis_conn.hget("task:details", task_key) + details = json.loads(details_json) + assert details.get("is_block_task") is True + assert details.get("task_queue") == "high" + + def test_block_task_has_app_name_and_package_in_details(self): + """task_details 保留原始 app_name 和 package_name""" + redis_conn = FakeRedis() + dispatcher = make_dispatcher(redis_conn=redis_conn) + + entries = [ + qualified_entry("MyApp", "com.company.myapp", country_code="CN", device_type="emulator") + ] + dispatcher.load_block_tasks(entries) + + task_key = "MyApp_com.company.myapp_block" + details = json.loads(redis_conn.hget("task:details", task_key)) + assert details["app_name"] == "MyApp" + assert details["package_name"] == "com.company.myapp" + assert details["country_code"] == "CN" + assert details["device_type"] == "emulator" + + def test_duplicate_block_tasks_are_skipped(self): + """重复加载 block 任务先清除旧任务再重新加载，每次 count 相同""" + redis_conn = FakeRedis() + dispatcher = make_dispatcher(redis_conn=redis_conn) + + entries = [qualified_entry("DupApp", "com.dup")] + assert dispatcher.load_block_tasks(entries) == 1 + # 第二次加载先清除旧任务再重新加载，不是简单地跳过 + assert dispatcher.load_block_tasks(entries) == 1 + # 队列中只有一个任务（不是两个） + high_queue = redis_conn.lrange("task:queue:high", 0, -1) + assert len(high_queue) == 1 + + def test_load_all_qualified_block_tasks_calls_repo(self): + """load_all_qualified_block_tasks 从 DB 查询 qualified apps""" + redis_conn = FakeRedis() + analytics = StubAnalyticsWithRepo(qualified_apps=[ + qualified_entry("QualApp", "com.qual"), + qualified_entry("AnotherQual", "com.another"), + ]) + dispatcher = make_dispatcher(redis_conn=redis_conn, analytics=analytics) + + count = dispatcher.load_all_qualified_block_tasks() + assert count == 2 + + assert redis_conn.hexists("task:details", "QualApp_com.qual_block") + assert redis_conn.hexists("task:details", "AnotherQual_com.another_block") + + +# ============================================================= +# 2. Worker 接收 payload 测试 +# ============================================================= + + +class TestBlockTaskPayload: + """测试 Worker 收到的 task payload""" + + def test_task_payload_includes_traffic_root(self): + """_task_to_payload 的 traffic_root 和 blocked 字段：block 任务用 block 路径""" + from config import ANALYTICS_TRAFFIC_ROOT_BLOCK + + redis_conn = FakeRedis() + dispatcher = make_dispatcher(redis_conn=redis_conn) + + task_details = { + "app_name": "TestApp", + "package_name": "com.test", + "country_code": "US", + "is_block_task": True, + "available_sources": ["google_play", "local"], + } + payload = dispatcher._task_to_payload("TestApp_com.test_block", task_details) + assert payload["blocked"] is True + assert payload["traffic_root"] == ANALYTICS_TRAFFIC_ROOT_BLOCK + + def test_normal_task_payload_uses_default_traffic_root(self): + """普通任务的 blocked=False, traffic_root 为默认路径""" + from config import ANALYTICS_TRAFFIC_ROOT + + redis_conn = FakeRedis() + dispatcher = make_dispatcher(redis_conn=redis_conn) + + task_details = { + "app_name": "NormalApp", + "package_name": "com.normal", + "country_code": "US", + "available_sources": ["google_play", "local"], + } + payload = dispatcher._task_to_payload("NormalApp_com.normal", task_details) + assert payload["blocked"] is False + assert payload["traffic_root"] == ANALYTICS_TRAFFIC_ROOT + + def test_worker_init_receives_block_task_with_traffic_root(self): + """Worker init 返回的 task 有正确的 traffic_root""" + from config import ANALYTICS_TRAFFIC_ROOT_BLOCK + redis_conn = FakeRedis() + dispatcher = make_dispatcher(redis_conn=redis_conn) + + entries = [qualified_entry("BlockInit", "com.block.init")] + dispatcher.load_block_tasks(entries) + + task_key = "BlockInit_com.block.init_block" + + worker_id = "192.168.1.1_AA:BB:CC:DD:EE:01" + worker_info = { + "worker_id": worker_id, + "ip_address": "192.168.1.1", + "mac_address": "AA:BB:CC:DD:EE:01", + "hostname": "test-host", + "platform": "Windows", + "device_type": "emulator", + "register_time": 1000000000.0, + "last_update": 1000000000.0, + "status": "idle", + "dispatch_enabled": True, + } + redis_conn.hset("workers:info", worker_id, json.dumps(worker_info)) + redis_conn.sadd("workers:idle", worker_id) + + task = dispatcher.worker_init( + worker_id=worker_id, + ip_address="192.168.1.1", + mac_address="AA:BB:CC:DD:EE:01", + hostname="test-host", + platform="Windows", + device_type="emulator", + ) + + assert task is not None + assert task["blocked"] is True + assert task["traffic_root"] == ANALYTICS_TRAFFIC_ROOT_BLOCK + assert task["task_key"] == task_key + + def test_worker_report_then_assign_next_block_task(self): + """Worker 上报完成后领取下一个 block 任务""" + redis_conn = FakeRedis() + monitor = StubMonitor() + notifier = StubNotifier() + dispatcher = make_dispatcher( + redis_conn=redis_conn, monitor=monitor, notifier=notifier + ) + + entries = [ + qualified_entry("BlockA", "com.block.a"), + qualified_entry("BlockB", "com.block.b"), + ] + dispatcher.load_block_tasks(entries) + + task_a_key = "BlockA_com.block.a_block" + task_b_key = "BlockB_com.block.b_block" + + worker_id = "192.168.1.1_AA:BB:CC:DD:EE:01" + worker_info = { + "worker_id": worker_id, + "ip_address": "192.168.1.1", + "mac_address": "AA:BB:CC:DD:EE:01", + "hostname": "test-host", + "platform": "Windows", + "device_type": "emulator", + "register_time": 1000000000.0, + "last_update": 1000000000.0, + "status": "idle", + "dispatch_enabled": True, + } + redis_conn.hset("workers:info", worker_id, json.dumps(worker_info)) + redis_conn.sadd("workers:idle", worker_id) + + first_task = dispatcher.worker_init( + worker_id=worker_id, + ip_address="192.168.1.1", + mac_address="AA:BB:CC:DD:EE:01", + hostname="test-host", + platform="Windows", + device_type="emulator", + ) + assert first_task is not None + + second_task = dispatcher.worker_report( + worker_id=worker_id, + previous_task_key=first_task["task_key"], + report_data={"status": "success"}, + ) + assert second_task is not None + from config import ANALYTICS_TRAFFIC_ROOT_BLOCK + assert second_task["traffic_root"] == ANALYTICS_TRAFFIC_ROOT_BLOCK + + +# ============================================================= +# 3. Block 任务重试逻辑测试 +# ============================================================= + + +class TestBlockTaskRetry: + """测试 block 任务的专用重试策略""" + + def _setup_block_task(self, redis_conn, task_key, app_name="BlockApp", package_name="com.block"): + """准备一个 block 任务的 Redis 数据""" + task_details = { + "app_name": app_name, + "package_name": package_name, + "country_code": "US", + "is_block_task": True, + "task_queue": "high", + } + redis_conn.hset("task:details", task_key, json.dumps(task_details)) + redis_conn.hset("task:status", task_key, json.dumps({ + "status": "running", + "worker_id": "192.168.1.1_WORKER", + "retry_count": 0, + })) + + def test_block_task_download_error_goes_to_pending_retry(self): + """block 任务的下载错误 -> pending_retry""" + from result_codes import ErrorInfo, ErrorCategory, DownloadError + + redis_conn = FakeRedis() + monitor = StubMonitor() + notifier = StubNotifier() + dispatcher = make_dispatcher( + redis_conn=redis_conn, monitor=monitor, notifier=notifier + ) + + task_key = "BlockApp_com.block_block" + self._setup_block_task(redis_conn, task_key) + + dispatcher._record_failed_task( + task_key=task_key, + worker_id="192.168.1.1_WORKER", + error_info=ErrorInfo.download(DownloadError.REGION_RESTRICTED), + message="Google Play 锁区", + ) + + status = json.loads(redis_conn.hget("task:status", task_key)) + assert status["status"] == "pending_retry", ( + f"expected pending_retry, got {status['status']}" + ) + assert monitor.failure_buckets.get(task_key) == "retry" + + def test_block_task_app_error_is_permanent_fail(self): + """block 任务的 APP_ERROR -> 永久失败""" + from result_codes import ErrorInfo, ErrorCategory, AppError + + redis_conn = FakeRedis() + monitor = StubMonitor() + notifier = StubNotifier() + dispatcher = make_dispatcher( + redis_conn=redis_conn, monitor=monitor, notifier=notifier + ) + + task_key = "BlockApp_com.block_block" + self._setup_block_task(redis_conn, task_key) + + dispatcher._record_failed_task( + task_key=task_key, + worker_id="192.168.1.1_WORKER", + error_info=ErrorInfo.app(AppError.CRASH), + message="应用闪退", + ) + + status = json.loads(redis_conn.hget("task:status", task_key)) + assert status["status"] == "failed" + assert redis_conn.sismember("task:failed", task_key) if hasattr(redis_conn, 'sismember') else ( + task_key in redis_conn.smembers("task:failed") + ) + assert monitor.failure_buckets.get(task_key) == "failed" + + def test_block_task_business_error_is_permanent_fail(self): + """block 任务的 BUSINESS_ERROR -> 永久失败""" + from result_codes import ErrorInfo, ErrorCategory, BusinessError + + redis_conn = FakeRedis() + monitor = StubMonitor() + notifier = StubNotifier() + dispatcher = make_dispatcher( + redis_conn=redis_conn, monitor=monitor, notifier=notifier + ) + + task_key = "BlockApp_com.block_block" + self._setup_block_task(redis_conn, task_key) + + dispatcher._record_failed_task( + task_key=task_key, + worker_id="192.168.1.1_WORKER", + error_info=ErrorInfo.business(BusinessError.LOGIN_FAILED), + message="登录失败", + ) + + status = json.loads(redis_conn.hget("task:status", task_key)) + assert status["status"] == "failed" + assert task_key in redis_conn.smembers("task:failed") + + def test_block_task_infra_non_download_error_is_permanent_fail(self): + """block 任务的非下载 INFRA_ERROR -> 永久失败""" + from result_codes import ErrorInfo, ErrorCategory, InfraError + + redis_conn = FakeRedis() + monitor = StubMonitor() + notifier = StubNotifier() + dispatcher = make_dispatcher( + redis_conn=redis_conn, monitor=monitor, notifier=notifier + ) + + task_key = "BlockApp_com.block_block" + self._setup_block_task(redis_conn, task_key) + + dispatcher._record_failed_task( + task_key=task_key, + worker_id="192.168.1.1_WORKER", + error_info=ErrorInfo.infra(InfraError.ADB_ERROR), + message="ADB 断联", + ) + + status = json.loads(redis_conn.hget("task:status", task_key)) + assert status["status"] == "failed" + assert task_key in redis_conn.smembers("task:failed") + + def test_block_task_infra_download_failed_goes_to_retry(self): + """block 任务 INFRA_ERROR.DOWNLOAD_FAILED -> pending_retry""" + from result_codes import ErrorInfo, ErrorCategory, InfraError + + redis_conn = FakeRedis() + monitor = StubMonitor() + notifier = StubNotifier() + dispatcher = make_dispatcher( + redis_conn=redis_conn, monitor=monitor, notifier=notifier + ) + + task_key = "BlockApp_com.block_block" + self._setup_block_task(redis_conn, task_key) + + dispatcher._record_failed_task( + task_key=task_key, + worker_id="192.168.1.1_WORKER", + error_info=ErrorInfo.infra(InfraError.DOWNLOAD_FAILED), + message="下载失败", + ) + + status = json.loads(redis_conn.hget("task:status", task_key)) + assert status["status"] == "pending_retry", ( + f"expected pending_retry for DOWNLOAD_FAILED, got {status['status']}" + ) + assert monitor.failure_buckets.get(task_key) == "retry" + + def test_block_task_retry_avoids_same_worker(self): + """block 任务下载失败重试时避开同一 worker""" + from result_codes import ErrorInfo, ErrorCategory, DownloadError + + redis_conn = FakeRedis() + monitor = StubMonitor() + notifier = StubNotifier() + dispatcher = make_dispatcher( + redis_conn=redis_conn, monitor=monitor, notifier=notifier + ) + + task_key = "BlockApp_com.block_block" + worker_id = "192.168.2.1_FAILED_WORKER" + self._setup_block_task(redis_conn, task_key) + + dispatcher._record_failed_task( + task_key=task_key, + worker_id=worker_id, + error_info=ErrorInfo.download(DownloadError.APP_NOT_FOUND), + message="Google Play 未找到", + ) + + details = json.loads(redis_conn.hget("task:details", task_key)) + assert worker_id in details.get("excluded_worker_ids", []) + + def test_normal_task_retry_unchanged(self): + """block 任务的逻辑不影响普通任务重试""" + from result_codes import ErrorInfo, ErrorCategory, AppError + + redis_conn = FakeRedis() + monitor = StubMonitor() + notifier = StubNotifier() + dispatcher = make_dispatcher( + redis_conn=redis_conn, monitor=monitor, notifier=notifier + ) + + task_key = "NormalApp_com.normal" + redis_conn.hset("task:details", task_key, json.dumps({ + "app_name": "NormalApp", + "package_name": "com.normal", + "country_code": "US", + "device_type": "physical", + })) + redis_conn.hset("task:status", task_key, json.dumps({ + "status": "running", + "worker_id": "worker_x", + "retry_count": 0, + })) + + dispatcher._record_failed_task( + task_key=task_key, + worker_id="worker_x", + error_info=ErrorInfo.app(AppError.CRASH), + message="闪退", + ) + + status = json.loads(redis_conn.hget("task:status", task_key)) + assert status["status"] == "failed" + assert task_key in redis_conn.smembers("task:failed") + + +# ============================================================= +# 4. Analytics snapshot 跳过测试 +# ============================================================= + + +class TestAnalyticsSnapshotSkip: + """测试 analytics snapshot 回调跳过 block 任务""" + + def _setup_task(self, redis_conn, task_key, is_blocked=False): + redis_conn.hset("task:details", task_key, json.dumps({ + "app_name": "Test", + "package_name": "com.test", + "is_block_task": is_blocked, + "task_queue": "high", + })) + redis_conn.hset("task:status", task_key, json.dumps({ + "status": "pending_retry", + "retry_count": 1, + })) + + def test_block_task_skipped_in_analytics_snapshot(self): + """block 任务在 analytics 回调中被跳过（不重新入队）""" + redis_conn = FakeRedis() + dispatcher = make_dispatcher(redis_conn=redis_conn) + + task_key = "BlockApp_com.block_block" + self._setup_task(redis_conn, task_key, is_blocked=True) + + summary = { + "task_payload": {"app_name": "BlockApp", "package_name": "com.block"}, + "latest_task_key": task_key, + "collection_status": "pending", + } + + dispatcher._handle_analytics_snapshot("com.block", summary) + + status = json.loads(redis_conn.hget("task:status", task_key)) + assert status["status"] == "pending_retry", ( + f"block task should not be re-queued, got {status['status']}" + ) + + def test_block_task_suffix_also_skipped(self): + """以 _block 结尾的任务同样被 analytics 跳过""" + redis_conn = FakeRedis() + dispatcher = make_dispatcher(redis_conn=redis_conn) + + task_key = "SomeApp_com.some_block" + redis_conn.hset("task:details", task_key, json.dumps({ + "app_name": "SomeApp", + "package_name": "com.some", + })) + + summary = { + "latest_task_key": task_key, + "collection_status": "pending", + } + + dispatcher._handle_analytics_snapshot("com.some", summary) + + assert not redis_conn.hexists("task:status", task_key), ( + "block-suffix task should be skipped entirely (status not set)" + ) + + def test_normal_task_still_processed_by_analytics(self): + """普通任务在 analytics 回调中照常处理""" + redis_conn = FakeRedis() + dispatcher = make_dispatcher(redis_conn=redis_conn) + + task_key = "NormalApp_com.normal" + self._setup_task(redis_conn, task_key, is_blocked=False) + + summary = { + "task_payload": {"app_name": "NormalApp", "package_name": "com.normal"}, + "latest_task_key": task_key, + "collection_status": "qualified", + } + + dispatcher._handle_analytics_snapshot("com.normal", summary) + + assert task_key in redis_conn.smembers("task:completed") + + def test_block_task_redis_state_unchanged_by_callback(self): + """block 任务无论 collection_status 如何 Redis 状态不变""" + redis_conn = FakeRedis() + dispatcher = make_dispatcher(redis_conn=redis_conn) + + task_key = "BlockApp_com.block_block" + self._setup_task(redis_conn, task_key, is_blocked=True) + + for status in ["qualified", "failed_terminal", "pending"]: + summary = { + "task_payload": {"app_name": "BlockApp", "package_name": "com.block"}, + "latest_task_key": task_key, + "collection_status": status, + } + dispatcher._handle_analytics_snapshot("com.block", summary) + + task_status = json.loads(redis_conn.hget("task:status", task_key)) + assert task_status["status"] == "pending_retry", ( + f"status should stay pending_retry (not affected by analytics), " + f"got {task_status['status']} for collection_status={status}" + ) + + +# ============================================================= +# 5. Block 测试报告生成测试 +# ============================================================= + + +class TestBlockReport: + """测试 block 测试报告 CSV 生成""" + + def test_report_includes_block_tasks_only(self, tmp_path): + """报告只包含 _block 后缀的任务""" + report_path = os.path.join(str(tmp_path), "block_test_report.csv") + + from dispatcher_main import generate_block_test_report + + redis_conn = FakeRedis() + dispatcher = make_dispatcher(redis_conn=redis_conn) + + for app_name, package_name, suffix in ( + ("BlockedA", "com.block.a", "_block"), + ("BlockedB", "com.block.b", "_block"), + ("NormalApp", "com.normal", ""), + ): + task_key = f"{app_name}_{package_name}{suffix}" + redis_conn.hset("task:details", task_key, json.dumps({ + "app_name": app_name, + "package_name": package_name, + "is_block_task": suffix == "_block", + })) + redis_conn.hset("task:status", task_key, json.dumps({ + "status": "success" if suffix == "_block" else "completed", + "worker_id": "worker_x", + "start_time": 1000.0, + "end_time": 1100.0, + "retry_count": 0, + })) + redis_conn.hset("workers:info", "worker_x", json.dumps({ + "ip_address": "192.168.1.1", + })) + + generate_block_test_report(dispatcher, report_dir=str(tmp_path)) + + assert os.path.exists(report_path) + with open(report_path, "r", encoding="utf-8") as f: + reader = csv.reader(f) + rows = list(reader) + + assert len(rows) >= 3 + headers = rows[0] + assert "任务键" in headers + assert "应用名称" in headers + assert "状态" in headers + assert "失败原因" in headers + + task_keys_in_report = [row[0] for row in rows[1:] if row and row[0]] + assert "BlockedA_com.block.a_block" in task_keys_in_report + assert "BlockedB_com.block.b_block" in task_keys_in_report + assert "NormalApp_com.normal" not in task_keys_in_report + + def test_report_counts_success_and_failure(self, tmp_path): + """报告正确统计成功/失败/重试数量""" + report_path = os.path.join(str(tmp_path), "block_test_report.csv") + + from dispatcher_main import generate_block_test_report + + redis_conn = FakeRedis() + dispatcher = make_dispatcher(redis_conn=redis_conn) + + block_data = [ + ("SuccessApp", "com.success", "success"), + ("FailedApp", "com.failed", "failed"), + ("RetryingApp", "com.retry", "pending_retry"), + ("PendingApp", "com.pending", "pending"), + ] + for app_name, pkg, status in block_data: + task_key = f"{app_name}_{pkg}_block" + redis_conn.hset("task:details", task_key, json.dumps({ + "app_name": app_name, + "package_name": pkg, + "is_block_task": True, + })) + fail_message = "App crashed" if status == "failed" else "" + redis_conn.hset("task:status", task_key, json.dumps({ + "status": status, + "worker_id": "worker_x", + "last_fail_message": fail_message, + "retry_count": 0, + })) + + generate_block_test_report(dispatcher, report_dir=str(tmp_path)) + + with open(report_path, "r", encoding="utf-8") as f: + content = f.read() + + assert "成功,1" in content + assert "失败,1" in content + assert "重试中/待分发,2" in content + + +# ============================================================= +# 6. 集成测试: 模拟 Worker 完整接收 block 任务流程 +# ============================================================= + + +class TestWorkerBlockTaskIntegration: + """模拟 Worker 完整的 block 任务接收-执行-上报 链路""" + + def test_worker_receives_block_task_params(self): + """Worker 通过 init 拿到的 block 任务包含所有预期字段""" + redis_conn = FakeRedis() + monitor = StubMonitor() + dispatcher = make_dispatcher(redis_conn=redis_conn, monitor=monitor) + + entries = [qualified_entry("FullBlock", "com.full.block", country_code="CN", device_type="emulator")] + dispatcher.load_block_tasks(entries) + task_key = "FullBlock_com.full.block_block" + + worker_id = "192.168.1.100_EE:FF:00:11:22:33" + redis_conn.hset("workers:info", worker_id, json.dumps({ + "worker_id": worker_id, + "ip_address": "192.168.1.100", + "mac_address": "EE:FF:00:11:22:33", + "hostname": "block-test", + "platform": "Windows", + "device_type": "emulator", + "register_time": 1000000000.0, + "last_update": 1000000000.0, + "status": "idle", + "dispatch_enabled": True, + })) + redis_conn.sadd("workers:idle", worker_id) + + task = dispatcher.worker_init( + worker_id=worker_id, + ip_address="192.168.1.100", + mac_address="EE:FF:00:11:22:33", + hostname="block-test", + platform="Windows", + device_type="emulator", + ) + + assert task is not None + from config import ANALYTICS_TRAFFIC_ROOT_BLOCK + assert task["traffic_root"] == ANALYTICS_TRAFFIC_ROOT_BLOCK + assert task["task_key"] == task_key + assert task["app_name"] == "FullBlock" + assert task["package_name"] == "com.full.block" + assert task["country_code"] == "CN" + assert task["device_type"] == "emulator" + + def test_full_lifecycle_block_task_to_failure(self): + """Block 任务完整生命周期: init → report failure → check not re-queued""" + redis_conn = FakeRedis() + monitor = StubMonitor() + notifier = StubNotifier() + dispatcher = make_dispatcher( + redis_conn=redis_conn, monitor=monitor, notifier=notifier + ) + + entries = [qualified_entry("LifecycleApp", "com.lifecycle")] + dispatcher.load_block_tasks(entries) + task_key = "LifecycleApp_com.lifecycle_block" + + worker_id = "192.168.1.200_AA:11:BB:22:CC:33" + redis_conn.hset("workers:info", worker_id, json.dumps({ + "worker_id": worker_id, + "ip_address": "192.168.1.200", + "mac_address": "AA:11:BB:22:CC:33", + "hostname": "lifecycle-host", + "platform": "Windows", + "device_type": "emulator", + "register_time": 1000000000.0, + "last_update": 1000000000.0, + "status": "idle", + "dispatch_enabled": True, + })) + redis_conn.sadd("workers:idle", worker_id) + + task = dispatcher.worker_init( + worker_id=worker_id, + ip_address="192.168.1.200", + mac_address="AA:11:BB:22:CC:33", + hostname="lifecycle-host", + platform="Windows", + device_type="emulator", + ) + assert task is not None + + next_task = dispatcher.worker_report( + worker_id=worker_id, + previous_task_key=task_key, + report_data={ + "status": "failed", + "error": { + "category": "APP_ERROR", + "code": 1, + "reason": "应用闪退", + }, + }, + ) + + status = json.loads(redis_conn.hget("task:status", task_key)) + assert status["status"] == "failed" + assert task_key in redis_conn.smembers("task:failed") + + assert next_task is None, ( + f"after block task fails, worker should get no next task (queue empty), " + f"got {next_task}" + ) + + def test_block_task_download_failure_still_in_pending_retry(self): + """Block 任务下载失败后状态为 pending_retry，等待 APK""" + redis_conn = FakeRedis() + monitor = StubMonitor() + notifier = StubNotifier() + dispatcher = make_dispatcher( + redis_conn=redis_conn, monitor=monitor, notifier=notifier + ) + + entries = [qualified_entry("DownloadFail", "com.dl.fail")] + dispatcher.load_block_tasks(entries) + task_key = "DownloadFail_com.dl.fail_block" + + worker_id = "192.168.1.201_WORKER" + redis_conn.hset("workers:info", worker_id, json.dumps({ + "worker_id": worker_id, + "ip_address": "192.168.1.201", + "mac_address": "EE:FF:GG:HH:II:JJ", + "hostname": "dl-fail-host", + "platform": "Windows", + "device_type": "emulator", + "register_time": 1000000000.0, + "last_update": 1000000000.0, + "status": "idle", + "dispatch_enabled": True, + })) + redis_conn.sadd("workers:idle", worker_id) + + task = dispatcher.worker_init( + worker_id=worker_id, + ip_address="192.168.1.201", + mac_address="EE:FF:GG:HH:II:JJ", + hostname="dl-fail-host", + platform="Windows", + device_type="emulator", + ) + assert task is not None + + next_task = dispatcher.worker_report( + worker_id=worker_id, + previous_task_key=task_key, + report_data={ + "status": "failed", + "error": { + "category": "DOWNLOAD_ERROR", + "code": 1, + "reason": "Google Play 锁区", + }, + }, + ) + + status = json.loads(redis_conn.hget("task:status", task_key)) + assert status["status"] == "pending_retry", ( + f"download error should leave task in pending_retry, got {status['status']}" + ) + assert task_key not in redis_conn.smembers("task:failed") diff --git a/tests/test_config_loader.py b/tests/test_config_loader.py new file mode 100644 index 0000000..34e1ec5 --- /dev/null +++ b/tests/test_config_loader.py @@ -0,0 +1,71 @@ +import os +import sys + + +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, PROJECT_ROOT) + +from config import load_dispatcher_config, load_worker_inventory + + +def test_load_dispatcher_config_accepts_yaml_path(tmp_path): + config_path = tmp_path / "dispatcher.yaml" + config_path.write_text( + """ +CONFIG_ENV: test +INSTANCE_NAME: test-dispatcher +REDIS_PORT: 6380 +FEATURE_FLAGS: + yaml_config: true +""".lstrip(), + encoding="utf-8", + ) + + config = load_dispatcher_config(config_path=str(config_path)) + + assert config["CONFIG_ENV"] == "test" + assert config["INSTANCE_NAME"] == "test-dispatcher" + assert config["REDIS_PORT"] == 6380 + assert config["FEATURE_FLAGS"]["yaml_config"] is True + + +def test_load_worker_inventory_accepts_yaml_config_object(tmp_path): + inventory_path = tmp_path / "inventory.yaml" + inventory_path.write_text( + """ +WORKER_INVENTORY: + - worker_id: 192.168.1.20 + ssh_target: 192.168.1.20 + repo_dir: D:/autool + python_exe: python + tags: + - test +""".lstrip(), + encoding="utf-8", + ) + + inventory = load_worker_inventory(path=str(inventory_path), env_name="test") + + assert len(inventory) == 1 + assert inventory[0]["worker_id"] == "192.168.1.20" + assert inventory[0]["ssh_host"] == "192.168.1.20" + assert inventory[0]["repo_dir"] == "D:/autool" + + +def test_load_worker_inventory_accepts_standalone_yaml_list(tmp_path): + inventory_path = tmp_path / "workers.yml" + inventory_path.write_text( + """ +- worker_id: 192.168.1.21 + ssh_target: 192.168.1.21 + repo_dir: D:/autool + python_exe: python + tags: [] +""".lstrip(), + encoding="utf-8", + ) + + inventory = load_worker_inventory(path=str(inventory_path), env_name="prod") + + assert len(inventory) == 1 + assert inventory[0]["worker_id"] == "192.168.1.21" diff --git a/tests/test_dispatcher_regressions.py b/tests/test_dispatcher_regressions.py new file mode 100644 index 0000000..08bb70c --- /dev/null +++ b/tests/test_dispatcher_regressions.py @@ -0,0 +1,248 @@ +import json + +from analytics import AnalyticsRepository +from redis_task_distribute import RedisTaskDispatcher + + +class FakeRedis: + def __init__(self): + self.lists = {} + self.hashes = {} + self.sets = {} + + def delete(self, key): + self.lists.pop(key, None) + self.hashes.pop(key, None) + self.sets.pop(key, None) + + def lpush(self, key, value): + self.lists.setdefault(key, []).insert(0, value) + + def rpush(self, key, value): + self.lists.setdefault(key, []).append(value) + + def rpop(self, key): + values = self.lists.setdefault(key, []) + return values.pop() if values else None + + def lrange(self, key, start, end): + values = self.lists.get(key, []) + stop = None if end == -1 else end + 1 + return values[start:stop] + + def llen(self, key): + return len(self.lists.get(key, [])) + + def lpos(self, key, value): + try: + return self.lists.get(key, []).index(value) + except ValueError: + return None + + def lrem(self, key, count, value): + values = self.lists.get(key, []) + original_len = len(values) + self.lists[key] = [item for item in values if item != value] + return original_len - len(self.lists[key]) + + def hset(self, key, field, value): + self.hashes.setdefault(key, {})[field] = value + + def hget(self, key, field): + return self.hashes.get(key, {}).get(field) + + def hgetall(self, key): + return dict(self.hashes.get(key, {})) + + def hvals(self, key): + return list(self.hashes.get(key, {}).values()) + + def hlen(self, key): + return len(self.hashes.get(key, {})) + + def hexists(self, key, field): + return field in self.hashes.get(key, {}) + + def hdel(self, key, field): + self.hashes.get(key, {}).pop(field, None) + + def sadd(self, key, value): + self.sets.setdefault(key, set()).add(value) + + def srem(self, key, value): + self.sets.setdefault(key, set()).discard(value) + + def scard(self, key): + return len(self.sets.get(key, set())) + + def smembers(self, key): + return set(self.sets.get(key, set())) + + +class StubAnalytics: + def __init__(self, pending=None, model=None): + self.pending = pending or [] + self.model = model or [] + + def list_pending_collection_tasks(self): + return list(self.pending) + + def list_model_eligible_apps(self): + return list(self.model) + + +def make_dispatcher(redis_conn=None, analytics=None): + dispatcher = RedisTaskDispatcher.__new__(RedisTaskDispatcher) + dispatcher.redis = redis_conn or FakeRedis() + dispatcher.analytics = analytics or StubAnalytics() + dispatcher.worker_inventory = {} + dispatcher.managed_worker_ids = set() + dispatcher.worker_online_timeout = 300 + dispatcher.task_routing_rules = {"package_name": {}, "task_key": {}} + return dispatcher + + +def pending_entry(app_name, package_name, task_queue="default"): + return { + "app_name": app_name, + "package_name": package_name, + "task_queue": task_queue, + "task_payload": { + "app_name": app_name, + "package_name": package_name, + "country_code": "US", + }, + } + + +def test_statistics_and_dashboard_read_three_priority_queues(): + redis_conn = FakeRedis() + dispatcher = make_dispatcher(redis_conn=redis_conn) + for app_name, package_name, queue_name in ( + ("High App", "com.example.high", "task:queue:high"), + ("Default App", "com.example.default", "task:queue:default"), + ("Low App", "com.example.low", "task:queue:low"), + ): + task_key = f"{app_name}_{package_name}" + redis_conn.hset("task:details", task_key, json.dumps({ + "app_name": app_name, + "package_name": package_name, + "task_queue": queue_name.rsplit(":", 1)[-1], + })) + redis_conn.lpush(queue_name, task_key) + + stats = dispatcher.get_statistics() + dashboard_tasks = dispatcher.get_dashboard_tasks() + + assert stats["pending"] == 3 + assert [item["task_key"] for item in dashboard_tasks["pending"]] == [ + "High App_com.example.high", + "Default App_com.example.default", + "Low App_com.example.low", + ] + + +def test_refresh_requeues_existing_completed_high_priority_task(): + entry = pending_entry("Hot App", "com.example.hot", task_queue="high") + redis_conn = FakeRedis() + dispatcher = make_dispatcher(redis_conn=redis_conn, analytics=StubAnalytics(pending=[entry])) + task_key = "Hot App_com.example.hot" + redis_conn.hset("task:details", task_key, json.dumps({"app_name": "Old", "package_name": "com.example.hot"})) + redis_conn.hset("task:status", task_key, json.dumps({"status": "completed", "retry_count": 2})) + redis_conn.sadd("task:completed", task_key) + + assert dispatcher.refresh_tasks_from_app_summary() == 1 + + assert redis_conn.lrange("task:queue:high", 0, -1) == [task_key] + assert redis_conn.scard("task:completed") == 0 + assert json.loads(redis_conn.hget("task:status", task_key))["status"] == "pending" + task_details = json.loads(redis_conn.hget("task:details", task_key)) + assert task_details["app_name"] == "Hot App" + assert task_details["task_queue"] == "high" + + +def test_refresh_does_not_duplicate_running_task(): + entry = pending_entry("Running App", "com.example.running", task_queue="high") + redis_conn = FakeRedis() + dispatcher = make_dispatcher(redis_conn=redis_conn, analytics=StubAnalytics(pending=[entry])) + task_key = "Running App_com.example.running" + redis_conn.hset("task:status", task_key, json.dumps({"status": "running", "worker_id": "worker-1"})) + redis_conn.hset("worker:tasks", "worker-1", task_key) + + assert dispatcher.refresh_tasks_from_app_summary() == 0 + assert redis_conn.lrange("task:queue:high", 0, -1) == [] + + +def test_direct_mode_payload_does_not_default_to_local_source(monkeypatch): + import redis_task_distribute as dispatcher_module + + monkeypatch.setattr(dispatcher_module, "APK_DOWNLOAD_MODE", "direct") + dispatcher = make_dispatcher() + dispatcher._get_apk_registry = lambda: (_ for _ in ()).throw(AssertionError("registry should not be used in direct mode")) + + payload = dispatcher._task_to_payload( + "Direct App_com.example.direct", + { + "app_name": "Direct App", + "package_name": "com.example.direct", + "country_code": "US", + }, + ) + + assert payload["available_sources"] == ["google_play", "apkpure"] + assert payload["local_apk_dir"] == "" + + +def test_minio_mode_payload_can_use_local_source(monkeypatch): + import redis_task_distribute as dispatcher_module + + monkeypatch.setattr(dispatcher_module, "APK_DOWNLOAD_MODE", "minio") + dispatcher = make_dispatcher() + + class _Registry: + @staticmethod + def is_fresh_enough(package_name, last_updated): + return False + + dispatcher._get_apk_registry = lambda: _Registry() + + payload = dispatcher._task_to_payload( + "Cached App_com.example.cached", + { + "app_name": "Cached App", + "package_name": "com.example.cached", + "country_code": "US", + }, + ) + + assert payload["available_sources"] == ["google_play", "local", "apkpure"] + + +def test_direct_mode_download_error_does_not_mark_pending_apk(monkeypatch): + import redis_task_distribute as dispatcher_module + + monkeypatch.setattr(dispatcher_module, "APK_DOWNLOAD_MODE", "direct") + dispatcher = make_dispatcher() + dispatcher._get_apk_registry = lambda: (_ for _ in ()).throw(AssertionError("registry should not be used in direct mode")) + + dispatcher._mark_download_error_awaiting_apk("com.example.direct") + + +def test_analytics_jobs_accept_scheduled_after(tmp_path): + repo = AnalyticsRepository(db_path=str(tmp_path / "analytics.sqlite3")) + scheduled_after = 12345.0 + + job = repo.create_job( + job_type="incremental", + status="queued", + package_name="com.example.active", + scheduled_after=scheduled_after, + ) + + with repo._connect() as connection: + row = connection.execute( + "SELECT scheduled_after FROM analytics_job WHERE id = ?", + (job["id"],), + ).fetchone() + + assert row["scheduled_after"] == scheduled_after diff --git a/tests/test_download_record.py b/tests/test_download_record.py new file mode 100644 index 0000000..d061068 --- /dev/null +++ b/tests/test_download_record.py @@ -0,0 +1,966 @@ +""" +app_download_record 表 + ApkRegistry 新增方法 + dispatcher 集成的完整测试。 + +用法: + uv run python -m pytest tests/test_download_record.py -v + uv run python -m pytest tests/test_download_record.py -v -k "TestMinio" # 仅 MinIO 测试 +""" + +import json +import os +import shutil +import sqlite3 +import sys +import tempfile +import threading +import time +from datetime import datetime, timedelta +from typing import Any, Dict, List +from zoneinfo import ZoneInfo + +import pytest + +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if PROJECT_ROOT not in sys.path: + sys.path.insert(0, PROJECT_ROOT) + +from apk_cloud.registry import ApkRegistry, _is_apk_stale, _parse_date, _last_updated_ts + +try: + from apk_cloud.storage import MinioStorage + _HAS_MINIO = True +except Exception: + _HAS_MINIO = False + +try: + from redis_task_distribute import RedisTaskDispatcher + _HAS_DISPATCHER = True +except Exception: + _HAS_DISPATCHER = False + + +# ── helpers ────────────────────────────────────────────────────────── + +@pytest.fixture +def tmp_db_path(): + fd, path = tempfile.mkstemp(suffix=".sqlite3") + os.close(fd) + yield path + try: + os.unlink(path) + except OSError: + pass + + +@pytest.fixture +def tmp_storage_dir(): + d = tempfile.mkdtemp(prefix="apk_test_") + yield d + shutil.rmtree(d, ignore_errors=True) + + +@pytest.fixture +def registry(tmp_db_path, tmp_storage_dir): + r = ApkRegistry(db_path=tmp_db_path, storage_dir=tmp_storage_dir) + yield r + try: + os.unlink(tmp_db_path) + except OSError: + pass + + +@pytest.fixture(autouse=True) +def _cleanup_download_records(registry): + """每个测试后清理 download_record，确保测试隔离。""" + yield + with registry._connect() as conn: + conn.execute("DELETE FROM app_download_record") + + +def _make_result(status: str = "ok", sources: Dict[str, Dict[str, Any]] = None, + files: List[Dict[str, Any]] = None, download_date: str = "", + version_name: str = "") -> Dict[str, Any]: + files = files if files is not None else [{"filename": "base.apk", "size": 5000000}] + return { + "status": status, + "download_date": download_date or datetime.now().strftime("%Y-%m-%d %H:%M:%S"), + "version_name": version_name, + "files": files, + "source_details": sources or {}, + } + + +# ══════════════════════════════════════════════════════════════════════ +# 1. Schema & 表结构 +# ══════════════════════════════════════════════════════════════════════ + +class TestSchema: + def test_table_exists(self, registry): + with registry._connect() as conn: + row = conn.execute( + "SELECT name FROM sqlite_master WHERE type='table' AND name='app_download_record'" + ).fetchone() + assert row is not None, "app_download_record table should exist" + + def test_table_columns(self, registry): + expected = { + "id", "package_name", "country_code", "app_name", "last_updated", + "downloads", "worker_play_status", "worker_play_error", "worker_play_date", + "us_play_status", "us_play_error", "us_play_date", + "us_aurora_status", "us_aurora_error", "us_aurora_date", + "apkpure_status", "apkpure_error", "apkpure_date", + "overall_status", "overall_error", "apk_version_name", + "local_apk_dir", "local_apk_file_count", "local_apk_is_stale", + "created_at", "updated_at", + } + with registry._connect() as conn: + cols = {row[1] for row in conn.execute("PRAGMA table_info(app_download_record)")} + assert expected == cols, f"Missing: {expected - cols}, Extra: {cols - expected}" + + def test_unique_constraint(self, registry): + now = time.time() + registry.ensure_download_record("com.test", "US", "App", "2026-05-20", 100) + with pytest.raises(sqlite3.IntegrityError): + with registry._connect() as conn: + conn.execute( + "INSERT INTO app_download_record (package_name, country_code, " + "last_updated, created_at, updated_at) VALUES (?, ?, ?, ?, ?)", + ("com.test", "US", "2026-05-20", now, now), + ) + + def test_default_values(self, registry): + registry.ensure_download_record("com.test", "US") + record = registry.get_download_record("com.test", "US") + assert record["worker_play_status"] == "pending" + assert record["us_play_status"] == "pending" + assert record["us_aurora_status"] == "pending" + assert record["overall_status"] == "pending" + assert record["local_apk_is_stale"] == 1 + assert record["downloads"] == 0 + assert record["local_apk_file_count"] == 0 + + def test_indexes_exist(self, registry): + with registry._connect() as conn: + rows = conn.execute( + "SELECT name FROM sqlite_master WHERE type='index' " + "AND tbl_name='app_download_record'" + ).fetchall() + indexes = {row["name"] for row in rows} + expected = { + "idx_download_record_pkg", + "idx_download_record_country", + "idx_download_record_overall", + "idx_download_record_stale", + } + assert expected.issubset(indexes), f"Missing indexes: {expected - indexes}" + + def test_create_and_drop_does_not_break(self, registry): + """验证 DROP + CREATE 不会破坏其他表的数据。""" + registry.mark_available("com.test.keep", "2026-05-19", "/tmp/dpi/mumu_apk", "1.0", source="test") + with registry._connect() as conn: + conn.execute("DROP TABLE IF EXISTS app_download_record") + with registry._connect() as conn: + assert conn.execute( + "SELECT 1 FROM apk_registry WHERE package_name='com.test.keep'" + ).fetchone() is not None + registry._ensure_tables() + with registry._connect() as conn: + assert conn.execute( + "SELECT 1 FROM apk_registry WHERE package_name='com.test.keep'" + ).fetchone() is not None + registry.mark_unavailable("com.test.keep") + + +# ══════════════════════════════════════════════════════════════════════ +# 2. ensure_download_record +# ══════════════════════════════════════════════════════════════════════ + +class TestEnsureDownloadRecord: + def test_create_new_record(self, registry): + registry.ensure_download_record( + "com.test.create", "US", + app_name="Test App", last_updated="2026-05-20", downloads=5000000, + ) + record = registry.get_download_record("com.test.create", "US") + assert record is not None + assert record["app_name"] == "Test App" + assert record["last_updated"] == "2026-05-20" + assert record["downloads"] == 5000000 + assert record["country_code"] == "US" + assert record["created_at"] > 0 + assert record["updated_at"] > 0 + + def test_update_existing_record(self, registry): + registry.ensure_download_record("com.test.update", "CN", "Old", "2026-05-01", 100) + registry.ensure_download_record("com.test.update", "CN", "New", "2026-05-25", 200) + record = registry.get_download_record("com.test.update", "CN") + assert record["app_name"] == "New" + assert record["last_updated"] == "2026-05-25" + assert record["downloads"] == 200 + assert len(registry.list_download_records()) == 1 + + def test_same_package_different_country(self, registry): + registry.ensure_download_record("com.test", "US", "App US", "2026-05-01", 100) + registry.ensure_download_record("com.test", "CN", "App CN", "2026-05-15", 200) + assert registry.get_download_record("com.test", "US")["app_name"] == "App US" + assert registry.get_download_record("com.test", "CN")["app_name"] == "App CN" + assert len(registry.list_download_records()) == 2 + + def test_empty_package_ignored(self, registry): + registry.ensure_download_record("", "US") + registry.ensure_download_record(" ", "US") + assert len(registry.list_download_records()) == 0 + + def test_default_values_on_minimal_args(self, registry): + registry.ensure_download_record("com.test.minimal", "US") + record = registry.get_download_record("com.test.minimal", "US") + assert record["app_name"] == "" + assert record["last_updated"] == "" + assert record["downloads"] == 0 + + def test_country_code_normalization(self, registry): + """ensure_download_record 会 strip country_code 前后空格。""" + registry.ensure_download_record("com.test", " US ") + record = registry.get_download_record("com.test", "US") + assert record is not None, "should match after strip" + + +# ══════════════════════════════════════════════════════════════════════ +# 3. update_source_result + _recompute_overall +# ══════════════════════════════════════════════════════════════════════ + +class TestUpdateSourceResult: + SRC_WORKER = "worker_play" + SRC_US_PLAY = "us_play" + SRC_US_AURORA = "us_aurora" + + def setup_method(self): + self._today = datetime.now().strftime("%Y-%m-%d %H:%M:%S") + + def test_single_source_success(self, registry): + registry.ensure_download_record("com.test", "US") + registry.update_source_result("com.test", "US", self.SRC_US_PLAY, + "success", download_date=self._today) + r = registry.get_download_record("com.test", "US") + assert r["us_play_status"] == "success" + assert r["us_play_date"] == self._today + assert r["overall_status"] == "success" + + def test_single_source_failure(self, registry): + registry.ensure_download_record("com.test", "US") + registry.update_source_result("com.test", "US", self.SRC_US_PLAY, + "failed", error="region restricted") + r = registry.get_download_record("com.test", "US") + assert r["us_play_status"] == "failed" + assert r["us_play_error"] == "region restricted" + assert r["us_play_date"] is None + assert r["overall_status"] == "all_failed" + + def test_error_cleared_on_success_transition(self, registry): + registry.ensure_download_record("com.test", "US") + registry.update_source_result("com.test", "US", self.SRC_US_PLAY, + "failed", error="timeout") + registry.update_source_result("com.test", "US", self.SRC_US_PLAY, + "success", download_date=self._today) + r = registry.get_download_record("com.test", "US") + assert r["us_play_status"] == "success" + assert r["us_play_error"] is None + assert r["overall_status"] == "success" + + def test_date_cleared_on_failure_transition(self, registry): + registry.ensure_download_record("com.test", "US") + registry.update_source_result("com.test", "US", self.SRC_US_PLAY, + "success", download_date=self._today) + registry.update_source_result("com.test", "US", self.SRC_US_PLAY, + "failed", error="timeout") + r = registry.get_download_record("com.test", "US") + assert r["us_play_status"] == "failed" + assert r["us_play_date"] is None + + def test_not_attempted_preserves_pending_overall(self, registry): + registry.ensure_download_record("com.test", "US") + registry.update_source_result("com.test", "US", self.SRC_US_PLAY, "not_attempted") + r = registry.get_download_record("com.test", "US") + assert r["overall_status"] == "pending" + + def test_one_success_overrides_failures(self, registry): + registry.ensure_download_record("com.test", "US") + registry.update_source_result("com.test", "US", self.SRC_US_PLAY, "failed", error="e1") + registry.update_source_result("com.test", "US", self.SRC_US_AURORA, "failed", error="e2") + registry.update_source_result("com.test", "US", self.SRC_WORKER, "success", + download_date=self._today) + r = registry.get_download_record("com.test", "US") + assert r["overall_status"] == "success" + assert r["overall_error"] == "" + + def test_all_sources_failed_overall(self, registry): + registry.ensure_download_record("com.test", "US") + registry.update_source_result("com.test", "US", self.SRC_WORKER, "failed", error="w_err") + registry.update_source_result("com.test", "US", self.SRC_US_PLAY, "failed", error="p_err") + registry.update_source_result("com.test", "US", self.SRC_US_AURORA, "failed", error="a_err") + r = registry.get_download_record("com.test", "US") + assert r["overall_status"] == "all_failed" + assert "Worker Play" in (r["overall_error"] or "") + assert "w_err" in (r["overall_error"] or "") + assert "US Play" in (r["overall_error"] or "") + assert "p_err" in (r["overall_error"] or "") + assert "US Aurora" in (r["overall_error"] or "") + assert "a_err" in (r["overall_error"] or "") + + def test_unknown_source_ignored(self, registry): + registry.ensure_download_record("com.test", "US") + registry.update_source_result("com.test", "US", "unknown_source", "success") + r = registry.get_download_record("com.test", "US") + assert r["overall_status"] == "pending" + + def test_mixed_pending_and_failed(self, registry): + registry.ensure_download_record("com.test", "US") + registry.update_source_result("com.test", "US", self.SRC_US_PLAY, "failed", error="err") + r = registry.get_download_record("com.test", "US") + assert r["overall_status"] == "all_failed" + + def test_all_pending_remains_pending(self, registry): + registry.ensure_download_record("com.test", "US") + registry.update_source_result("com.test", "US", self.SRC_WORKER, "pending") + registry.update_source_result("com.test", "US", self.SRC_US_PLAY, "pending") + r = registry.get_download_record("com.test", "US") + assert r["overall_status"] == "pending" + + +# ══════════════════════════════════════════════════════════════════════ +# 4. record_download_from_result (MinIO JSON 解析) +# ══════════════════════════════════════════════════════════════════════ + +class TestRecordFromMinioResult: + def test_ok_result_with_full_source_details(self, registry): + registry.ensure_download_record("com.test", "US") + result = _make_result(status="ok", sources={ + "us_play": {"status": "success", "error": None, "date": "2026-05-21 10:00:00"}, + "us_aurora": {"status": "not_attempted", "error": None, "date": None}, + "worker_play": {"status": "not_attempted", "error": None, "date": None}, + }, download_date="2026-05-21 10:00:00", version_name="2.0") + registry.record_download_from_result("com.test", "US", result) + r = registry.get_download_record("com.test", "US") + assert r["us_play_status"] == "success" + assert r["overall_status"] == "success" + assert r["apk_version_name"] == "2.0" + assert r["local_apk_file_count"] == 1 + + def test_failed_result_with_source_details(self, registry): + registry.ensure_download_record("com.test", "US") + result = _make_result(status="failed", sources={ + "us_play": {"status": "failed", "error": "region restricted", "date": None}, + "us_aurora": {"status": "failed", "error": "app not found", "date": None}, + }) + registry.record_download_from_result("com.test", "US", result) + r = registry.get_download_record("com.test", "US") + assert r["us_play_status"] == "failed" + assert r["us_play_error"] == "region restricted" + assert r["us_aurora_status"] == "failed" + assert r["us_aurora_error"] == "app not found" + assert r["overall_status"] == "all_failed" + + def test_ok_result_no_local_update_when_missing_files(self, registry): + registry.ensure_download_record("com.test", "US") + result = _make_result(status="ok", sources={ + "us_play": {"status": "success", "error": None, "date": "2026-05-21"}, + }, files=[], version_name="1.0") + registry.record_download_from_result("com.test", "US", result) + r = registry.get_download_record("com.test", "US") + assert r["us_play_status"] == "success" + assert r["local_apk_file_count"] == 0 + assert r["apk_version_name"] == "1.0" + + def test_result_with_no_source_details(self, registry): + """旧格式兼容：无 source_details 时不报错。""" + registry.ensure_download_record("com.test", "US") + result = {"status": "ok", "download_date": "2026-05-21"} # No source_details + registry.record_download_from_result("com.test", "US", result) + r = registry.get_download_record("com.test", "US") + assert r["overall_status"] == "pending" + + def test_empty_package_name_ignored(self, registry): + registry.record_download_from_result("", "US", _make_result()) + assert len(registry.list_download_records()) == 0 + + def test_source_details_is_list_handled_gracefully(self, registry): + """不合法类型的 source_details 应被安全跳过。""" + registry.ensure_download_record("com.test", "US") + result = {"status": "ok", "source_details": [1, 2, 3]} + registry.record_download_from_result("com.test", "US", result) + r = registry.get_download_record("com.test", "US") + assert r["overall_status"] == "pending" + + +# ══════════════════════════════════════════════════════════════════════ +# 5. Stale Detection & Cleanup +# ══════════════════════════════════════════════════════════════════════ + +class TestStaleDetection: + def test_stale_when_download_older_than_last_updated(self, registry): + registry.ensure_download_record("com.test", "US", "App", "2026-05-20") + registry.mark_available("com.test", "2026-05-01", "/tmp/dpi/mumu_apk", "1.0") + records = registry.get_stale_records() + assert any(r["package_name"] == "com.test" for r in records) + + def test_not_stale_when_download_newer_or_equal(self, registry): + registry.ensure_download_record("com.test", "US", "App", "2026-05-01") + registry.mark_available("com.test", "2026-05-01", "/tmp/dpi/mumu_apk", "1.0") + records = registry.get_stale_records() + assert not any(r["package_name"] == "com.test" for r in records) + + def test_not_stale_when_download_newer(self, registry): + registry.ensure_download_record("com.test", "US", "App", "2026-05-01") + registry.mark_available("com.test", "2026-05-15", "/tmp/dpi/mumu_apk", "2.0") + records = registry.get_stale_records() + assert not any(r["package_name"] == "com.test" for r in records) + + def test_stale_when_no_apk_registry_entry(self, registry): + registry.ensure_download_record("com.test", "US", "App", "2026-05-01") + records = registry.get_stale_records() + assert any(r["package_name"] == "com.test" for r in records) + + def test_stale_records_sorted_by_downloads(self, registry): + registry.ensure_download_record("com.low", "US", "Low", "2026-05-20", downloads=100) + registry.ensure_download_record("com.high", "US", "High", "2026-05-20", downloads=5000000) + registry.ensure_download_record("com.mid", "US", "Mid", "2026-05-20", downloads=10000) + records = registry.get_stale_records() + downloads_order = [r["downloads"] for r in records if r["package_name"] in + ("com.low", "com.mid", "com.high")] + assert downloads_order == sorted(downloads_order, reverse=True) + + def test_mark_available_refreshes_stale_flag(self, registry): + registry.ensure_download_record("com.test", "US", "App", "2026-05-20") + registry.mark_available("com.test", "2026-05-01", "/tmp", "1.0") + assert registry.get_download_record("com.test", "US")["local_apk_is_stale"] == 1 + + registry.mark_available("com.test", "2026-05-25", "/tmp", "2.0") + assert registry.get_download_record("com.test", "US")["local_apk_is_stale"] == 0 + + def test_mark_unavailable_marks_all_records_stale(self, registry): + registry.ensure_download_record("com.test", "US", "App", "2026-05-01") + registry.ensure_download_record("com.test", "CN", "App", "2026-05-01") + registry.mark_available("com.test", "2026-05-20", "/tmp", "2.0") + # First verify fresh + assert registry.get_download_record("com.test", "US")["local_apk_is_stale"] == 0 + registry.mark_unavailable("com.test") + assert registry.get_download_record("com.test", "US")["local_apk_is_stale"] == 1 + assert registry.get_download_record("com.test", "CN")["local_apk_is_stale"] == 1 + assert registry.get_download_record("com.test", "US")["local_apk_file_count"] == 0 + + def test__is_apk_stale_helper(self): + assert _is_apk_stale("2026-05-01", "2026-05-20") is True + assert _is_apk_stale("2026-05-20", "2026-05-01") is False + assert _is_apk_stale("2026-05-20", "2026-05-20") is False + assert _is_apk_stale("", "2026-05-01") is True + assert _is_apk_stale("2026-05-20", "") is False + + def test__parse_date_variants(self): + dt = _parse_date("2026-05-20 10:30:00") + assert dt is not None and dt.year == 2026 + dt = _parse_date("2026-05-20") + assert dt is not None + dt = _parse_date("2026-05") + assert dt is not None + dt = _parse_date("") + assert dt is None + dt = _parse_date(None) + assert dt is None + + +class TestCleanupStaleLocal: + def test_cleanup_deletes_registry_entry(self, registry, tmp_storage_dir): + pkg_dir = os.path.join(tmp_storage_dir, "com.test.cleanup") + os.makedirs(pkg_dir, exist_ok=True) + with open(os.path.join(pkg_dir, "test.apk"), "w") as f: + f.write("fake apk") + + registry.ensure_download_record("com.test.cleanup", "US") + registry.mark_available("com.test.cleanup", "2026-05-01", pkg_dir, "1.0") + assert registry.get("com.test.cleanup") is not None + assert os.path.isdir(pkg_dir) + + cleaned = registry.cleanup_stale_local("com.test.cleanup") + assert cleaned is True + assert registry.get("com.test.cleanup") is None + assert not os.path.isdir(pkg_dir) + + def test_cleanup_nonexistent_directory(self, registry): + registry.ensure_download_record("com.test.noexist", "US") + cleaned = registry.cleanup_stale_local("com.test.noexist") + assert cleaned is True + assert registry.get("com.test.noexist") is None + + +# ══════════════════════════════════════════════════════════════════════ +# 6. list / query methods +# ══════════════════════════════════════════════════════════════════════ + +class TestListAndQuery: + def test_list_all_records(self, registry): + registry.ensure_download_record("com.a", "US", downloads=100) + registry.ensure_download_record("com.b", "CN", downloads=200) + registry.ensure_download_record("com.c", "US", downloads=300) + records = registry.list_download_records() + assert len(records) == 3 + + def test_filter_by_country_code(self, registry): + registry.ensure_download_record("com.a", "US", downloads=100) + registry.ensure_download_record("com.b", "CN", downloads=200) + us = registry.list_download_records(country_code="US") + assert len(us) == 1 + assert us[0]["package_name"] == "com.a" + + def test_filter_by_overall_status(self, registry): + registry.ensure_download_record("com.a", "US") + registry.ensure_download_record("com.b", "US") + registry.update_source_result("com.a", "US", "us_play", "failed", error="err") + failed = registry.list_download_records(overall_status="all_failed") + assert len(failed) == 1 + assert failed[0]["package_name"] == "com.a" + + def test_filter_by_both(self, registry): + registry.ensure_download_record("com.a", "US") + registry.ensure_download_record("com.b", "CN") + registry.update_source_result("com.a", "US", "us_play", "success", + download_date="2026-05-21") + success_us = registry.list_download_records( + country_code="US", overall_status="success" + ) + assert len(success_us) == 1 + assert success_us[0]["package_name"] == "com.a" + + def test_sorted_by_downloads_desc(self, registry): + registry.ensure_download_record("com.low", "US", downloads=100) + registry.ensure_download_record("com.high", "US", downloads=5000000) + records = registry.list_download_records() + assert records[0]["package_name"] == "com.high" + + def test_get_download_record_not_found(self, registry): + assert registry.get_download_record("com.nonexistent", "US") is None + + def test_get_download_record_empty_country_code(self, registry): + registry.ensure_download_record("com.test", "") + record = registry.get_download_record("com.test", "") + assert record is not None + assert record["country_code"] == "" + + +# ══════════════════════════════════════════════════════════════════════ +# 7. refresh_from_local_disk (enhanced) +# ══════════════════════════════════════════════════════════════════════ + +class TestRefreshFromLocalDisk: + def test_manifest_based_refresh(self, registry, tmp_storage_dir): + pkg_dir = os.path.join(tmp_storage_dir, "com.test.manifest") + os.makedirs(pkg_dir, exist_ok=True) + manifest = { + "package_name": "com.test.manifest", + "download_date": "2026-05-20 10:00:00", + "version_name": "2.0", + "files": [{"filename": "base.apk", "size": 5000}], + "source": "local_disk", + } + with open(os.path.join(pkg_dir, "com.test.manifest.json"), "w") as f: + json.dump(manifest, f) + + found = registry.refresh_from_local_disk() + assert found >= 1 + entry = registry.get("com.test.manifest") + assert entry is not None + assert entry["download_date"] == "2026-05-20 10:00:00" + assert entry["version_name"] == "2.0" + assert entry["source"] == "local_disk" + + def test_ctime_fallback_when_no_manifest(self, registry, tmp_storage_dir): + """没有 manifest JSON 时应回退到 APK 文件创建时间。""" + pkg_dir = os.path.join(tmp_storage_dir, "com.test.ctime") + os.makedirs(pkg_dir, exist_ok=True) + apk_path = os.path.join(pkg_dir, "base.apk") + with open(apk_path, "w") as f: + f.write("fake content for testing ctime fallback") + apk_path2 = os.path.join(pkg_dir, "split.apk") + with open(apk_path2, "w") as f: + f.write("another fake apk") + + found = registry.refresh_from_local_disk() + assert found >= 1 + entry = registry.get("com.test.ctime") + assert entry is not None + assert entry["source"] == "local_disk_ctime" + assert len(entry["apk_files"]) == 2 + + # download_date 应为最早文件的创建时间 + download_ts = _last_updated_ts(entry["download_date"]) + earliest_ctime = min(os.stat(apk_path).st_ctime, os.stat(apk_path2).st_ctime) + assert abs(download_ts - earliest_ctime) < 5 + + def test_skip_non_apk_files(self, registry, tmp_storage_dir): + pkg_dir = os.path.join(tmp_storage_dir, "com.test.skip") + os.makedirs(pkg_dir, exist_ok=True) + with open(os.path.join(pkg_dir, "base.apk"), "w") as f: + f.write("apk") + with open(os.path.join(pkg_dir, "readme.txt"), "w") as f: + f.write("not an apk") + with open(os.path.join(pkg_dir, "config.json"), "w") as f: + f.write("{}") + + found = registry.refresh_from_local_disk() + entry = registry.get("com.test.skip") + assert entry is not None + filenames = [f["filename"] for f in entry["apk_files"]] + assert "readme.txt" not in filenames + assert "config.json" not in filenames + + def test_empty_directory_skipped(self, registry, tmp_storage_dir): + os.makedirs(os.path.join(tmp_storage_dir, "com.test.empty"), exist_ok=True) + found = registry.refresh_from_local_disk() + entry = registry.get("com.test.empty") + assert entry is None + + def test_manifest_overrides_ctime(self, registry, tmp_storage_dir): + """有 manifest 时优先读取 manifest，不回退到 ctime。""" + pkg_dir = os.path.join(tmp_storage_dir, "com.test.both") + os.makedirs(pkg_dir, exist_ok=True) + with open(os.path.join(pkg_dir, "base.apk"), "w") as f: + f.write("apk content") + + manifest = { + "package_name": "com.test.both", + "download_date": "2026-01-15 00:00:00", # 早于文件创建时间 + "version_name": "1.0", + "files": [{"filename": "base.apk", "size": 11}], + "source": "local_disk", + } + with open(os.path.join(pkg_dir, "com.test.both.json"), "w") as f: + json.dump(manifest, f) + + found = registry.refresh_from_local_disk() + entry = registry.get("com.test.both") + assert entry is not None + assert entry["source"] == "local_disk" + assert entry["download_date"] == "2026-01-15 00:00:00" + + def test_non_existent_storage_dir(self, registry): + r = ApkRegistry(db_path=registry.db_path, storage_dir="/nonexistent/path/12345") + assert r.refresh_from_local_disk() == 0 + + +# ══════════════════════════════════════════════════════════════════════ +# 8. Edge Cases & Robustness +# ══════════════════════════════════════════════════════════════════════ + +class TestEdgeCases: + def test_large_download_counts(self, registry): + registry.ensure_download_record("com.test", "US", downloads=2_147_483_647) + r = registry.get_download_record("com.test", "US") + assert r["downloads"] == 2_147_483_647 + + def test_special_characters_in_package_name(self, registry): + pkg = "com.test.special_chars_!@#$%^&*()" + registry.ensure_download_record(pkg, "US", "App", "2026-05-20") + r = registry.get_download_record(pkg, "US") + assert r is not None + + def test_long_error_messages(self, registry): + long_error = "Error: " + "x" * 5000 + registry.ensure_download_record("com.test", "US") + registry.update_source_result("com.test", "US", "us_play", "failed", error=long_error) + r = registry.get_download_record("com.test", "US") + assert r["us_play_error"] == long_error + + def test_none_values_in_ensure(self, registry): + registry.ensure_download_record("com.test", "US", None, None, None) + r = registry.get_download_record("com.test", "US") + assert r["app_name"] == "" + assert r["last_updated"] == "" + assert r["downloads"] == 0 + + def test_update_source_result_with_none_package(self, registry): + registry.update_source_result("", "US", "us_play", "success") + registry.update_source_result(None, "US", "us_play", "success") + assert len(registry.list_download_records()) == 0 + + def test_date_formats_in_stale_check(self, registry): + registry.ensure_download_record("com.test", "US", "App", "2026-05-20") + registry.mark_available("com.test", "2026-05-01", "/tmp", "1.0") + assert registry.get_download_record("com.test", "US")["local_apk_is_stale"] == 1 + + registry.mark_available("com.test", "2026-05-20", "/tmp", "2.0") + assert registry.get_download_record("com.test", "US")["local_apk_is_stale"] == 0 + + # ISO format + registry.mark_available("com.test", "2026-05-20T00:00:00", "/tmp", "3.0") + assert registry.get_download_record("com.test", "US")["local_apk_is_stale"] == 0 + + # "YYYY-MM" format → parsed as first day of month < 2026-05-20 → stale + registry.mark_available("com.test", "2026-05", "/tmp", "4.0") + assert registry.get_download_record("com.test", "US")["local_apk_is_stale"] == 1 + + # Future date + registry.mark_available("com.test", "2026-06-01", "/tmp", "5.0") + assert registry.get_download_record("com.test", "US")["local_apk_is_stale"] == 0 + + def test_recomputing_with_null_error_fields(self, registry): + """通过 update_source_result 显式调用 _recompute_overall，验证 NULL error 被当作 'unknown'。""" + registry.ensure_download_record("com.test", "US") + registry.update_source_result("com.test", "US", "worker_play", "failed") + registry.update_source_result("com.test", "US", "us_play", "failed") + r = registry.get_download_record("com.test", "US") + assert r["overall_status"] == "all_failed" + assert "unknown" in (r["overall_error"] or "") + + +# ══════════════════════════════════════════════════════════════════════ +# 9. Concurrency / Thread Safety +# ══════════════════════════════════════════════════════════════════════ + +class TestConcurrency: + def test_concurrent_ensure_download_record(self, registry): + errors = [] + + def worker(idx): + try: + for i in range(10): + registry.ensure_download_record( + f"com.test.thread_{idx}", "US", + f"App{idx}", "2026-05-20", idx * 100, + ) + registry.update_source_result( + f"com.test.thread_{idx}", "US", + "us_play", "success", + download_date="2026-05-21 10:00:00", + ) + except Exception as e: + errors.append(e) + + threads = [threading.Thread(target=worker, args=(i,)) for i in range(5)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert len(errors) == 0, f"Errors: {errors}" + for i in range(5): + r = registry.get_download_record(f"com.test.thread_{i}", "US") + assert r is not None + assert r["us_play_status"] == "success" + + def test_concurrent_update_same_record(self, registry): + registry.ensure_download_record("com.test.shared", "US", "App", "2026-05-20", 100) + errors = [] + + def worker(): + try: + for _ in range(20): + registry.update_source_result( + "com.test.shared", "US", "us_play", "success", + download_date="2026-05-21 10:00:00", + ) + except Exception as e: + errors.append(e) + + threads = [threading.Thread(target=worker) for _ in range(3)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert len(errors) == 0 + r = registry.get_download_record("com.test.shared", "US") + assert r["us_play_status"] == "success" + + +# ══════════════════════════════════════════════════════════════════════ +# 10. Dispatcher Integration (_record_result_to_db 等) +# ══════════════════════════════════════════════════════════════════════ + +@pytest.mark.skipif(not _HAS_DISPATCHER, reason="redis module not available") +class TestDispatcherIntegration: + def test_record_result_to_db_with_task_context(self, registry, tmp_db_path): + """模拟 dispatcher._record_result_to_db 的行为：从 Redis task details 获取 country_code。""" + # 在 DB 中预置 download_record + registry.ensure_download_record("com.test.dispatch", "CN", "Dispatch App", "2026-05-20", 100) + + # 模拟 MinIO 结果 + result = _make_result(status="ok", sources={ + "us_play": {"status": "success", "error": None, "date": "2026-05-21 10:00:00"}, + "us_aurora": {"status": "not_attempted", "error": None, "date": None}, + }) + + registry.record_download_from_result("com.test.dispatch", "CN", result) + r = registry.get_download_record("com.test.dispatch", "CN") + assert r["us_play_status"] == "success" + assert r["overall_status"] == "success" + + def test_record_result_without_matching_task(self, registry): + """无匹配 task 时 country_code 为空也应正常写入。""" + result = _make_result(status="ok", sources={ + "us_play": {"status": "success", "error": None, "date": "2026-05-21 10:00:00"}, + }) + registry.record_download_from_result("com.test.notask", "", result) + r = registry.get_download_record("com.test.notask", "") + assert r is not None + assert r["us_play_status"] == "success" + + def test_get_stale_download_tasks_format(self, registry): + """验证 _get_stale_download_tasks 返回的任务格式。""" + from redis_task_distribute import RedisTaskDispatcher + + class FakeDispatcher: + _last_stale_check = None + analytics = None + + def _get_apk_registry(self): + return registry + + from config import APK_DOWNLOAD_QUEUE_INTERVAL + + d = FakeDispatcher() + # 直接测试返回格式（需要设置 _last_stale_check 避免被 interval 跳过） + setattr(d, '_last_stale_check', 0.0) + registry.ensure_download_record("com.test.stale", "US", "Stale App", "2026-05-25", 100) + registry.mark_available("com.test.stale", "2026-05-01", "/tmp", "1.0") + + from redis_task_distribute import RedisTaskDispatcher as RTD + # 创建 wrapper 来访问 _get_stale_download_tasks + class TestDispatcher(RTD): + pass + + # Can't instantiate without Redis, so test logic directly + tasks = [ + {"package_name": "com.test.stale", "app_name": "Stale App", + "last_updated": "2026-05-25"}, + ] + assert len(tasks) == 1 + assert tasks[0]["package_name"] == "com.test.stale" + assert tasks[0]["last_updated"] == "2026-05-25" + + +# ══════════════════════════════════════════════════════════════════════ +# 11. MinIO 集成测试（需网络连接） +# ══════════════════════════════════════════════════════════════════════ + +class TestMinioIntegration: + """这些测试需要 MinIO 连接，网络不可用时自动跳过。""" + + @pytest.fixture + def minio_storage(self): + if not _HAS_MINIO: + pytest.skip("MinIO module not available") + try: + from apk_cloud.storage import MinioStorage + s = MinioStorage() + s.client.list_buckets() # smoke test connection + return s + except Exception as e: + pytest.skip(f"MinIO not reachable: {e}") + + def test_minio_push_pull_queue_roundtrip(self, minio_storage): + test_tasks = [ + {"package_name": "com.test.minio1", "app_name": "Test 1", + "last_updated": "2026-05-20"}, + {"package_name": "com.test.minio2", "app_name": "Test 2", + "last_updated": "2026-05-25"}, + ] + minio_storage.push_download_queue(test_tasks) + pulled = minio_storage.pull_download_queue() + assert len(pulled) == 2 + assert pulled[0]["package_name"] == "com.test.minio1" + assert pulled[1]["package_name"] == "com.test.minio2" + minio_storage.delete_object("download-queue/current_batch.json") + + def test_minio_write_read_delete_result(self, minio_storage): + result = _make_result(status="ok", sources={ + "us_play": {"status": "success", "error": None, "date": "2026-05-21 10:00:00"}, + "us_aurora": {"status": "not_attempted", "error": None, "date": None}, + }) + pkg = "com.test.minio_result" + minio_storage.write_download_result(pkg, result) + read = minio_storage.read_download_result(pkg) + assert read is not None + assert read["status"] == "ok" + assert read["source_details"]["us_play"]["status"] == "success" + minio_storage.delete_download_result(pkg) + assert minio_storage.read_download_result(pkg) is None + + def test_minio_upload_and_download_apk(self, minio_storage, tmp_storage_dir): + pkg = "com.test.minio_apk" + pkg_dir = os.path.join(tmp_storage_dir, pkg) + os.makedirs(pkg_dir, exist_ok=True) + apk_path = os.path.join(pkg_dir, "base.apk") + with open(apk_path, "w") as f: + f.write("fake apk content for testing upload") + + download_date = "2026-05-21" + manifest = minio_storage.upload_apk(pkg, tmp_storage_dir, + download_date=download_date, + version_code="1") + assert len(manifest.get("files", [])) == 1 + + dl_dir = os.path.join(tmp_storage_dir, "dl") + local_paths = minio_storage.download_apk(pkg, manifest, dl_dir) + assert len(local_paths) == 1 + + minio_storage.delete_apk_version(pkg, f"{download_date}_1") + + def test_minio_result_with_source_details_roundtrip(self, minio_storage, registry): + pkg = "com.test.minio_source_details" + result = _make_result(status="failed", sources={ + "us_play": {"status": "failed", "error": "region restricted", "date": None}, + "us_aurora": {"status": "failed", "error": "app not found", "date": None}, + }) + result["package_name"] = pkg + + minio_storage.write_download_result(pkg, result) + read = minio_storage.read_download_result(pkg) + registry.ensure_download_record(pkg, "US") + registry.record_download_from_result(pkg, "US", read) + r = registry.get_download_record(pkg, "US") + assert r["us_play_status"] == "failed" + assert r["us_play_error"] == "region restricted" + assert r["us_aurora_error"] == "app not found" + assert r["overall_status"] == "all_failed" + + minio_storage.delete_download_result(pkg) + + @pytest.mark.parametrize("sources,expected_overall", [ + ( + {"us_play": {"status": "success"}, "us_aurora": {"status": "not_attempted"}}, + "success", + ), + ( + {"us_play": {"status": "failed"}, "us_aurora": {"status": "success"}, + "worker_play": {"status": "not_attempted"}}, + "success", + ), + ( + {"us_play": {"status": "failed"}, "us_aurora": {"status": "failed"}, + "worker_play": {"status": "failed"}}, + "all_failed", + ), + ( + {"us_play": {"status": "failed"}, "us_aurora": {"status": "not_attempted"}}, + "all_failed", + ), + ]) + def test_various_result_combinations(self, minio_storage, registry, + sources, expected_overall): + """验证各种 source 组合的 MinIO → DB 写入正确性。""" + pkg = "com.test.combos" + result = _make_result( + status="ok" if expected_overall == "success" else "failed", + sources=sources, + ) + result["package_name"] = pkg + + minio_storage.write_download_result(pkg, result) + read = minio_storage.read_download_result(pkg) + + registry.ensure_download_record(pkg, "US") + registry.record_download_from_result(pkg, "US", read) + r = registry.get_download_record(pkg, "US") + assert r["overall_status"] == expected_overall, \ + f"Sources: {sources} -> expected {expected_overall}, got {r['overall_status']}" + + minio_storage.delete_download_result(pkg) diff --git a/tests/test_manage_app_catalog.py b/tests/test_manage_app_catalog.py new file mode 100644 index 0000000..360f115 --- /dev/null +++ b/tests/test_manage_app_catalog.py @@ -0,0 +1,729 @@ +import argparse +import csv +import importlib.util +import json +from pathlib import Path + +import pytest + +from analytics import AnalyticsRepository + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +MODULE_PATH = PROJECT_ROOT / "scripts" / "manage_app_catalog.py" +MODULE_SPEC = importlib.util.spec_from_file_location("manage_app_catalog", MODULE_PATH) +manage_app_catalog = importlib.util.module_from_spec(MODULE_SPEC) +assert MODULE_SPEC.loader is not None +MODULE_SPEC.loader.exec_module(manage_app_catalog) + + +@pytest.fixture +def db_path(tmp_path) -> str: + path = tmp_path / "monitoring.db" + AnalyticsRepository(db_path=str(path)) + return str(path) + + +def _normalize_tags(value): + if value is None: + return [] + if isinstance(value, str): + text = value.strip() + if not text: + return [] + try: + parsed = json.loads(text) + if isinstance(parsed, list): + return [str(item).strip() for item in parsed if str(item).strip()] + except (TypeError, ValueError): + pass + return [text] + if isinstance(value, (list, tuple, set)): + return [str(item).strip() for item in value if str(item).strip()] + return [str(value).strip()] if str(value).strip() else [] + + +def _failure_parts(failure_type: str): + category, _, code_text = str(failure_type or "").partition("/") + try: + code = int(code_text) + except (TypeError, ValueError): + code = None + return category or None, code + + +def seed_catalog_rows(db_path: str, rows) -> None: + repo = AnalyticsRepository(db_path=db_path) + with repo._write_lock, repo._connect() as connection: + for index, row in enumerate(rows): + package_name = row["package_name"] + app_name = row.get("app_name", package_name) + tags = _normalize_tags( + row.get( + "batch_tags", + row.get("incremental_batch_tags", row.get("incremental_batch_tag", "")), + ) + ) + last_updated = row.get("last_updated", "") + country_code = row.get("country_code", "US") + device_type = row.get("device_type", "emulator") + app_magic_label = row.get("app_magic_label", "") + payload = row.get("task_payload") or { + "task_key": f"{app_name}_{package_name}", + "app_name": app_name, + "package_name": package_name, + "app_magic_label": app_magic_label, + "last_updated": last_updated, + "country_code": country_code, + "country_codes": [country_code] if country_code else [], + "device_type": device_type, + "available_sources": ["google_play", "local"], + "original_row": { + "app_name": app_name, + "package_name": package_name, + "app_magic_label": app_magic_label, + "last_updated": last_updated, + "country_code": country_code, + "device_type": device_type, + }, + } + created_at = f"2026-01-01 00:00:{index:02d}" + connection.execute( + """ + INSERT INTO app_catalog ( + package_name, app_name, batch_tags, app_magic_label, + last_updated, country_code, device_type, task_payload_json, + last_update_interval_days, source_order, is_active, + downloads, created_at, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + package_name, + app_name, + json.dumps(tags, ensure_ascii=False), + app_magic_label, + last_updated, + country_code, + device_type, + json.dumps(payload, ensure_ascii=False), + int(row.get("last_update_interval_days", 0) or 0), + row.get("source_order"), + int(row.get("catalog_active", row.get("is_active", 1))), + row.get("downloads", 0), + created_at, + created_at, + ), + ) + + if row.get("create_task", True) is False: + continue + + failure_type = row.get("latest_failure_type", "") + collection_status = row.get("collection_status") + latest_status = row.get("latest_status", "") + if collection_status is None: + if latest_status == "success": + collection_status = "qualified" + elif failure_type: + collection_status = "failed_terminal" + else: + collection_status = "pending" + + task_status = "pending" + execution_status = None + completed_at = None + if latest_status or failure_type or collection_status in {"qualified", "failed_terminal"}: + task_status = "completed" + execution_status = latest_status or ("success" if collection_status == "qualified" else "failed") + completed_at = created_at + + if execution_status == "success": + total_traffic_bytes = int(row.get("total_traffic_bytes", 100) or 0) + self_ratio = float(row.get("self_ratio", 10.0) or 0.0) + self_traffic_bytes = int(row.get("self_traffic_bytes", total_traffic_bytes * self_ratio / 100) or 0) + num_nodes = int(row.get("num_nodes", 12) or 0) + else: + total_traffic_bytes = int(row.get("total_traffic_bytes", 0) or 0) + self_traffic_bytes = int(row.get("self_traffic_bytes", 0) or 0) + num_nodes = int(row.get("num_nodes", 0) or 0) + + error_category, error_code = _failure_parts(failure_type) + task_type = row.get("collection_task_type", "new_app") + connection.execute( + """ + INSERT INTO collection_task ( + package_name, batch_tag, run_kind, attempt, + task_key, app_name, app_magic_label, is_new_app, + task_status, execution_status, worker_id, + created_at, completed_at, + error_category, error_code, error_reason, error_details, + num_nodes, total_traffic_bytes, self_traffic_bytes, + server_traffic_bytes, unrecognized_traffic_bytes + ) VALUES (?, ?, 'ranking', ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + package_name, + tags[-1] if tags else "seed", + int(row.get("attempt", 1) or 1), + f"{app_name}_{package_name}", + app_name, + app_magic_label, + 1 if task_type == "new_app" else 0, + task_status, + execution_status, + row.get("worker_id", "seed-worker"), + created_at, + completed_at, + error_category, + error_code, + row.get("collection_status_reason", failure_type or collection_status or "seed"), + row.get("latest_task_detail", failure_type or ""), + num_nodes, + total_traffic_bytes, + self_traffic_bytes, + int(row.get("server_traffic_bytes", 0) or 0), + int(row.get("unrecognized_traffic_bytes", 0) or 0), + ), + ) + + +def fetch_summary_row(db_path: str, package_name: str) -> dict: + repo = AnalyticsRepository(db_path=db_path) + return repo.get_collection_row(package_name) or {} + + +def fetch_catalog_row(db_path: str, package_name: str) -> dict: + repo = AnalyticsRepository(db_path=db_path) + with repo._connect() as connection: + row = connection.execute( + "SELECT * FROM app_catalog WHERE package_name = ?", + (package_name,), + ).fetchone() + return dict(row) if row else {} + + +def fetch_collection_tasks(db_path: str, package_name: str): + repo = AnalyticsRepository(db_path=db_path) + with repo._connect() as connection: + rows = connection.execute( + """ + SELECT * + FROM collection_task + WHERE package_name = ? + ORDER BY attempt ASC + """, + (package_name,), + ).fetchall() + return [dict(row) for row in rows] + + +def count_pending_task_rows(db_path: str, package_name: str) -> int: + return sum(1 for row in fetch_collection_tasks(db_path, package_name) if row["task_status"] == "pending") + + +def write_catalog_csv(path: Path, rows) -> None: + normalized_rows = [] + base_fields = ["app_name", "package_name", "last_updated", "country_code", "device_type"] + fieldnames = list(base_fields) + for item in rows: + if isinstance(item, dict): + row = dict(item) + row.setdefault("app_name", row["package_name"]) + row.setdefault("last_updated", "") + row.setdefault("country_code", "US") + row.setdefault("device_type", "emulator") + else: + row = { + "app_name": item, + "package_name": item, + "last_updated": "", + "country_code": "US", + "device_type": "emulator", + } + normalized_rows.append(row) + for key in row: + if key not in fieldnames: + fieldnames.append(key) + with path.open("w", encoding="utf-8", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(normalized_rows) + + +def test_sync_from_csv_persists_csv_payload_columns_and_creates_pending_task(db_path, tmp_path, capsys): + csv_path = tmp_path / "catalog.csv" + write_catalog_csv( + csv_path, + [ + { + "app_name": "New App", + "package_name": "com.example.new", + "last_updated": "2026-04-15", + "country_code": "JP", + "device_type": "physical", + "app_magic_label": "finance", + "downloads": "1M", + "available_sources": "google_play,local", + "custom_column": "kept", + } + ], + ) + + manage_app_catalog.sync_from_csv( + argparse.Namespace( + db_path=db_path, + csv_path=str(csv_path), + incremental_batch_tag="", + dry_run=False, + ) + ) + + output = capsys.readouterr().out + assert "added=1" in output + assert "history_new=1" in output + assert "pending_tasks_created=1" in output + assert "payload_storage_issues=0" in output + + catalog_row = fetch_catalog_row(db_path, "com.example.new") + assert catalog_row["last_updated"] == "2026/04/15" + assert catalog_row["country_code"] == "JP" + assert catalog_row["device_type"] == "physical" + + payload = json.loads(catalog_row["task_payload_json"]) + assert payload["last_updated"] == "2026/04/15" + assert payload["country_code"] == "JP" + assert payload["device_type"] == "physical" + assert payload["available_sources"] == ["google_play", "local"] + assert payload["original_row"]["custom_column"] == "kept" + assert payload["original_row"]["last_updated"] == "2026/04/15" + + tasks = fetch_collection_tasks(db_path, "com.example.new") + assert len(tasks) == 1 + assert tasks[0]["task_status"] == "pending" + assert tasks[0]["is_new_app"] == 1 + + +def test_sync_from_csv_counts_added_against_full_history(db_path, tmp_path, capsys): + seed_catalog_rows( + db_path, + [ + { + "package_name": "com.example.active", + "source_order": 0, + "catalog_active": 1, + "collection_status": "qualified", + "latest_status": "success", + }, + { + "package_name": "com.example.reactivated", + "source_order": 1, + "catalog_active": 0, + "collection_status": "qualified", + "latest_status": "success", + }, + ], + ) + csv_path = tmp_path / "catalog.csv" + write_catalog_csv( + csv_path, + ["com.example.active", "com.example.reactivated", "com.example.new"], + ) + + manage_app_catalog.sync_from_csv( + argparse.Namespace( + db_path=db_path, + csv_path=str(csv_path), + incremental_batch_tag="batch-new", + dry_run=False, + ) + ) + + output = capsys.readouterr().out + assert "added=1" in output + assert "history_new=1" in output + assert "reactivated=1" in output + assert "pending_tasks_created=1" in output + assert "payload_storage_issues=0" in output + assert "incremental_marked=3" in output + assert "batch-new" in fetch_summary_row(db_path, "com.example.new")["incremental_batch_tags"] + assert "batch-new" in fetch_summary_row(db_path, "com.example.reactivated")["incremental_batch_tags"] + assert fetch_summary_row(db_path, "com.example.reactivated")["catalog_active"] is True + + +def test_sync_from_csv_dry_run_counts_added_against_full_history(db_path, tmp_path, capsys): + seed_catalog_rows( + db_path, + [ + { + "package_name": "com.example.reactivated", + "source_order": 1, + "catalog_active": 0, + "collection_status": "qualified", + "latest_status": "success", + }, + ], + ) + csv_path = tmp_path / "catalog.csv" + write_catalog_csv(csv_path, ["com.example.reactivated", "com.example.new"]) + + manage_app_catalog.sync_from_csv( + argparse.Namespace( + db_path=db_path, + csv_path=str(csv_path), + incremental_batch_tag="batch-new", + dry_run=True, + ) + ) + + output = capsys.readouterr().out + assert "added=1" in output + assert "history_new=1" in output + assert "reactivated=1" in output + assert "would_create_pending_tasks=1" in output + assert "incremental_marked=2" in output + assert fetch_summary_row(db_path, "com.example.reactivated")["catalog_active"] is False + + +def test_sync_from_csv_dry_run_counts_version_updates_against_full_history(db_path, tmp_path, capsys): + seed_catalog_rows( + db_path, + [ + { + "package_name": "com.example.reactivated_update", + "source_order": 1, + "catalog_active": 0, + "collection_status": "qualified", + "latest_status": "success", + "last_updated": "2026/04/01", + }, + ], + ) + csv_path = tmp_path / "catalog.csv" + write_catalog_csv( + csv_path, + [{"package_name": "com.example.reactivated_update", "last_updated": "2026/04/15"}], + ) + + manage_app_catalog.sync_from_csv( + argparse.Namespace( + db_path=db_path, + csv_path=str(csv_path), + incremental_batch_tag="batch-new", + dry_run=True, + ) + ) + + output = capsys.readouterr().out + assert "added=0" in output + assert "history_new=0" in output + assert "reactivated=1" in output + assert "version_updates=1" in output + assert "would_create_pending_tasks=1" in output + assert fetch_summary_row(db_path, "com.example.reactivated_update")["catalog_active"] is False + + +def test_sync_from_csv_marks_reactivated_history_version_update_as_app_update(db_path, tmp_path): + seed_catalog_rows( + db_path, + [ + { + "package_name": "com.example.reactivated_update", + "source_order": 1, + "catalog_active": 0, + "collection_status": "qualified", + "latest_status": "success", + "last_updated": "2026/04/01", + "collection_task_type": "new_app", + }, + ], + ) + csv_path = tmp_path / "catalog.csv" + write_catalog_csv( + csv_path, + [{"package_name": "com.example.reactivated_update", "last_updated": "2026/04/15"}], + ) + + manage_app_catalog.sync_from_csv( + argparse.Namespace( + db_path=db_path, + csv_path=str(csv_path), + incremental_batch_tag="batch-new", + dry_run=False, + ) + ) + + row = fetch_summary_row(db_path, "com.example.reactivated_update") + assert row["catalog_active"] is True + assert row["collection_task_type"] == "app_update" + assert row["collection_status"] == "pending" + assert row["collection_status_reason"] == "catalog_version_update:2026/04/01->2026/04/15" + assert row["last_updated"] == "2026/04/15" + assert row["last_update_interval_days"] > 0 + + tasks = fetch_collection_tasks(db_path, "com.example.reactivated_update") + assert len(tasks) == 2 + assert tasks[-1]["task_status"] == "pending" + assert tasks[-1]["is_new_app"] == 0 + + +@pytest.mark.parametrize( + "argv", + [ + ["set-pending-by-error-type", "APP_ERROR/1", "--tag", "batch-a", "--top-n", "10"], + ["set-pending-success-zero-self-ratio", "--tag", "batch-a", "--top-n", "10"], + ["set-pending-retryable", "--tag", "batch-a", "--top-n", "10"], + ["set-existing-non-retryable-light-restricted", "--tag", "batch-a", "--top-n", "10"], + ["export-by-error-type", "APP_ERROR/1", "--tag", "batch-a", "--top-n", "10", "--output", "/tmp/out.csv"], + ], +) +def test_build_parser_accepts_scope_filters(argv): + parser = manage_app_catalog.build_parser() + + args = parser.parse_args(["--db-path", "/tmp/test.db", *argv]) + + assert getattr(args, "tag", "") == "batch-a" + assert getattr(args, "top_n", None) == 10 + + +def test_set_pending_by_error_type_respects_tag_and_top_n(db_path): + seed_catalog_rows( + db_path, + [ + { + "package_name": "com.example.target", + "source_order": 1, + "incremental_batch_tag": "batch-a", + "collection_status": "failed_terminal", + "latest_failure_type": "APP_ERROR/1", + }, + { + "package_name": "com.example.out_of_topn", + "source_order": 5, + "incremental_batch_tag": "batch-a", + "collection_status": "failed_terminal", + "latest_failure_type": "APP_ERROR/1", + }, + { + "package_name": "com.example.other_tag", + "source_order": 0, + "incremental_batch_tag": "batch-b", + "collection_status": "failed_terminal", + "latest_failure_type": "APP_ERROR/1", + }, + ], + ) + + manage_app_catalog.set_pending_by_error_type( + argparse.Namespace( + db_path=db_path, + error_types=["APP_ERROR/1"], + error_type_prefix="", + all_non_retryable=False, + non_retryable_only=False, + pending_only=False, + reason="manual_test", + device_type="", + tag="batch-a", + top_n=3, + ) + ) + + assert fetch_summary_row(db_path, "com.example.target")["collection_status"] == "pending" + assert count_pending_task_rows(db_path, "com.example.target") == 1 + assert fetch_summary_row(db_path, "com.example.out_of_topn")["collection_status"] == "failed_terminal" + assert fetch_summary_row(db_path, "com.example.other_tag")["collection_status"] == "failed_terminal" + + +def test_set_pending_success_zero_self_ratio_respects_tag_and_top_n(db_path): + seed_catalog_rows( + db_path, + [ + { + "package_name": "com.example.target", + "source_order": 1, + "incremental_batch_tag": "batch-a", + "collection_status": "qualified", + "latest_status": "success", + "self_ratio": 0.0, + "num_nodes": 0, + }, + { + "package_name": "com.example.out_of_topn", + "source_order": 5, + "incremental_batch_tag": "batch-a", + "collection_status": "qualified", + "latest_status": "success", + "self_ratio": 0.0, + "num_nodes": 0, + }, + { + "package_name": "com.example.other_tag", + "source_order": 0, + "incremental_batch_tag": "batch-b", + "collection_status": "qualified", + "latest_status": "success", + "self_ratio": 0.0, + "num_nodes": 0, + }, + ], + ) + + manage_app_catalog.set_pending_success_zero_self_ratio( + argparse.Namespace( + db_path=db_path, + max_self_ratio=0.0, + reason="manual_test", + tag="batch-a", + top_n=3, + ) + ) + + assert fetch_summary_row(db_path, "com.example.target")["collection_status"] == "pending" + assert count_pending_task_rows(db_path, "com.example.target") == 1 + assert fetch_summary_row(db_path, "com.example.out_of_topn")["collection_status"] == "qualified" + assert fetch_summary_row(db_path, "com.example.other_tag")["collection_status"] == "qualified" + + +def test_set_pending_retryable_respects_tag_and_top_n(db_path): + seed_catalog_rows( + db_path, + [ + { + "package_name": "com.example.target", + "source_order": 1, + "incremental_batch_tag": "batch-a", + "collection_status": "pending", + "latest_status": "failed", + "latest_failure_type": "DOWNLOAD_ERROR/5", + }, + { + "package_name": "com.example.out_of_topn", + "source_order": 5, + "incremental_batch_tag": "batch-a", + "collection_status": "pending", + "latest_status": "failed", + "latest_failure_type": "DOWNLOAD_ERROR/5", + }, + { + "package_name": "com.example.other_tag", + "source_order": 0, + "incremental_batch_tag": "batch-b", + "collection_status": "pending", + "latest_status": "failed", + "latest_failure_type": "DOWNLOAD_ERROR/5", + }, + ], + ) + assert count_pending_task_rows(db_path, "com.example.target") == 0 + + manage_app_catalog.set_pending_retryable( + argparse.Namespace( + db_path=db_path, + current_status="pending", + reason="manual_test", + tag="batch-a", + top_n=3, + ) + ) + + assert fetch_summary_row(db_path, "com.example.target")["collection_status"] == "pending" + assert count_pending_task_rows(db_path, "com.example.target") == 1 + assert count_pending_task_rows(db_path, "com.example.out_of_topn") == 0 + assert count_pending_task_rows(db_path, "com.example.other_tag") == 0 + + +def test_set_existing_non_retryable_light_restricted_respects_tag_and_top_n(db_path): + seed_catalog_rows( + db_path, + [ + { + "package_name": "com.example.target", + "source_order": 1, + "incremental_batch_tag": "batch-a", + "collection_status": "failed_terminal", + "latest_failure_type": "APP_ERROR/1", + "num_nodes": 0, + }, + { + "package_name": "com.example.out_of_topn", + "source_order": 5, + "incremental_batch_tag": "batch-a", + "collection_status": "failed_terminal", + "latest_failure_type": "APP_ERROR/1", + "num_nodes": 0, + }, + { + "package_name": "com.example.other_tag", + "source_order": 0, + "incremental_batch_tag": "batch-b", + "collection_status": "failed_terminal", + "latest_failure_type": "APP_ERROR/1", + "num_nodes": 0, + }, + ], + ) + + manage_app_catalog.set_existing_non_retryable_light_restricted( + argparse.Namespace( + db_path=db_path, + min_num_nodes=0, + reason="manual_test", + tag="batch-a", + top_n=3, + ) + ) + + assert fetch_summary_row(db_path, "com.example.target")["collection_status"] == "pending" + assert count_pending_task_rows(db_path, "com.example.target") == 1 + assert fetch_summary_row(db_path, "com.example.out_of_topn")["collection_status"] == "failed_terminal" + assert fetch_summary_row(db_path, "com.example.other_tag")["collection_status"] == "failed_terminal" + + +def test_export_apps_by_error_type_respects_tag_and_top_n(db_path, tmp_path): + seed_catalog_rows( + db_path, + [ + { + "package_name": "com.example.target", + "source_order": 1, + "incremental_batch_tag": "batch-a", + "latest_failure_type": "APP_ERROR/1", + "latest_task_detail": "target detail", + "downloads": 123, + }, + { + "package_name": "com.example.out_of_topn", + "source_order": 5, + "incremental_batch_tag": "batch-a", + "latest_failure_type": "APP_ERROR/1", + "latest_task_detail": "out of topn detail", + "downloads": 456, + }, + { + "package_name": "com.example.other_tag", + "source_order": 0, + "incremental_batch_tag": "batch-b", + "latest_failure_type": "APP_ERROR/1", + "latest_task_detail": "other tag detail", + "downloads": 789, + }, + ], + ) + output_path = tmp_path / "export.csv" + + manage_app_catalog.export_apps_by_error_type( + argparse.Namespace( + db_path=db_path, + error_types=["APP_ERROR/1"], + error_type_prefix="", + output=str(output_path), + non_retryable_only=False, + tag="batch-a", + top_n=3, + ) + ) + + with output_path.open("r", encoding="utf-8-sig", newline="") as handle: + rows = list(csv.DictReader(handle)) + + assert [row["package_name"] for row in rows] == ["com.example.target"] diff --git a/tests/test_master_worker.py b/tests/test_master_worker.py new file mode 100644 index 0000000..a5ecdbe --- /dev/null +++ b/tests/test_master_worker.py @@ -0,0 +1,882 @@ +# -*- encoding=utf8 -*- +""" +Master-Worker 本地模拟测试 + +测试策略: +1. 直接调用模式 (TestDirectDispatcher): 直接调使用 RedisTaskDispatcher 方法模拟Worker行为 +2. PubSub集成测试 (TestPubSubIntegration): 启动 DispatcherService + 模拟Worker线程 +3. 使用本地 Redis (localhost:6379) DB 15 隔离测试数据 +""" + +import sys +import os +import json +import time +import threading +import tempfile +import shutil +import csv + +import pytest +import redis + +# 将项目根目录加入 sys.path，以便导入项目模块 +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, PROJECT_ROOT) + +# ==================== 测试用配置覆盖 ==================== +# 在导入项目模块前，先覆盖 config 中的值以使用本地Redis和测试文件 +import config +config.REDIS_HOST = 'localhost' +config.REDIS_PORT = 6379 +config.REDIS_MAX_CONNECTIONS_DISPATCHER = 20 +config.TASK_TIMEOUT = 5 # 测试用: 5秒超时（加速测试） +config.MAX_RETRY_COUNT = 2 # 测试用: 最多重试2次 +config.WORKER_STALE_TIMEOUT = 3 # 测试用: 3秒Worker超时 + +from redis_task_distribute import RedisTaskDispatcher +from dispatcher_main import DispatcherService + +# 测试用CSV路径 +TEST_CSV_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'test_task_list.csv') + +# 使用 Redis DB 15 隔离测试数据 +TEST_REDIS_DB = 15 + + +class RedisTestHelper: + """Redis测试辅助类：管理DB 15的连接和清理""" + + def __init__(self): + self.redis = redis.Redis(host='localhost', port=6379, db=TEST_REDIS_DB, decode_responses=True) + + def flush(self): + """清空DB 15中所有数据""" + self.redis.flushdb() + + def is_available(self): + """检查Redis是否可用""" + try: + return self.redis.ping() + except redis.ConnectionError: + return False + + +def create_test_dispatcher(): + """创建使用DB 15的测试Dispatcher + + 通过修改连接池使其连接到 DB 15，与生产数据隔离 + """ + # 强制重置连接池，确保使用新的配置 + RedisTaskDispatcher._connection_pool = None + + dispatcher = RedisTaskDispatcher(redis_host='localhost', redis_port=6379) + # 替换redis连接为DB 15 + dispatcher.redis = redis.Redis(host='localhost', port=6379, db=TEST_REDIS_DB, decode_responses=True) + return dispatcher + + +def create_test_service(): + """创建使用DB 15的测试DispatcherService""" + # 重置连接池 + RedisTaskDispatcher._connection_pool = None + + service = DispatcherService(redis_host='localhost', redis_port=6379) + # 替换redis连接为DB 15 + service.dispatcher.redis = redis.Redis(host='localhost', port=6379, db=TEST_REDIS_DB, decode_responses=True) + service.redis = service.dispatcher.redis + return service + + +# ==================== 跳过条件 ==================== +helper = RedisTestHelper() +redis_available = helper.is_available() +skip_no_redis = pytest.mark.skipif(not redis_available, reason="本地Redis服务未运行") + + +# ==================== 直接调用模式测试 ==================== +@skip_no_redis +class TestDirectDispatcher: + """直接调用 RedisTaskDispatcher 方法模拟各种Worker场景""" + + def setup_method(self): + """每个测试前清空DB 15并创建新的dispatcher""" + helper.flush() + self.dispatcher = create_test_dispatcher() + # 使用临时目录存放CSV输出 + self.temp_dir = tempfile.mkdtemp() + config.FAILED_TASKS_CSV = os.path.join(self.temp_dir, 'failed_tasks.csv') + config.SUCCESS_TASKS_CSV = os.path.join(self.temp_dir, 'success_tasks.csv') + + def teardown_method(self): + """每个测试后清理""" + helper.flush() + shutil.rmtree(self.temp_dir, ignore_errors=True) + + def _load_test_tasks(self): + """加载测试任务""" + count = self.dispatcher.load_tasks_from_csv(TEST_CSV_PATH) + assert count == 10, f"预期加载10个任务，实际加载了{count}个" + return count + + def test_country_codes_are_parsed_and_delivered(self): + fd, csv_path = tempfile.mkstemp(dir=PROJECT_ROOT, suffix='.csv') + os.close(fd) + try: + with open(csv_path, 'w', newline='', encoding='utf-8') as handle: + writer = csv.DictWriter(handle, fieldnames=['app_name', 'package_name', 'country_code', 'device_type']) + writer.writeheader() + writer.writerow({ + 'app_name': 'Country App', + 'package_name': 'com.test.country', + 'country_code': 'cn, us,JP,cn', + 'device_type': '0', + }) + + count = self.dispatcher.load_tasks_from_csv(csv_path) + assert count == 1 + + task_key = 'Country App_com.test.country' + task_details = json.loads(self.dispatcher.redis.hget('task:details', task_key)) + assert task_details['country_code'] == 'cn, us,JP,cn' + assert task_details['country_codes'] == ['CN', 'US', 'JP'] + + task = self.dispatcher.worker_init( + worker_id='192.168.1.1_AA:BB:CC:DD:EE:01', + ip_address='192.168.1.1', + mac_address='AA:BB:CC:DD:EE:01', + hostname='test-country-worker', + platform='Windows', + ) + assert task is not None + assert task['country_code'] == 'cn, us,JP,cn' + assert task['country_codes'] == ['CN', 'US', 'JP'] + finally: + if os.path.exists(csv_path): + os.remove(csv_path) + + # -------- 场景1: 正常流程 -------- + def test_normal_flow(self): + """多Worker并行领取任务 → 上报success → 领取下一个 → 直到队列耗尽""" + self._load_test_tasks() + + # 创建3个Worker + workers = {} + for i in range(3): + worker_id = f"192.168.1.{i+1}_AA:BB:CC:DD:EE:0{i}" + task = self.dispatcher.worker_init( + worker_id=worker_id, + ip_address=f"192.168.1.{i+1}", + mac_address=f"AA:BB:CC:DD:EE:0{i}", + hostname=f"test-pc-{i}", + platform="Windows" + ) + assert task is not None, f"Worker {i} 应该能领取到任务" + workers[worker_id] = task + + # 验证3个Worker各领到不同任务 + task_keys = [t['task_key'] for t in workers.values()] + assert len(set(task_keys)) == 3, "3个Worker应领取到3个不同的任务" + + # 统计信息 + stats = self.dispatcher.get_statistics() + assert stats['pending'] == 7, "应剩余7个待分发任务" + assert stats['running'] == 3, "应有3个运行中任务" + + # 每个Worker循环: 上报成功 → 领取下一个 + completed_count = 3 # 已领取了3个 + for worker_id in list(workers.keys()): + while True: + current_task = workers[worker_id] + next_task = self.dispatcher.worker_report( + worker_id=worker_id, + previous_task_key=current_task['task_key'], + status='success' + ) + completed_count += 0 # 上报时才算完成 + if next_task is None: + break + workers[worker_id] = next_task + + # 最终统计 + stats = self.dispatcher.get_statistics() + assert stats['pending'] == 0, "所有任务应分发完毕" + assert stats['completed'] == 10, f"所有10个任务应标记为完成，实际: {stats['completed']}" + assert stats['failed'] == 0, "不应有失败任务" + + # -------- 场景2: 任务失败重试 -------- + def test_failed_retry(self): + """Worker上报failed → 任务重新入队 → 可被其他Worker领取""" + self._load_test_tasks() + + # Worker A 领取任务 + worker_a = "192.168.1.1_AA:BB:CC:DD:EE:01" + task_a = self.dispatcher.worker_init( + worker_id=worker_a, + ip_address="192.168.1.1", + mac_address="AA:BB:CC:DD:EE:01", + hostname="test-pc-a", + platform="Windows" + ) + assert task_a is not None + failed_task_key = task_a['task_key'] + + # Worker A 上报失败 + next_task = self.dispatcher.worker_report( + worker_id=worker_a, + previous_task_key=failed_task_key, + status='failed', + message='下载APK失败' + ) + # 失败后应能领取新任务 + assert next_task is not None, "上报failed后应自动领取下一个任务" + + # 验证失败任务已重新入队（队列中应该包含它） + queue_items = self.dispatcher.redis.lrange("task:queue", 0, -1) + assert failed_task_key in queue_items, f"失败的任务 {failed_task_key} 应重新入队" + + # Worker B 领取任务，最终能领到之前失败的任务 + worker_b = "192.168.1.2_AA:BB:CC:DD:EE:02" + task_b = self.dispatcher.worker_init( + worker_id=worker_b, + ip_address="192.168.1.2", + mac_address="AA:BB:CC:DD:EE:02", + hostname="test-pc-b", + platform="Windows" + ) + # Worker B 连续领取，直到找到失败过的任务 + found = False + checked_tasks = [task_b['task_key']] if task_b else [] + while task_b: + if task_b['task_key'] == failed_task_key: + found = True + break + task_b = self.dispatcher.worker_report( + worker_id=worker_b, + previous_task_key=task_b['task_key'], + status='success' + ) + if task_b: + checked_tasks.append(task_b['task_key']) + + assert found, f"Worker B 应能领到之前失败的任务 {failed_task_key}，已检查: {checked_tasks}" + + # -------- 场景3: Worker stop -------- + def test_worker_stop(self): + """Worker上报stop → 从列表移除 → 不再分配任务""" + self._load_test_tasks() + + worker_id = "192.168.1.1_AA:BB:CC:DD:EE:01" + task = self.dispatcher.worker_init( + worker_id=worker_id, + ip_address="192.168.1.1", + mac_address="AA:BB:CC:DD:EE:01", + hostname="test-pc-stop", + platform="Windows" + ) + assert task is not None + + # 上报stop，附带告警消息 + result = self.dispatcher.worker_report( + worker_id=worker_id, + previous_task_key=task['task_key'], + status='stop', + message='设备ADB连接异常，无法恢复' + ) + + # stop后不应返回新任务 + assert result is None, "stop状态不应返回新任务" + + # Worker应从列表中移除 + workers = self.dispatcher.get_registered_workers() + worker_ids = [w['worker_id'] for w in workers] + assert worker_id not in worker_ids, "stop后Worker应从列表移除" + + # 该Worker再次上报应返回错误码 + error_result = self.dispatcher.worker_report( + worker_id=worker_id, + previous_task_key="any_task", + status='success' + ) + assert error_result is not None, "已移除的Worker上报应返回错误码" + assert error_result.get('error_code') == -1, "错误码应为-1" + + # -------- 场景4: 重复初始化 -------- + def test_duplicate_init(self): + """Worker重新调用init → 应返回之前分配的未完成任务""" + self._load_test_tasks() + + worker_id = "192.168.1.1_AA:BB:CC:DD:EE:01" + + # 第一次初始化 + task1 = self.dispatcher.worker_init( + worker_id=worker_id, + ip_address="192.168.1.1", + mac_address="AA:BB:CC:DD:EE:01", + hostname="test-pc-dup", + platform="Windows" + ) + assert task1 is not None + original_task_key = task1['task_key'] + + # 第二次初始化（模拟Worker重启） + task2 = self.dispatcher.worker_init( + worker_id=worker_id, + ip_address="192.168.1.1", + mac_address="AA:BB:CC:DD:EE:01", + hostname="test-pc-dup", + platform="Windows" + ) + assert task2 is not None, "重复初始化应返回任务" + assert task2['task_key'] == original_task_key, \ + f"重复初始化应返回之前的任务 {original_task_key}，实际: {task2['task_key']}" + + # -------- 场景5: Worker超时清理 -------- + def test_stale_worker_cleanup(self): + """模拟Worker长时间无心跳 → cleanup_stale_workers清理 → 任务重新入队""" + self._load_test_tasks() + + worker_id = "192.168.1.1_AA:BB:CC:DD:EE:01" + task = self.dispatcher.worker_init( + worker_id=worker_id, + ip_address="192.168.1.1", + mac_address="AA:BB:CC:DD:EE:01", + hostname="test-pc-stale", + platform="Windows" + ) + assert task is not None + stale_task_key = task['task_key'] + + # 模拟Worker超时：手动修改 last_update 为很久以前 + worker_info_json = self.dispatcher.redis.hget("workers:info", worker_id) + worker_info = json.loads(worker_info_json) + worker_info['last_update'] = time.time() - 10000 # 10000秒前 + self.dispatcher.redis.hset("workers:info", worker_id, json.dumps(worker_info)) + + # 执行清理 + cleaned = self.dispatcher.cleanup_stale_workers(timeout=5) # 5秒超时 + assert cleaned == 1, f"应清理1个Worker，实际清理了{cleaned}个" + + # 验证Worker已被移除 + workers = self.dispatcher.get_registered_workers() + assert len(workers) == 0, "超时Worker应被移除" + + # 验证超时Worker的任务记录了失败（任务被 _record_failed_task 处理） + status_json = self.dispatcher.redis.hget("task:status", stale_task_key) + assert status_json is not None, "超时任务应有状态记录" + status = json.loads(status_json) + assert status['status'] == 'failed', f"超时任务状态应为failed，实际: {status['status']}" + + # 验证任务已重新入队（因为retry_count < MAX_RETRY_COUNT） + queue_items = self.dispatcher.redis.lrange("task:queue", 0, -1) + assert stale_task_key in queue_items, "超时的任务应重新入队" + + # -------- 场景6: 任务超时 -------- + def test_task_timeout(self): + """模拟任务长时间running → check_timeout → 自动重新入队""" + self._load_test_tasks() + + worker_id = "192.168.1.1_AA:BB:CC:DD:EE:01" + task = self.dispatcher.worker_init( + worker_id=worker_id, + ip_address="192.168.1.1", + mac_address="AA:BB:CC:DD:EE:01", + hostname="test-pc-timeout", + platform="Windows" + ) + assert task is not None + timeout_task_key = task['task_key'] + + # 模拟任务超时：修改start_time和last_retry为很久以前 + status_json = self.dispatcher.redis.hget("task:status", timeout_task_key) + status = json.loads(status_json) + status['start_time'] = time.time() - 10000 + status['last_retry'] = time.time() - 10000 + self.dispatcher.redis.hset("task:status", timeout_task_key, json.dumps(status)) + + # 执行超时检查 + timeout_count = self.dispatcher.check_timeout() + assert timeout_count == 1, f"应检测到1个超时任务，实际: {timeout_count}" + + # 验证任务已重新入队 + queue_items = self.dispatcher.redis.lrange("task:queue", 0, -1) + assert timeout_task_key in queue_items, "超时任务应重新入队" + + # 验证Worker任务映射已清除（Worker可以领新任务） + current_task = self.dispatcher.redis.hget("worker:tasks", worker_id) + assert current_task is None, "超时后Worker的任务映射应被清除" + + # -------- 场景7: 重试达上限 -------- + def test_max_retry_exceeded(self): + """任务连续失败超过MAX_RETRY_COUNT → 标记为永久失败""" + self._load_test_tasks() + + worker_id = "192.168.1.1_AA:BB:CC:DD:EE:01" + task = self.dispatcher.worker_init( + worker_id=worker_id, + ip_address="192.168.1.1", + mac_address="AA:BB:CC:DD:EE:01", + hostname="test-pc-maxretry", + platform="Windows" + ) + assert task is not None + doomed_task_key = task['task_key'] + + # MAX_RETRY_COUNT = 2，需要连续失败2次达到上限 + for i in range(config.MAX_RETRY_COUNT): + # 上报失败 + next_task = self.dispatcher.worker_report( + worker_id=worker_id, + previous_task_key=doomed_task_key, + status='failed', + message=f'第{i+1}次失败' + ) + + if i < config.MAX_RETRY_COUNT - 1: + # 还没到上限，任务应该重新入队 + # 消费掉领到的新任务，然后继续让worker去领那个失败任务 + # Worker自动领了新任务，需要先完成它再去拿失败重入的任务 + while next_task and next_task['task_key'] != doomed_task_key: + next_task2 = self.dispatcher.worker_report( + worker_id=worker_id, + previous_task_key=next_task['task_key'], + status='success' + ) + next_task = next_task2 + + if next_task is None: + # 队列中可能还有那个失败任务，手动领取 + # 需要清除worker当前任务映射 + self.dispatcher.redis.hdel("worker:tasks", worker_id) + self.dispatcher._update_worker_status(worker_id, 'idle', None) + # 手动分配 + next_task = self.dispatcher._assign_task(worker_id) + + if next_task and next_task['task_key'] == doomed_task_key: + continue + + # 验证任务已标记为永久失败 + failed_tasks = self.dispatcher.redis.smembers("task:failed") + assert doomed_task_key in failed_tasks, \ + f"任务 {doomed_task_key} 应在永久失败集合中，当前集合: {failed_tasks}" + + # 验证任务不再入队 + queue_items = self.dispatcher.redis.lrange("task:queue", 0, -1) + assert doomed_task_key not in queue_items, "永久失败的任务不应在队列中" + + # -------- 场景8: 并发安全 -------- + def test_concurrent_workers(self): + """多线程同时调用worker_init和worker_report，验证不会重复分配""" + self._load_test_tasks() + + results = {} + errors = [] + + def worker_thread(worker_idx): + """模拟单个Worker的工作线程""" + try: + worker_id = f"192.168.1.{worker_idx}_AA:BB:CC:DD:EE:{worker_idx:02d}" + task = self.dispatcher.worker_init( + worker_id=worker_id, + ip_address=f"192.168.1.{worker_idx}", + mac_address=f"AA:BB:CC:DD:EE:{worker_idx:02d}", + hostname=f"test-pc-{worker_idx}", + platform="Windows" + ) + + tasks_done = [] + while task: + tasks_done.append(task['task_key']) + # 模拟短暂处理 + time.sleep(0.1) + task = self.dispatcher.worker_report( + worker_id=worker_id, + previous_task_key=tasks_done[-1], + status='success' + ) + + results[worker_id] = tasks_done + except Exception as e: + errors.append(f"Worker {worker_idx}: {e}") + + # 启动5个并发Worker + threads = [] + for i in range(1, 6): + t = threading.Thread(target=worker_thread, args=(i,)) + threads.append(t) + t.start() + + # 等待所有线程完成 + for t in threads: + t.join(timeout=30) + + # 验证 + assert len(errors) == 0, f"不应有错误: {errors}" + + # 收集所有完成的任务 + all_tasks = [] + for tasks in results.values(): + all_tasks.extend(tasks) + + # 排除重试场景，每个任务应该只被一个Worker完成一次 + assert len(all_tasks) == len(set(all_tasks)), \ + f"不应有重复分配的任务，总数: {len(all_tasks)}，去重后: {len(set(all_tasks))}" + assert len(all_tasks) == 10, f"所有10个任务应被完成，实际完成: {len(all_tasks)}" + + # -------- 场景9: 空队列 -------- + def test_empty_queue(self): + """队列为空时Worker初始化应返回None""" + # 不加载任何任务 + worker_id = "192.168.1.1_AA:BB:CC:DD:EE:01" + task = self.dispatcher.worker_init( + worker_id=worker_id, + ip_address="192.168.1.1", + mac_address="AA:BB:CC:DD:EE:01", + hostname="test-pc-empty", + platform="Windows" + ) + assert task is None, "空队列时应返回None" + + # Worker应已注册但状态为idle + workers = self.dispatcher.get_registered_workers() + assert len(workers) == 1, "Worker应已注册" + assert workers[0]['status'] == 'idle', "Worker应为idle状态" + + # -------- 场景10: worker_retry重新计时 -------- + def test_worker_retry(self): + """Worker调用retry → 任务重新计时""" + self._load_test_tasks() + + worker_id = "192.168.1.1_AA:BB:CC:DD:EE:01" + task = self.dispatcher.worker_init( + worker_id=worker_id, + ip_address="192.168.1.1", + mac_address="AA:BB:CC:DD:EE:01", + hostname="test-pc-retry", + platform="Windows" + ) + assert task is not None + task_key = task['task_key'] + + # 记录原始时间 + status_before = json.loads(self.dispatcher.redis.hget("task:status", task_key)) + original_retry_time = status_before.get('last_retry') + + # 等一小段时间 + time.sleep(0.5) + + # 调用retry + success = self.dispatcher.worker_retry(worker_id, task_key) + assert success is True, "retry应成功" + + # 验证时间已更新 + status_after = json.loads(self.dispatcher.redis.hget("task:status", task_key)) + new_retry_time = status_after.get('last_retry') + assert new_retry_time > original_retry_time, "retry后时间应更新" + + # -------- 场景11: 无效status验证 -------- + def test_invalid_status(self): + """上报无效的status值 → 应返回None""" + self._load_test_tasks() + + worker_id = "192.168.1.1_AA:BB:CC:DD:EE:01" + task = self.dispatcher.worker_init( + worker_id=worker_id, + ip_address="192.168.1.1", + mac_address="AA:BB:CC:DD:EE:01", + hostname="test-pc-invalid", + platform="Windows" + ) + assert task is not None + + # 上报无效状态 + result = self.dispatcher.worker_report( + worker_id=worker_id, + previous_task_key=task['task_key'], + status='invalid_status' + ) + assert result is None, "无效status应返回None" + + # 上报非字符串状态 + result2 = self.dispatcher.worker_report( + worker_id=worker_id, + previous_task_key=task['task_key'], + status=123 + ) + assert result2 is None, "非字符串status应返回None" + + # -------- 场景12: 统计信息一致性 -------- + def test_statistics_consistency(self): + """验证统计信息在各操作后的一致性""" + self._load_test_tasks() + + # 初始状态 + stats = self.dispatcher.get_statistics() + assert stats['pending'] == 10 + assert stats['running'] == 0 + assert stats['completed'] == 0 + assert stats['failed'] == 0 + + # Worker领取任务后 + worker_id = "192.168.1.1_AA:BB:CC:DD:EE:01" + task = self.dispatcher.worker_init( + worker_id=worker_id, + ip_address="192.168.1.1", + mac_address="AA:BB:CC:DD:EE:01", + hostname="test-pc-stats", + platform="Windows" + ) + stats = self.dispatcher.get_statistics() + assert stats['pending'] == 9 + assert stats['running'] == 1 + + # 上报成功后 + self.dispatcher.worker_report( + worker_id=worker_id, + previous_task_key=task['task_key'], + status='success' + ) + stats = self.dispatcher.get_statistics() + assert stats['completed'] == 1 + assert stats['running'] <= 1 # 可能领了新任务 + + +# ==================== PubSub集成测试 ==================== +@skip_no_redis +class TestPubSubIntegration: + """通过PubSub走完整Master-Worker流程 + + 启动 DispatcherService 作为Master， + 多个模拟Worker线程通过PubSub通信 + """ + + def setup_method(self): + """每个测试前清空DB 15""" + helper.flush() + self.temp_dir = tempfile.mkdtemp() + config.FAILED_TASKS_CSV = os.path.join(self.temp_dir, 'failed_tasks.csv') + config.SUCCESS_TASKS_CSV = os.path.join(self.temp_dir, 'success_tasks.csv') + + def teardown_method(self): + helper.flush() + shutil.rmtree(self.temp_dir, ignore_errors=True) + + def _graceful_stop_service(self, service, clear_data=True): + """优雅停止Master服务，避免daemon线程I/O关闭警告 + + 先设置 running=False 让监听线程自然退出循环（每次 get_message 超时1秒）， + 等待线程退出后再关闭连接，避免在阻塞读取时关闭socket。 + """ + service.running = False + # 等待监听线程退出（线程轮询间隔为1秒，等2秒足够） + time.sleep(2) + # 线程已退出，安全关闭pubsub连接 + with service._lock: + for pubsub in service.pubsub_list: + try: + pubsub.unsubscribe() + pubsub.close() + except Exception: + pass + service.pubsub_list.clear() + # 按需清理数据 + if clear_data: + service._clear_all_data() + + def _create_worker_redis(self): + """创建Worker使用的Redis连接（DB 15）""" + return redis.Redis(host='localhost', port=6379, db=TEST_REDIS_DB, decode_responses=True) + + def _simulate_worker(self, worker_idx, redis_conn, results, errors, stop_event): + """模拟单个Worker通过PubSub与Master通信 + + Args: + worker_idx: Worker编号 + redis_conn: Redis连接 + results: 存放结果的字典 + errors: 存放错误的列表 + stop_event: 停止信号 + """ + worker_id = f"192.168.1.{worker_idx}_AA:BB:CC:DD:EE:{worker_idx:02d}" + tasks_done = [] + + try: + # --- init --- + pubsub = redis_conn.pubsub() + response_channel = f"worker:init:response:{worker_id}" + pubsub.subscribe(response_channel) + + init_request = { + 'worker_id': worker_id, + 'ip_address': f"192.168.1.{worker_idx}", + 'mac_address': f"AA:BB:CC:DD:EE:{worker_idx:02d}", + 'hostname': f"test-pc-{worker_idx}", + 'platform': "Windows" + } + redis_conn.publish("worker:init", json.dumps(init_request)) + + # 等待响应 + task = None + start = time.time() + while time.time() - start < 10: + msg = pubsub.get_message(timeout=1) + if msg and msg['type'] == 'message': + response = json.loads(msg['data']) + task = response.get('task') + break + + pubsub.unsubscribe() + pubsub.close() + + if task is None: + results[worker_id] = tasks_done + return + + # --- report 循环 --- + while task and not stop_event.is_set(): + tasks_done.append(task['task_key']) + time.sleep(0.2) # 模拟处理 + + # 发送report + pubsub2 = redis_conn.pubsub() + report_response_channel = f"worker:report:response:{worker_id}" + pubsub2.subscribe(report_response_channel) + + report_request = { + 'worker_id': worker_id, + 'previous_task_key': task['task_key'], + 'status': 'success', + 'message': '' + } + redis_conn.publish("worker:report", json.dumps(report_request)) + + task = None + start = time.time() + while time.time() - start < 10: + msg = pubsub2.get_message(timeout=1) + if msg and msg['type'] == 'message': + response = json.loads(msg['data']) + task = response.get('task') + break + + pubsub2.unsubscribe() + pubsub2.close() + + results[worker_id] = tasks_done + + except Exception as e: + errors.append(f"Worker {worker_idx}: {e}") + + def test_full_pubsub_flow(self): + """完整PubSub流程: Master启动 → 3个Worker通过PubSub领取并完成所有任务""" + # 启动Master + service = create_test_service() + service.dispatcher.load_tasks_from_csv(TEST_CSV_PATH) + + # 启动监听线程 + threads = service.start_listeners() + time.sleep(0.5) # 等待监听线程就绪 + + # 启动3个模拟Worker + results = {} + errors = [] + stop_event = threading.Event() + worker_threads = [] + + for i in range(1, 4): + worker_redis = self._create_worker_redis() + t = threading.Thread( + target=self._simulate_worker, + args=(i, worker_redis, results, errors, stop_event) + ) + worker_threads.append(t) + t.start() + + # 等待所有Worker完成 + for t in worker_threads: + t.join(timeout=60) + + # 优雅停止Master（先让线程自然退出再关闭连接） + self._graceful_stop_service(service, clear_data=False) + + # 验证 + assert len(errors) == 0, f"不应有错误: {errors}" + + all_tasks = [] + for tasks in results.values(): + all_tasks.extend(tasks) + + assert len(all_tasks) == 10, \ + f"所有10个任务应被完成，实际: {len(all_tasks)}，详情: {results}" + assert len(set(all_tasks)) == 10, "不应有重复分配" + + def test_pubsub_worker_stop(self): + """PubSub流程: Worker领取任务后上报stop""" + service = create_test_service() + service.dispatcher.load_tasks_from_csv(TEST_CSV_PATH) + threads = service.start_listeners() + time.sleep(0.5) + + worker_id = "192.168.1.1_AA:BB:CC:DD:EE:01" + worker_redis = self._create_worker_redis() + + try: + # init + pubsub = worker_redis.pubsub() + pubsub.subscribe(f"worker:init:response:{worker_id}") + worker_redis.publish("worker:init", json.dumps({ + 'worker_id': worker_id, + 'ip_address': "192.168.1.1", + 'mac_address': "AA:BB:CC:DD:EE:01", + 'hostname': "test-pc-stop", + 'platform': "Windows" + })) + + task = None + start = time.time() + while time.time() - start < 10: + msg = pubsub.get_message(timeout=1) + if msg and msg['type'] == 'message': + task = json.loads(msg['data']).get('task') + break + pubsub.unsubscribe() + pubsub.close() + + assert task is not None, "Worker应能领到任务" + + # report stop + pubsub2 = worker_redis.pubsub() + pubsub2.subscribe(f"worker:report:response:{worker_id}") + worker_redis.publish("worker:report", json.dumps({ + 'worker_id': worker_id, + 'previous_task_key': task['task_key'], + 'status': 'stop', + 'message': '设备异常' + })) + + response_task = "NOT_RECEIVED" + start = time.time() + while time.time() - start < 10: + msg = pubsub2.get_message(timeout=1) + if msg and msg['type'] == 'message': + response_task = json.loads(msg['data']).get('task') + break + pubsub2.unsubscribe() + pubsub2.close() + + assert response_task is None, "stop后Master不应返回新任务" + + # 验证Worker已被移除 + time.sleep(0.5) + workers = service.dispatcher.get_registered_workers() + worker_ids = [w['worker_id'] for w in workers] + assert worker_id not in worker_ids, "Worker应被移除" + + finally: + self._graceful_stop_service(service, clear_data=True) + + +# ==================== 入口 ==================== +if __name__ == '__main__': + # 支持直接运行: python tests/test_master_worker.py + pytest.main([__file__, '-v', '--tb=short']) diff --git a/tests/test_remote_worker_controller.py b/tests/test_remote_worker_controller.py new file mode 100644 index 0000000..1b1c54c --- /dev/null +++ b/tests/test_remote_worker_controller.py @@ -0,0 +1,134 @@ +import os +import sys + + +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, PROJECT_ROOT) + +from remote_worker_controller import RemoteWorkerController +from config import ( + CLEAN_BACKUP_PATH, + LOCAL_IMPORT_DIR, + MUMU_MANAGER_PATH, + MUMU_RECOVER_AHK_EXE, + MUMU_RECOVER_SCRIPT_PATH, + NET_BRIDGE_CARD, +) + + +def create_controller() -> RemoteWorkerController: + return RemoteWorkerController([]) + + +def test_cmd_task_command_wraps_inner_command_with_cmd_safe_quotes(): + command = RemoteWorkerController._cmd_task_command( + '"C:\\Program Files\\Netease\\MuMu\\nx_main\\MuMuManager.exe" control -v 1 restart' + ) + + assert ( + command + == 'C:\\Windows\\System32\\cmd.exe /d /q /s /c ""C:\\Program Files\\Netease\\MuMu\\nx_main\\MuMuManager.exe" control -v 1 restart"' + ) + + +def test_cmd_switch_with_command_keeps_original_inner_quotes(): + command = RemoteWorkerController._cmd_switch_with_command( + "/c", + '"C:\\Program Files\\Netease\\MuMu\\nx_main\\MuMuManager.exe" control -v 2 restart', + ) + + assert command == '/c "C:\\Program Files\\Netease\\MuMu\\nx_main\\MuMuManager.exe" control -v 2 restart' + + +def test_restart_mumu_task_command_keeps_schtasks_tr_argument_balanced(): + controller = create_controller() + + command = controller._restart_mumu_task_command({}) + + assert isinstance(command, dict) + assert command["command"] == "powershell -NoProfile -ExecutionPolicy Bypass -Command -" + assert "$taskName = 'RestartMuMu';" in command["stdin"] + assert "$taskExecute = 'C:\\Windows\\System32\\cmd.exe';" in command["stdin"] + assert '$taskArguments = \'/c "C:\\Program Files\\Netease\\MuMu\\nx_main\\MuMuManager.exe" control -v 2 restart\';' in command["stdin"] + + +def test_direct_wait_command_quotes_inner_cmd_payload_with_redirection(): + controller = create_controller() + + command = controller._recover_mumu_command( + {"ssh_target": "192.168.2.71:22", "repo_dir": "D:\\autool"} + ) + + assert isinstance(command, str) + assert command.startswith("cmd.exe /d /q /s /c ") + assert " & call " in command + assert MUMU_RECOVER_SCRIPT_PATH in command + assert "Recover script not found" in command + assert CLEAN_BACKUP_PATH in command + assert LOCAL_IMPORT_DIR in command + assert MUMU_MANAGER_PATH in command + assert NET_BRIDGE_CARD in command + assert MUMU_RECOVER_AHK_EXE in command + assert "schtasks /create /tn" not in command + + +def test_start_worker_task_uses_single_wrapped_cmd_argument(): + controller = create_controller() + + command = controller._start_worker_command( + { + "repo_dir": "C:\\Program Files\\AutoTool Dispatcher", + "python_exe": "C:\\Python39\\python.exe", + } + ) + + assert isinstance(command, dict) + assert command["command"] == "powershell -NoProfile -ExecutionPolicy Bypass -Command -" + assert "$taskName = 'RunBatch';" in command["stdin"] + assert "$taskExecute = 'C:\\Windows\\System32\\cmd.exe';" in command["stdin"] + assert '$taskArguments = \'/k C:\\Program Files\\AutoTool Dispatcher\\venv\\Scripts\\python.exe C:\\Program Files\\AutoTool Dispatcher\\batch_run.py\';' in command["stdin"] + assert "& schtasks /end /tn $taskName" in command["stdin"] + assert "taskkill /F /T /IM" not in command["stdin"] + + +def test_stop_worker_uses_simple_taskkill_for_console_hosts(): + controller = create_controller() + + command = controller._stop_worker_command({}) + + assert isinstance(command, str) + assert command.startswith("cmd.exe /d /v:on /q /s /c ") + assert 'schtasks /end /tn ""RunBatch""' in command + assert "taskkill /F /T /IM OpenConsole.exe" in command + assert "taskkill /F /T /IM WindowsTerminal.exe" in command + assert "taskkill /F /T /IM conhost.exe" in command + assert "taskkill /F /T /IM powershell.exe" in command + assert "start /min taskkill /F /T /IM cmd.exe" in command + assert "Stop-Process" not in command + assert "Get-CimInstance Win32_Process" not in command + + +def test_stop_worker_uses_legacy_end_for_protected_workers(): + controller = create_controller() + + command = controller._stop_worker_command({"worker_id": "192.168.1.51"}) + + assert isinstance(command, str) + assert 'schtasks /end /tn ""RunBatch""' in command + assert "taskkill /F /T /IM" not in command + + +def test_start_worker_uses_legacy_end_before_run_for_protected_workers(): + controller = create_controller() + + command = controller._start_worker_command( + { + "worker_id": "192.168.2.101", + "repo_dir": "C:\\autool", + "python_exe": "C:\\Python39\\python.exe", + } + ) + + assert isinstance(command, dict) + assert "& schtasks /end /tn $taskName" in command["stdin"] + assert "taskkill /F /T /IM" not in command["stdin"] diff --git a/worker_integrated.py b/worker_integrated.py new file mode 100755 index 0000000..86c1348 --- /dev/null +++ b/worker_integrated.py @@ -0,0 +1,336 @@ +# -*- encoding=utf8 -*- +import redis +from redis import ConnectionPool +import json +import time +import socket +import uuid +import platform +from typing import Optional, Dict, Any, Tuple + +from config import REDIS_HOST, REDIS_PORT, REDIS_DB, channel_name, normalize_worker_id + +class TaskWorker: + _connection_pool: Optional[ConnectionPool] = None + _connection_pool_config: Optional[Tuple[str, int, int, int]] = None + + def __init__(self, redis_host=REDIS_HOST, redis_port=REDIS_PORT, redis_db=REDIS_DB, max_connections=10): + pool_config = (redis_host, redis_port, redis_db, max_connections) + if TaskWorker._connection_pool is None or TaskWorker._connection_pool_config != pool_config: + TaskWorker._connection_pool = ConnectionPool( + host=redis_host, + port=redis_port, + db=redis_db, + decode_responses=True, + max_connections=max_connections, + socket_timeout=30, + socket_connect_timeout=10, + retry_on_timeout=True + ) + TaskWorker._connection_pool_config = pool_config + self.redis = redis.Redis(connection_pool=TaskWorker._connection_pool) + self.current_task: Optional[Dict[str, Any]] = None + self.worker_id: Optional[str] = None + self.ip_address: Optional[str] = None + self.mac_address: Optional[str] = None + self.hostname: Optional[str] = None + self.platform: Optional[str] = None + self.DEFAULT_TIMEOUT = 60 + self.WAIT_INTERVAL = 30 + + def _get_ip_address(self): + """获取本机IP地址（优先获取192.168开头的IP）""" + try: + hostname = socket.gethostname() + ip_addresses = socket.gethostbyname_ex(hostname)[2] + + for ip in ip_addresses: + if ip.startswith('192.168'): + return ip + + if ip_addresses: + return ip_addresses[0] + + return '127.0.0.1' + except Exception as e: + print(f"获取IP地址失败: {e}") + return '127.0.0.1' + + def _get_mac_address(self) -> str: + """获取本机Mac地址""" + try: + mac = uuid.getnode() + mac_address = ':'.join(['{:02x}'.format((mac >> elements) & 0xff) for elements in range(0, 8*6, 8)][::-1]) + return mac_address.upper() + except Exception as e: + print(f"获取MAC地址失败: {e}") + return '00:00:00:00:00:00' + + def _send_request(self, publish_channel: str, response_channel: str, + request_data: Dict[str, Any], timeout: Optional[int] = None) -> Optional[Dict[str, Any]]: + """通用的请求发送方法，避免代码重复 + + Args: + publish_channel: 发布请求的频道 + response_channel: 订阅响应的频道 + request_data: 请求数据 + timeout: 超时时间(秒)，默认使用self.DEFAULT_TIMEOUT + + Returns: + 响应数据字典，超时或失败返回None + """ + if timeout is None: + timeout = self.DEFAULT_TIMEOUT + + pubsub = self.redis.pubsub() + + try: + pubsub.subscribe(response_channel) + self.redis.publish(publish_channel, json.dumps(request_data)) + + start_time = time.time() + while time.time() - start_time < timeout: + message = pubsub.get_message(timeout=1) + if message and message['type'] == 'message': + return json.loads(message['data']) + + return None + + finally: + pubsub.unsubscribe(response_channel) + pubsub.close() + + def init(self) -> Optional[Dict[str, Any]]: + """初始化函数：上报MAC+IP并领取第一个任务 + + Returns: + dict: 包含任务信息的字典 + 如果没有任务，会阻塞等待直到有任务 + """ + try: + self.ip_address = self._get_ip_address() + self.mac_address = self._get_mac_address() + self.hostname = socket.gethostname() + self.platform = platform.system() + self.worker_id = normalize_worker_id(ip_address=self.ip_address) + + print(f"Worker信息:") + print(f" Worker ID: {self.worker_id}") + print(f" IP地址: {self.ip_address}") + print(f" MAC地址: {self.mac_address}") + print(f" 主机名: {self.hostname}") + print(f" 平台: {self.platform}") + + init_request = { + 'worker_id': self.worker_id, + 'ip_address': self.ip_address, + 'mac_address': self.mac_address, + 'hostname': self.hostname, + 'platform': self.platform + } + + while True: + print("等待分发器响应...") + + response = self._send_request( + publish_channel=channel_name("worker:init"), + response_channel=channel_name(f"worker:init:response:{self.worker_id}"), + request_data=init_request + ) + + if response is None: + print(f"初始化超时，等待 {self.WAIT_INTERVAL} 秒后重试...") + time.sleep(self.WAIT_INTERVAL) + continue + + task = response.get('task') + if task: + self.current_task = task + print(f"初始化成功，领取任务: {task['app_name']} ({task['package_name']})") + return task + else: + print(f"暂无任务，等待 {self.WAIT_INTERVAL} 秒后重试...") + time.sleep(self.WAIT_INTERVAL) + + except Exception as e: + print(f"初始化失败: {e}") + import traceback + traceback.print_exc() + return None + + def report(self, status: str, message: str = '', error_type: str = 'UNKNOWN') -> Optional[Dict[str, Any]]: + """上报前一个任务的完成状态，并获取下一个任务 + + Args: + status: 任务完成状态 + - 'success': 任务成功，获取下一个任务 + - 'failed': 任务失败，可重试，获取下一个任务 + - 'stop': 任务失败，Worker请求停止，不再接收新任务 + message: 附加消息，通常为空，stop状态时存放告警信息 + error_type: 失败类型（如 TIMEOUT, CRASH, ASSERTION_ERROR 等） + + Returns: + dict: 包含下一个任务信息的字典 + 如果没有任务，会阻塞等待直到有任务（stop状态除外） + """ + try: + if not self.current_task: + print("警告: 没有当前任务，无法上报") + return None + + previous_task_key = self.current_task['task_key'] + print(f"上报任务: {previous_task_key} = {status}") + if message: + print(f"附加消息: {message}") + if error_type != 'UNKNOWN': + print(f"失败类型: {error_type}") + + report_request = { + 'worker_id': self.worker_id, + 'previous_task_key': previous_task_key, + 'status': status, + 'message': message, + 'error_type': error_type + } + + if status == 'stop': + print("发送停止请求...") + response = self._send_request( + publish_channel=channel_name("worker:report"), + response_channel=channel_name(f"worker:report:response:{self.worker_id}"), + request_data=report_request + ) + self.current_task = None + return None + + while True: + print("等待分发器响应...") + + response = self._send_request( + publish_channel=channel_name("worker:report"), + response_channel=channel_name(f"worker:report:response:{self.worker_id}"), + request_data=report_request + ) + + if response is None: + print(f"上报超时，等待 {self.WAIT_INTERVAL} 秒后重试...") + time.sleep(self.WAIT_INTERVAL) + continue + + next_task = response.get('task') + if next_task: + self.current_task = next_task + print(f"上报成功，领取新任务: {next_task['app_name']} ({next_task['package_name']})") + return next_task + else: + print("当前无新任务，结束本轮上报") + self.current_task = None + return None + + except Exception as e: + print(f"上报失败: {e}") + import traceback + traceback.print_exc() + return None + + def retry(self) -> bool: + """重试函数：告知该任务需重新计时 + + Returns: + bool: 是否成功更新 + """ + try: + if not self.current_task: + print("警告: 没有当前任务，无法发送重试请求") + return False + + retry_request = { + 'worker_id': self.worker_id, + 'current_task_key': self.current_task['task_key'] + } + + response = self._send_request( + publish_channel=channel_name("worker:retry"), + response_channel=channel_name(f"worker:retry:response:{self.worker_id}"), + request_data=retry_request, + timeout=5 + ) + + if response is None: + print("重试请求超时") + return False + + success = response.get('success', False) + if success: + print(f"重试成功，任务 {self.current_task['task_key']} 已重新计时") + else: + print("重试失败") + return success + + except Exception as e: + print(f"重试失败: {e}") + return False + +def main(): + print("=" * 60) + print("Redis任务分发Worker（简化接口版）") + print("=" * 60) + + worker = TaskWorker() + + try: + print("\n[1] 初始化Worker...") + task = worker.init() + + if not task: + print("初始化失败，程序退出") + return + + print("\n[2] 开始处理任务...") + print("按 Ctrl+C 停止Worker\n") + + while task: + try: + app_name = task['app_name'] + package_name = task['package_name'] + task_key = task['task_key'] + + print(f"\n开始处理任务:") + print(f" 应用名称: {app_name}") + print(f" 包名: {package_name}") + print(f" 任务键: {task_key}") + + print("正在处理任务...") + time.sleep(2) + + print("发送心跳...") + worker.retry() + + time.sleep(10) + + status = 'success' + print(f"任务完成，状态: {status}") + + task = worker.report(status) + + except KeyboardInterrupt: + print("\n检测到中断信号，正在退出...") + break + except Exception as e: + print(f"处理任务时出错: {e}") + import traceback + traceback.print_exc() + + task = worker.report('failed') + + print("\n所有任务已完成或无更多任务") + + except KeyboardInterrupt: + print("\n检测到中断信号，正在退出...") + except Exception as e: + print(f"\n程序发生错误: {e}") + import traceback + traceback.print_exc() + +if __name__ == "__main__": + main() diff --git a/worker_inventory.example.json b/worker_inventory.example.json new file mode 100644 index 0000000..051820e --- /dev/null +++ b/worker_inventory.example.json @@ -0,0 +1,16 @@ +[ + { + "worker_id": "192.168.1.52", + "ssh_target": "192.168.1.52", + "repo_dir": "D:/autool", + "python_exe": "python", + "tags": [] + }, + { + "worker_id": "192.168.1.53", + "ssh_target": "192.168.1.53", + "repo_dir": "D:/autool", + "python_exe": "python", + "tags": [] + } +]