Initial draft for new image API design
This commit is contained in:
parent
a98601e2e8
commit
41068ae763
3 changed files with 314 additions and 263 deletions
173
api/index.js
173
api/index.js
|
|
@ -4,9 +4,32 @@ import { cpus } from "os";
|
|||
import { Worker } from "worker_threads";
|
||||
import { join } from "path";
|
||||
import { createServer } from "http";
|
||||
import ws from "ws";
|
||||
import { WebSocketServer } from "ws";
|
||||
import { fileURLToPath } from "url";
|
||||
import { dirname } from "path";
|
||||
import { createRequire } from "module";
|
||||
|
||||
const nodeRequire = createRequire(import.meta.url);
|
||||
const magick = nodeRequire(`../build/${process.env.DEBUG && process.env.DEBUG === "true" ? "Debug" : "Release"}/image.node`);
|
||||
|
||||
/*
|
||||
Rerror 0x01
|
||||
Tqueue 0x02
|
||||
Rqueue 0x03
|
||||
Tcancel 0x04
|
||||
Rcancel 0x05
|
||||
Twait 0x06
|
||||
Rwait 0x07
|
||||
Rinit 0x08
|
||||
*/
|
||||
const Rerror = 0x01;
|
||||
const Tqueue = 0x02;
|
||||
const Rqueue = 0x03;
|
||||
const Tcancel = 0x04;
|
||||
const Rcancel = 0x05;
|
||||
const Twait = 0x06;
|
||||
const Rwait = 0x07;
|
||||
const Rinit = 0x08;
|
||||
|
||||
const start = process.hrtime();
|
||||
const log = (msg, jobNum) => {
|
||||
|
|
@ -23,30 +46,34 @@ class JobCache extends Map {
|
|||
}
|
||||
|
||||
const jobs = new JobCache();
|
||||
// Should look like UUID : { msg: "request", num: <job number> }
|
||||
// Should look like ID : { msg: "request", num: <job number> }
|
||||
const queue = [];
|
||||
// Array of UUIDs
|
||||
|
||||
import { v4 as uuidv4 } from "uuid";
|
||||
// Array of IDs
|
||||
|
||||
const MAX_JOBS = process.env.JOBS !== "" && process.env.JOBS !== undefined ? parseInt(process.env.JOBS) : cpus().length * 4; // Completely arbitrary, should usually be some multiple of your amount of cores
|
||||
const PASS = process.env.PASS !== "" && process.env.PASS !== undefined ? process.env.PASS : undefined;
|
||||
let jobAmount = 0;
|
||||
|
||||
const acceptJob = (uuid, sock) => {
|
||||
const acceptJob = (id, sock) => {
|
||||
jobAmount++;
|
||||
queue.shift();
|
||||
const job = jobs.get(uuid);
|
||||
const job = jobs.get(id);
|
||||
return runJob({
|
||||
uuid: uuid,
|
||||
id: id,
|
||||
msg: job.msg,
|
||||
num: job.num
|
||||
}, sock).then(() => {
|
||||
log(`Job ${uuid} has finished`);
|
||||
log(`Job ${id} has finished`);
|
||||
}).catch((err) => {
|
||||
console.error(`Error on job ${uuid}:`, err);
|
||||
jobs.delete(uuid);
|
||||
sock.send(Buffer.concat([Buffer.from([0x2]), Buffer.from(uuid), Buffer.from(err.message)]));
|
||||
console.error(`Error on job ${id}:`, err);
|
||||
const newJob = jobs.get(id);
|
||||
if (!newJob.tag) {
|
||||
newJob.error = err.message;
|
||||
jobs.set(id, newJob);
|
||||
return;
|
||||
}
|
||||
jobs.delete(id);
|
||||
sock.send(Buffer.concat([Buffer.from([Rerror]), newJob.tag, Buffer.from(err.message)]));
|
||||
}).finally(() => {
|
||||
jobAmount--;
|
||||
if (queue.length > 0) {
|
||||
|
|
@ -55,42 +82,63 @@ const acceptJob = (uuid, sock) => {
|
|||
});
|
||||
};
|
||||
|
||||
const wss = new ws.Server({ clientTracking: true, noServer: true });
|
||||
const wss = new WebSocketServer({ clientTracking: true, noServer: true });
|
||||
|
||||
wss.on("connection", (ws, request) => {
|
||||
log(`WS client ${request.socket.remoteAddress}:${request.socket.remotePort} has connected`);
|
||||
|
||||
ws.on("open", () => {
|
||||
const init = Buffer.concat([Buffer.from([Rinit]), Buffer.from(MAX_JOBS), Buffer.from(JSON.stringify(Object.keys(magick)))]);
|
||||
ws.send(init);
|
||||
});
|
||||
|
||||
ws.on("error", (err) => {
|
||||
console.error(err);
|
||||
});
|
||||
|
||||
ws.on("message", (msg) => {
|
||||
const opcode = msg.readUint8(0);
|
||||
const req = msg.toString().slice(1,msg.length);
|
||||
console.log(req);
|
||||
// 0x00 == Cancel job
|
||||
// 0x01 == Queue job
|
||||
if (opcode == 0x00) {
|
||||
delete queue[queue.indexOf(req) - 1];
|
||||
jobs.delete(req);
|
||||
} else if (opcode == 0x01) {
|
||||
const length = parseInt(req.slice(0, 1));
|
||||
const num = req.slice(1, length + 1);
|
||||
const obj = req.slice(length + 1);
|
||||
const tag = msg.slice(1, 5);
|
||||
const req = msg.toString().slice(5, msg.length);
|
||||
if (opcode == Tqueue) {
|
||||
const id = msg.readUInt32LE(5);
|
||||
const obj = msg.slice(9, msg.length);
|
||||
const job = { msg: obj, num: jobAmount };
|
||||
const uuid = uuidv4();
|
||||
jobs.set(uuid, job);
|
||||
queue.push(uuid);
|
||||
jobs.set(id, job);
|
||||
queue.push(id);
|
||||
|
||||
const newBuffer = Buffer.concat([Buffer.from([0x00]), Buffer.from(uuid), Buffer.from(num)]);
|
||||
const newBuffer = Buffer.concat([Buffer.from([Rqueue]), tag]);
|
||||
ws.send(newBuffer);
|
||||
|
||||
if (jobAmount < MAX_JOBS) {
|
||||
log(`Got WS request for job ${job.msg} with id ${uuid}`, job.num);
|
||||
acceptJob(uuid, ws);
|
||||
log(`Got WS request for job ${job.msg} with id ${id}`, job.num);
|
||||
acceptJob(id, ws);
|
||||
} else {
|
||||
log(`Got WS request for job ${job.msg} with id ${uuid}, queued in position ${queue.indexOf(uuid)}`, job.num);
|
||||
log(`Got WS request for job ${job.msg} with id ${id}, queued in position ${queue.indexOf(id)}`, job.num);
|
||||
}
|
||||
} else if (opcode == Tcancel) {
|
||||
delete queue[queue.indexOf(req) - 1];
|
||||
jobs.delete(req);
|
||||
const cancelResponse = Buffer.concat([Buffer.from([Rcancel]), tag]);
|
||||
ws.send(cancelResponse);
|
||||
} else if (opcode == Twait) {
|
||||
const id = msg.readUInt32LE(5);
|
||||
const job = jobs.get(id);
|
||||
if (!job) {
|
||||
const errorResponse = Buffer.concat([Buffer.from([Rerror]), tag, Buffer.from("Invalid job ID")]);
|
||||
ws.send(errorResponse);
|
||||
return;
|
||||
}
|
||||
if (job.error) {
|
||||
jobs.delete(id);
|
||||
const errorResponse = Buffer.concat([Buffer.from([Rerror]), tag, Buffer.from(job.error)]);
|
||||
ws.send(errorResponse);
|
||||
return;
|
||||
}
|
||||
job.tag = tag;
|
||||
jobs.set(id, job);
|
||||
//const waitResponse = Buffer.concat([Buffer.from([Rwait]), tag]);
|
||||
//ws.send(waitResponse);
|
||||
} else {
|
||||
log("Could not parse WS message");
|
||||
}
|
||||
|
|
@ -117,41 +165,35 @@ httpServer.on("request", async (req, res) => {
|
|||
return res.end("401 Unauthorized");
|
||||
}
|
||||
const reqUrl = new URL(req.url, `http://${req.headers.host}`);
|
||||
if (reqUrl.pathname === "/status" && req.method === "GET") {
|
||||
log(`Sending server status to ${req.socket.remoteAddress}:${req.socket.remotePort} via HTTP`);
|
||||
const statusObject = {
|
||||
load: MAX_JOBS - jobAmount,
|
||||
queued: queue.length
|
||||
};
|
||||
return res.end(JSON.stringify(statusObject));
|
||||
} else if (reqUrl.pathname === "/running" && req.method === "GET") {
|
||||
log(`Sending currently running jobs to ${req.socket.remoteAddress}:${req.socket.remotePort} via HTTP`);
|
||||
const keys = jobs.keys();
|
||||
const newObject = { queued: queue.length, runningJobs: jobAmount, max: MAX_JOBS };
|
||||
for (const key of keys) {
|
||||
const validKeys = Object.keys(jobs.get(key)).filter((value) => value !== "addr" && value !== "port" && value !== "data" && value !== "ext");
|
||||
newObject[key] = {};
|
||||
for (const validKey of validKeys) {
|
||||
if (validKey === "msg") {
|
||||
newObject[key][validKey] = JSON.parse(jobs.get(key)[validKey]);
|
||||
} else {
|
||||
newObject[key][validKey] = jobs.get(key)[validKey];
|
||||
}
|
||||
}
|
||||
}
|
||||
return res.end(JSON.stringify(newObject));
|
||||
} else if (reqUrl.pathname === "/image" && req.method === "GET") {
|
||||
if (reqUrl.pathname === "/image" && req.method === "GET") {
|
||||
if (!reqUrl.searchParams.has("id")) {
|
||||
res.statusCode = 400;
|
||||
return res.end("400 Bad Request");
|
||||
}
|
||||
const id = reqUrl.searchParams.get("id");
|
||||
const id = parseInt(reqUrl.searchParams.get("id"));
|
||||
if (!jobs.has(id)) {
|
||||
res.statusCode = 410;
|
||||
return res.end("410 Gone");
|
||||
}
|
||||
log(`Sending image data for job ${id} to ${req.socket.remoteAddress}:${req.socket.remotePort} via HTTP`);
|
||||
res.setHeader("ext", jobs.get(id).ext);
|
||||
const ext = jobs.get(id).ext;
|
||||
let contentType;
|
||||
switch (ext) {
|
||||
case "gif":
|
||||
contentType = "image/gif";
|
||||
break;
|
||||
case "png":
|
||||
contentType = "image/png";
|
||||
break;
|
||||
case "jpeg":
|
||||
case "jpg":
|
||||
contentType = "image/jpeg";
|
||||
break;
|
||||
case "webp":
|
||||
contentType = "image/webp";
|
||||
break;
|
||||
}
|
||||
res.setHeader("Content-Type", contentType);
|
||||
const data = jobs.get(id).data;
|
||||
jobs.delete(id);
|
||||
return res.end(data, (err) => {
|
||||
|
|
@ -189,9 +231,9 @@ httpServer.listen(8080, () => {
|
|||
log("HTTP and WS listening on port 8080");
|
||||
});
|
||||
|
||||
const runJob = (job, sock) => {
|
||||
const runJob = (job, ws) => {
|
||||
return new Promise((resolve, reject) => {
|
||||
log(`Job ${job.uuid} starting...`, job.num);
|
||||
log(`Job ${job.id} starting...`, job.num);
|
||||
|
||||
const object = JSON.parse(job.msg);
|
||||
// If the image has a path, it must also have a type
|
||||
|
|
@ -206,17 +248,16 @@ const runJob = (job, sock) => {
|
|||
worker.terminate();
|
||||
reject(new Error("Job timed out"));
|
||||
}, 900000);
|
||||
log(`Job ${job.uuid} started`, job.num);
|
||||
log(`Job ${job.id} started`, job.num);
|
||||
worker.once("message", (data) => {
|
||||
clearTimeout(timeout);
|
||||
log(`Sending result of job ${job.uuid} back to the bot`, job.num);
|
||||
const jobObject = jobs.get(job.uuid);
|
||||
log(`Sending result of job ${job.id} back to the bot`, job.num);
|
||||
const jobObject = jobs.get(job.id);
|
||||
jobObject.data = data.buffer;
|
||||
jobObject.ext = data.fileExtension;
|
||||
jobs.set(job.uuid, jobObject);
|
||||
sock.send(Buffer.concat([Buffer.from([0x1]), Buffer.from(job.uuid)]), () => {
|
||||
return resolve();
|
||||
});
|
||||
jobs.set(job.id, jobObject);
|
||||
const waitResponse = Buffer.concat([Buffer.from([Rwait]), jobObject.tag]);
|
||||
ws.send(waitResponse);
|
||||
});
|
||||
worker.once("error", (e) => {
|
||||
clearTimeout(timeout);
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue