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
163
utils/imageConnection.js
Normal file
163
utils/imageConnection.js
Normal file
|
|
@ -0,0 +1,163 @@
|
|||
import fetch from "node-fetch";
|
||||
import WebSocket from "ws";
|
||||
import * as logger from "./logger.js";
|
||||
import { setTimeout } from "timers/promises";
|
||||
|
||||
/*
|
||||
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;
|
||||
|
||||
class ImageConnection {
|
||||
constructor(host, auth, tls = false) {
|
||||
this.requests = new Map();
|
||||
this.host = host;
|
||||
this.auth = auth;
|
||||
this.tag = null;
|
||||
this.disconnected = false;
|
||||
this.njobs = 0;
|
||||
this.max = 0;
|
||||
this.formats = {};
|
||||
this.wsproto = null;
|
||||
if (tls) {
|
||||
this.wsproto = "wss";
|
||||
} else {
|
||||
this.wsproto = "ws";
|
||||
}
|
||||
this.sockurl = `${this.wsproto}://${host}/sock`;
|
||||
this.conn = new WebSocket(this.sockurl, {
|
||||
headers: {
|
||||
"Authentication": auth && auth !== "" ? auth : undefined
|
||||
}
|
||||
});
|
||||
let httpproto;
|
||||
if (tls) {
|
||||
httpproto = "https";
|
||||
} else {
|
||||
httpproto = "http";
|
||||
}
|
||||
this.httpurl = `${httpproto}://${host}/image`;
|
||||
this.conn.on("message", (msg) => this.onMessage(msg));
|
||||
this.conn.once("error", (err) => this.onError(err));
|
||||
this.conn.once("close", () => this.onClose());
|
||||
}
|
||||
|
||||
onMessage(msg) {
|
||||
const op = msg.readUint8(0);
|
||||
if (op === Rinit) {
|
||||
this.max = msg.readUint16LE(1);
|
||||
this.formats = JSON.parse(msg.toString("utf8", 3));
|
||||
return;
|
||||
}
|
||||
const tag = msg.readUint32LE(1);
|
||||
const promise = this.requests.get(tag);
|
||||
this.requests.delete(tag);
|
||||
if (op === Rerror) {
|
||||
promise.reject(new Error(msg.slice(5, msg.length).toString()));
|
||||
return;
|
||||
}
|
||||
promise.resolve();
|
||||
}
|
||||
|
||||
onError(e) {
|
||||
logger.error(e.toString());
|
||||
}
|
||||
|
||||
async onClose() {
|
||||
for (const promise of this.requests.values()) {
|
||||
promise.reject(new Error("Request ended prematurely due to a closed connection"));
|
||||
}
|
||||
this.requests.clear();
|
||||
if (!this.disconnected) {
|
||||
logger.warn(`Lost connection to ${this.host}, attempting to reconnect in 5 seconds...`);
|
||||
await setTimeout(5000);
|
||||
this.conn = new WebSocket(this.sockurl, {
|
||||
headers: {
|
||||
"Authentication": this.auth
|
||||
}
|
||||
});
|
||||
this.conn.on("message", (msg) => this.onMessage(msg));
|
||||
this.conn.once("error", (err) => this.onError(err));
|
||||
this.conn.once("close", () => this.onClose());
|
||||
}
|
||||
this.disconnected = false;
|
||||
}
|
||||
|
||||
close() {
|
||||
this.disconnected = true;
|
||||
this.conn.close();
|
||||
}
|
||||
|
||||
queue(jobid, jobobj) {
|
||||
const str = JSON.stringify(jobobj);
|
||||
const buf = Buffer.alloc(4 + str.length);
|
||||
buf.writeUint32LE(jobid);
|
||||
buf.write(str, 4);
|
||||
return this.do(Tqueue, buf);
|
||||
}
|
||||
|
||||
wait(jobid) {
|
||||
const buf = Buffer.alloc(4);
|
||||
buf.writeUint32LE(jobid);
|
||||
return this.do(Twait, buf);
|
||||
}
|
||||
|
||||
cancel(jobid) {
|
||||
const buf = Buffer.alloc(4);
|
||||
buf.writeUint32LE(jobid);
|
||||
return this.do(Tcancel, buf);
|
||||
}
|
||||
|
||||
async getOutput(jobid) {
|
||||
const req = await fetch(`${this.httpurl}?id=${jobid}`, {
|
||||
headers: {
|
||||
"Authentication": this.auth && this.auth !== "" ? this.auth : undefined
|
||||
}
|
||||
});
|
||||
const contentType = req.headers.get("Content-Type");
|
||||
let type;
|
||||
switch (contentType) {
|
||||
case "image/gif":
|
||||
type = "gif";
|
||||
break;
|
||||
case "image/png":
|
||||
type = "png";
|
||||
break;
|
||||
case "image/jpeg":
|
||||
type = "jpg";
|
||||
break;
|
||||
case "image/webp":
|
||||
type = "webp";
|
||||
break;
|
||||
}
|
||||
return { buffer: Buffer.from(await req.arrayBuffer()), type };
|
||||
}
|
||||
|
||||
async do(op, data) {
|
||||
const buf = Buffer.alloc(1 + 4);
|
||||
const tag = this.tag++;
|
||||
buf.writeUint8(op);
|
||||
buf.writeUint32LE(tag, 1);
|
||||
this.conn.send(Buffer.concat([buf, data]));
|
||||
const promise = new Promise((resolve, reject) => {
|
||||
this.requests.set(tag, { resolve, reject });
|
||||
});
|
||||
return promise;
|
||||
}
|
||||
}
|
||||
|
||||
export default ImageConnection;
|
||||
|
|
@ -1,12 +1,11 @@
|
|||
import { BaseServiceWorker } from "eris-fleet";
|
||||
import * as logger from "../logger.js";
|
||||
import fetch from "node-fetch";
|
||||
import WebSocket from "ws";
|
||||
import fs from "fs";
|
||||
import path from "path";
|
||||
import { fileURLToPath } from "url";
|
||||
import { Worker } from "worker_threads";
|
||||
import { EventEmitter } from "events";
|
||||
|
||||
import ImageConnection from "../imageConnection.js";
|
||||
|
||||
class ImageWorker extends BaseServiceWorker {
|
||||
constructor(setup) {
|
||||
|
|
@ -16,6 +15,7 @@ class ImageWorker extends BaseServiceWorker {
|
|||
this.jobs = {};
|
||||
this.connections = new Map();
|
||||
this.servers = JSON.parse(fs.readFileSync("./servers.json", { encoding: "utf8" })).image;
|
||||
this.nextID = 0;
|
||||
}
|
||||
|
||||
this.begin().then(() => this.serviceReady());
|
||||
|
|
@ -41,46 +41,18 @@ class ImageWorker extends BaseServiceWorker {
|
|||
}
|
||||
|
||||
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--;
|
||||
for (const [address, connection] of this.connections) {
|
||||
if (connection.conn.readyState !== 0 && connection.conn.readyState !== 1) {
|
||||
continue;
|
||||
}
|
||||
const controller = new AbortController(); // eslint-disable-line no-undef
|
||||
const timeout = setTimeout(() => {
|
||||
controller.abort();
|
||||
}, 2000);
|
||||
try {
|
||||
const auth = this.servers.filter((val) => val.server === address)[0].auth;
|
||||
const statusRequest = await fetch(`http://${address}:8080/running`, {
|
||||
signal: controller.signal,
|
||||
headers: {
|
||||
"Authentication": auth && auth !== "" ? auth : undefined
|
||||
}
|
||||
});
|
||||
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");
|
||||
statuses.push({
|
||||
address,
|
||||
runningJobs: connection.njobs,
|
||||
max: connection.max
|
||||
});
|
||||
}
|
||||
return statuses;
|
||||
}
|
||||
|
||||
async chooseServer(ideal) {
|
||||
|
|
@ -89,119 +61,27 @@ class ImageWorker extends BaseServiceWorker {
|
|||
return b.load - a.load;
|
||||
}).filter((e, i, array) => {
|
||||
return !(e.load < array[0].load);
|
||||
}).sort((a, b) => {
|
||||
return a.queued - b.queued;
|
||||
});
|
||||
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.server)) await this.connect(server.server, server.auth);
|
||||
} 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--;
|
||||
for (const [address, connection] of this.connections) {
|
||||
if (connection.conn.readyState !== 0 && connection.conn.readyState !== 1) {
|
||||
continue;
|
||||
}
|
||||
const controller = new AbortController(); // eslint-disable-line no-undef
|
||||
const timeout = setTimeout(() => {
|
||||
controller.abort();
|
||||
}, 5000);
|
||||
try {
|
||||
const auth = this.servers.filter((val) => val.server === address)[0].auth;
|
||||
const statusRequest = await fetch(`http://${address}:8080/status`, {
|
||||
signal: controller.signal,
|
||||
headers: {
|
||||
"Authentication": auth && auth !== "" ? auth : undefined
|
||||
}
|
||||
});
|
||||
clearTimeout(timeout);
|
||||
const status = await statusRequest.json();
|
||||
serversLeft--;
|
||||
idealServers.push({
|
||||
addr: address,
|
||||
load: status.load,
|
||||
queued: status.queued
|
||||
});
|
||||
} 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");
|
||||
idealServers.push({
|
||||
addr: address,
|
||||
load: connection.njobs / connection.max
|
||||
});
|
||||
}
|
||||
const server = await this.chooseServer(idealServers);
|
||||
return this.connections.get(server.addr);
|
||||
}
|
||||
|
||||
async connect(server, auth) {
|
||||
const connection = new WebSocket(`ws://${server}:8080/sock`, {
|
||||
headers: {
|
||||
"Authentication": auth && auth !== "" ? auth : undefined
|
||||
}
|
||||
});
|
||||
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}`, {
|
||||
headers: {
|
||||
"Authentication": auth && auth !== "" ? auth : undefined
|
||||
}
|
||||
});
|
||||
const image = Buffer.from(await imageReq.arrayBuffer());
|
||||
// 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);
|
||||
});
|
||||
const connection = new ImageConnection(`${server}:8080`, auth);
|
||||
this.connections.set(server, connection);
|
||||
}
|
||||
|
||||
|
|
@ -209,70 +89,37 @@ class ImageWorker extends BaseServiceWorker {
|
|||
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);
|
||||
waitForWorker(worker) {
|
||||
return new Promise((resolve, reject) => {
|
||||
worker.once("message", (data) => {
|
||||
resolve({
|
||||
buffer: Buffer.from([...data.buffer]),
|
||||
type: data.fileExtension
|
||||
});
|
||||
});
|
||||
worker.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).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(path.dirname(fileURLToPath(import.meta.url)), "../image-runner.js"), {
|
||||
workerData: object
|
||||
});
|
||||
worker.once("message", (data) => {
|
||||
resolve({
|
||||
buffer: Buffer.from([...data.buffer]),
|
||||
type: data.fileExtension
|
||||
});
|
||||
});
|
||||
worker.once("error", reject);
|
||||
}
|
||||
});
|
||||
async run(object) {
|
||||
if (process.env.API === "true") {
|
||||
const num = this.nextID++;
|
||||
const currentServer = await this.getIdeal();
|
||||
await currentServer.queue(num, object);
|
||||
await currentServer.wait(num);
|
||||
const output = await currentServer.getOutput(num);
|
||||
return output;
|
||||
} else {
|
||||
// Called from command (not using image API)
|
||||
const worker = new Worker(path.join(path.dirname(fileURLToPath(import.meta.url)), "../image-runner.js"), {
|
||||
workerData: object
|
||||
});
|
||||
return await this.waitForWorker(worker);
|
||||
}
|
||||
}
|
||||
|
||||
async handleCommand(data) {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue