murm/mrmBot-Matrix
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Attempt to move image processing to an eris-fleet service

Browse source
This commit is contained in:
Essem 2021-07-05 19:20:21 -05:00
parent 736207075b
commit 97050f0cf1
No known key found for this signature in database
GPG key ID: 2502A99EDC3F6FB9
8 changed files with 302 additions and 306 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
api/index.js
View file

@ -1,6 +1,5 @@
require("dotenv").config();
const os = require("os");
//const { run } = require("../utils/image-runner.js");
const { Worker } = require("worker_threads");
const path = require("path");
const http = require("http");

16
app.js
View file

@ -51,6 +51,9 @@ const Admiral = new Fleet({
name: "Starting esmBot..."
}
},
whatToLog: {
blacklist: ["stats_update"]
},
clientOptions: {
disableEvents: {
CHANNEL_DELETE: true,
@ -82,7 +85,11 @@ const Admiral = new Fleet({
stats: {
requestTimeout: 30000
}
}
},
services: [
//{ name: "prometheus", path: path.join(__dirname, "./utils/services/prometheus.js") },
{ name: "image", path: path.join(__dirname, "./utils/services/image.js")}
]
});
if (isMaster) {
@ -119,12 +126,9 @@ if (isMaster) {
warn: "yellow",
error: "red"
});
Admiral.on("log", (m) => logger.main(m));
Admiral.on("info", (m) => {
console.log("hi");
logger.info(m);
});
Admiral.on("info", (m) => logger.info(m));
Admiral.on("debug", (m) => logger.debug(m));
Admiral.on("warn", (m) => logger.warn(m));
Admiral.on("error", (m) => logger.error(m));

5
classes/imageCommand.js
View file

@ -1,5 +1,4 @@
const Command = require("./command.js");
const magick = require("../utils/image.js");
const imageDetect = require("../utils/imagedetect.js");
const collections = require("../utils/collections.js");
const { emotes } = require("../messages.json");
@ -99,12 +98,12 @@ class ImageCommand extends Command {
}
try {
const { buffer, type } = await magick.run(magickParams).catch(e => {
const { buffer, type } = await this.ipc.command("image", { type: "run", obj: magickParams }, true).catch(e => {
throw e;
});
if (type === "nogif" && this.constructor.requiresGIF) return "That isn't a GIF!";
return {
file: buffer,
file: Buffer.from(buffer.data),
name: `${this.constructor.command}.${type}`
};
} catch (e) {

14
commands/general/imagereload.js
View file

@ -1,21 +1,9 @@
const image = require("../../utils/image.js");
const logger = require("../../utils/logger.js");
const Command = require("../../classes/command.js");
class ImageReloadCommand extends Command {
async run() {
if (this.message.author.id !== process.env.OWNER) return "Only the bot owner can reload the image servers!";
await image.disconnect();
await image.repopulate();
let amount = 0;
for (const server of image.servers) {
try {
await image.connect(server);
amount += 1;
} catch (e) {
logger.error(e);
}
}
const amount = await this.ipc.command("image", { type: "reload" }, true);
if (amount > 0) {
return `Successfully connected to ${amount} image servers.`;
} else {

5
commands/general/imagestats.js
View file

@ -1,8 +1,8 @@
const image = require("../../utils/image.js");
const Command = require("../../classes/command.js");
class ImageStatsCommand extends Command {
async run() {
const servers = await this.ipc.command("image", { type: "stats" }, true);
const embed = {
embed: {
"author": {
@ -10,11 +10,10 @@ class ImageStatsCommand extends Command {
"icon_url": this.client.user.avatarURL
},
"color": 16711680,
"description": `The bot is currently connected to ${image.connections.size} image server(s).`,
"description": `The bot is currently connected to ${servers.length} image server(s).`,
"fields": []
}
};
const servers = await image.getRunning();
for (let i = 0; i < servers.length; i++) {
embed.embed.fields.push({
name: `Server ${i + 1}`,

51
shard.js
View file

@ -8,8 +8,6 @@ const logger = require("./utils/logger.js");
const handler = require("./utils/handler.js");
// lavalink stuff
const sound = require("./utils/soundplayer.js");
// image processing stuff
const image = require("./utils/image.js");
// database stuff
const database = require("./utils/database.js");
// command collections
@ -21,7 +19,6 @@ const misc = require("./utils/misc.js");
// generate help page
const helpGenerator =
process.env.OUTPUT !== "" ? require("./utils/help.js") : null;
const http = require("http");
class Shard extends BaseClusterWorker {
constructor(bot) {
@ -53,17 +50,6 @@ class Shard extends BaseClusterWorker {
this.bot.on(eventName, event.bind(null, this.bot, this.clusterID, this.workerID, this.ipc));
}
// connect to image api if enabled
if (process.env.API === "true") {
for (const server of image.servers) {
try {
await image.connect(server);
} catch (e) {
logger.error(e);
}
}
}
// generate docs
if (helpGenerator) {
await helpGenerator.generateList();
@ -73,41 +59,6 @@ class Shard extends BaseClusterWorker {
}
}
if (process.env.METRICS !== "" && process.env.METRICS !== undefined) {
const httpServer = http.createServer(async (req, res) => {
if (req.method !== "GET") {
res.statusCode = 405;
return res.end("GET only");
}
res.write(`# HELP connected_workers Number of workers connected
# TYPE connected_workers gauge
connected_workers ${image.connections.size}
# HELP running_jobs Number of running jobs on this worker
# TYPE running_jobs gauge
# HELP queued_jobs Number of queued jobs on this worker
# TYPE queued_jobs gauge
# HELP max_jobs Number of max allowed jobs on this worker
# TYPE max_jobs gauge
# HELP command_count Number of times a command has been run
# TYPE command_count counter
`);
const servers = await image.getRunning();
for (const [i, w] of servers.entries()) {
res.write(`running_jobs{worker="${i}"} ${w.runningJobs}\n`);
res.write(`queued_jobs{worker="${i}"} ${w.queued}\n`);
res.write(`max_jobs{worker="${i}"} ${w.max}\n`);
}
const counts = await database.getCounts();
for (const [i, w] of Object.entries(counts)) {
res.write(`command_count{command="${i}"} ${w}\n`);
}
res.end();
});
httpServer.listen(process.env.METRICS, () => {
logger.log("info", `Serving metrics at ${process.env.METRICS}`);
});
}
this.ipc.register("reload", async (message) => {
const result = await handler.unload(message.cmd);
if (result) return this.ipc.broadcast("reloadFail", { result: result });
@ -156,7 +107,7 @@ connected_workers ${image.connections.size}
}
shutdown(done) {
logger.log("warn", "SIGINT detected, shutting down...");
logger.log("warn", "Shutting down...");
this.bot.editStatus("dnd", {
name: "Restarting/shutting down..."
});

230
utils/image.js
View file

@ -1,12 +1,6 @@
const magick = require("../build/Release/image.node");
const { Worker } = require("worker_threads");
const fetch = require("node-fetch");
const fs = require("fs");
const WebSocket = require("ws");
const fileType = require("file-type");
const path = require("path");
const { EventEmitter } = require("events");
const logger = require("./logger.js");
const formats = ["image/jpeg", "image/png", "image/webp", "image/gif", "video/mp4", "video/webm", "video/mov"];
@ -16,183 +10,6 @@ exports.connections = new Map();
exports.servers = JSON.parse(fs.readFileSync("./servers.json", { encoding: "utf8" })).image;
const chooseServer = async (ideal) => {
if (ideal.length === 0) throw "No available servers";
const sorted = ideal.sort((a, b) => {
return b.load - a.load;
});
return sorted[0];
};
exports.repopulate = async () => {
const data = await fs.promises.readFile("./servers.json", { encoding: "utf8" });
this.servers = JSON.parse(data).image;
return;
};
exports.getRunning = async () => {
let serversLeft = this.connections.size;
const statuses = [];
for (const address of this.connections.keys()) {
const connection = this.connections.get(address);
if (connection.readyState !== 0 && connection.readyState !== 1) {
serversLeft--;
continue;
}
const controller = new AbortController(); // eslint-disable-line no-undef
const timeout = setTimeout(() => {
controller.abort();
}, 2000);
try {
const statusRequest = await fetch(`http://${address}:8080/running`, { signal: controller.signal });
clearTimeout(timeout);
const status = await statusRequest.json();
serversLeft--;
statuses.push(status);
} catch (e) {
if (e.name === "AbortError") {
serversLeft--;
continue;
} else if (e.code === "ECONNREFUSED") {
serversLeft--;
continue;
}
throw e;
}
}
if (!serversLeft) {
return statuses;
} else {
throw new Error("Loop ended before all servers could be checked");
}
};
exports.connect = async (server) => {
const connection = new WebSocket(`ws://${server}:8080/sock`);
connection.on("message", async (msg) => {
const opcode = msg.readUint8(0);
const req = msg.slice(37, msg.length);
const uuid = msg.slice(1, 37).toString();
if (opcode === 0x00) { // Job queued
if (this.jobs[req]) {
this.jobs[req].event.emit("uuid", uuid);
}
} else if (opcode === 0x01) { // Job completed successfully
// the image API sends all job responses over the same socket; make sure this is ours
if (this.jobs[uuid]) {
const imageReq = await fetch(`http://${server}:8080/image?id=${uuid}`);
const image = await imageReq.buffer();
// The response data is given as the file extension/ImageMagick type of the image (e.g. "png"), followed
// by a newline, followed by the image data.
this.jobs[uuid].event.emit("image", image, imageReq.headers.get("ext"));
}
} else if (opcode === 0x02) { // Job errored
if (this.jobs[uuid]) {
this.jobs[uuid].event.emit("error", new Error(req));
}
}
});
connection.on("error", (e) => {
logger.error(e.toString());
});
connection.once("close", () => {
for (const uuid of Object.keys(this.jobs)) {
if (this.jobs[uuid].addr === server) {
this.jobs[uuid].event.emit("error", "Job ended prematurely due to a closed connection; please run your image job again");
delete this.jobs[uuid];
}
}
//logger.log(`Lost connection to ${server}, attempting to reconnect...`);
this.connections.delete(server);
});
this.connections.set(server, connection);
};
exports.disconnect = async () => {
for (const connection of this.connections.values()) {
connection.close();
}
for (const uuid of Object.keys(this.jobs)) {
this.jobs[uuid].event.emit("error", "Job ended prematurely (not really an error; just run your image job again)");
delete this.jobs[uuid];
}
this.connections.clear();
return;
};
const getIdeal = async () => {
let serversLeft = this.connections.size;
if (serversLeft < this.servers.length) {
for (const server of this.servers) {
try {
if (!this.connections.has(server)) await this.connect(server);
} catch (e) {
logger.error(e);
}
}
serversLeft = this.connections.size;
}
const idealServers = [];
for (const address of this.connections.keys()) {
const connection = this.connections.get(address);
if (connection.readyState !== 0 && connection.readyState !== 1) {
serversLeft--;
continue;
}
const controller = new AbortController(); // eslint-disable-line no-undef
const timeout = setTimeout(() => {
controller.abort();
}, 5000);
try {
const statusRequest = await fetch(`http://${address}:8080/status`, { signal: controller.signal });
clearTimeout(timeout);
const status = await statusRequest.text();
serversLeft--;
idealServers.push({
addr: address,
load: parseInt(status)
});
} catch (e) {
if (e.name === "AbortError") {
serversLeft--;
continue;
} else if (e.code === "ECONNREFUSED") {
serversLeft--;
continue;
}
throw e;
} finally {
clearTimeout(timeout);
}
}
if (!serversLeft) {
const server = await chooseServer(idealServers);
return { addr: server.addr, sock: this.connections.get(server.addr) };
} else {
throw new Error("Loop ended before all servers could be checked");
}
};
const start = async (object, num) => {
const currentServer = await getIdeal();
const data = Buffer.concat([Buffer.from([0x01 /* queue job */]), Buffer.from(num.length.toString()), Buffer.from(num), Buffer.from(JSON.stringify(object))]);
currentServer.sock.send(data);
const event = new EventEmitter();
this.jobs[num] = { event, addr: currentServer.addr };
const uuid = await new Promise((resolve, reject) => {
event.once("uuid", (uuid) => resolve(uuid));
event.once("error", reject);
});
delete this.jobs[num];
this.jobs[uuid] = { event: event, addr: currentServer.addr };
return { uuid: uuid, event: event };
};
exports.check = (cmd) => {
return magick[cmd] ? true : false;
};
exports.getType = async (image, extraReturnTypes) => {
if (!image.startsWith("http")) {
const imageType = await fileType.fromFile(image);
@ -234,50 +51,3 @@ exports.getType = async (image, extraReturnTypes) => {
}
return type;
};
exports.run = (object) => {
return new Promise((resolve, reject) => {
if (process.env.API === "true") {
// Connect to best image server
const num = Math.floor(Math.random() * 100000).toString().slice(0, 5);
const timeout = setTimeout(() => {
if (this.jobs[num]) delete this.jobs[num];
reject("the image request timed out after 25 seconds. Try uploading your image elsewhere.");
}, 25000);
start(object, num).catch(err => { // incredibly hacky code incoming
clearTimeout(timeout);
if (err instanceof Error) return reject(err);
return err;
}).then((data) => {
clearTimeout(timeout);
if (!data.event) reject("Not connected to image server");
data.event.once("image", (image, type) => {
delete this.jobs[data.uuid];
const payload = {
// Take just the image data
buffer: image,
type: type
};
resolve(payload);
});
data.event.once("error", (err) => {
delete this.jobs[data.uuid];
reject(err);
});
return;
}).catch(err => reject(err));
} else {
// Called from command (not using image API)
const worker = new Worker(path.join(__dirname, "image-runner.js"), {
workerData: object
});
worker.once("message", (data) => {
resolve({
buffer: Buffer.from([...data.buffer]),
type: data.fileExtension
});
});
worker.once("error", reject);
}
});
};

286
utils/services/image.js Normal file
View file

@ -0,0 +1,286 @@
const { BaseServiceWorker } = require("eris-fleet");
const logger = require("../logger.js");
const fetch = require("node-fetch");
const WebSocket = require("ws");
const fs = require("fs");
const path = require("path");
const { Worker } = require("worker_threads");
const { EventEmitter } = require("events");
class ImageWorker extends BaseServiceWorker {
constructor(setup) {
super(setup);
if (process.env.API === "true") {
this.jobs = {};
this.connections = new Map();
this.servers = JSON.parse(fs.readFileSync("./servers.json", { encoding: "utf8" })).image;
}
this.begin().then(() => this.serviceReady());
}
async begin() {
// connect to image api if enabled
if (process.env.API === "true") {
for (const server of this.servers) {
try {
await this.connect(server);
} catch (e) {
logger.error(e);
}
}
}
}
async repopulate() {
const data = await fs.promises.readFile("./servers.json", { encoding: "utf8" });
this.servers = JSON.parse(data).image;
return;
}
async getRunning() {
let serversLeft = this.connections.size;
const statuses = [];
for (const address of this.connections.keys()) {
const connection = this.connections.get(address);
if (connection.readyState !== 0 && connection.readyState !== 1) {
serversLeft--;
continue;
}
const controller = new AbortController(); // eslint-disable-line no-undef
const timeout = setTimeout(() => {
controller.abort();
}, 2000);
try {
const statusRequest = await fetch(`http://${address}:8080/running`, { signal: controller.signal });
clearTimeout(timeout);
const status = await statusRequest.json();
serversLeft--;
statuses.push(status);
} catch (e) {
if (e.name === "AbortError") {
serversLeft--;
continue;
} else if (e.code === "ECONNREFUSED") {
serversLeft--;
continue;
}
throw e;
}
}
if (!serversLeft) {
return statuses;
} else {
throw new Error("Loop ended before all servers could be checked");
}
}
async chooseServer(ideal) {
if (ideal.length === 0) throw "No available servers";
const sorted = ideal.sort((a, b) => {
return b.load - a.load;
});
return sorted[0];
}
async getIdeal() {
let serversLeft = this.connections.size;
if (serversLeft < this.servers.length) {
for (const server of this.servers) {
try {
if (!this.connections.has(server)) await this.connect(server);
} catch (e) {
logger.error(e);
}
}
serversLeft = this.connections.size;
}
const idealServers = [];
for (const address of this.connections.keys()) {
const connection = this.connections.get(address);
if (connection.readyState !== 0 && connection.readyState !== 1) {
serversLeft--;
continue;
}
const controller = new AbortController(); // eslint-disable-line no-undef
const timeout = setTimeout(() => {
controller.abort();
}, 5000);
try {
const statusRequest = await fetch(`http://${address}:8080/status`, { signal: controller.signal });
clearTimeout(timeout);
const status = await statusRequest.text();
serversLeft--;
idealServers.push({
addr: address,
load: parseInt(status)
});
} catch (e) {
if (e.name === "AbortError") {
serversLeft--;
continue;
} else if (e.code === "ECONNREFUSED") {
serversLeft--;
continue;
}
throw e;
} finally {
clearTimeout(timeout);
}
}
if (!serversLeft) {
const server = await this.chooseServer(idealServers);
return { addr: server.addr, sock: this.connections.get(server.addr) };
} else {
throw new Error("Loop ended before all servers could be checked");
}
}
async connect(server) {
const connection = new WebSocket(`ws://${server}:8080/sock`);
connection.on("message", async (msg) => {
const opcode = msg.readUint8(0);
const req = msg.slice(37, msg.length);
const uuid = msg.slice(1, 37).toString();
if (opcode === 0x00) { // Job queued
if (this.jobs[req]) {
this.jobs[req].event.emit("uuid", uuid);
}
} else if (opcode === 0x01) { // Job completed successfully
// the image API sends all job responses over the same socket; make sure this is ours
if (this.jobs[uuid]) {
const imageReq = await fetch(`http://${server}:8080/image?id=${uuid}`);
const image = await imageReq.buffer();
// The response data is given as the file extension/ImageMagick type of the image (e.g. "png"), followed
// by a newline, followed by the image data.
this.jobs[uuid].event.emit("image", image, imageReq.headers.get("ext"));
}
} else if (opcode === 0x02) { // Job errored
if (this.jobs[uuid]) {
this.jobs[uuid].event.emit("error", new Error(req));
}
}
});
connection.on("error", (e) => {
logger.error(e.toString());
});
connection.once("close", () => {
for (const uuid of Object.keys(this.jobs)) {
if (this.jobs[uuid].addr === server) {
this.jobs[uuid].event.emit("error", "Job ended prematurely due to a closed connection; please run your image job again");
delete this.jobs[uuid];
}
}
//logger.log(`Lost connection to ${server}, attempting to reconnect...`);
this.connections.delete(server);
});
this.connections.set(server, connection);
}
async disconnect() {
for (const connection of this.connections.values()) {
connection.close();
}
for (const uuid of Object.keys(this.jobs)) {
this.jobs[uuid].event.emit("error", "Job ended prematurely (not really an error; just run your image job again)");
delete this.jobs[uuid];
}
this.connections.clear();
return;
}
async start(object, num) {
const currentServer = await this.getIdeal();
const data = Buffer.concat([Buffer.from([0x01 /* queue job */]), Buffer.from(num.length.toString()), Buffer.from(num), Buffer.from(JSON.stringify(object))]);
currentServer.sock.send(data);
const event = new EventEmitter();
this.jobs[num] = { event, addr: currentServer.addr };
const uuid = await new Promise((resolve, reject) => {
event.once("uuid", (uuid) => resolve(uuid));
event.once("error", reject);
});
delete this.jobs[num];
this.jobs[uuid] = { event: event, addr: currentServer.addr };
return { uuid: uuid, event: event };
}
run(object) {
return new Promise((resolve, reject) => {
if (process.env.API === "true") {
// Connect to best image server
const num = Math.floor(Math.random() * 100000).toString().slice(0, 5);
const timeout = setTimeout(() => {
if (this.jobs[num]) delete this.jobs[num];
reject("the image request timed out after 25 seconds. Try uploading your image elsewhere.");
}, 25000);
this.start(object, num).catch(err => { // incredibly hacky code incoming
clearTimeout(timeout);
if (err instanceof Error) return reject(err);
return err;
}).then((data) => {
clearTimeout(timeout);
if (!data.event) reject("Not connected to image server");
data.event.once("image", (image, type) => {
delete this.jobs[data.uuid];
const payload = {
// Take just the image data
buffer: image,
type: type
};
resolve(payload);
});
data.event.once("error", (err) => {
delete this.jobs[data.uuid];
reject(err);
});
return;
}).catch(err => reject(err));
} else {
// Called from command (not using image API)
const worker = new Worker(path.join(__dirname, "../image-runner.js"), {
workerData: object
});
worker.once("message", (data) => {
resolve({
buffer: Buffer.from([...data.buffer]),
type: data.fileExtension
});
});
worker.once("error", reject);
}
});
}
async handleCommand(data) {
try {
if (data.type === "run") {
return await this.run(data.obj);
} else if (data.type === "reload") {
await this.disconnect();
await this.repopulate();
let amount = 0;
for (const server of this.servers) {
try {
await this.connect(server);
amount += 1;
} catch (e) {
logger.error(e);
}
}
return amount;
} else if (data.type === "stats") {
return await this.getRunning();
}
} catch (err) {
return { err: err.message };
}
}
shutdown(done) {
done();
}
}
module.exports = ImageWorker;