wakaq
Version:
Background task queue for Node backed by Redis, a super minimal Celery
318 lines (317 loc) • 13.7 kB
JavaScript
"use strict";
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
var desc = Object.getOwnPropertyDescriptor(m, k);
if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
desc = { enumerable: true, get: function() { return m[k]; } };
}
Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
o["default"] = v;
});
var __importStar = (this && this.__importStar) || function (mod) {
if (mod && mod.__esModule) return mod;
var result = {};
if (mod != null) for (var k in mod) if (k !== "default" && Object.prototype.hasOwnProperty.call(mod, k)) __createBinding(result, mod, k);
__setModuleDefault(result, mod);
return result;
};
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.WakaQWorker = void 0;
const node_child_process_1 = require("node:child_process");
const node_process_1 = __importDefault(require("node:process"));
const os = __importStar(require("os"));
const pidusage_1 = __importDefault(require("pidusage"));
const ts_duration_1 = require("ts-duration");
const child_js_1 = require("./child.js");
const logger_js_1 = require("./logger.js");
const serializer_js_1 = require("./serializer.js");
class WakaQWorker {
constructor(wakaq, childWorkerCommand) {
this.children = [];
this._stopProcessing = false;
this._numTasksProcessed = 0;
this.wakaq = wakaq;
if (childWorkerCommand.length == 0)
throw Error('Missing child worker command.');
this.childWorkerCommand = childWorkerCommand.shift() ?? '';
this.childWorkerArgs = childWorkerCommand;
this.logger = (0, logger_js_1.setupLogging)(this.wakaq);
this.wakaq.logger = this.logger;
}
async start() {
if (!this.wakaq.singleProcess) {
this.logger.info(`concurrency=${this.wakaq.concurrency}`);
this.logger.info(`soft_timeout=${this.wakaq.softTimeout.seconds}`);
this.logger.info(`hard_timeout=${this.wakaq.hardTimeout.seconds}`);
}
this.logger.info(`wait_timeout=${this.wakaq.waitTimeout.seconds}`);
this.logger.info(`exclude_queues=${this.wakaq.excludeQueues}`);
this.logger.info(`max_retries=${this.wakaq.maxRetries}`);
this.logger.info(`max_mem_percent=${this.wakaq.maxMemPercent}`);
this.logger.info(`max_tasks_per_worker=${this.wakaq.maxTasksPerWorker}`);
this.logger.info(`worker_log_file=${this.wakaq.workerLogFile}`);
this.logger.info(`worker_log_level=${this.wakaq.workerLogLevel}`);
if (this.wakaq.singleProcess) {
this.logger.info('running in single process mode...');
}
else {
this.logger.info(`starting ${this.wakaq.concurrency} workers...`);
}
const _this = this;
node_process_1.default.on('SIGINT', () => _this._onExitParent());
node_process_1.default.on('SIGTERM', () => _this._onExitParent());
node_process_1.default.on('SIGQUIT', () => _this._onExitParent());
if (!this.wakaq.singleProcess) {
for (let i = 0; i < this.wakaq.concurrency; i++) {
this._spawnChild();
}
this.logger.info('finished spawning all workers');
}
if (this.wakaq.singleProcess && this.wakaq.afterWorkerStartedCallback)
await this.wakaq.afterWorkerStartedCallback();
try {
await this.wakaq.connect();
const pubsub = await this.wakaq.pubsub();
pubsub.on('message', this._handleBroadcastTask);
await pubsub.subscribe(this.wakaq.broadcastKey, (err) => {
if (err)
this.logger.error(`Failed to subscribe to broadcast tasks: ${err.message}`);
});
this._numTasksProcessed = 0;
while (!this._stopProcessing) {
this._respawnMissingChildren();
await this._enqueueReadyEtaTasks();
if (this.wakaq.singleProcess)
await this._processTasksSingleProcessMode();
await this._checkChildMemoryUsages();
this._checkChildRuntimes();
await this.wakaq.sleep(ts_duration_1.Duration.millisecond(500));
}
this.logger.info('shutting down...');
if (this.children.length > 0) {
while (this.children.length > 0) {
this._stop();
await this._checkChildMemoryUsages();
this._checkChildRuntimes();
await this.wakaq.sleep(ts_duration_1.Duration.millisecond(500));
}
}
}
catch (error) {
this.logger.error(error);
this._stop();
}
finally {
this.wakaq.disconnect();
}
}
_spawnChild() {
const t = this;
this.logger.debug(`spawning child worker: ${this.childWorkerCommand} ${this.childWorkerArgs.join(' ')}`);
const p = (0, node_child_process_1.spawn)(this.childWorkerCommand, this.childWorkerArgs, {
stdio: [null, null, null, 'ipc'],
});
const child = new child_js_1.Child(this.wakaq, p);
p.on('close', (code) => {
t._onChildExited(child, code);
});
p.on('message', (message) => {
t._onMessageReceivedFromChild(child, message);
});
p.stdout?.on('data', (data) => {
t._onOutputReceivedFromChild(child, data);
});
p.stderr?.on('data', (data) => {
t._onOutputReceivedFromChild(child, data);
});
this.children.push(child);
}
async _processTasksSingleProcessMode() {
const { queueBrokerKey, payload } = await this.wakaq.blockingDequeue();
if (queueBrokerKey !== undefined && payload !== undefined) {
const task = this.wakaq.tasks.get(payload.name);
if (!task && payload.name)
this.logger.error(`Task not found: ${payload.name}`);
if (task) {
const queue = this.wakaq.queuesByKey.get(queueBrokerKey);
this.wakaq.currentTask = task;
const retry = payload.retry ?? 0;
this.logger.debug(`working on task ${task.name}`);
try {
await this._executeTask(task, payload.args, queue);
}
catch (error) {
const maxRetries = task.maxRetries ?? queue?.maxRetries ?? this.wakaq.maxRetries;
if (retry + 1 > maxRetries) {
this.logger.error(error);
}
else {
this.logger.warning(error);
this.wakaq.enqueueAtEnd(task.name, payload.args, queue, retry);
}
}
finally {
this.wakaq.currentTask = undefined;
}
}
}
if (this.wakaq.maxTasksPerWorker && this._numTasksProcessed >= this.wakaq.maxTasksPerWorker) {
this.logger.info(`exiting single process worker after ${this._numTasksProcessed} tasks`);
this._stopProcessing = true;
}
}
_stop() {
this._stopProcessing = true;
this._stopAllChildren();
}
_stopAllChildren() {
this.children.forEach((child) => {
child.sigterm();
});
}
_onExitParent() {
this._stop();
}
_onChildExited(child, code) {
this.logger.debug(`child process ${child.process.pid} exited: ${code}`);
this.children = this.children.filter((c) => c !== child);
}
_onOutputReceivedFromChild(child, data) {
if (data instanceof Buffer)
data = data.toString();
if (!data)
return;
child.outputBuffer = `${child.outputBuffer}${data}`;
let i = -1;
while ((i = child.outputBuffer.indexOf('\n')) && i > -1) {
const payload = child.outputBuffer.slice(0, i);
child.outputBuffer = child.outputBuffer.slice(i + 1);
if (payload.length > 0) {
try {
const parsed = JSON.parse(payload);
if (parsed?.level) {
this.logger.log(parsed.level, parsed.message, { worker: child.process.pid, payload: parsed.payload });
}
else {
this.logger.info(payload);
}
}
catch (e) {
this.logger.info(payload);
}
}
}
}
_onMessageReceivedFromChild(child, message) {
if (typeof message !== 'object')
return;
const payload = message;
if (payload.type !== 'wakaq-ping')
return;
child.lastPing = Math.round(Date.now() / 1000);
const taskName = payload.task;
const task = taskName ? this.wakaq.tasks.get(taskName) : undefined;
const queueName = payload.queue;
const queue = this.wakaq.queues.find((q) => {
return q.name === queueName;
});
child.setTimeouts(this.wakaq, task, queue);
child.softTimeoutReached = false;
}
async _enqueueReadyEtaTasks() {
await Promise.all(this.wakaq.queues.map(async (q) => {
const results = await this.wakaq.broker.getetatasks(q.brokerEtaKey, String(Math.round(Date.now() / 1000)));
await Promise.all(results.map(async (result) => {
const payload = (0, serializer_js_1.deserialize)(result);
const taskName = payload.name;
const args = payload.args;
await this.wakaq.enqueueAtFront(taskName, args, q);
}));
}));
}
async _checkChildMemoryUsages() {
if (!this.wakaq.maxMemPercent)
return;
const totalMem = os.totalmem();
const percent = ((totalMem - os.freemem()) / totalMem) * 100;
if (percent < this.wakaq.maxMemPercent)
return;
if (this.wakaq.singleProcess) {
this.logger.info('stopping single process worker from too much ram usage');
this._stop();
return;
}
const usages = await Promise.all(this.children.map(async (child) => (await (0, pidusage_1.default)(child.process.pid ?? 0)).memory || 0));
const maxIndex = usages.reduce((iMax, x, i, arr) => (x > (arr[iMax] ?? 0) ? i : iMax), 0);
const child = this.children.at(maxIndex);
if (!child)
return;
this.logger.info(`killing child ${child.process.pid} from too much ram usage`);
child.sigterm();
}
_checkChildRuntimes() {
this.children.forEach((child) => {
const softTimeout = child.softTimeout || this.wakaq.softTimeout;
const hardTimeout = child.hardTimeout || this.wakaq.hardTimeout;
if (softTimeout || hardTimeout) {
const now = Math.round(Date.now() / 1000);
const runtime = ts_duration_1.Duration.second(now - child.lastPing);
if (hardTimeout && runtime.seconds > hardTimeout.seconds) {
this.logger.debug(`child process ${child.process.pid} runtime ${runtime} reached hard timeout, sending sigkill`);
child.sigkill();
}
else if (!child.softTimeoutReached && softTimeout && runtime.seconds > softTimeout.seconds) {
this.logger.debug(`child process ${child.process.pid} runtime ${runtime} reached soft timeout, sending sigquit`);
child.softTimeoutReached = true;
child.sigquit();
}
}
});
}
_handleBroadcastTask(channel, message) {
const child = this.children.at(0);
if (!child) {
this.logger.error(`Unable to run broadcast task because no available child workers: ${message}`);
return;
}
this.logger.debug(`run broadcast task: ${message}`);
child.process.stdin?.write(`${message}\n`, (err) => {
if (err)
this.logger.error(`Unable to run broadcast task because writing to child stdin encountered an error: ${err}`);
});
}
_respawnMissingChildren() {
if (this._stopProcessing)
return;
if (this.wakaq.singleProcess)
return;
for (let i = this.wakaq.concurrency - this.children.length; i > 0; i--) {
this.logger.debug('restarting a crashed worker');
this._spawnChild();
}
}
async _executeTask(task, variables, queue) {
this.logger.debug(`running with args ${variables}`);
if (this.wakaq.beforeTaskStartedCallback)
this.wakaq.beforeTaskStartedCallback(task);
try {
await task.fn(variables);
}
finally {
this._numTasksProcessed += 1;
if (this.wakaq.afterTaskFinishedCallback)
this.wakaq.afterTaskFinishedCallback(task);
}
}
}
exports.WakaQWorker = WakaQWorker;