267 lines
8.4 KiB
JavaScript
267 lines
8.4 KiB
JavaScript
import { BaseServiceWorker } from "eris-fleet";
|
|
import * as logger from "../logger.js";
|
|
import fs from "fs";
|
|
import path from "path";
|
|
import { fileURLToPath } from "url";
|
|
import { Worker } from "worker_threads";
|
|
import { createRequire } from "module";
|
|
import { createServer } from "http";
|
|
import fetch from "node-fetch";
|
|
import EventEmitter from "events";
|
|
|
|
// only requiring this to work around an issue regarding worker threads
|
|
const nodeRequire = createRequire(import.meta.url);
|
|
nodeRequire(`../../build/${process.env.DEBUG && process.env.DEBUG === "true" ? "Debug" : "Release"}/image.node`);
|
|
|
|
import ImageConnection from "../imageConnection.js";
|
|
|
|
class ImageWorker extends BaseServiceWorker {
|
|
constructor(setup) {
|
|
super(setup);
|
|
|
|
console.info = (str) => this.ipc.sendToAdmiral("info", str);
|
|
|
|
if (process.env.API_TYPE === "ws") {
|
|
this.connections = new Map();
|
|
this.servers = JSON.parse(fs.readFileSync(new URL("../../servers.json", import.meta.url), { encoding: "utf8" })).image;
|
|
this.nextID = 0;
|
|
} else if (process.env.API_TYPE === "azure") {
|
|
this.jobs = new Map();
|
|
this.webhook = createServer();
|
|
this.port = parseInt(process.env.WEBHOOK_PORT) || 3763;
|
|
}
|
|
|
|
this.begin().then(() => this.serviceReady());
|
|
}
|
|
|
|
async begin() {
|
|
// connect to image api if enabled
|
|
if (process.env.API_TYPE === "ws") {
|
|
for (const server of this.servers) {
|
|
try {
|
|
await this.connect(server.server, server.auth);
|
|
} catch (e) {
|
|
logger.error(e);
|
|
}
|
|
}
|
|
} else if (process.env.API_TYPE === "azure") {
|
|
this.webhook.on("request", async (req, res) => {
|
|
if (req.method !== "POST") {
|
|
res.statusCode = 405;
|
|
return res.end("405 Method Not Allowed");
|
|
}
|
|
if (process.env.AZURE_PASS && req.headers.authorization !== process.env.AZURE_PASS) {
|
|
res.statusCode = 401;
|
|
return res.end("401 Unauthorized");
|
|
}
|
|
const reqUrl = new URL(req.url, `http://${req.headers.host}`);
|
|
if (reqUrl.pathname === "/callback") {
|
|
try {
|
|
const chunks = [];
|
|
req.on("data", (data) => {
|
|
chunks.push(data);
|
|
});
|
|
req.once("end", () => {
|
|
if (this.jobs.has(req.headers["x-azure-id"])) {
|
|
try {
|
|
const error = JSON.parse(Buffer.concat(chunks).toString());
|
|
if (error.error) this.jobs.get(req.headers["x-azure-id"]).emit("error", new Error(error.message));
|
|
} catch {
|
|
// no-op
|
|
}
|
|
const contentType = req.headers["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;
|
|
default:
|
|
type = contentType;
|
|
break;
|
|
}
|
|
this.jobs.get(req.headers["x-azure-id"]).emit("image", { buffer: Buffer.concat(chunks), type });
|
|
return res.end("OK");
|
|
} else {
|
|
res.statusCode = 409;
|
|
return res.end("409 Conflict");
|
|
}
|
|
});
|
|
} catch (e) {
|
|
logger.error("An error occurred while processing a webhook request: ", e);
|
|
res.statusCode = 500;
|
|
return res.end("500 Internal Server Error");
|
|
}
|
|
} else {
|
|
res.statusCode = 404;
|
|
return res.end("404 Not Found");
|
|
}
|
|
});
|
|
this.webhook.on("error", (e) => {
|
|
logger.error("An error occurred on the Azure webhook: ", e);
|
|
});
|
|
this.webhook.listen(this.port, () => {
|
|
logger.log(`Azure HTTP webhook listening on port ${this.port}`);
|
|
});
|
|
}
|
|
}
|
|
|
|
async repopulate() {
|
|
const data = await fs.promises.readFile(new URL("../../servers.json", import.meta.url), { encoding: "utf8" });
|
|
this.servers = JSON.parse(data).image;
|
|
return;
|
|
}
|
|
|
|
async getRunning() {
|
|
const statuses = [];
|
|
if (process.env.API_TYPE === "ws") {
|
|
for (const [address, connection] of this.connections) {
|
|
if (connection.conn.readyState !== 0 && connection.conn.readyState !== 1) {
|
|
continue;
|
|
}
|
|
statuses.push({
|
|
address,
|
|
runningJobs: connection.njobs,
|
|
max: connection.max
|
|
});
|
|
}
|
|
}
|
|
return statuses;
|
|
}
|
|
|
|
async chooseServer(ideal) {
|
|
if (ideal.length === 0) throw "No available servers";
|
|
const sorted = ideal.sort((a, b) => {
|
|
return a.load - b.load;
|
|
}).filter((e, i, array) => {
|
|
return !(e.load < array[0].load);
|
|
});
|
|
return sorted[0];
|
|
}
|
|
|
|
async getIdeal(object) {
|
|
const idealServers = [];
|
|
for (const [address, connection] of this.connections) {
|
|
if (connection.conn.readyState !== 0 && connection.conn.readyState !== 1) {
|
|
continue;
|
|
}
|
|
if (object.params.type && connection.formats[object.cmd] && !connection.formats[object.cmd].includes(object.params.type)) continue;
|
|
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 ImageConnection(server, auth);
|
|
this.connections.set(server, connection);
|
|
}
|
|
|
|
async disconnect() {
|
|
for (const connection of this.connections.values()) {
|
|
connection.close();
|
|
}
|
|
this.connections.clear();
|
|
return;
|
|
}
|
|
|
|
waitForWorker(worker) {
|
|
return new Promise((resolve, reject) => {
|
|
worker.once("message", (data) => {
|
|
resolve({
|
|
buffer: Buffer.from([...data.buffer]),
|
|
type: data.fileExtension
|
|
});
|
|
});
|
|
worker.once("error", reject);
|
|
});
|
|
}
|
|
|
|
waitForAzure(event) {
|
|
return new Promise((resolve, reject) => {
|
|
event.once("image", (data) => {
|
|
resolve(data);
|
|
});
|
|
event.once("error", reject);
|
|
});
|
|
}
|
|
|
|
async run(object) {
|
|
if (process.env.API_TYPE === "ws") {
|
|
let num = this.nextID++;
|
|
if (num > 4294967295) num = this.nextID = 0;
|
|
for (let i = 0; i < 3; i++) {
|
|
const currentServer = await this.getIdeal(object);
|
|
try {
|
|
await currentServer.queue(num, object);
|
|
await currentServer.wait(num);
|
|
const output = await currentServer.getOutput(num);
|
|
return output;
|
|
} catch (e) {
|
|
if (i < 2 && e === "Request ended prematurely due to a closed connection") {
|
|
continue;
|
|
} else {
|
|
if (e === "No available servers" && i >= 2) throw "Request ended prematurely due to a closed connection";
|
|
throw e;
|
|
}
|
|
}
|
|
}
|
|
} else if (process.env.API_TYPE === "azure") {
|
|
object.callback = `${process.env.AZURE_CALLBACK_URL}:${this.port}/callback`;
|
|
const response = await fetch(`${process.env.AZURE_URL}/api/orchestrators/ImageOrchestrator`, { method: "POST", body: JSON.stringify(object) }).then(r => r.json());
|
|
const event = new EventEmitter();
|
|
this.jobs.set(response.id, event);
|
|
return await this.waitForAzure(event);
|
|
} 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) {
|
|
try {
|
|
if (data.type === "run") {
|
|
const result = await this.run(data.obj);
|
|
return result;
|
|
} 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.server, server.auth);
|
|
amount += 1;
|
|
} catch (e) {
|
|
logger.error(e);
|
|
}
|
|
}
|
|
return amount;
|
|
} else if (data.type === "stats") {
|
|
return await this.getRunning();
|
|
}
|
|
} catch (err) {
|
|
return { err: typeof err === "string" ? err : err.message };
|
|
}
|
|
}
|
|
|
|
shutdown(done) {
|
|
done();
|
|
}
|
|
}
|
|
|
|
export default ImageWorker;
|