图书下载修改多线程并发下载分配策略，统一由主线程分配给子线程下载图书信息

lyg

2024-06-11 35a49bbd1b9c131a3a2db734f1351837022930a5

 src/main.mjs

@@ -4,7 +4,8 @@
import proxy from "selenium-webdriver/proxy.js";
import axios from "axios";
import * as fs from "fs";
import { Worker, isMainThread, parentPort, workerData } from 'worker_threads';
import path from "path";
import { Worker, isMainThread, parentPort, workerData, threadId } from 'worker_threads';
import { HttpsProxyAgent } from "https-proxy-agent";
import { resolve } from "path";

@@ -18,7 +19,7 @@
  if (!fs.existsSync('./logs')) {
    fs.mkdirSync('./logs', { recursive: true });
  }
  logFile = fs.createWriteStream(`./logs/logs-${config.startRow}-${config.endRow}.log`, { flags: 'a', encoding: 'utf8' });
  logFile = fs.createWriteStream(`./logs/logs-${config.startRow}-${config.endRow}-thread${threadId}.log`, { flags: 'a', encoding: 'utf8' });
  console.log = function (...text) {
    text = `${new Date().toLocaleString()} ${text.join(' ') ?? ''}`;
    _log(text);
@@ -307,14 +308,43 @@
  });
}

function isAlreadyDownloaded(book) {
  const id = `${book.id} ${book.isbn}`;
  return alreadyDownloadedBooks.includes(id);
}

function nextBook() {
  return new Promise(resolve => {
    const cb = (message) => {
      if (message.type === 'book') {
        resolve(message.data);
        parentPort.removeListener('message', cb);
      }
    };
    parentPort.on('message', cb);
    parentPort.postMessage({ type: 'get-book', threadId });

  });
}

async function downloadBooks(books) {
  driver = await createDriver();
  for (const book of books) {

  for (; ;) {
    const book = await nextBook();
    if (!book) {
      break;
    }
    books.push(book);
    if (config.endOfTime && Date.now() - startTime > 1000 * 60 * config.endOfTime) {
      // 定时退出
      break;
    }
    bookCount++;
    if (isAlreadyDownloaded(book)) {
      skipCount++;
      continue;
    }
    if (book.state && (book.state === "没有搜索结果" || book.state === "没有pdf或text文件" || book.state === "下载完成")) {
      // 跳过没有搜索结果或没有pdf或text文件的书籍
      skipCount++;
@@ -343,14 +373,14 @@
      continue;
    }
    // 等一段时间再打开详情页
    sleep(getRandomNumber(3000, 10000));
    sleep(getRandomNumber(1000, 30000));
    // 打开详情页
    await openBookDetailPage(book, detailPageUrl);
    // 获取下载链接
    const url = await getDownloadUrl(book);
    if (!url) { continue; }
    // 等待一段时间再下载
    await sleep(getRandomNumber(3000, 10000));
    await sleep(getRandomNumber(1000, 30000));
    // 下载文件
    try {
      await downloadFile(book, url);
@@ -358,7 +388,7 @@
    } catch (e) { }
    successCount++;
    // 等一段时间再下一个
    sleep(getRandomNumber(3000, 10000));
    sleep(getRandomNumber(1000, 30000));
  }
}

@@ -378,7 +408,7 @@
  }

  const buffer = xlsx.build([{ name: "Sheet1", data }]);
  fs.writeFile("./【第二批二次处理后】交付清单.xlsx", buffer, (err) => { });
  fs.writeFileSync("./【第二批二次处理后】交付清单.xlsx", buffer, (err) => { });
  console.log("保存完成: ./【第二批二次处理后】交付清单.xlsx");
}

@@ -416,9 +446,19 @@
let skipCount = 0;
// chrome驱动
let driver;
let alreadyDownloadedBooks = [];

function getAlreadyDownloadedBooks() {
  const text = fs.readFileSync('./alreadyDownloadedBooks.txt', 'utf-8');
  const books = text.replace(/\r/g, '').split('\n').map(it => it.trim()).filter(it => it);
  const files = fs.readdirSync('./downloads');
  books.push(...files);
  return books.map(it => path.basename(it, path.extname(it)).trim());
}

function main() {
  initLogger();
  const books = getBooksFromExcel(config.startRow, config.endRow);
  const books = [];
  downloadBooks(books)
    .then(() => {
      console.log(`全部完成，共下载${bookCount}本，成功下载${successCount}本，跳过${skipCount}本，失败${bookCount - skipCount - successCount}本，耗时： ${msFormat(Date.now() - startTime)}。`);
@@ -440,32 +480,34 @@
// 多进程执行
if (isMainThread) {
  initLogger();
  console.log(`线程数：${config.threadSize}, 开始行：${config.startRow}, 结束行：${config.endRow}`);
  let startRow = config.startRow;
  let endRow = config.endRow;
  const alreadyDownloadedBooks = getAlreadyDownloadedBooks();
  const { startRow, endRow, threadSize } = config;
  console.log(`线程数：${threadSize}, 开始行：${startRow}, 结束行：${endRow}`);
  let finishCnt = 0;
  const finishBooks = [];
  const threadSize = config.threadSize;
  const thBookSize = endRow - startRow / threadSize
  const thBookSize = (endRow - startRow) / threadSize;
  const books = getBooksFromExcel(startRow, endRow);

  for (let sr = startRow; sr < endRow; sr += thBookSize) {
    let er = sr + thBookSize;
    if (er > endRow) {
      er = endRow;
    }
    const worker = new Worker("./src/main.mjs", { workerData: { startRow: sr, endRow: er } });
    const worker = new Worker("./src/main.mjs", { workerData: { startRow: sr, endRow: er, alreadyDownloadedBooks } });
    worker.on("message", (message) => {
      if (message.type === 'books') {
        finishBooks.push(...message.data);
        finishCnt++;
        if (finishCnt >= config.threadSize) {
        if (finishCnt >= threadSize) {
          saveBooks(finishBooks);
        }
      } else if (message.type === 'get-book') {
        worker.postMessage({ type: "book", data: books.shift() });
      }
    });
  }
} else {
  config.startRow = workerData.startRow;
  config.endRow = workerData.endRow;
  alreadyDownloadedBooks = workerData.alreadyDownloadedBooks;
  main();
}