任务调度 RabbitQ

任务调度demo

Producer

发任务测试用，支持命令行参数队列名：

node producer.js gpu_task 10

node producer.js cpu_task 20

javascript 复制代码

#!/usr/bin/env node
// producer.js
const amqp = require('amqplib');

const [,, queue = 'task_task', count = '1'] = process.argv;
const total = parseInt(count, 10);

if (!Number.isInteger(total) || total <= 0) {
  console.error('用法：node producer.js <queue> <数量>');
  process.exit(1);
}

(async () => {
  const conn = await amqp.connect('amqp://localhost');
  const ch   = await conn.createConfirmChannel(); // 开启 confirm 模式
  await ch.assertQueue(queue, { durable: true });

  console.log(`[Producer] 开始向队列 "${queue}" 发送 ${total} 条任务...`);

  for (let i = 1; i <= total; i++) {
    const task = `task-${i}`;
    ch.sendToQueue(queue, Buffer.from(task), { persistent: true }, (err) => {
      if (err) console.error(`发送失败：${task}`, err);
    });
  }

  // 等待所有消息落盘再退出
  await ch.waitForConfirms();
  console.log('[Producer] 全部发送完成，关闭连接');
  await ch.close();
  await conn.close();
})();

worker

也可以单独起

启动单个 Worker

set WORKER_CFG={"name":"GPU-1","queue":"gpu_task","concurrency":2} && node worker.js
启动第二个 Worker（另一窗口）

set WORKER_CFG={"name":"GPU-2","queue":"gpu_task","concurrency":2} && node worker.js

javascript 复制代码

#!/usr/bin/env node
// worker.js
const amqp      = require('amqplib');
const fs        = require('fs');
const path      = require('path');
const pool      = require('tiny-async-pool');

// 1. 读配置
const cfg = JSON.parse(process.env.WORKER_CFG || '{}');
const {
  name        = 'Worker',
  queue       = 'task_queue',
  prefetch    = 1,
  heartbeat   = 5000,
  concurrency = 1,
  logFile     = `${name}.log`,
  beatFile    = `${name}.beat`,
  taskTimeout = 60000            // 单任务超时 60 s
} = cfg;

// 2. 日志函数（同时写文件+控制台）
function log(msg) {
  const line = `[${name}] ${new Date().toISOString()}  ${msg}\n`;
  console.log(line.slice(0, -1));
  fs.mkdirSync(path.dirname(logFile), { recursive: true });
  fs.appendFileSync(logFile, line);
}

// 3. 心跳
fs.mkdirSync(path.dirname(beatFile), { recursive: true });
const heartTimer = setInterval(() => {
  fs.writeFileSync(beatFile, Date.now().toString());
}, heartbeat);

// 4. 伪模型转换（替换成真正脚本即可）
async function fakeConvert(taskId) {
  const sec = 3 + Math.floor(Math.random() * 5); // 3~7 秒
  log(`>>> 开始任务 ${taskId}，预计 ${sec}s`);
  await new Promise(r => setTimeout(r, sec * 1000));
  log(`<<< 完成任务 ${taskId}`);
}

// 5. 建立连接 & 通道
let conn, ch;
async function connect() {
  conn = await amqp.connect('amqp://localhost');
  ch   = await conn.createChannel();
  await ch.assertQueue(queue, { durable: true });
  await ch.prefetch(prefetch);
  log('已连接 RabbitMQ，等待任务...');
}

// 6. 消费逻辑（带并发池 + 超时）
async function consume() {
  await ch.consume(queue, async (msg) => {
    if (msg === null) return;
    const taskId = msg.content.toString();
    try {
      // 用 async 池同时跑 concurrency 个
      await pool(concurrency, [taskId], (id) =>
        Promise.race([
          fakeConvert(id),
          new Promise((_, rej) =>
            setTimeout(() => rej(new Error('TASK_TIMEOUT')), taskTimeout)
          )
        ])
      );
      ch.ack(msg);
      log(`任务 ${taskId} ack 成功`);
    } catch (e) {
      log(`任务 ${taskId} 失败：${e.message}，重回队列`);
      ch.nack(msg, false, true); // requeue=true
    }
  }, { noAck: false });
}

// 7. 优雅退出
async function gracefulShutdown() {
  log('收到 SIGINT，开始优雅退出...');
  clearInterval(heartTimer);
  if (ch)  await ch.close();
  if (conn) await conn.close();
  process.exit(0);
}
process.on('SIGINT', gracefulShutdown);

// 8. 启动
(async () => {
  try {
    await connect();
    await consume();
  } catch (e) {
    log('启动失败：' + e.message);
    process.exit(1);
  }
})();

launcher

用来控制worker示例的启动，停止，重启

javascript 复制代码

#!/usr/bin/env node
// launcher.js
const fs   = require('fs');
const path = require('path');
const { spawn } = require('child_process');

const CFG_FILE = './config.json';
const PID_FILE = './launcher.pid';

const action = process.argv[2] || 'start';

function readCfg() {
  if (!fs.existsSync(CFG_FILE)) throw 'config.json 不存在';
  return JSON.parse(fs.readFileSync(CFG_FILE));
}

function start() {
  if (fs.existsSync(PID_FILE)) {
    console.log('已有 launcher 在运行，请先 stop');
    return;
  }
  const cfg = readCfg();
  const children = [];
  cfg.forEach(c => {
    // 确保目录存在
    fs.mkdirSync(path.dirname(c.logFile), { recursive: true });
    fs.mkdirSync(path.dirname(c.beatFile), { recursive: true });
    const env = { ...process.env, WORKER_CFG: JSON.stringify(c) };
    const child = spawn('node', ['worker.js'], { env, stdio: 'inherit' });
    children.push(child);
  });
  // 把子进程 pid 写文件，方便 stop
  fs.writeFileSync(PID_FILE, JSON.stringify(children.map(c => c.pid)));
  console.log(`已启动 ${children.length} 个 Worker`);
}

function stop() {
  if (!fs.existsSync(PID_FILE)) return;
  const pids = JSON.parse(fs.readFileSync(PID_FILE));
  pids.forEach(pid => {
    try { process.kill(pid, 'SIGINT'); } catch (e) {}
  });
  fs.unlinkSync(PID_FILE);
  console.log('已发送停止信号');
}

function restart() {
  stop();
  setTimeout(() => start(), 1000);
}

switch (action) {
  case 'start':  start();  break;
  case 'stop':   stop();   break;
  case 'restart':restart();break;
  default: console.log('用法: node launcher.js [start|stop|restart]');
}

cfg配置文件

可配置多个实例

javascript 复制代码

[
  {
    "name": "GPU-Worker-1",
    "queue": "gpu_task",
    "prefetch": 1,
    "heartbeat": 5000,
    "concurrency": 2,
    "logFile": "logs/GPU-1.log",
    "beatFile": "beats/GPU-1.beat"
  },
  {
    "name": "GPU-Worker-2",
    "queue": "gpu_task",
    "prefetch": 1,
    "heartbeat": 5000,
    "concurrency": 2,
    "logFile": "logs/GPU-2.log",
    "beatFile": "beats/GPU-2.beat"
  },
  {
    "name": "CPU-Worker-1",
    "queue": "cpu_task",
    "prefetch": 3,
    "heartbeat": 10000,
    "concurrency": 4,
    "logFile": "logs/CPU-1.log",
    "beatFile": "beats/CPU-1.beat"
  }
]

演示

思考

任务的状态，进度存储--redis

worker的日志-上传S3，方便管理模块查看下载

各个worker的控制-launcher

多台服务器搭局域网互通