A little bit more image work

This commit is contained in:
TheEssem 2021-06-18 19:10:13 -05:00
parent a0a32c31b9
commit 4d3e1655e0
No known key found for this signature in database
GPG Key ID: A3F9F02129092FCA
5 changed files with 123 additions and 145 deletions

View File

@ -169,7 +169,7 @@ httpServer.on("error", (e) => {
});
httpServer.listen(8080, () => {
log("HTTP listening on port 8080");
log("HTTP and WS listening on port 8080");
});
const runJob = (job, sock) => {

View File

@ -10,11 +10,11 @@ class ImageStatsCommand extends Command {
"icon_url": this.client.user.avatarURL
},
"color": 16711680,
"description": `The bot is currently connected to ${image.connections.length} image server(s).`,
"description": `The bot is currently connected to ${image.connections.size} image server(s).`,
"fields": []
}
};
const servers = await image.getStatus();
const servers = await image.getRunning();
for (let i = 0; i < servers.length; i++) {
embed.embed.fields.push({
name: `Server ${i + 1}`,

View File

@ -79,7 +79,7 @@ class Shard extends Base {
}
res.write(`# HELP connected_workers Number of workers connected
# TYPE connected_workers gauge
connected_workers ${image.connections.length}
connected_workers ${image.connections.size}
# HELP running_jobs Number of running jobs on this worker
# TYPE running_jobs gauge
# HELP queued_jobs Number of queued jobs on this worker
@ -89,7 +89,7 @@ connected_workers ${image.connections.length}
# HELP command_count Number of times a command has been run
# TYPE command_count counter
`);
const servers = await image.getStatus();
const servers = await image.getRunning();
for (const [i, w] of servers.entries()) {
res.write(`running_jobs{worker="${i}"} ${w.runningJobs}\n`);
res.write(`queued_jobs{worker="${i}"} ${w.queued}\n`);

View File

@ -30,90 +30,80 @@ exports.repopulate = async () => {
return;
};
exports.getStatus = () => {
return new Promise((resolve, reject) => {
let serversLeft = this.connections.size;
const statuses = [];
exports.getRunning = async () => {
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--;
continue;
}
const controller = new AbortController(); // eslint-disable-line no-undef
const timeout = setTimeout(() => {
resolve(statuses);
}, 5000);
for (const address of this.connections.keys()) {
const connection = this.connections.get(address);
if (connection.readyState !== 1) {
controller.abort();
}, 2000);
try {
const statusRequest = await fetch(`http://${address}:8080/running`, { signal: controller.signal });
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;
}
fetch(`http://${address}:8080/running`).then(statusRequest => statusRequest.json()).then((status) => {
serversLeft--;
statuses.push(status);
if (!serversLeft) {
clearTimeout(timeout);
resolve(statuses);
}
return;
}).catch(e => {
if (e.code === "ECONNREFUSED") {
serversLeft--;
return;
}
reject(e);
});
throw e;
}
if (!serversLeft) {
clearTimeout(timeout);
resolve(statuses);
}
});
}
if (!serversLeft) {
return statuses;
} else {
throw new Error("Loop ended before all servers could be checked");
}
};
exports.connect = (server) => {
return new Promise((resolve, reject) => {
const connection = new WebSocket(`ws://${server}:8080/sock`);
const timeout = setTimeout(() => {
this.connections.delete(server);
reject(`Failed to connect to ${server}`);
}, 5000);
connection.once("open", () => {
clearTimeout(timeout);
});
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}`);
const image = await imageReq.buffer();
// 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.
const connection = new WebSocket(`ws://${server}:8080/sock`);
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}`);
const image = await imageReq.buffer();
// 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));
}
this.jobs[uuid].event.emit("image", image, imageReq.headers.get("ext"));
}
});
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");
} else if (opcode === 0x02) { // Job errored
if (this.jobs[uuid]) {
this.jobs[uuid].event.emit("error", new Error(req));
}
logger.log(`Lost connection to ${server}, attempting to reconnect...`);
this.connections.delete(server);
//this.connect(server);
});
this.connections.set(server, connection);
resolve();
}
});
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");
}
//logger.log(`Lost connection to ${server}, attempting to reconnect...`);
this.connections.delete(server);
});
this.connections.set(server, connection);
};
exports.disconnect = async () => {
@ -128,84 +118,72 @@ exports.disconnect = async () => {
return;
};
const getIdeal = () => {
return new Promise((resolve, reject) => {
let serversLeft = this.connections.size;
if (serversLeft === 0) {
for (const server of this.servers) {
this.connect(server).catch(e => {
logger.error(e);
});
}
serversLeft = this.connections.size;
}
const idealServers = [];
const timeout = setTimeout(async () => {
const getIdeal = async () => {
let serversLeft = this.connections.size;
if (serversLeft === 0) {
for (const server of this.servers) {
try {
const server = await chooseServer(idealServers);
resolve({ addr: server.addr, sock: this.connections.get(server.addr) });
await this.connect(server);
} catch (e) {
reject(e);
logger.error(e);
}
}, 5000);
for (const address of this.connections.keys()) {
const connection = this.connections.get(address);
if (connection.readyState !== 1) {
}
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--;
continue;
}
const controller = new AbortController(); // eslint-disable-line no-undef
const timeout = setTimeout(() => {
controller.abort();
}, 2000);
try {
const statusRequest = await fetch(`http://${address}:8080/status`, { signal: controller.signal });
clearTimeout(timeout);
const status = await statusRequest.text();
serversLeft--;
idealServers.push({
addr: address,
load: parseInt(status)
});
} catch (e) {
if (e.name === "AbortError") {
serversLeft--;
continue;
} else if (e.code === "ECONNREFUSED") {
serversLeft--;
continue;
}
fetch(`http://${address}:8080/status`).then(statusRequest => statusRequest.text()).then((status) => {
serversLeft--;
idealServers.push({
addr: address,
load: parseInt(status)
});
return;
}).then(async () => {
if (!serversLeft) {
clearTimeout(timeout);
const server = await chooseServer(idealServers);
resolve({ addr: server.addr, sock: this.connections.get(server.addr) });
}
}).catch(e => {
if (e.code === "ECONNREFUSED") {
serversLeft--;
return;
}
reject(e);
});
}
if (!serversLeft) {
throw e;
} finally {
clearTimeout(timeout);
chooseServer(idealServers).then(server => {
resolve({ addr: server.addr, sock: this.connections.get(server.addr) });
}).catch(e => reject(e));
}
});
}
if (!serversLeft) {
const server = await chooseServer(idealServers);
return { addr: server.addr, sock: this.connections.get(server.addr) };
} else {
throw new Error("Loop ended before all servers could be checked");
}
};
const start = (object, num) => {
return getIdeal().then(async (currentServer) => {
const data = Buffer.concat([Buffer.from([0x01 /* queue job */]), Buffer.from(num.length.toString()), Buffer.from(num), Buffer.from(JSON.stringify(object))]);
return new Promise((resolve) => {
currentServer.sock.send(data, () => {
resolve(currentServer.addr);
});
});
}).catch((result) => {
throw result;
}).then((addr) => {
const event = new EventEmitter();
return new Promise((resolve, reject) => {
event.once("uuid", (uuid) => resolve({ event, uuid, addr }));
event.once("error", reject);
this.jobs[num] = { event, addr };
});
}).then(data => {
delete this.jobs[num];
this.jobs[data.uuid] = { event: data.event, addr: data.addr };
return { uuid: data.uuid, event: data.event };
const start = async (object, num) => {
const currentServer = await 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);
});
delete this.jobs[num];
this.jobs[uuid] = { event: event, addr: currentServer.addr };
return { uuid: uuid, event: event };
};
exports.check = (cmd) => {

View File

@ -7,7 +7,7 @@ module.exports = (input, allowed) => {
const a = input[i];
if (a.startsWith("--") && !curr && isAllowed(a.slice(2).split("=")[0], allowed)) {
if (a.includes("=")) {
const [ arg, value ] = a.slice(2).split("=");
const [arg, value] = a.slice(2).split("=");
let ended = true;
if (value.startsWith("\"")) {
if (!value.endsWith("\"")) {