Switched to an image API implementation by Terradice

This commit is contained in:
TheEssem 2020-11-17 08:52:12 -06:00
parent 6761e39d70
commit ad149156ad
6 changed files with 209 additions and 165 deletions

View file

@ -1,98 +1,136 @@
require("dotenv").config();
// code provided by terradice/tzlil
const os = require("os");
const { Worker, isMainThread, parentPort } = require("worker_threads");
const magick = require("../utils/image.js");
const Job = require("./job.js");
const { version } = require("../package.json");
const express = require("express");
const execPromise = require("util").promisify(require("child_process").exec);
const app = express();
const port = 3000;
const net = require("net");
const dgram = require("dgram"); // for UDP servers
const socket = dgram.createSocket("udp4"); // Our universal UDP socket, this might cause issues and we may have to use a seperate socket for each connection
const jobs = new Map();
var start = process.hrtime();
const log = (msg, thread) => {
console.log(`[${process.hrtime(start)[1] / 1000000}${(thread)?`:${thread}`:""}]\t ${msg}`);
};
app.get("/", (req, res) => {
res.send(`esmBot v${version}`);
});
const jobs = {};
// Should look like UUID : { addr : "someaddr", port: someport msg: "request" }
const queue = [];
// Array of UUIDs
app.post("/run", express.json(), async (req, res, next) => {
const object = req.body;
if (!magick.check(object.cmd)) return res.sendStatus(400);
if (isMainThread) {
const { v4: uuidv4 } = require("uuid");
try {
let type;
if (object.path) {
type = object.type;
if (!object.type) {
type = await magick.getType(object.path);
const MAX_WORKERS = process.env.WORKERS === "" ? parseInt(process.env.WORKERS) : os.cpus().length * 4; // Completely arbitrary, should usually be some multiple of your amount of cores
let workingWorkers = 0;
const acceptJob = (uuid) => {
workingWorkers++;
const worker = new Worker(__filename);
queue.shift();
worker.once("message", (uuid) => {
// This means the worker is finished
workingWorkers--;
if (queue.length > 0) {
acceptJob(queue[0]); // Get the next job UUID in queue and remove it from the original queue
delete jobs[uuid];
}
if (!type) {
return res.sendStatus(400);
});
worker.on("error", err => console.error("worker error:", err));
worker.on("exit", (code) => {
if (code !== 0)
console.error(`Worker stopped with exit code ${code}`);
});
worker.postMessage({
uuid: uuid,
msg: jobs[uuid].msg,
addr: jobs[uuid].addr,
port: jobs[uuid].port,
threadNum: workingWorkers
});
};
const server = dgram.createSocket("udp4"); //Create a UDP server for listening to requests, we dont need tcp
server.on("message", (msg, rinfo) => {
const opcode = msg.readUint8(0);
const req = msg.toString().slice(1,msg.length);
// 0x0 == Cancel job
// 0x1 == Queue job
if (opcode == 0x0) {
queue.shift();
delete jobs[req];
} else if (opcode == 0x1) {
const job = { addr: rinfo.address, port: rinfo.port, msg: req };
const uuid = uuidv4();
jobs[uuid] = job;
queue.push(uuid);
if (workingWorkers < MAX_WORKERS) {
acceptJob(uuid);
}
object.type = type.split("/")[1];
if (object.type !== "gif" && object.onlyGIF) return res.send({
status: "nogif"
const newBuffer = Buffer.concat([Buffer.from([0x0]), Buffer.from(uuid)]);
socket.send(newBuffer, rinfo.port, rinfo.address);
} else {
log("Could not parse message");
}
});
server.on("listening", () => {
const address = server.address();
log(`server listening ${address.address}:${address.port}`);
});
server.bind(8080); // ATTENTION: Always going to be bound to 0.0.0.0 !!!
} else {
parentPort.once("message", async (job) => {
log(`${job.uuid} worker got: ${job.msg}`, job.threadNum);
try {
const object = JSON.parse(job.msg);
let type;
if (object.path) {
type = object.type;
if (!object.type) {
type = await magick.getType(object.path);
}
if (!type) {
throw new TypeError("Unknown image type");
}
object.type = type.split("/")[1];
if (object.type !== "gif" && object.onlyGIF) throw new TypeError(`Expected a GIF, got ${object.type}`);
object.delay = object.delay ? object.delay : 0;
}
if (object.type === "gif" && !object.delay) {
const delay = (await execPromise(`ffprobe -v 0 -of csv=p=0 -select_streams v:0 -show_entries stream=r_frame_rate ${object.path}`)).stdout.replace("\n", "");
object.delay = (100 / delay.split("/")[0]) * delay.split("/")[1];
}
const data = await magick.run(object, true);
log(`${job.uuid} is done`, job.threadNum);
const server = net.createServer(function(socket) {
socket.write(Buffer.concat([Buffer.from(type ? type : "image/png"), Buffer.from("\n"), data]));
});
object.delay = object.delay ? object.delay : 0;
server.listen(job.port, job.addr);
// handle address in use errors
server.on("error", (e) => {
if (e.code === "EADDRINUSE") {
console.log("Address in use, retrying...");
setTimeout(() => {
server.close();
server.listen(job.port, job.addr);
}, 500);
}
});
socket.send(Buffer.concat([Buffer.from([0x1]), Buffer.from(job.port.toString())]), job.port, job.addr);
parentPort.postMessage(job.uuid); //Inform main thread about this worker freeing up
} catch (e) {
socket.send(Buffer.concat([Buffer.from([0x2]), Buffer.from(e.toString())]), job.port, job.address);
parentPort.postMessage(job.uuid);
}
const id = Math.random().toString(36).substring(2, 15);
if (object.type === "gif" && !object.delay) {
const delay = (await execPromise(`ffprobe -v 0 -of csv=p=0 -select_streams v:0 -show_entries stream=r_frame_rate ${object.path}`)).stdout.replace("\n", "");
object.delay = (100 / delay.split("/")[0]) * delay.split("/")[1];
}
const job = new Job(object);
jobs.set(id, job);
res.send({
id: id,
status: "queued"
});
job.run();
} catch (e) {
next(e);
}
});
app.get("/status", (req, res) => {
if (!req.query.id) return res.sendStatus(400);
const job = jobs.get(req.query.id);
if (!job) return res.sendStatus(400);
const timeout = setTimeout(function() {
job.removeAllListeners();
return res.send({
id: req.query.id,
status: job.status
});
}, 10000);
job.once("data", function() {
clearTimeout(timeout);
res.send({
id: req.query.id,
status: job.status
});
//jobs.delete(req.query.id);
});
job.on("error", function(e) {
clearTimeout(timeout);
res.status(500);
res.send({
id: req.query.id,
status: job.status,
error: e
});
jobs.delete(req.query.id);
});
});
app.get("/image", (req, res) => {
if (!req.query.id) return res.sendStatus(400);
const job = jobs.get(req.query.id);
if (!job) return res.sendStatus(400);
if (!job.data) return res.sendStatus(400);
if (job.error) return;
jobs.delete(req.query.id);
res.contentType(job.options.type ? job.options.type : "png");
return res.send(job.data);
});
app.listen(port, () => {
console.log(`Started image API on port ${port}.`);
});
}

View file

@ -1,28 +0,0 @@
const { EventEmitter } = require("events");
const magick = require("../utils/image.js");
class Job extends EventEmitter {
constructor(options) {
super();
this.options = options;
this.status = "queued";
this.data = null;
this.error = null;
}
run() {
this.status = "processing";
magick.run(this.options, true).then(data => {
this.status = "success";
this.data = data;
return this.emit("data", data, this.options.type);
}).catch(e => {
this.status = "error";
this.error = e;
return this.emit("error", e);
});
return;
}
}
module.exports = Job;