Switched image API from raw TCP to websockets, removed selective image reading code

This commit is contained in:
TheEssem 2021-06-18 00:10:11 -05:00
parent c2e13c26b6
commit a0a32c31b9
No known key found for this signature in database
GPG key ID: A3F9F02129092FCA
7 changed files with 133 additions and 138 deletions

View file

@ -1,12 +1,10 @@
// code originally provided by tzlil
require("dotenv").config();
const os = require("os");
//const { run } = require("../utils/image-runner.js");
const { Worker } = require("worker_threads");
const path = require("path");
const net = require("net");
const http = require("http");
const WebSocket = require("ws");
const start = process.hrtime();
const log = (msg, jobNum) => {
@ -23,7 +21,7 @@ class JobCache extends Map {
}
const jobs = new JobCache();
// Should look like UUID : { addr : "someaddr", port: someport, msg: "request" }
// Should look like UUID : { msg: "request", num: <job number> }
const queue = [];
// Array of UUIDs
@ -45,7 +43,7 @@ const acceptJob = (uuid, sock) => {
}).catch((err) => {
console.error(`Error on job ${uuid}:`, err);
jobs.delete(uuid);
sock.write(Buffer.concat([Buffer.from([0x2]), Buffer.from(uuid), Buffer.from(err.message)]));
sock.send(Buffer.concat([Buffer.from([0x2]), Buffer.from(uuid), Buffer.from(err.message)]));
}).finally(() => {
jobAmount--;
if (queue.length > 0) {
@ -54,16 +52,68 @@ const acceptJob = (uuid, sock) => {
});
};
const httpServer = http.createServer((req, res) => {
if (req.method !== "GET") {
const wss = new WebSocket.Server({ clientTracking: true, noServer: true });
wss.on("connection", (ws, request) => {
log(`WS client ${request.socket.remoteAddress}:${request.socket.remotePort} has connected`);
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 job = { msg: obj, num: jobAmount };
const uuid = uuidv4();
jobs.set(uuid, job);
queue.push(uuid);
const newBuffer = Buffer.concat([Buffer.from([0x00]), Buffer.from(uuid), Buffer.from(num)]);
ws.send(newBuffer);
if (jobAmount < MAX_JOBS) {
log(`Got WS request for job ${job.msg} with id ${uuid}`, job.num);
acceptJob(uuid, ws);
} else {
log(`Got WS request for job ${job.msg} with id ${uuid}, queued in position ${queue.indexOf(uuid)}`, job.num);
}
} else {
log("Could not parse WS message");
}
});
ws.on("close", () => {
log(`WS client ${request.socket.remoteAddress}:${request.socket.remotePort} has disconnected`);
});
});
wss.on("error", (err) => {
console.error("A WS error occurred: ", err);
});
const httpServer = http.createServer();
httpServer.on("request", async (req, res) => {
if (req.method !== "GET" && req.method !== "POST") {
res.statusCode = 405;
return res.end("405 Method Not Allowed");
}
const reqUrl = new URL(req.url, `http://${req.headers.host}`);
if (reqUrl.pathname === "/status") {
if (reqUrl.pathname === "/status" && req.method === "GET") {
log(`Sending server status to ${req.socket.remoteAddress}:${req.socket.remotePort} via HTTP`);
return res.end(Buffer.from((MAX_JOBS - jobAmount).toString()));
} else if (reqUrl.pathname === "/running") {
} 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 };
@ -79,7 +129,7 @@ const httpServer = http.createServer((req, res) => {
}
}
return res.end(JSON.stringify(newObject));
} else if (reqUrl.pathname === "/image") {
} else if (reqUrl.pathname === "/image" && req.method === "GET") {
if (!reqUrl.searchParams.has("id")) {
res.statusCode = 400;
return res.end("400 Bad Request");
@ -102,64 +152,24 @@ const httpServer = http.createServer((req, res) => {
}
});
httpServer.on("upgrade", (req, sock, head) => {
const reqUrl = new URL(req.url, `http://${req.headers.host}`);
if (reqUrl.pathname === "/sock") {
wss.handleUpgrade(req, sock, head, (ws) => {
wss.emit("connection", ws, req);
});
} else {
sock.destroy();
}
});
httpServer.on("error", (e) => {
console.error("An HTTP error occurred: ", e);
});
httpServer.listen(8081, () => {
log("HTTP listening on port 8081");
});
const server = net.createServer((sock) => { // Create a TCP socket/server to listen to requests
log(`TCP client ${sock.remoteAddress}:${sock.remotePort} has connected`);
sock.on("error", (e) => {
console.error(e);
});
sock.on("data", (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 job = { addr: sock.remoteAddress, port: sock.remotePort, msg: obj, num: jobAmount };
const uuid = uuidv4();
jobs.set(uuid, job);
queue.push(uuid);
const newBuffer = Buffer.concat([Buffer.from([0x00]), Buffer.from(uuid), Buffer.from(num)]);
sock.write(newBuffer);
if (jobAmount < MAX_JOBS) {
log(`Got TCP request for job ${job.msg} with id ${uuid}`, job.num);
acceptJob(uuid, sock);
} else {
log(`Got TCP request for job ${job.msg} with id ${uuid}, queued in position ${queue.indexOf(uuid)}`, job.num);
}
} else {
log("Could not parse TCP message");
}
});
sock.on("end", () => {
log(`TCP client ${sock.remoteAddress}:${sock.remotePort} has disconnected`);
});
});
server.on("error", (e) => {
console.error("A TCP error occurred: ", e);
});
server.listen(8080, () => {
log("TCP listening on port 8080");
httpServer.listen(8080, () => {
log("HTTP listening on port 8080");
});
const runJob = (job, sock) => {
@ -187,8 +197,7 @@ const runJob = (job, sock) => {
jobObject.data = data.buffer;
jobObject.ext = data.fileExtension;
jobs.set(job.uuid, jobObject);
sock.write(Buffer.concat([Buffer.from([0x1]), Buffer.from(job.uuid)]), (e) => {
if (e) return reject(e);
sock.send(Buffer.concat([Buffer.from([0x1]), Buffer.from(job.uuid)]), () => {
return resolve();
});
});