John
102 lines
2.9 KiB
JavaScript
Executable File
102 lines
2.9 KiB
JavaScript
Executable File
#!/usr/bin/env node
|
|
import fs from 'node:fs';
|
|
import path from 'node:path';
|
|
import { createClient } from 'redis';
|
|
|
|
function loadEnvFile(filePath) {
|
|
if (!fs.existsSync(filePath)) return;
|
|
for (const line of fs.readFileSync(filePath, 'utf8').split('\n')) {
|
|
const trimmed = line.trim();
|
|
if (!trimmed || trimmed.startsWith('#')) continue;
|
|
const idx = trimmed.indexOf('=');
|
|
if (idx < 0) continue;
|
|
const key = trimmed.slice(0, idx).trim();
|
|
const value = trimmed.slice(idx + 1).trim();
|
|
if (!process.env[key]) process.env[key] = value;
|
|
}
|
|
}
|
|
|
|
loadEnvFile(path.join(process.cwd(), '.env'));
|
|
|
|
const redisUrl = process.env.MEMIND_RUNTIME_REDIS_URL || 'redis://127.0.0.1:6379/0';
|
|
const namespace = process.env.MEMIND_RUNTIME_REDIS_NAMESPACE || 'memind:runtime';
|
|
const configuredWorkers = (process.env.TKMIND_API_TARGETS || process.env.TKMIND_API_TARGET || '')
|
|
.split(',')
|
|
.map((value) => value.trim())
|
|
.filter(Boolean)
|
|
.map((_, index) => `goosed-${index + 1}`);
|
|
const action = process.argv[2] || 'status';
|
|
const workerId = process.argv[3] || null;
|
|
|
|
function usage() {
|
|
console.error('Usage: node scripts/runtime-worker-drain.mjs <status|drain|undrain> [goosed-N]');
|
|
}
|
|
|
|
function workerKey(id, field) {
|
|
return [namespace, 'worker', id, field].join(':');
|
|
}
|
|
|
|
if (!['status', 'drain', 'undrain'].includes(action)) {
|
|
usage();
|
|
process.exit(2);
|
|
}
|
|
if (['drain', 'undrain'].includes(action) && !workerId) {
|
|
usage();
|
|
process.exit(2);
|
|
}
|
|
|
|
const client = createClient({ url: redisUrl });
|
|
client.on('error', (err) => {
|
|
console.error(`Redis error: ${err instanceof Error ? err.message : err}`);
|
|
});
|
|
await client.connect();
|
|
|
|
if (action === 'drain') {
|
|
await client.set(workerKey(workerId, 'drain'), '1');
|
|
}
|
|
if (action === 'undrain') {
|
|
await client.del(workerKey(workerId, 'drain'));
|
|
}
|
|
|
|
const keys = await client.keys(workerKey('*', 'active_streams'));
|
|
const workers = [
|
|
...configuredWorkers,
|
|
...keys
|
|
.map((key) => key.split(':').at(-2))
|
|
.filter(Boolean),
|
|
];
|
|
if (workerId && !workers.includes(workerId)) workers.push(workerId);
|
|
|
|
const rows = [];
|
|
for (const id of [...new Set(workers)].sort()) {
|
|
const values = await client.mGet([
|
|
workerKey(id, 'active_streams'),
|
|
workerKey(id, 'drain'),
|
|
workerKey(id, 'stream_open_count'),
|
|
workerKey(id, 'stream_abort_count'),
|
|
workerKey(id, 'stream_error_count'),
|
|
workerKey(id, 'last_stream_started_at'),
|
|
workerKey(id, 'last_stream_ended_at'),
|
|
]);
|
|
rows.push({
|
|
id,
|
|
activeStreams: Number(values[0] || 0),
|
|
drain: /^(1|true|yes)$/i.test(String(values[1] || '')),
|
|
streamOpenCount: Number(values[2] || 0),
|
|
streamAbortCount: Number(values[3] || 0),
|
|
streamErrorCount: Number(values[4] || 0),
|
|
lastStreamStartedAt: values[5] ? Number(values[5]) : null,
|
|
lastStreamEndedAt: values[6] ? Number(values[6]) : null,
|
|
});
|
|
}
|
|
|
|
await client.quit();
|
|
|
|
console.log(JSON.stringify({
|
|
ok: true,
|
|
action,
|
|
workerId,
|
|
namespace,
|
|
workers: rows,
|
|
}, null, 2));
|