mrmBot-Matrix/utils/services/image.js

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 { request } from "undici";
import EventEmitter from "events";

// only requiring this to work around an issue regarding worker threads
const nodeRequire = createRequire(import.meta.url);
if (process.env.API_TYPE === "none") {
  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("../../config/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("../../config/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]?.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 request(`${process.env.AZURE_URL}/api/orchestrators/ImageOrchestrator`, { method: "POST", body: JSON.stringify(object) }).then(r => r.body.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;