import queue import threading import uuid from dataclasses import dataclass, field from datetime import UTC, datetime from pathlib import Path from app.db.session import SessionLocal from app.models import Certificate, CertificateAccessToken, Learner, ProjectCourse from app.services.pdf import cached_pdf_is_fresh, certificate_template_path, pdf_cache_path, render_certificate_pdf MAX_PDF_PREGENERATION_JOBS = 30 def _utc_now() -> datetime: return datetime.now(UTC).replace(tzinfo=None) @dataclass class PdfPregenerationJob: id: str certificate_ids: list[int] concurrency_limit: int status: str = "running" total: int = 0 completed: int = 0 generated: int = 0 cached: int = 0 skipped: int = 0 failed: int = 0 errors: list[str] = field(default_factory=list) created_at: datetime = field(default_factory=_utc_now) finished_at: datetime | None = None def to_dict(self) -> dict[str, object]: percent = round((self.completed / self.total) * 100) if self.total else 0 return { "id": self.id, "status": self.status, "total": self.total, "completed": self.completed, "generated": self.generated, "cached": self.cached, "skipped": self.skipped, "failed": self.failed, "percent": percent, "concurrency_limit": self.concurrency_limit, "errors": self.errors[-10:], "created_at": self.created_at, "finished_at": self.finished_at, } class PdfPregenerationManager: def __init__(self) -> None: self._jobs: dict[str, PdfPregenerationJob] = {} self._lock = threading.Lock() def start(self, certificate_ids: list[int], concurrency_limit: int) -> PdfPregenerationJob: unique_ids = _unique_ids(certificate_ids) if not unique_ids: raise ValueError("请选择需要预生成PDF的证书") job = PdfPregenerationJob( id=uuid.uuid4().hex, certificate_ids=unique_ids, concurrency_limit=max(1, concurrency_limit), total=len(unique_ids), ) with self._lock: self._jobs[job.id] = job self._trim_jobs() thread = threading.Thread(target=self._run_job, args=(job.id,), daemon=True) thread.start() return job def get(self, job_id: str) -> PdfPregenerationJob | None: with self._lock: return self._jobs.get(job_id) def list(self) -> list[PdfPregenerationJob]: with self._lock: return sorted(self._jobs.values(), key=lambda item: item.created_at, reverse=True) def _run_job(self, job_id: str) -> None: job = self.get(job_id) if not job: return work_queue: queue.Queue[int] = queue.Queue() for certificate_id in job.certificate_ids: work_queue.put(certificate_id) workers = [ threading.Thread(target=self._worker, args=(job_id, work_queue), daemon=True) for _ in range(min(job.concurrency_limit, job.total)) ] for worker in workers: worker.start() for worker in workers: worker.join() with self._lock: job.status = "completed_with_errors" if job.failed else "completed" job.finished_at = _utc_now() def _worker(self, job_id: str, work_queue: queue.Queue[int]) -> None: while True: try: certificate_id = work_queue.get_nowait() except queue.Empty: return try: job = self.get(job_id) if not job: return result = pre_generate_certificate_pdf(certificate_id, job.concurrency_limit) self._record(job_id, result) except Exception as exc: self._record(job_id, "failed", f"证书ID {certificate_id}:{exc}") finally: work_queue.task_done() def _record(self, job_id: str, result: str, error: str | None = None) -> None: job = self.get(job_id) if not job: return with self._lock: job.completed += 1 if result == "generated": job.generated += 1 elif result == "cached": job.cached += 1 elif result == "skipped": job.skipped += 1 else: job.failed += 1 if error: job.errors.append(error) def _trim_jobs(self) -> None: if len(self._jobs) <= MAX_PDF_PREGENERATION_JOBS: return old_jobs = sorted(self._jobs.values(), key=lambda item: item.created_at) for job in old_jobs[: len(self._jobs) - MAX_PDF_PREGENERATION_JOBS]: self._jobs.pop(job.id, None) def pre_generate_certificate_pdf(certificate_id: int, concurrency_limit: int) -> str: with SessionLocal() as db: certificate = db.get(Certificate, certificate_id) if not certificate or certificate.status != "valid": return "skipped" learner = db.get(Learner, certificate.learner_id) token = db.get(CertificateAccessToken, certificate.qr_token_id or certificate.public_token_id) project = db.query(ProjectCourse).filter(ProjectCourse.code == certificate.project_code).first() if not learner or not token: raise ValueError("证书数据不完整") was_cached = certificate_pdf_is_ready(certificate) render_certificate_pdf(certificate, learner, project, token, concurrency_limit=concurrency_limit) db.commit() return "cached" if was_cached else "generated" def certificate_pdf_is_ready(certificate: Certificate) -> bool: output_path = pdf_cache_path(certificate.certificate_no) template_path = certificate_template_path() latest_renderer_mtime = max(template_path.stat().st_mtime, Path(__file__).with_name("pdf.py").stat().st_mtime) return cached_pdf_is_fresh(output_path) and output_path.stat().st_mtime >= latest_renderer_mtime def _unique_ids(values: list[int]) -> list[int]: seen: set[int] = set() result: list[int] = [] for value in values: if value not in seen: seen.add(value) result.append(value) return result pdf_pregeneration_manager = PdfPregenerationManager()