文档解析异步化

package nexus.io.mosskb.controller;

import java.util.ArrayList;
import java.util.List;

import com.jfinal.kit.Kv;

import nexus.io.annotation.Get;
import nexus.io.annotation.Post;
import nexus.io.annotation.RequestPath;
import nexus.io.jfinal.aop.Aop;
import nexus.io.mosskb.service.SystemFileService;
import nexus.io.mosskb.service.kb.MossKbDocumentSplitService;
import nexus.io.mosskb.service.kb.MossKbDocumentSplitTaskService;
import nexus.io.mosskb.vo.MossKbDocumentSplitTaskVo;
import nexus.io.model.result.ResultVo;
import nexus.io.model.upload.UploadFile;
import nexus.io.model.upload.UploadResult;
import nexus.io.tio.boot.http.TioRequestContext;

@RequestPath("/api/dataset/document")
public class ApiDatasetDocumentController {

  /**
   * 分段预览：只负责上传与建任务，解析在后台线程逐页推进。
   *
   * <p>大文档解析要几分钟，请求线程不再等待，前端拿 task_id 轮询 {@link #splitTask(Long)} 取结果。
   */
  @Post("/split")
  public ResultVo split(nexus.io.tio.http.common.HttpRequest request) {
    Object[] files = request.getParams().get("file");
    if (files == null || files.length == 0) {
      return ResultVo.fail("请求体中未找到文件");
    }
    Long userId = TioRequestContext.getUserIdLong();
    List<Long> taskIds = new ArrayList<>();
    List<Kv> failures = new ArrayList<>();
    for (Object item : files) {
      if (!(item instanceof UploadFile)) {
        return ResultVo.fail("文件格式不正确");
      }
      UploadFile file = (UploadFile) item;
      UploadResult uploaded = Aop.get(SystemFileService.class).upload(file, "default", "default");
      if (uploaded == null) {
        failures.add(Kv.by("name", file.getName()).set("message", "文件保存失败"));
        continue;
      }
      try {
        taskIds.add(Aop.get(MossKbDocumentSplitService.class).splitAsync(file.getData(), uploaded, userId));
      } catch (Exception e) {
        failures.add(Kv.by("name", file.getName()).set("message", e.getMessage()));
      }
    }
    if (taskIds.isEmpty()) {
      return ResultVo.fail(failures.isEmpty() ? "没有可解析的文件" : String.valueOf(failures.get(0).get("message")));
    }
    return ResultVo.ok(Kv.by("task_id_list", taskIds).set("failures", failures));
  }
}

import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

public class ExecutorServiceUtils {

  /**
   * 文档解析专用执行器。
   *
   * <p>解析一份大文档要几分钟且几乎全在等远程 OCR，用每任务一个虚拟线程的线程池，避免占用固定池的
   * 100 个工作线程，也避免多份大文档互相顶掉。
   */
  private static final ExecutorService documentParseExecutor =
      Executors.newThreadPerTaskExecutor(Thread.ofVirtual().name("doc-parse-", 0).factory());

  public static ExecutorService getDocumentParseExecutor() {
    return documentParseExecutor;
  }
}

  /** 记录已完成页数。 */
  public void markProgress(Long taskId, int completed, int total) {
    if (total <= 0) {
      return;
    }
    int percent = (int) Math.min(99, Math.round(completed * 100.0 / total));
    boolean lastPage = completed >= total;
    // 页数多时按页步进节流，页数少时按百分比节流，两种情况都能看到推进。
    boolean reachStep = completed % PROGRESS_STEP == 0
        || percent >= currentPercent(taskId) + PROGRESS_PERCENT_STEP;
    if (!lastPage && !reachStep) {
      return;
    }
    Db.update("update moss_kb_document_split_task set progress=?, update_time=now() where id=?", percent, taskId);
  }

  @Get("/split/task/{taskId}")
  public ResultVo splitTask(Long taskId) {
    Long userId = TioRequestContext.getUserIdLong();
    MossKbDocumentSplitTaskVo task = Aop.get(MossKbDocumentSplitTaskService.class).get(userId, taskId);
    if (task == null) {
      return ResultVo.fail("任务不存在或无权访问");
    }
    if (MossKbDocumentSplitTaskService.STATUS_FAILED.equals(task.getStatus())) {
      return ResultVo.fail(task.getErrorMessage() == null ? "文档解析失败" : task.getErrorMessage());
    }
    if (MossKbDocumentSplitTaskService.STATUS_SUCCESS.equals(task.getStatus())) {
      return ResultVo.ok(task.getResult() == null ? new ArrayList<>() : task.getResult());
    }
    return ResultVo.ok();
  }

  protected String recognize(byte[] data, String filename) throws Exception {
    GiteeClient client = new GiteeClient();
    GiteeDocumentParseRequest request = new GiteeDocumentParseRequest();
    request.setModel(OCR_MODEL);
    request.setPrompt(GiteePromptConst.pdf_to_markdown_prompt);
    request.setInclude_image(false);
    request.setInclude_image_base64(false);
    GiteeTaskResponse task = client.parseDocument(data, filename, request);
    long deadline = System.nanoTime() + java.util.concurrent.TimeUnit.MINUTES.toNanos(8);
    while (!Arrays.asList("success", "succeeded", "completed").contains(String.valueOf(task.getStatus()).toLowerCase(Locale.ROOT))) {
      if (task.getStatus() == null || task.getTask_id() == null) {
        throw new IOException("OCR服务未返回有效任务状态");
      }
      if (Arrays.asList("failed", "failure", "cancelled", "canceled").contains(task.getStatus().toLowerCase(Locale.ROOT))) {
        throw new IOException("OCR任务失败: " + task.getTask_id());
      }
      if (System.nanoTime() > deadline) {
        throw new IOException("OCR任务超时: " + task.getTask_id());
      }
      Thread.sleep(2000); task = client.getTask(task.getTask_id());
    }
    // 扫描页可能是空白页或整页图片，服务会正常返回空正文，这里按空内容处理，由调用方决定是否记录。
    String markdown = textOnly(GiteeSimpleMarkdownUtils.toMarkdown(task.getOutput(), null));
    if (markdown == null || markdown.isBlank()) {
      log.info("OCR returned no text for {}", filename);
      return "";
    }
    return markdown;
  }

  /**
   * OCR 重试：失败时按 1s、2s 退避重试，最后一次仍失败就记为失败页。
   *
   * <p>“未识别到内容”是正常返回，不是异常，因此不在这里重试。
   *
   * @return 是否拿到结果
   */
  private boolean withOcrRetry(OcrCall call, String[] result, boolean[] failed, String describe) {
    for (int attempt = 1; attempt <= OCR_MAX_ATTEMPTS; attempt++) {
      try {
        result[0] = call.recognize();
        return true;
      } catch (Exception e) {
        log.warn("OCR attempt {}/{} failed for {}: {}", attempt, OCR_MAX_ATTEMPTS, describe, e.getMessage());
        if (attempt < OCR_MAX_ATTEMPTS) {
          try {
            Thread.sleep(1000L * attempt);
          } catch (InterruptedException interrupted) {
            Thread.currentThread().interrupt();
            failed[0] = true;
            return false;
          }
        }
      }
    }
    log.warn("OCR gave up for {}", describe);
    failed[0] = true;
    return false;
  }

    if (result.text().isBlank()) {
      // 全部走 OCR 的文档如果每一页都失败，原因是远程服务不可用，不是文档本身没有内容。
      if (result.ocrPages() > 0 && result.ocrPages() == result.failedPages()) {
        throw new IOException("扫描页OCR识别失败，请稍后重试");
      }
      throw new IllegalArgumentException("未提取到内容，请检查文档或使用扫描PDF");
    }

  /** 并发提取的页面数：OCR 是网络调用，页与页之间并发即可，不需要按 CPU 核数拉满。 */
  private static final int OCR_PAGE_PARALLELISM = Math.max(2, Math.min(4, Runtime.getRuntime().availableProcessors()));

  /** 逐页提取 PDF：文本层优先，扫描页走 OCR，页面之间并发处理但按原页序输出。 */
  private Parsed parsePdf(byte[] data, Consumer<Progress> progress) throws Exception {
    try (PDDocument pdf = PDDocument.load(data)) {
      int totalPages = pdf.getNumberOfPages();
      PDFTextStripper stripper = new PDFTextStripper();
      stripper.setSortByPosition(true);
      String sourceHash = Md5Utils.md5Hex(data);
      PageResult[] results = new PageResult[totalPages];
      AtomicInteger completed = new AtomicInteger();
      ForkJoinPool pool = new ForkJoinPool(OCR_PAGE_PARALLELISM);
      try {
        List<ForkJoinTask<?>> tasks = new ArrayList<>(totalPages);
        for (int page = 1; page <= totalPages; page++) {
          int pageNumber = page;
          tasks.add(pool.submit(() -> {
            try {
              results[pageNumber - 1] = extractPdfPage(pdf, stripper, sourceHash, pageNumber);
            } catch (IOException e) {
              throw new UncheckedIOException("解析第 " + pageNumber + " 页失败", e);
            }
            if (progress != null) {
              progress.accept(new Progress(completed.incrementAndGet(), totalPages));
            }
          }));
        }
        for (ForkJoinTask<?> task : tasks) {
          task.join();
        }
      } finally {
        pool.shutdown();
      }

    // 缓存键固定用“原文件内容 + 页码”：序列化单页时 PDFBox 会写入新的内部文档 ID，
    // 同一页重复序列化的字节并不相同，用这些字节做键会让缓存永远不命中。
    boolean[] failed = new boolean[1];
    String[] recognized = new String[1];
    boolean ok = withOcrRetry(() -> ocrPdfPage(sourceHash, pageNumber, pageData), recognized, failed,
        sourceHash + ":page:" + pageNumber + ":" + OCR_CACHE_VERSION);
    return new PageResult(pageNumber, ok ? recognized[0] : "", true, failed[0]);

  private String ocrWithCache(byte[] data, String filename, String hash) throws Exception {
    String cached = readOcrCache(hash);
    if (cached != null && !cached.isBlank()) {
      return textOnly(cached);
    }
    synchronized (OCR_LOCK_STRIPES[Math.floorMod(hash.hashCode(), OCR_LOCKS)]) {
      // 同一份内容的第一个请求可能已经写好缓存，这里再查一次就能省掉重复的付费调用。
      cached = readOcrCache(hash);
      if (cached != null && !cached.isBlank()) {
        return textOnly(cached);
      }
      String markdown = recognize(data, filename);
      if (markdown != null && !markdown.isBlank()) {
        writeOcrCache(hash, markdown);
      }
      return markdown;
    }
  }